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

Such batches are rejected by the backend so we should prevent them to avoid stuckness. Parallel requests for the same work item are unexpected but could be caused by bugs in the harness or by incorrect parallel state fetching from a single bundle.

Author: scwhittle

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;
       }
@@ -482,17 +482,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) {
@@ -532,7 +530,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
-    }
-  }
 }