/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.beam.runners.dataflow.worker.windmill.client.grpc;

import static org.junit.Assert.*;

import java.io.InputStream;
import java.io.SequenceInputStream;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1ImplBase;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.CommitStatus;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.ComputationGetDataRequest;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.ComputationHeartbeatRequest;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.ComputationWorkItemMetadata;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GetWorkRequest;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalData;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalDataId;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalDataRequest;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.HeartbeatRequest;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.KeyedGetDataRequest;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.KeyedGetDataResponse;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.LatencyAttribution;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.StreamingCommitRequestChunk;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.StreamingCommitResponse;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.StreamingCommitWorkRequest;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.StreamingGetDataRequest;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.StreamingGetDataResponse;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.StreamingGetWorkRequest;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.StreamingGetWorkResponseChunk;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.TagValue;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.Value;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItem;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItemCommitRequest;
import org.apache.beam.runners.dataflow.worker.windmill.WindmillApplianceGrpc;
import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.CommitWorkStream;
import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.GetDataStream;
import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.GetWorkStream;
import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStreamShutdownException;
import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillChannels;
import org.apache.beam.runners.dataflow.worker.windmill.testing.FakeWindmillStubFactory;
import org.apache.beam.runners.dataflow.worker.windmill.testing.FakeWindmillStubFactoryFactory;
import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.CallOptions;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Channel;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ClientCall;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ClientInterceptor;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ClientInterceptors;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Deadline;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ManagedChannel;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.MethodDescriptor;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Server;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.inprocess.InProcessChannelBuilder;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.inprocess.InProcessServerBuilder;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.testing.GrpcCleanupRule;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.util.MutableHandlerRegistry;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.hamcrest.Matchers;
import org.joda.time.Instant;
import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ErrorCollector;
import org.junit.rules.Timeout;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** Unit tests for {@link GrpcWindmillServer}. */
@RunWith(JUnit4.class)
@SuppressWarnings({
  "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
})
public class GrpcWindmillServerTest {

  private static final Logger LOG = LoggerFactory.getLogger(GrpcWindmillServerTest.class);
  private static final int STREAM_CHUNK_SIZE = 2 << 20;
  private final long clientId = 10L;
  private final Set<ManagedChannel> openedChannels = new HashSet<>();
  private final MutableHandlerRegistry serviceRegistry = new MutableHandlerRegistry();

  @Rule
  public transient Timeout globalTimeout =
      Timeout.builder().withTimeout(10, TimeUnit.MINUTES).withLookingForStuckThread(true).build();

  @Rule public GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
  @Rule public ErrorCollector errorCollector = new ErrorCollector();
  private Server server;
  private GrpcWindmillServer client;
  private int remainingErrors = 20;

  @Before
  public void setUp() throws Exception {
    startServerAndClient(new ArrayList<>());
  }

  @After
  public void tearDown() throws Exception {
    server.shutdownNow();
    openedChannels.forEach(ManagedChannel::shutdownNow);
  }

  private void startServerAndClient(List<String> experiments) throws Exception {
    String name = "Fake server for " + getClass();
    this.server =
        grpcCleanup.register(
            InProcessServerBuilder.forName(name)
                .fallbackHandlerRegistry(serviceRegistry)
                .executor(Executors.newFixedThreadPool(1))
                .build()
                .start());

    this.client =
        GrpcWindmillServer.newTestInstance(
            name,
            experiments,
            clientId,
            new FakeWindmillStubFactoryFactory(
                new FakeWindmillStubFactory(
                    () -> {
                      ManagedChannel channel =
                          grpcCleanup.register(WindmillChannels.inProcessChannel(name));
                      openedChannels.add(channel);
                      return channel;
                    })));
  }

  private <Stream extends StreamObserver> void maybeInjectError(Stream stream) {
    if (remainingErrors > 0 && ThreadLocalRandom.current().nextInt(20) == 0) {
      try {
        stream.onError(new RuntimeException("oops"));
        remainingErrors--;
      } catch (IllegalStateException e) {
        // The stream is already closed.
      }
    }
  }

  @Test
  public void testWindmillApplianceGrpcDeadlineInterceptor() {
    Windmill.GetWorkResponse getWorkResponse =
        Windmill.GetWorkResponse.newBuilder()
            .addWork(
                Windmill.ComputationWorkItems.newBuilder()
                    .setComputationId("someComputation1")
                    .build())
            .build();

    serviceRegistry.addService(
        new WindmillApplianceGrpc.WindmillApplianceImplBase() {
          @Override
          public void getWork(
              Windmill.GetWorkRequest request,
              StreamObserver<Windmill.GetWorkResponse> responseObserver) {
            responseObserver.onNext(getWorkResponse);
            responseObserver.onCompleted();
          }
        });

    ClientInterceptor testInterceptor =
        new ClientInterceptor() {
          Deadline deadline;

          @Override
          public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
              MethodDescriptor<ReqT, RespT> methodDescriptor,
              CallOptions callOptions,
              Channel channel) {
            assertNotNull(callOptions.getDeadline());
            if (deadline != null) {
              // We want to assert that a new deadline is being created on every call since
              // deadlines are absolute.
              assertNotEquals(deadline, callOptions.getDeadline());
            }

            deadline = callOptions.getDeadline();
            // Just forward the call once we record the deadline
            return channel.newCall(methodDescriptor, callOptions);
          }
        };

    Channel inprocessChannel =
        ClientInterceptors.intercept(
            InProcessChannelBuilder.forName("Fake server for " + getClass())
                .directExecutor()
                .build(),
            testInterceptor);

