[Dataflow Streaming] Enforce that get data requests for the same work item are not batched.

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStream.java

@@ -193,7 +193,7 @@ class GetDataPhysicalStreamHandler extends PhysicalStreamHandler {
     public void sendBatch(QueuedBatch batch) throws WindmillStreamShutdownException {
       // Synchronization of pending inserts is necessary with send to ensure duplicates are not
       // sent on stream reconnect.
-      for (QueuedRequest request : batch.requestsReadOnly()) {
+      for (QueuedRequest request : batch.requestsView()) {
         boolean alreadyPresent = pending.put(request.id(), request.getResponseStream()) != null;
         verify(!alreadyPresent, "Request already sent, id: %s", request.id());
       }
@@ -277,7 +277,7 @@ protected synchronized void onFlushPending(boolean isNewStream)
     }
     while (!batches.isEmpty()) {
       QueuedBatch batch = checkNotNull(batches.peekFirst());
-      verify(!batch.isEmpty());
+      verify(batch.requestsCount() > 0);
       if (!batch.isFinalized()) {
         break;
       }
@@ -481,17 +481,15 @@ private void queueRequestAndWait(QueuedRequest request)
 
       batch = batches.isEmpty() ? null : batches.getLast();
       if (batch == null
-          || batch.isFinalized()
-          || batch.requestsCount() >= streamingRpcBatchLimit
-          || batch.byteSize() + request.byteSize() > AbstractWindmillStream.RPC_STREAM_CHUNK_SIZE) {
-        if (batch != null) {
-          prevBatch = batch;
-        }
+          || !batch.tryAddRequest(
+              request, streamingRpcBatchLimit, AbstractWindmillStream.RPC_STREAM_CHUNK_SIZE)) {
+        // We need a new batch.
+        prevBatch = batch; // may be null
         batch = new QueuedBatch();
         batches.addLast(batch);
         responsibleForSend = true;
+        verify(batch.tryAddRequest(request, Integer.MAX_VALUE, Long.MAX_VALUE));
       }
-      batch.addRequest(request);
     }
     if (responsibleForSend) {
       if (prevBatch == null) {
@@ -531,7 +529,7 @@ private synchronized void trySendBatch(QueuedBatch batch) throws WindmillStreamS
       // an error and will
       // resend requests (possibly with new batching).
       verify(batch == batches.pollFirst());
-      verify(!batch.isEmpty());
+      verify(batch.requestsCount() > 0);
       currentGetDataPhysicalStream.sendBatch(batch);
       // Notify all waiters with requests in this batch as well as the sender
       // of the next batch (if one exists).

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStreamRequests.java

@@ -19,18 +19,18 @@
 
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList.toImmutableList;
 
-import com.google.auto.value.AutoOneOf;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.HashSet;
 import java.util.List;
 import java.util.concurrent.CountDownLatch;
-import java.util.stream.Stream;
+import javax.annotation.Nullable;
 import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
-import org.apache.beam.runners.dataflow.worker.windmill.Windmill.ComputationGetDataRequest;
 import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalDataRequest;
 import org.apache.beam.runners.dataflow.worker.windmill.Windmill.KeyedGetDataRequest;
 import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStreamShutdownException;
+import org.apache.beam.sdk.util.Preconditions;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -46,15 +46,42 @@ private static String debugFormat(long value) {
     return String.format("%016x", value);
   }
 
+  static class ComputationAndKeyRequest {
+    private final String computation;
+    private final KeyedGetDataRequest request;
+
+    ComputationAndKeyRequest(String computation, KeyedGetDataRequest request) {
+      this.computation = computation;
+      this.request = request;
+    }
+
+    String getComputation() {
+      return computation;
+    }
+
+    KeyedGetDataRequest getKeyedGetDataRequest() {
+      return request;
+    }
+  }
+
   static class QueuedRequest {
     private final long id;
-    private final ComputationOrGlobalDataRequest dataRequest;
+    private final @Nullable ComputationAndKeyRequest computationAndKeyRequest;
+    private final @Nullable GlobalDataRequest globalDataRequest;
     private AppendableInputStream responseStream;
 
+    private QueuedRequest(long id, GlobalDataRequest globalDataRequest, long deadlineSeconds) {
+      this.id = id;
+      this.computationAndKeyRequest = null;
+      this.globalDataRequest = globalDataRequest;
+      responseStream = new AppendableInputStream(deadlineSeconds);
+    }
+
     private QueuedRequest(
-        long id, ComputationOrGlobalDataRequest dataRequest, long deadlineSeconds) {
+        long id, ComputationAndKeyRequest computationAndKeyRequest, long deadlineSeconds) {
       this.id = id;
-      this.dataRequest = dataRequest;
+      this.computationAndKeyRequest = computationAndKeyRequest;
+      this.globalDataRequest = null;
       responseStream = new AppendableInputStream(deadlineSeconds);
     }
 
@@ -63,27 +90,19 @@ static QueuedRequest forComputation(
         String computation,
         KeyedGetDataRequest keyedGetDataRequest,
         long deadlineSeconds) {
-      ComputationGetDataRequest computationGetDataRequest =
-          ComputationGetDataRequest.newBuilder()
-              .setComputationId(computation)
-              .addRequests(keyedGetDataRequest)
-              .build();
       return new QueuedRequest(
-          id,
-          ComputationOrGlobalDataRequest.computation(computationGetDataRequest),
-          deadlineSeconds);
+          id, new ComputationAndKeyRequest(computation, keyedGetDataRequest), deadlineSeconds);
     }
 
     static QueuedRequest global(
         long id, GlobalDataRequest globalDataRequest, long deadlineSeconds) {
-      return new QueuedRequest(
-          id, ComputationOrGlobalDataRequest.global(globalDataRequest), deadlineSeconds);
+      return new QueuedRequest(id, globalDataRequest, deadlineSeconds);
     }
 
     static Comparator<QueuedRequest> globalRequestsFirst() {
       return (QueuedRequest r1, QueuedRequest r2) -> {
-        boolean r1gd = r1.dataRequest.isGlobal();
-        boolean r2gd = r2.dataRequest.isGlobal();
+        boolean r1gd = r1.getKind() == Kind.GLOBAL;
+        boolean r2gd = r2.getKind() == Kind.GLOBAL;
         return r1gd == r2gd ? 0 : (r1gd ? -1 : 1);
       };
     }
@@ -93,7 +112,13 @@ long id() {
     }
 
     long byteSize() {
-      return dataRequest.serializedSize();
+      if (globalDataRequest != null) {
+        return globalDataRequest.getSerializedSize();
+      }
+      Preconditions.checkStateNotNull(computationAndKeyRequest);
+      return 10L
+          + computationAndKeyRequest.request.getSerializedSize()
+          + computationAndKeyRequest.getComputation().length();
     }
 
     AppendableInputStream getResponseStream() {
@@ -104,22 +129,56 @@ void resetResponseStream() {
       this.responseStream = new AppendableInputStream(responseStream.getDeadlineSeconds());
     }
 
-    public ComputationOrGlobalDataRequest getDataRequest() {
-      return dataRequest;
+    enum Kind {
+      COMPUTATION_AND_KEY_REQUEST,
+      GLOBAL
+    }
+
+    Kind getKind() {
+      return computationAndKeyRequest != null ? Kind.COMPUTATION_AND_KEY_REQUEST : Kind.GLOBAL;
+    }
+
+    ComputationAndKeyRequest getComputationAndKeyRequest() {
+      return Preconditions.checkStateNotNull(computationAndKeyRequest);
+    }
+
+    GlobalDataRequest getGlobalDataRequest() {
+      return Preconditions.checkStateNotNull(globalDataRequest);
     }
 
     void addToStreamingGetDataRequest(Windmill.StreamingGetDataRequest.Builder builder) {
       builder.addRequestId(id);
-      if (dataRequest.isForComputation()) {
-        builder.addStateRequest(dataRequest.computation());
-      } else {
-        builder.addGlobalDataRequest(dataRequest.global());
+      switch (getKind()) {
+        case COMPUTATION_AND_KEY_REQUEST:
+          ComputationAndKeyRequest request = getComputationAndKeyRequest();
+          builder
+              .addStateRequestBuilder()
+              .setComputationId(request.getComputation())
+              .addRequests(request.request);
+          break;
+        case GLOBAL:
+          builder.addGlobalDataRequest(getGlobalDataRequest());
+          break;
       }
     }
 
     @Override
     public final String toString() {
-      return "QueuedRequest{" + "dataRequest=" + dataRequest + ", id=" + id + '}';
+      StringBuilder result = new StringBuilder("QueuedRequest{id=").append(id).append(", ");
+      if (getKind() == Kind.GLOBAL) {
+        result.append("GetSideInput=").append(getGlobalDataRequest());
+      } else {
+        KeyedGetDataRequest key = getComputationAndKeyRequest().request;
+        result
+            .append("KeyedGetState=[shardingKey=")
+            .append(debugFormat(key.getShardingKey()))
+            .append("cacheToken=")
+            .append(debugFormat(key.getCacheToken()))
+            .append("workToken")
+            .append(debugFormat(key.getWorkToken()))
+            .append("]");
+      }
+      return result.append('}').toString();
     }
   }
 
@@ -128,13 +187,14 @@ public final String toString() {
    */
   static class QueuedBatch {
     private final List<QueuedRequest> requests = new ArrayList<>();
+    private final HashSet<Long> workTokens = new HashSet<>();
     private final CountDownLatch sent = new CountDownLatch(1);
     private long byteSize = 0;
     private volatile boolean finalized = false;
     private volatile boolean failed = false;
 
     /** Returns a read-only view of requests. */
-    List<QueuedRequest> requestsReadOnly() {
+    List<QueuedRequest> requestsView() {
       return Collections.unmodifiableList(requests);
     }
 
@@ -155,18 +215,10 @@ Windmill.StreamingGetDataRequest asGetDataRequest() {
       return builder.build();
     }
 
-    boolean isEmpty() {
-      return requests.isEmpty();
-    }
-
     int requestsCount() {
       return requests.size();
     }
 
-    long byteSize() {
-      return byteSize;
-    }
-
     boolean isFinalized() {
       return finalized;
     }
@@ -176,9 +228,26 @@ void markFinalized() {
     }
 
     /** Adds a request to the batch. */
-    void addRequest(QueuedRequest request) {
+    boolean tryAddRequest(QueuedRequest request, int countLimit, long byteLimit) {
+      if (finalized) {
+        return false;
+      }
+      if (requests.size() >= countLimit) {
+        return false;
+      }
+      long estimatedBytes = request.byteSize();
+      if (byteSize + estimatedBytes >= byteLimit) {
+        return false;
+      }
+
+      if (request.getKind() == QueuedRequest.Kind.COMPUTATION_AND_KEY_REQUEST
+          && !workTokens.add(request.getComputationAndKeyRequest().request.getWorkToken())) {
+        return false;
+      }
+      // At this point we have added to work items so we must accept the item.
       requests.add(request);
-      byteSize += request.byteSize();
+      byteSize += estimatedBytes;
+      return true;
     }
 
     /**
@@ -227,75 +296,9 @@ void waitForSendOrFailNotification()
 
     private ImmutableList<String> createStreamCancelledErrorMessages() {
       return requests.stream()
-          .flatMap(
-              request -> {
-                switch (request.getDataRequest().getKind()) {
-                  case GLOBAL:
-                    return Stream.of("GetSideInput=" + request.getDataRequest().global());
-                  case COMPUTATION:
-                    return request.getDataRequest().computation().getRequestsList().stream()
-                        .map(
-                            keyedRequest ->
-                                "KeyedGetState=["
-                                    + "shardingKey="
-                                    + debugFormat(keyedRequest.getShardingKey())
-                                    + "cacheToken="
-                                    + debugFormat(keyedRequest.getCacheToken())
-                                    + "workToken"
-                                    + debugFormat(keyedRequest.getWorkToken())
-                                    + "]");
-                  default:
-                    // Will never happen switch is exhaustive.
-                    throw new IllegalStateException();
-                }
-              })
+          .map(QueuedRequest::toString)
           .limit(STREAM_CANCELLED_ERROR_LOG_LIMIT)
           .collect(toImmutableList());
     }
   }
-
-  @AutoOneOf(ComputationOrGlobalDataRequest.Kind.class)
-  abstract static class ComputationOrGlobalDataRequest {
-    static ComputationOrGlobalDataRequest computation(
-        ComputationGetDataRequest computationGetDataRequest) {
-      return AutoOneOf_GrpcGetDataStreamRequests_ComputationOrGlobalDataRequest.computation(
-          computationGetDataRequest);
-    }
-
-    static ComputationOrGlobalDataRequest global(GlobalDataRequest globalDataRequest) {
-      return AutoOneOf_GrpcGetDataStreamRequests_ComputationOrGlobalDataRequest.global(
-          globalDataRequest);
-    }
-
-    abstract Kind getKind();
-
-    abstract ComputationGetDataRequest computation();
-
-    abstract GlobalDataRequest global();
-
-    boolean isGlobal() {
-      return getKind() == Kind.GLOBAL;
-    }
-
-    boolean isForComputation() {
-      return getKind() == Kind.COMPUTATION;
-    }
-
-    long serializedSize() {
-      switch (getKind()) {
-        case GLOBAL:
-          return global().getSerializedSize();
-        case COMPUTATION:
-          return computation().getSerializedSize();
-          // this will never happen since the switch is exhaustive.
-        default:
-          throw new UnsupportedOperationException("unknown dataRequest type.");
-      }
-    }
-
-    enum Kind {
-      COMPUTATION,
-      GLOBAL
-    }
-  }
 }