diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStream.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStream.java index 375c94e2156d..d91303c5267e 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStream.java +++ b/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). diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStreamRequests.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStreamRequests.java index 7d51350571d2..d27b42d5a353 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStreamRequests.java +++ b/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 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 requests = new ArrayList<>(); + private final HashSet 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 requestsReadOnly() { + List 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 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 - } - } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStreamRequestsTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStreamRequestsTest.java index 150db4ed4815..c7bef43a4542 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStreamRequestsTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStreamRequestsTest.java @@ -18,6 +18,7 @@ package org.apache.beam.runners.dataflow.worker.windmill.client.grpc; import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; @@ -80,7 +81,7 @@ public void testQueuedRequest_globalRequestsFirstComparator() { requests.sort(GrpcGetDataStreamRequests.QueuedRequest.globalRequestsFirst()); // First one should be the global request. - assertTrue(requests.get(0).getDataRequest().isGlobal()); + assertTrue(requests.get(0).getKind() == GrpcGetDataStreamRequests.QueuedRequest.Kind.GLOBAL); } @Test @@ -95,9 +96,12 @@ public void testQueuedBatch_asGetDataRequest() { .setWorkToken(1L) .setMaxBytes(Long.MAX_VALUE) .build(); - queuedBatch.addRequest( - GrpcGetDataStreamRequests.QueuedRequest.forComputation( - 1, "computation1", keyedGetDataRequest1, DEADLINE_SECONDS)); + assertTrue( + queuedBatch.tryAddRequest( + GrpcGetDataStreamRequests.QueuedRequest.forComputation( + 1, "computation1", keyedGetDataRequest1, DEADLINE_SECONDS), + Integer.MAX_VALUE, + Long.MAX_VALUE)); Windmill.KeyedGetDataRequest keyedGetDataRequest2 = Windmill.KeyedGetDataRequest.newBuilder() @@ -107,9 +111,12 @@ public void testQueuedBatch_asGetDataRequest() { .setWorkToken(2L) .setMaxBytes(Long.MAX_VALUE) .build(); - queuedBatch.addRequest( - GrpcGetDataStreamRequests.QueuedRequest.forComputation( - 2, "computation2", keyedGetDataRequest2, DEADLINE_SECONDS)); + assertTrue( + queuedBatch.tryAddRequest( + GrpcGetDataStreamRequests.QueuedRequest.forComputation( + 2, "computation2", keyedGetDataRequest2, DEADLINE_SECONDS), + Integer.MAX_VALUE, + Long.MAX_VALUE)); Windmill.GlobalDataRequest globalDataRequest = Windmill.GlobalDataRequest.newBuilder() @@ -120,12 +127,15 @@ public void testQueuedBatch_asGetDataRequest() { .build()) .setComputationId("computation1") .build(); - queuedBatch.addRequest( - GrpcGetDataStreamRequests.QueuedRequest.global(3, globalDataRequest, DEADLINE_SECONDS)); + assertTrue( + queuedBatch.tryAddRequest( + GrpcGetDataStreamRequests.QueuedRequest.global(3, globalDataRequest, DEADLINE_SECONDS), + Integer.MAX_VALUE, + Long.MAX_VALUE)); Windmill.StreamingGetDataRequest getDataRequest = queuedBatch.asGetDataRequest(); - assertThat(getDataRequest.getRequestIdCount()).isEqualTo(3); + assertThat(getDataRequest.getRequestIdList()).containsExactly(3L, 1L, 2L); assertThat(getDataRequest.getGlobalDataRequestList()).containsExactly(globalDataRequest); assertThat(getDataRequest.getStateRequestList()) .containsExactly( @@ -153,4 +163,134 @@ public void testQueuedBatch_notifyFailed_throwsWindmillStreamShutdownExceptionOn queuedBatch.notifyFailed(); waitFuture.join(); } + + @Test + public void testQueuedBatch_tryAddRequest_exceedsMaxCount() { + GrpcGetDataStreamRequests.QueuedBatch queuedBatch = new GrpcGetDataStreamRequests.QueuedBatch(); + Windmill.KeyedGetDataRequest keyedGetDataRequest = + Windmill.KeyedGetDataRequest.newBuilder() + .setKey(ByteString.EMPTY) + .setCacheToken(1L) + .setShardingKey(1L) + .setWorkToken(1L) + .build(); + + // Add one request successfully. + assertTrue( + queuedBatch.tryAddRequest( + GrpcGetDataStreamRequests.QueuedRequest.forComputation( + 1, "computation1", keyedGetDataRequest, DEADLINE_SECONDS), + 1, + Long.MAX_VALUE)); + + // Adding another request should fail due to max count. + assertFalse( + queuedBatch.tryAddRequest( + GrpcGetDataStreamRequests.QueuedRequest.forComputation( + 2, "computation1", keyedGetDataRequest, DEADLINE_SECONDS), + 1, + Long.MAX_VALUE)); + } + + @Test + public void testQueuedBatch_tryAddRequest_exceedsMaxBytes() { + GrpcGetDataStreamRequests.QueuedBatch queuedBatch = new GrpcGetDataStreamRequests.QueuedBatch(); + Windmill.KeyedGetDataRequest keyedGetDataRequest = + Windmill.KeyedGetDataRequest.newBuilder() + .setKey(ByteString.EMPTY) + .setCacheToken(1L) + .setShardingKey(1L) + .setWorkToken(1L) + .build(); + + // Add one request successfully. + assertTrue( + queuedBatch.tryAddRequest( + GrpcGetDataStreamRequests.QueuedRequest.forComputation( + 1, "computation1", keyedGetDataRequest, DEADLINE_SECONDS), + Integer.MAX_VALUE, + 80L)); + + // Adding another request should fail due to max bytes. + assertFalse( + queuedBatch.tryAddRequest( + GrpcGetDataStreamRequests.QueuedRequest.forComputation( + 2, "computation1", keyedGetDataRequest, DEADLINE_SECONDS), + Integer.MAX_VALUE, + 80L)); + + Windmill.GlobalDataRequest globalDataRequest = + Windmill.GlobalDataRequest.newBuilder() + .setDataId( + Windmill.GlobalDataId.newBuilder() + .setTag("globalData") + .setVersion(ByteString.EMPTY) + .build()) + .setComputationId("computation1") + .build(); + assertFalse( + queuedBatch.tryAddRequest( + GrpcGetDataStreamRequests.QueuedRequest.global(3, globalDataRequest, DEADLINE_SECONDS), + Integer.MAX_VALUE, + 80)); + } + + @Test + public void testQueuedBatch_tryAddRequest_duplicateWorkToken() { + GrpcGetDataStreamRequests.QueuedBatch queuedBatch = new GrpcGetDataStreamRequests.QueuedBatch(); + Windmill.KeyedGetDataRequest keyedGetDataRequest1 = + Windmill.KeyedGetDataRequest.newBuilder() + .setKey(ByteString.EMPTY) + .setCacheToken(1L) + .setShardingKey(1L) + .setWorkToken(1L) + .build(); + + Windmill.KeyedGetDataRequest keyedGetDataRequest2 = + Windmill.KeyedGetDataRequest.newBuilder() + .setKey(ByteString.EMPTY) + .setCacheToken(2L) + .setShardingKey(2L) + .setWorkToken(1L) + .build(); + + // Add one request successfully. + assertTrue( + queuedBatch.tryAddRequest( + GrpcGetDataStreamRequests.QueuedRequest.forComputation( + 1, "computation1", keyedGetDataRequest1, DEADLINE_SECONDS), + Integer.MAX_VALUE, + Long.MAX_VALUE)); + + // Adding another request with same work token should fail. + assertFalse( + queuedBatch.tryAddRequest( + GrpcGetDataStreamRequests.QueuedRequest.forComputation( + 2, "computation1", keyedGetDataRequest2, DEADLINE_SECONDS), + Integer.MAX_VALUE, + Long.MAX_VALUE)); + } + + @Test + public void testQueuedBatch_tryAddRequest_afterFinalized() { + GrpcGetDataStreamRequests.QueuedBatch queuedBatch = new GrpcGetDataStreamRequests.QueuedBatch(); + Windmill.KeyedGetDataRequest keyedGetDataRequest = + Windmill.KeyedGetDataRequest.newBuilder() + .setKey(ByteString.EMPTY) + .setCacheToken(1L) + .setShardingKey(1L) + .setWorkToken(1L) + .setMaxBytes(Long.MAX_VALUE) + .build(); + + queuedBatch.markFinalized(); + + // Adding request after finalization should fail. + assertFalse( + queuedBatch.tryAddRequest( + GrpcGetDataStreamRequests.QueuedRequest.forComputation( + 1, "computation1", keyedGetDataRequest, DEADLINE_SECONDS), + Integer.MAX_VALUE, + Long.MAX_VALUE)); + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStreamTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStreamTest.java index 849b2612cecf..fccc32af4c7d 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStreamTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStreamTest.java @@ -27,6 +27,7 @@ import java.io.IOException; import java.time.Duration; +import java.util.ArrayList; import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; @@ -235,6 +236,56 @@ public void testRequestKeyedData_sendOnShutdownStreamThrowsWindmillStreamShutdow } } + @Test + public void testRequestKeyedData_multipleRequestsSameWorkItemSeparateBatches() + throws InterruptedException { + GrpcGetDataStream getDataStream = createGetDataStream(); + FakeWindmillGrpcService.GetDataStreamInfo streamInfo = waitForConnectionAndConsumeHeader(); + + final CountDownLatch requestStarter = new CountDownLatch(1); + + // Get a bunch of threads ready to send a request with the same work token. These should racily + // attempt to batch but be prevented due to work token separation logic. + // These will block until they are successfully sent. + List> futures = new ArrayList<>(); + final Windmill.KeyedGetDataRequest keyedGetDataRequest = createTestRequest(1); + for (int i = 0; i < 10; ++i) { + futures.add( + CompletableFuture.supplyAsync( + () -> { + try { + requestStarter.await(); + return getDataStream.requestKeyedData("computationId", keyedGetDataRequest); + } catch (Exception e) { + throw new RuntimeException(e); + } + })); + } + + // Unblock and verify that 10 requests are made and not batched. + requestStarter.countDown(); + for (int i = 0; i < 10; ++i) { + Windmill.StreamingGetDataRequest request = streamInfo.requests.take(); + assertEquals(1, request.getRequestIdCount()); + assertEquals(keyedGetDataRequest, request.getStateRequest(0).getRequests(0)); + } + + // Send the responses. + Windmill.KeyedGetDataResponse keyedGetDataResponse = createTestResponse(1); + for (int i = 0; i < 10; ++i) { + streamInfo.responseObserver.onNext( + Windmill.StreamingGetDataResponse.newBuilder() + .addRequestId(i + 1) + .addSerializedResponse(keyedGetDataResponse.toByteString()) + .build()); + } + + for (CompletableFuture future : futures) { + assertThat(future.join()).isEqualTo(keyedGetDataResponse); + } + getDataStream.shutdown(); + } + @Test public void testRequestKeyedData_reconnectOnStreamError() throws InterruptedException { GrpcGetDataStream getDataStream = createGetDataStream();