    this.client =
        GrpcWindmillServer.newApplianceTestInstance(
            inprocessChannel,
            new FakeWindmillStubFactoryFactory(
                new FakeWindmillStubFactory(() -> (ManagedChannel) inprocessChannel)));

    Windmill.GetWorkResponse response1 = client.getWork(GetWorkRequest.getDefaultInstance());
    Windmill.GetWorkResponse response2 = client.getWork(GetWorkRequest.getDefaultInstance());

    // Assert that the request/responses are being forwarded correctly.
    assertEquals(response1, getWorkResponse);
    assertEquals(response2, getWorkResponse);
  }

  @Test
  public void testStreamingGetWork() throws Exception {
    // This fake server returns an infinite stream of identical WorkItems, obeying the request size
    // limits set by the client.
    serviceRegistry.addService(
        new CloudWindmillServiceV1Alpha1ImplBase() {
          @Override
          public StreamObserver<StreamingGetWorkRequest> getWorkStream(
              StreamObserver<StreamingGetWorkResponseChunk> responseObserver) {
            return new StreamObserver<StreamingGetWorkRequest>() {
              final ResponseErrorInjector injector = new ResponseErrorInjector(responseObserver);
              boolean sawHeader = false;

              @Override
              public void onNext(StreamingGetWorkRequest request) {
                maybeInjectError(responseObserver);

                try {
                  long maxItems;
                  if (!sawHeader) {
                    errorCollector.checkThat(
                        request.getRequest(),
                        Matchers.equalTo(
                            GetWorkRequest.newBuilder()
                                .setClientId(10)
                                .setJobId("job")
                                .setProjectId("project")
                                .setWorkerId("worker")
                                .setMaxItems(3)
                                .setMaxBytes(10000)
                                .build()));
                    sawHeader = true;
                    maxItems = request.getRequest().getMaxItems();
                  } else {
                    maxItems = request.getRequestExtension().getMaxItems();
                  }

                  for (int item = 0; item < maxItems; item++) {
                    long id = ThreadLocalRandom.current().nextLong();
                    ByteString serializedResponse =
                        WorkItem.newBuilder()
                            .setKey(ByteString.copyFromUtf8("somewhat_long_key"))
                            .setWorkToken(id)
                            .setShardingKey(id)
                            .build()
                            .toByteString();

                    // Break the WorkItem into smaller chunks to test chunking code.
                    for (int i = 0; i < serializedResponse.size(); i += 10) {
                      int end = Math.min(serializedResponse.size(), i + 10);
                      StreamingGetWorkResponseChunk.Builder builder =
                          StreamingGetWorkResponseChunk.newBuilder()
                              .setStreamId(id)
                              .addSerializedWorkItem(serializedResponse.substring(i, end))
                              .setRemainingBytesForWorkItem(serializedResponse.size() - end);

                      if (i == 0) {
                        builder.setComputationMetadata(
                            ComputationWorkItemMetadata.newBuilder()
                                .setComputationId("comp")
                                .setDependentRealtimeInputWatermark(17000)
                                .setInputDataWatermark(18000));
                      }

                      try {
                        responseObserver.onNext(builder.build());
                      } catch (IllegalStateException e) {
                        // Client closed stream, we're done.
                        return;
                      }
                    }
                  }
                } catch (Exception e) {
                  errorCollector.addError(e);
                }
              }

              @Override
              public void onError(Throwable throwable) {}

              @Override
              public void onCompleted() {
                injector.cancel();
                responseObserver.onCompleted();
              }
            };
          }
        });

    // Read the stream of WorkItems until 100 of them are received.
    CountDownLatch latch = new CountDownLatch(100);
    GetWorkStream stream =
        client.getWorkStream(
            GetWorkRequest.newBuilder().setClientId(10).setMaxItems(3).setMaxBytes(10000).build(),
            (String computation,
                @Nullable Instant inputDataWatermark,
                Instant synchronizedProcessingTime,
                boolean drainMode,
                WorkItem workItem,
                long serializedWorkItemSize,
                ImmutableList<Long> appliedFinalizeIds,
                ImmutableList<LatencyAttribution> getWorkStreamLatencies) -> {
              latch.countDown();
              assertEquals(inputDataWatermark, new Instant(18));
              assertEquals(synchronizedProcessingTime, new Instant(17));
              assertEquals(workItem.getKey(), ByteString.copyFromUtf8("somewhat_long_key"));
            });
    assertTrue(latch.await(30, TimeUnit.SECONDS));

    stream.halfClose();
    assertTrue(stream.awaitTermination(30, TimeUnit.SECONDS));
  }

  @Test
  public void testStreamingGetWorkBatchedGetWorkResponse() throws Exception {
    ConcurrentHashMap<Long, Boolean> sentResponseIds = new ConcurrentHashMap<>();
    // This fake server returns an infinite stream of identical WorkItems, obeying the request size
    // limits set by the client.
    serviceRegistry.addService(
        new CloudWindmillServiceV1Alpha1ImplBase() {
          @Override
          public StreamObserver<StreamingGetWorkRequest> getWorkStream(
              StreamObserver<StreamingGetWorkResponseChunk> responseObserver) {
            return new StreamObserver<StreamingGetWorkRequest>() {
              final ResponseErrorInjector injector = new ResponseErrorInjector(responseObserver);
              boolean sawHeader = false;

              @Override
              public void onNext(StreamingGetWorkRequest request) {
                maybeInjectError(responseObserver);

                try {
                  long maxItems;
                  if (!sawHeader) {
                    errorCollector.checkThat(
                        request.getRequest(),
                        Matchers.equalTo(
                            GetWorkRequest.newBuilder()
                                .setClientId(10)
                                .setJobId("job")
                                .setProjectId("project")
                                .setWorkerId("worker")
                                .setMaxItems(10)
                                .setMaxBytes(10000)
                                .build()));
                    sawHeader = true;
                    maxItems = request.getRequest().getMaxItems();
                  } else {
                    maxItems = request.getRequestExtension().getMaxItems();
                  }

                  ArrayDeque<ByteString> serializedResponses = new ArrayDeque<>((int) maxItems);
                  for (int item = 0; item < maxItems; item++) {
                    long id = ThreadLocalRandom.current().nextLong();
                    ByteString serializedResponse =
                        WorkItem.newBuilder()
                            .setKey(ByteString.copyFromUtf8("somewhat_long_key"))
                            .setWorkToken(id)
                            .setShardingKey(id)
                            .build()
                            .toByteString();
                    serializedResponses.push(serializedResponse);
                    sentResponseIds.put(id, true);
                  }

                  int loop = 0;
                  // Send a mix of chunks with single messages and partial messages
                  while (serializedResponses.size() > 0) {
                    StreamingGetWorkResponseChunk.Builder responseChunk =
                        StreamingGetWorkResponseChunk.newBuilder()
                            .setStreamId(ThreadLocalRandom.current().nextLong())
                            .setComputationMetadata(
                                ComputationWorkItemMetadata.newBuilder()
                                    .setComputationId("comp")
                                    .setDependentRealtimeInputWatermark(17000)
                                    .setInputDataWatermark(18000)
                                    .setDrainMode(true));
                    int loopVariant = loop % 3;
                    if (loopVariant < 1) {
                      responseChunk.addSerializedWorkItem(serializedResponses.pop());
                    }
                    if (serializedResponses.size() > 0 && loopVariant < 2) {
                      ByteString serializedResponse = serializedResponses.pop();
                      int end = serializedResponse.size() / 2;
                      responseChunk.addSerializedWorkItem(serializedResponse.substring(0, end));
                      responseChunk.setRemainingBytesForWorkItem(serializedResponse.size() - end);
                      try {
                        responseObserver.onNext(responseChunk.build());
                      } catch (IllegalStateException e) {
                        // Client closed stream, we're done.
                        return;
                      }
                      responseChunk.clearSerializedWorkItem();
                      responseChunk.clearRemainingBytesForWorkItem();

                      responseChunk.addSerializedWorkItem(serializedResponse.substring(end));
                    }
                    if (serializedResponses.size() > 0 && loopVariant < 3) {
                      responseChunk.addSerializedWorkItem(serializedResponses.pop());
                    }
                    try {
                      responseObserver.onNext(responseChunk.build());
                    } catch (IllegalStateException e) {
                      // Client closed stream, we're done.
                      return;
                    }
                    ++loop;
                  }
                } catch (Exception e) {
                  errorCollector.addError(e);
                }
              }

              @Override
              public void onError(Throwable throwable) {}

              @Override
              public void onCompleted() {
                injector.cancel();
                responseObserver.onCompleted();
              }
            };
          }
        });

    // Read the stream of WorkItems until 100 of them are received.
    CountDownLatch latch = new CountDownLatch(100);
    GetWorkStream stream =
        client.getWorkStream(
            GetWorkRequest.newBuilder().setClientId(10).setMaxItems(10).setMaxBytes(10000).build(),
            (String computation,
                @Nullable Instant inputDataWatermark,
                Instant synchronizedProcessingTime,
                boolean drainMode,
                WorkItem workItem,
                long serializedWorkItemSize,
                ImmutableList<Long> appliedFinalizeIds,
                ImmutableList<LatencyAttribution> getWorkStreamLatencies) -> {
              assertEquals(inputDataWatermark, new Instant(18));
              assertEquals(synchronizedProcessingTime, new Instant(17));
              assertEquals(workItem.getKey(), ByteString.copyFromUtf8("somewhat_long_key"));
              assertTrue(drainMode);
              assertTrue(sentResponseIds.containsKey(workItem.getWorkToken()));
              sentResponseIds.remove(workItem.getWorkToken());
              latch.countDown();
            });
    assertTrue(latch.await(30, TimeUnit.SECONDS));

    stream.halfClose();
    assertTrue(stream.awaitTermination(30, TimeUnit.SECONDS));
  }

  @Test
  public void testStreamingGetData() throws Exception {
    // This server responds to GetDataRequests with responses that mirror the requests.
    serviceRegistry.addService(
        new CloudWindmillServiceV1Alpha1ImplBase() {
          @Override
          public StreamObserver<StreamingGetDataRequest> getDataStream(
              StreamObserver<StreamingGetDataResponse> responseObserver) {
            return new StreamObserver<StreamingGetDataRequest>() {
              final HashSet<Long> seenIds = new HashSet<>();
              final ResponseErrorInjector injector = new ResponseErrorInjector(responseObserver);
              final StreamingGetDataResponse.Builder responseBuilder =
                  StreamingGetDataResponse.newBuilder();
              boolean sawHeader = false;

              @Override
              public void onNext(StreamingGetDataRequest chunk) {
                maybeInjectError(responseObserver);
                try {
                  if (!sawHeader) {
                    LOG.info("Received header");
                    errorCollector.checkThat(
                        chunk.getHeader(),
                        Matchers.equalTo(
                            JobHeader.newBuilder()
                                .setJobId("job")
                                .setProjectId("project")
                                .setWorkerId("worker")
                                .setClientId(clientId)
                                .build()));
                    sawHeader = true;
                  } else {
                    LOG.info(
                        "Received get data of {} global data, {} data requests",
                        chunk.getGlobalDataRequestCount(),
                        chunk.getStateRequestCount());
                    errorCollector.checkThat(
                        chunk.getSerializedSize(), Matchers.lessThanOrEqualTo(STREAM_CHUNK_SIZE));

                    int i = 0;
                    for (GlobalDataRequest request : chunk.getGlobalDataRequestList()) {
                      long requestId = chunk.getRequestId(i++);
                      errorCollector.checkThat(seenIds.add(requestId), Matchers.is(true));
                      sendResponse(requestId, processGlobalDataRequest(request));
                    }

                    for (ComputationGetDataRequest request : chunk.getStateRequestList()) {
                      long requestId = chunk.getRequestId(i++);
                      errorCollector.checkThat(seenIds.add(requestId), Matchers.is(true));
                      sendResponse(requestId, processStateRequest(request));
                    }
                    flushResponse();
                  }
                } catch (Exception e) {
                  errorCollector.addError(e);
                }
              }

              @Override
              public void onError(Throwable throwable) {}

              @Override
              public void onCompleted() {
                injector.cancel();
                responseObserver.onCompleted();
              }

              private ByteString processGlobalDataRequest(GlobalDataRequest request) {
                errorCollector.checkThat(request.getStateFamily(), Matchers.is("family"));

                return GlobalData.newBuilder()
                    .setDataId(request.getDataId())
                    .setStateFamily("family")
                    .setData(ByteString.copyFromUtf8(request.getDataId().getTag()))
                    .build()
                    .toByteString();
              }

              private ByteString processStateRequest(ComputationGetDataRequest compRequest) {
                errorCollector.checkThat(compRequest.getRequestsCount(), Matchers.is(1));
                errorCollector.checkThat(
                    compRequest.getComputationId(), Matchers.is("computation"));
                KeyedGetDataRequest request = compRequest.getRequests(0);
                KeyedGetDataResponse response =
                    makeGetDataResponse(request.getValuesToFetch(0).getTag().toStringUtf8());
                return response.toByteString();
              }

              private void sendResponse(long id, ByteString serializedResponse) {
                if (ThreadLocalRandom.current().nextInt(4) == 0) {
                  sendChunkedResponse(id, serializedResponse);
                } else {
                  responseBuilder.addRequestId(id).addSerializedResponse(serializedResponse);
                  if (responseBuilder.getRequestIdCount() > 10) {
                    flushResponse();
                  }
                }
              }

              private void sendChunkedResponse(long id, ByteString serializedResponse) {
                LOG.info("Sending response with {} chunks", (serializedResponse.size() / 10) + 1);
                for (int i = 0; i < serializedResponse.size(); i += 10) {
                  int end = Math.min(serializedResponse.size(), i + 10);
                  try {
                    responseObserver.onNext(
                        StreamingGetDataResponse.newBuilder()
                            .addRequestId(id)
                            .addSerializedResponse(serializedResponse.substring(i, end))
                            .setRemainingBytesForResponse(serializedResponse.size() - end)
                            .build());
                  } catch (IllegalStateException e) {
                    // Stream is already closed.
                  }
                }
              }

              private void flushResponse() {
                if (responseBuilder.getRequestIdCount() > 0) {
                  LOG.info(
                      "Sending batched response of {} ids", responseBuilder.getRequestIdCount());
                  try {
                    responseObserver.onNext(responseBuilder.build());
                  } catch (Exception e) {
                    // Stream is already closed.
                    LOG.warn("{}", Arrays.toString(e.getStackTrace()));
                  }
                  responseBuilder.clear();
                }
              }
            };
          }
        });

    GetDataStream stream = client.getDataStream();

    // Make requests of varying sizes to test chunking, and verify the responses.
    ExecutorService executor = Executors.newFixedThreadPool(50);
    final CountDownLatch done = new CountDownLatch(200);
    for (int i = 0; i < 100; ++i) {
      final String key = "key" + i;
      final String s = i % 5 == 0 ? largeString(i) : "tag";
      executor.execute(
          () -> {
            try {
              errorCollector.checkThat(
                  stream.requestKeyedData("computation", makeGetDataRequest(key, s)),
                  Matchers.equalTo(makeGetDataResponse(s)));
            } catch (WindmillStreamShutdownException e) {
              throw new RuntimeException(e);
            }
            done.countDown();
          });
      executor.execute(
          () -> {
            try {
              errorCollector.checkThat(
                  stream.requestGlobalData(makeGlobalDataRequest(key)),
                  Matchers.equalTo(makeGlobalDataResponse(key)));
            } catch (WindmillStreamShutdownException e) {
              throw new RuntimeException(e);
            }
            done.countDown();
          });
    }
    done.await();
    stream.halfClose();
    assertTrue(stream.awaitTermination(60, TimeUnit.SECONDS));
    executor.shutdown();
  }

  private String largeString(int length) {
    return String.join("", Collections.nCopies(length, "."));
  }

  private KeyedGetDataRequest makeGetDataRequest(String key, String tag) {
    return KeyedGetDataRequest.newBuilder()
        .setKey(ByteString.copyFromUtf8(key))
        .setWorkToken(17)
        .addValuesToFetch(TagValue.newBuilder().setTag(ByteString.copyFromUtf8(tag)))
        .build();
  }

  private KeyedGetDataResponse makeGetDataResponse(String tag) {
    return KeyedGetDataResponse.newBuilder()
        .setKey(ByteString.copyFromUtf8("key"))
        .addValues(
            TagValue.newBuilder()
                .setTag(ByteString.copyFromUtf8("tag"))
                .setValue(
                    Value.newBuilder()
                        .setTimestamp(0)
                        .setData(ByteString.copyFromUtf8(tag + "-value"))))
        .build();
  }

  private GlobalDataRequest makeGlobalDataRequest(String key) {
    return GlobalDataRequest.newBuilder()
        .setStateFamily("family")
        .setDataId(GlobalDataId.newBuilder().setTag(key).setVersion(ByteString.EMPTY).build())
        .build();
  }

  private GlobalData makeGlobalDataResponse(String key) {
    return GlobalData.newBuilder()
        .setStateFamily("family")
        .setDataId(GlobalDataId.newBuilder().setTag(key).setVersion(ByteString.EMPTY))
        .setData(ByteString.copyFromUtf8(key))
        .build();
  }

  private WorkItemCommitRequest makeCommitRequest(int i, int size) {
    return WorkItemCommitRequest.newBuilder()
        .setKey(ByteString.copyFromUtf8("key" + i))
        .setWorkToken(i)
        .addValueUpdates(
            TagValue.newBuilder()
                .setTag(ByteString.copyFromUtf8("tag" + i))
                .setValue(
                    Value.newBuilder()
                        .setData(ByteString.copyFromUtf8(largeString(size)))
                        .setTimestamp(i)))
        .build();
  }

  // This server receives WorkItemCommitRequests, and verifies they are equal to the provided
  // commitRequest.
  private StreamObserver<StreamingCommitWorkRequest> getTestCommitStreamObserver(
      StreamObserver<StreamingCommitResponse> responseObserver,
      Map<Long, WorkItemCommitRequest> commitRequests) {
    return new StreamObserver<StreamingCommitWorkRequest>() {
      final ResponseErrorInjector injector = new ResponseErrorInjector(responseObserver);
      boolean sawHeader = false;
      InputStream buffer = null;
      long remainingBytes = 0;

      @Override
      public void onNext(StreamingCommitWorkRequest request) {
        maybeInjectError(responseObserver);

        if (!sawHeader) {
          errorCollector.checkThat(
              request.getHeader(),
              Matchers.equalTo(
                  JobHeader.newBuilder()
                      .setJobId("job")
                      .setProjectId("project")
                      .setWorkerId("worker")
                      .setClientId(clientId)
                      .build()));
          sawHeader = true;
          LOG.info("Received header");
        } else {
          boolean first = true;
          LOG.info("Received request with {} chunks", request.getCommitChunkCount());
          for (StreamingCommitRequestChunk chunk : request.getCommitChunkList()) {
            assertTrue(chunk.getSerializedWorkItemCommit().size() <= STREAM_CHUNK_SIZE);
            if (first || chunk.hasComputationId()) {
              errorCollector.checkThat(chunk.getComputationId(), Matchers.equalTo("computation"));
            }

            if (remainingBytes != 0) {
              errorCollector.checkThat(buffer, Matchers.notNullValue());
              errorCollector.checkThat(
                  remainingBytes,
                  Matchers.is(
                      chunk.getSerializedWorkItemCommit().size()
                          + chunk.getRemainingBytesForWorkItem()));
              buffer =
                  new SequenceInputStream(buffer, chunk.getSerializedWorkItemCommit().newInput());
            } else {
              errorCollector.checkThat(buffer, Matchers.nullValue());
              buffer = chunk.getSerializedWorkItemCommit().newInput();
            }
            remainingBytes = chunk.getRemainingBytesForWorkItem();
            if (remainingBytes == 0) {
              try {
                WorkItemCommitRequest received = WorkItemCommitRequest.parseFrom(buffer);
                errorCollector.checkThat(
                    received, Matchers.equalTo(commitRequests.get(received.getWorkToken())));
                try {
                  responseObserver.onNext(
                      StreamingCommitResponse.newBuilder()
                          .addRequestId(chunk.getRequestId())
                          .build());
                } catch (IllegalStateException e) {
                  // Stream is closed.
                }
              } catch (Exception e) {
                errorCollector.addError(e);
              }
              buffer = null;
            } else {
              errorCollector.checkThat(first, Matchers.is(true));
            }
            first = false;
          }
        }
      }

      @Override
      public void onError(Throwable throwable) {}

      @Override
      public void onCompleted() {
        injector.cancel();
        responseObserver.onCompleted();
      }
    };
  }

  private void commitWorkTestHelper(
      CommitWorkStream stream,
      ConcurrentHashMap<Long, WorkItemCommitRequest> requestsForService,
      int requestIdStart,
      int numRequests)
      throws InterruptedException {
    List<WorkItemCommitRequest> commitRequestList = new ArrayList<>();
    List<CountDownLatch> latches = new ArrayList<>();
    for (int i = 0; i < numRequests; ++i) {
      // Build some requests of varying size with a few big ones.
      WorkItemCommitRequest request =
          makeCommitRequest(i + requestIdStart, i * (i < numRequests * .9 ? 8 : 128));
      commitRequestList.add(request);
      requestsForService.put((long) i + requestIdStart, request);
      latches.add(new CountDownLatch(1));
    }
    Collections.shuffle(commitRequestList);
    try (CommitWorkStream.RequestBatcher batcher = stream.batcher()) {
      for (int i = 0; i < commitRequestList.size(); ) {
        final CountDownLatch latch = latches.get(i);
        if (batcher.commitWorkItem(
            "computation",
            commitRequestList.get(i),
            (CommitStatus status) -> {
              assertEquals(status, CommitStatus.OK);
              latch.countDown();
            })) {
          i++;
        } else {
          batcher.flush();
        }
      }
    }
    for (CountDownLatch latch : latches) {
      assertTrue(latch.await(1, TimeUnit.MINUTES));
    }
  }

  @Test
  public void testStreamingCommit() throws Exception {
    ConcurrentHashMap<Long, WorkItemCommitRequest> commitRequests = new ConcurrentHashMap<>();
    serviceRegistry.addService(
        new CloudWindmillServiceV1Alpha1ImplBase() {
          @Override
          public StreamObserver<StreamingCommitWorkRequest> commitWorkStream(
              StreamObserver<StreamingCommitResponse> responseObserver) {
            return getTestCommitStreamObserver(responseObserver, commitRequests);
          }
        });

    // Make the commit requests, waiting for each of them to be verified and acknowledged.
    CommitWorkStream stream = client.commitWorkStream();
    commitWorkTestHelper(stream, commitRequests, 0, 500);
    stream.halfClose();
    assertTrue(stream.awaitTermination(30, TimeUnit.SECONDS));
  }

  @Test
  public void testStreamingCommitManyThreads() throws Exception {
    ConcurrentHashMap<Long, WorkItemCommitRequest> commitRequests = new ConcurrentHashMap<>();
    serviceRegistry.addService(
        new CloudWindmillServiceV1Alpha1ImplBase() {
          @Override
          public StreamObserver<StreamingCommitWorkRequest> commitWorkStream(
              StreamObserver<StreamingCommitResponse> responseObserver) {
            return getTestCommitStreamObserver(responseObserver, commitRequests);
          }
        });
    ScheduledExecutorService executor = Executors.newScheduledThreadPool(10);
    // Make the commit requests, waiting for each of them to be verified and acknowledged.
    CommitWorkStream stream = client.commitWorkStream();
    List<Future<?>> futures = new ArrayList<>();
    for (int i = 0; i < 10; ++i) {
      final int startRequestId = i * 50;
      futures.add(
          executor.submit(
              () -> {
                try {
                  commitWorkTestHelper(stream, commitRequests, startRequestId, 50);
                } catch (InterruptedException e) {
                  throw new RuntimeException(e);
                }
              }));
    }
    // Surface any exceptions that might be thrown by submitting by blocking on the future.
    for (Future<?> f : futures) {
      f.get();
    }
    stream.halfClose();
    assertTrue(stream.awaitTermination(30, TimeUnit.SECONDS));
    executor.shutdown();
  }

  @Test
  // Tests stream retries on server errors before and after `close()`
  public void testStreamingCommitClosedStream() throws Exception {
    List<WorkItemCommitRequest> commitRequestList = new ArrayList<>();
    List<CountDownLatch> latches = new ArrayList<>();
    Map<Long, WorkItemCommitRequest> commitRequests = new ConcurrentHashMap<>();
    AtomicBoolean shouldServerReturnError = new AtomicBoolean(true);
    AtomicBoolean isClientClosed = new AtomicBoolean(false);
    AtomicInteger errorsBeforeClose = new AtomicInteger();
    AtomicInteger errorsAfterClose = new AtomicInteger();
    for (int i = 0; i < 500; ++i) {
      // Build some requests of varying size with a few big ones.
      WorkItemCommitRequest request = makeCommitRequest(i, i * (i < 480 ? 8 : 128));
      commitRequestList.add(request);
      commitRequests.put((long) i, request);
      latches.add(new CountDownLatch(1));
    }
    Collections.shuffle(commitRequestList);

    // This server returns errors if shouldServerReturnError is true, else returns valid responses.
    serviceRegistry.addService(
        new CloudWindmillServiceV1Alpha1ImplBase() {
          @Override
          public StreamObserver<StreamingCommitWorkRequest> commitWorkStream(
              StreamObserver<StreamingCommitResponse> responseObserver) {
            StreamObserver<StreamingCommitWorkRequest> testCommitStreamObserver =
                getTestCommitStreamObserver(responseObserver, commitRequests);
            return new StreamObserver<StreamingCommitWorkRequest>() {
              @Override
              public void onNext(StreamingCommitWorkRequest request) {
                if (shouldServerReturnError.get()) {
                  try {
                    responseObserver.onError(
                        new RuntimeException("shouldServerReturnError = true"));
                    if (isClientClosed.get()) {
                      errorsAfterClose.incrementAndGet();
                    } else {
                      errorsBeforeClose.incrementAndGet();
                    }
                  } catch (IllegalStateException e) {
                    // The stream is already closed.
                  }
                } else {
                  testCommitStreamObserver.onNext(request);
                }
              }

              @Override
              public void onError(Throwable throwable) {
                testCommitStreamObserver.onError(throwable);
              }

              @Override
              public void onCompleted() {
                testCommitStreamObserver.onCompleted();
              }
            };
          }
        });

    // Make the commit requests, waiting for each of them to be verified and acknowledged.
    CommitWorkStream stream = client.commitWorkStream();
    try (CommitWorkStream.RequestBatcher batcher = stream.batcher()) {
      for (int i = 0; i < commitRequestList.size(); ) {
        final CountDownLatch latch = latches.get(i);
        if (batcher.commitWorkItem(
            "computation",
            commitRequestList.get(i),
            (CommitStatus status) -> {
              assertEquals(status, CommitStatus.OK);
              latch.countDown();
            })) {
          i++;
        } else {
          batcher.flush();
        }
      }
    }

    long deadline = System.currentTimeMillis() + 60_000; // 1 min
    while (true) {
      Thread.sleep(100);
      int tmpErrorsBeforeClose = errorsBeforeClose.get();
      // wait for at least 1 error before close
      if (tmpErrorsBeforeClose > 0) {
        break;
      }
      if (System.currentTimeMillis() > deadline) {
        // Control should not reach here if the test is working as expected
        fail(
            String.format(
                "Expected errors not sent by server errorsBeforeClose: %s"
                    + " \n Should not reach here if the test is working as expected.",
                tmpErrorsBeforeClose));
      }
    }

    stream.halfClose();
    isClientClosed.set(true);

    deadline = System.currentTimeMillis() + 60_000; // 1 min
    while (true) {
      Thread.sleep(100);
      int tmpErrorsAfterClose = errorsAfterClose.get();
      // wait for at least 1 error after close
      if (tmpErrorsAfterClose > 0) {
        break;
      }
      if (System.currentTimeMillis() > deadline) {
        // Control should not reach here if the test is working as expected
        fail(
            String.format(
                "Expected errors not sent by server errorsAfterClose: %s"
                    + " \n Should not reach here if the test is working as expected.",
                tmpErrorsAfterClose));
      }
    }

    shouldServerReturnError.set(false);
    for (CountDownLatch latch : latches) {
      assertTrue(latch.await(1, TimeUnit.MINUTES));
    }
    assertTrue(stream.awaitTermination(30, TimeUnit.SECONDS));
  }

  private List<KeyedGetDataRequest> makeGetDataHeartbeatRequest(List<String> keys) {
    List<KeyedGetDataRequest> result = new ArrayList<>();
    for (String key : keys) {
      result.add(
          Windmill.KeyedGetDataRequest.newBuilder()
              .setShardingKey(key.hashCode())
              .setWorkToken(0)
              .setCacheToken(0)
              .build());
    }
    return result;
  }

  private List<HeartbeatRequest> makeHeartbeatRequest(List<String> keys) {
    List<HeartbeatRequest> result = new ArrayList<>();
    for (String key : keys) {
      result.add(
          Windmill.HeartbeatRequest.newBuilder()
              .setShardingKey(key.hashCode())
              .setWorkToken(0)
              .setCacheToken(0)
              .build());
    }
    return result;
  }

  @Test
  public void testStreamingGetDataHeartbeatsAsKeyedGetDataRequests() throws Exception {
    // This server records the heartbeats observed but doesn't respond.
    final Map<String, List<KeyedGetDataRequest>> getDataHeartbeats = new HashMap<>();
    // Create a client and server different from the one in SetUp so we can add an experiment to the
    // options passed in. This requires teardown and re-constructing the client and server
    tearDown();
    startServerAndClient(
        Collections.singletonList("streaming_engine_disable_new_heartbeat_requests"));

    serviceRegistry.addService(
        new CloudWindmillServiceV1Alpha1ImplBase() {
          @Override
          public StreamObserver<StreamingGetDataRequest> getDataStream(
              StreamObserver<StreamingGetDataResponse> responseObserver) {
            return new StreamObserver<StreamingGetDataRequest>() {
              boolean sawHeader = false;

              @Override
              public void onNext(StreamingGetDataRequest chunk) {
                try {
                  if (!sawHeader) {
                    LOG.info("Received header");
                    errorCollector.checkThat(
                        chunk.getHeader(),
                        Matchers.equalTo(
                            JobHeader.newBuilder()
                                .setJobId("job")
                                .setProjectId("project")
                                .setWorkerId("worker")
                                .setClientId(clientId)
                                .build()));
                    sawHeader = true;
                  } else {
                    LOG.info("Received {} getDataHeartbeats", chunk.getStateRequestCount());
                    errorCollector.checkThat(
                        chunk.getSerializedSize(), Matchers.lessThanOrEqualTo(STREAM_CHUNK_SIZE));
                    errorCollector.checkThat(chunk.getRequestIdCount(), Matchers.is(0));

                    synchronized (getDataHeartbeats) {
                      for (ComputationGetDataRequest request : chunk.getStateRequestList()) {
                        errorCollector.checkThat(request.getRequestsCount(), Matchers.is(1));
                        getDataHeartbeats.putIfAbsent(
                            request.getComputationId(), new ArrayList<>());
                        getDataHeartbeats
                            .get(request.getComputationId())
                            .add(request.getRequestsList().get(0));
                      }
                    }
                  }
                } catch (Exception e) {
                  errorCollector.addError(e);
                }
              }

              @Override
              public void onError(Throwable throwable) {}

              @Override
              public void onCompleted() {
                responseObserver.onCompleted();
              }
            };
          }
        });

    List<String> computation1Keys = new ArrayList<>();
    List<String> computation2Keys = new ArrayList<>();

    for (int i = 0; i < 100; ++i) {
      computation1Keys.add("Computation1Key" + i);
      computation2Keys.add("Computation2Key" + largeString(i * 20));
    }
    // We're adding HeartbeatRequests to refreshActiveWork, but expecting to get back
    // KeyedGetDataRequests, so make a Map of both types.
    Map<String, List<KeyedGetDataRequest>> expectedKeyedGetDataRequests = new HashMap<>();
    expectedKeyedGetDataRequests.put("Computation1", makeGetDataHeartbeatRequest(computation1Keys));
    expectedKeyedGetDataRequests.put("Computation2", makeGetDataHeartbeatRequest(computation2Keys));
    Map<String, Collection<HeartbeatRequest>> heartbeatsToRefresh = new HashMap<>();
    heartbeatsToRefresh.put("Computation1", makeHeartbeatRequest(computation1Keys));
    heartbeatsToRefresh.put("Computation2", makeHeartbeatRequest(computation2Keys));

    GetDataStream stream = client.getDataStream();
    stream.refreshActiveWork(heartbeatsToRefresh);
    stream.halfClose();
    assertTrue(stream.awaitTermination(60, TimeUnit.SECONDS));

    boolean receivedAllGetDataHeartbeats = false;
    while (!receivedAllGetDataHeartbeats) {
      Thread.sleep(100);
      synchronized (getDataHeartbeats) {
        if (getDataHeartbeats.size() != expectedKeyedGetDataRequests.size()) {
          continue;
        }
        assertEquals(expectedKeyedGetDataRequests, getDataHeartbeats);
        receivedAllGetDataHeartbeats = true;
      }
    }
  }

  @Test
  public void testStreamingGetDataHeartbeatsAsHeartbeatRequests() throws Exception {
    // This server records the heartbeats observed but doesn't respond.
    final List<ComputationHeartbeatRequest> receivedHeartbeats = new ArrayList<>();

    serviceRegistry.addService(
        new CloudWindmillServiceV1Alpha1ImplBase() {
          @Override
          public StreamObserver<StreamingGetDataRequest> getDataStream(
              StreamObserver<StreamingGetDataResponse> responseObserver) {
            return new StreamObserver<StreamingGetDataRequest>() {
              boolean sawHeader = false;

              @Override
              public void onNext(StreamingGetDataRequest chunk) {
                try {
                  if (!sawHeader) {
                    LOG.info("Received header");
                    errorCollector.checkThat(
                        chunk.getHeader(),
                        Matchers.equalTo(
                            JobHeader.newBuilder()
                                .setJobId("job")
                                .setProjectId("project")
                                .setWorkerId("worker")
                                .setClientId(clientId)
                                .build()));
                    sawHeader = true;
                  } else {
                    LOG.info(
                        "Received {} computationHeartbeatRequests",
                        chunk.getComputationHeartbeatRequestCount());
                    errorCollector.checkThat(
                        chunk.getSerializedSize(), Matchers.lessThanOrEqualTo(STREAM_CHUNK_SIZE));
                    errorCollector.checkThat(chunk.getRequestIdCount(), Matchers.is(0));

                    synchronized (receivedHeartbeats) {
                      receivedHeartbeats.addAll(chunk.getComputationHeartbeatRequestList());
                    }
                  }
                } catch (Exception e) {
                  errorCollector.addError(e);
                }
              }

              @Override
              public void onError(Throwable throwable) {}

              @Override
              public void onCompleted() {
                responseObserver.onCompleted();
              }
            };
          }
        });

    List<String> computation1Keys = new ArrayList<>();
    List<String> computation2Keys = new ArrayList<>();

    // When sending heartbeats as HeartbeatRequest protos, all keys for the same computation should
    // be batched into the same ComputationHeartbeatRequest. Compare to the KeyedGetDataRequest
    // version in the test above, which only sends one key per ComputationGetDataRequest.
    List<ComputationHeartbeatRequest> expectedHeartbeats = new ArrayList<>();
    ComputationHeartbeatRequest.Builder comp1Builder =
        ComputationHeartbeatRequest.newBuilder().setComputationId("Computation1");
    ComputationHeartbeatRequest.Builder comp2Builder =
        ComputationHeartbeatRequest.newBuilder().setComputationId("Computation2");
    for (int i = 0; i < 100; ++i) {
      String computation1Key = "Computation1Key" + i;
      computation1Keys.add(computation1Key);
      comp1Builder.addHeartbeatRequests(
          makeHeartbeatRequest(Collections.singletonList(computation1Key)).get(0));
      String computation2Key = "Computation2Key" + largeString(i * 20);
      computation2Keys.add(computation2Key);
      comp2Builder.addHeartbeatRequests(
          makeHeartbeatRequest(Collections.singletonList(computation2Key)).get(0));
    }
    expectedHeartbeats.add(comp1Builder.build());
    expectedHeartbeats.add(comp2Builder.build());
    Map<String, Collection<HeartbeatRequest>> heartbeatRequestMap = new HashMap<>();
    heartbeatRequestMap.put("Computation1", makeHeartbeatRequest(computation1Keys));
    heartbeatRequestMap.put("Computation2", makeHeartbeatRequest(computation2Keys));

    GetDataStream stream = client.getDataStream();
    stream.refreshActiveWork(heartbeatRequestMap);
    stream.halfClose();
    assertTrue(stream.awaitTermination(60, TimeUnit.SECONDS));

    boolean receivedAllHeartbeatRequests = false;
    while (!receivedAllHeartbeatRequests) {
      Thread.sleep(100);
      synchronized (receivedHeartbeats) {
        if (receivedHeartbeats.size() != expectedHeartbeats.size()) {
          continue;
        }
        assertEquals(expectedHeartbeats, receivedHeartbeats);
        receivedAllHeartbeatRequests = true;
      }
    }
  }

  class ResponseErrorInjector<Stream extends StreamObserver> {

    private final Stream stream;
    private final Thread errorThread;
    private boolean cancelled = false;

    public ResponseErrorInjector(Stream stream) {
      this.stream = stream;
      errorThread = new Thread(this::errorThreadBody);
      errorThread.start();
    }

    private void errorThreadBody() {
      int i = 0;
      while (true) {
        try {
          Thread.sleep(ThreadLocalRandom.current().nextInt(++i * 10));
        } catch (InterruptedException e) {
          Thread.currentThread().interrupt();
          break;
        }
        synchronized (this) {
          if (cancelled) {
            break;
          }
        }
        maybeInjectError(stream);
      }
    }

    public void cancel() {
      LOG.info("Starting cancel of error injector.");
      synchronized (this) {
        cancelled = true;
      }
      errorThread.interrupt();
      try {
        errorThread.join();
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
      }
      LOG.info("Done cancelling.");
    }
  }
}
