/*
 * 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.fn.control;

import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;

import java.io.Closeable;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import javax.annotation.Nullable;
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionRequest;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleProgressRequest;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleRequest;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.RegisterRequest;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateAppendResponse;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateClearResponse;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateGetResponse;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest.RequestCase;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateResponse;
import org.apache.beam.model.pipeline.v1.MetricsApi.MonitoringInfo;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
import org.apache.beam.runners.core.SideInputReader;
import org.apache.beam.runners.core.StateNamespaces;
import org.apache.beam.runners.core.StateTags;
import org.apache.beam.runners.core.metrics.MonitoringInfoConstants;
import org.apache.beam.runners.dataflow.worker.ByteStringCoder;
import org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext;
import org.apache.beam.runners.dataflow.worker.DataflowOperationContext;
import org.apache.beam.runners.dataflow.worker.counters.NameContext;
import org.apache.beam.runners.dataflow.worker.util.common.worker.Operation;
import org.apache.beam.runners.dataflow.worker.util.common.worker.OperationContext;
import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputReceiver;
import org.apache.beam.runners.fnexecution.control.InstructionRequestHandler;
import org.apache.beam.runners.fnexecution.state.StateDelegator;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.fn.IdGenerator;
import org.apache.beam.sdk.fn.data.RemoteGrpcPortRead;
import org.apache.beam.sdk.fn.data.RemoteGrpcPortWrite;
import org.apache.beam.sdk.state.BagState;
import org.apache.beam.sdk.transforms.Materializations;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.util.MoreFutures;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.TextFormat;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.MoreObjects;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Maps;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Table;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * This {@link Operation} is responsible for communicating with the SDK harness and asking it to
 * process a bundle of work. This operation registers the {@link
 * org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor} when executed the first
 * time. Afterwards, it only asks the SDK harness to process the bundle using the already registered
 * {@link org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor}.
 *
 * <p>This operation supports restart.
 */
public class RegisterAndProcessBundleOperation extends Operation {
  private static final Logger LOG =
      LoggerFactory.getLogger(RegisterAndProcessBundleOperation.class);

  private static final OutputReceiver[] EMPTY_RECEIVERS = new OutputReceiver[0];

  private final IdGenerator idGenerator;
  private final InstructionRequestHandler instructionRequestHandler;
  private final StateDelegator beamFnStateDelegator;
  private final RegisterRequest registerRequest;
  private final Map<String, DataflowStepContext> ptransformIdToUserStepContext;
  private final Map<String, SideInputReader> ptransformIdToSideInputReader;
  private final Table<String, String, PCollectionView<?>>
      ptransformIdToSideInputIdToPCollectionView;
  private final ConcurrentHashMap<StateKey, BagState<ByteString>> userStateData;
  private final Map<String, NameContext> pcollectionIdToNameContext;

  private @Nullable CompletionStage<InstructionResponse> registerFuture;
  private @Nullable CompletionStage<InstructionResponse> processBundleResponse;
  private volatile @Nullable String processBundleId = null;
  private StateDelegator.Registration deregisterStateHandler;

  private @Nullable String grpcReadTransformId = null;
  private String grpcReadTransformOutputName = null;
  private String grpcReadTransformOutputPCollectionName = null;
  private final Set<String> grpcReadTransformReadWritePCollectionNames;

  public RegisterAndProcessBundleOperation(
      IdGenerator idGenerator,
      InstructionRequestHandler instructionRequestHandler,
      StateDelegator beamFnStateDelegator,
      RegisterRequest registerRequest,
      Map<String, DataflowOperationContext> ptransformIdToOperationContext,
      Map<String, DataflowStepContext> ptransformIdToSystemStepContext,
      Map<String, SideInputReader> ptransformIdToSideInputReader,
      Table<String, String, PCollectionView<?>> ptransformIdToSideInputIdToPCollectionView,
      Map<String, NameContext> pcollectionIdToNameContext,
      OperationContext context) {
    super(EMPTY_RECEIVERS, context);
    this.idGenerator = idGenerator;
    this.instructionRequestHandler = instructionRequestHandler;
    this.beamFnStateDelegator = beamFnStateDelegator;
    this.registerRequest = registerRequest;
    this.ptransformIdToSideInputReader = ptransformIdToSideInputReader;
    this.ptransformIdToSideInputIdToPCollectionView = ptransformIdToSideInputIdToPCollectionView;
    this.pcollectionIdToNameContext = pcollectionIdToNameContext;
    ImmutableMap.Builder<String, DataflowStepContext> userStepContextsMap = ImmutableMap.builder();
    for (Map.Entry<String, DataflowStepContext> entry :
        ptransformIdToSystemStepContext.entrySet()) {
      userStepContextsMap.put(entry.getKey(), entry.getValue().namespacedToUser());
    }
    this.ptransformIdToUserStepContext = userStepContextsMap.build();
    this.userStateData = new ConcurrentHashMap<>();

    checkState(
        registerRequest.getProcessBundleDescriptorCount() == 1,
        "Only one bundle registration at a time currently supported.");
    if (LOG.isDebugEnabled()) {
      LOG.debug(
          "Process bundle descriptor {}", toDot(registerRequest.getProcessBundleDescriptor(0)));
    }

    for (Map.Entry<String, RunnerApi.PTransform> pTransform :
        registerRequest.getProcessBundleDescriptor(0).getTransformsMap().entrySet()) {
      if (pTransform.getValue().getSpec().getUrn().equals(RemoteGrpcPortRead.URN)) {
        if (grpcReadTransformId != null) {
          // TODO: Handle the case of more than one input.
          grpcReadTransformId = null;
          grpcReadTransformOutputName = null;
          grpcReadTransformOutputPCollectionName = null;
          break;
        }
        grpcReadTransformId = pTransform.getKey();
        grpcReadTransformOutputName =
            Iterables.getOnlyElement(pTransform.getValue().getOutputsMap().keySet());
        grpcReadTransformOutputPCollectionName =
            pTransform.getValue().getOutputsMap().get(grpcReadTransformOutputName);
      }
    }

    grpcReadTransformReadWritePCollectionNames =
        extractCrossBoundaryGrpcPCollectionNames(
            registerRequest.getProcessBundleDescriptor(0).getTransformsMap().entrySet());
  }

  private Set<String> extractCrossBoundaryGrpcPCollectionNames(
      final Set<Entry<String, PTransform>> ptransforms) {
    Set<String> result = new HashSet<>();

    // GRPC Read/Write expected to only have one Output/Input respectively.
    for (Map.Entry<String, RunnerApi.PTransform> pTransform : ptransforms) {
      if (pTransform.getValue().getSpec().getUrn().equals(RemoteGrpcPortRead.URN)) {
        String grpcReadTransformOutputName =
            Iterables.getOnlyElement(pTransform.getValue().getOutputsMap().keySet());
        String pcollectionName =
            pTransform.getValue().getOutputsMap().get(grpcReadTransformOutputName);
        result.add(pcollectionName);
      }

      if (pTransform.getValue().getSpec().getUrn().equals(RemoteGrpcPortWrite.URN)) {
        String grpcTransformInputName =
            Iterables.getOnlyElement(pTransform.getValue().getInputsMap().keySet());
        String pcollectionName = pTransform.getValue().getInputsMap().get(grpcTransformInputName);
        result.add(pcollectionName);
      }
    }

    return result;
  }

  /** Generates a dot description of the process bundle descriptor. */
  private static String toDot(ProcessBundleDescriptor processBundleDescriptor) {
    StringBuilder builder = new StringBuilder();
    builder.append("digraph network {\n");
    Map<String, String> nodeName = Maps.newHashMap();
    processBundleDescriptor
        .getPcollectionsMap()
        .forEach((key, node) -> nodeName.put("pc " + key, "n" + nodeName.size()));
    processBundleDescriptor
        .getTransformsMap()
        .forEach((key, node) -> nodeName.put("pt " + key, "n" + nodeName.size()));
    for (Entry<String, RunnerApi.PCollection> nodeEntry :
        processBundleDescriptor.getPcollectionsMap().entrySet()) {
      builder.append(
          String.format(
              "  %s [fontname=\"Courier New\" label=\"%s\"];%n",
              nodeName.get("pc " + nodeEntry.getKey()),
              escapeDot(nodeEntry.getKey() + ": " + nodeEntry.getValue().getUniqueName())));
    }
    for (Entry<String, RunnerApi.PTransform> nodeEntry :
        processBundleDescriptor.getTransformsMap().entrySet()) {
      builder.append(
          String.format(
              "  %s [fontname=\"Courier New\" label=\"%s\"];%n",
              nodeName.get("pt " + nodeEntry.getKey()),
              escapeDot(
                  nodeEntry.getKey()
                      + ": "
                      + nodeEntry.getValue().getSpec().getUrn()
                      + " "
                      + nodeEntry.getValue().getUniqueName())));
      for (Entry<String, String> inputEntry : nodeEntry.getValue().getInputsMap().entrySet()) {
        builder.append(
            String.format(
                "  %s -> %s [fontname=\"Courier New\" label=\"%s\"];%n",
                nodeName.get("pc " + inputEntry.getValue()),
                nodeName.get("pt " + nodeEntry.getKey()),
                escapeDot(inputEntry.getKey())));
      }
      for (Entry<String, String> outputEntry : nodeEntry.getValue().getOutputsMap().entrySet()) {
        builder.append(
            String.format(
                "  %s -> %s [fontname=\"Courier New\" label=\"%s\"];%n",
                nodeName.get("pt " + nodeEntry.getKey()),
                nodeName.get("pc " + outputEntry.getValue()),
                escapeDot(outputEntry.getKey())));
      }
    }
    builder.append("}");
    return builder.toString();
  }

  private static String escapeDot(String s) {
    return s.replace("\\", "\\\\")
        .replace("\"", "\\\"")
        // http://www.graphviz.org/doc/info/attrs.html#k:escString
        // The escape sequences "\n", "\l" and "\r" divide the label into lines, centered,
        // left-justified, and right-justified, respectively.
        .replace("\n", "\\l");
  }

  /**
   * Returns an id for the current bundle being processed.
   *
   * <p>Generates new id with idGenerator if no id is cached.
   *
   * <p><b>Note</b>: This operation could be used across multiple bundles, so a unique id is
   * generated for every bundle. {@link Operation Operations} accessing the bundle id should only
   * call this once per bundle and cache the id in the {@link Operation#start()} method and clear it
   * in the {@link Operation#finish()} method.
   */
  public synchronized String getProcessBundleInstructionId() {
    if (processBundleId == null) {
      processBundleId = idGenerator.getId();
    }
    return processBundleId;
  }

  public String getCurrentProcessBundleInstructionId() {
    return processBundleId;
  }

  @Override
  public void start() throws Exception {
    try (Closeable scope = context.enterStart()) {
      super.start();

      // Only register once by using the presence of the future as a signal.
      if (registerFuture == null) {
        InstructionRequest request =
            InstructionRequest.newBuilder()
                .setInstructionId(idGenerator.getId())
                .setRegister(registerRequest)
                .build();
        registerFuture = instructionRequestHandler.handle(request);
        getRegisterResponse(registerFuture);
      }

      checkState(
          registerRequest.getProcessBundleDescriptorCount() == 1,
          "Only one bundle registration at a time currently supported.");
      InstructionRequest processBundleRequest =
          InstructionRequest.newBuilder()
              .setInstructionId(getProcessBundleInstructionId())
              .setProcessBundle(
                  ProcessBundleRequest.newBuilder()
                      .setProcessBundleDescriptorReference(
                          registerRequest.getProcessBundleDescriptor(0).getId()))
              .build();

      deregisterStateHandler =
          beamFnStateDelegator.registerForProcessBundleInstructionId(
              getProcessBundleInstructionId(), this::delegateByStateKeyType);
      processBundleResponse = instructionRequestHandler.handle(processBundleRequest);
    }
  }

  @Override
  public void finish() throws Exception {
    // TODO: Once we have access to windowing strategy via the ParDoPayload, add support to garbage
    // collect any user state set. Also add support for consuming those garbage collection timers.

    try (Closeable scope = context.enterFinish()) {
      // Await completion or failure
      BeamFnApi.ProcessBundleResponse completedResponse =
          MoreFutures.get(getProcessBundleResponse(processBundleResponse));
      if (completedResponse.getResidualRootsCount() > 0) {
        throw new IllegalStateException(
            "TODO: [BEAM-2939] residual roots in process bundle response not yet supported.");
      }

      deregisterStateHandler.deregister();
      userStateData.clear();
      processBundleId = null;
      super.finish();
    }
  }

  @Override
  public void abort() throws Exception {
    try (Closeable scope = context.enterAbort()) {
      deregisterStateHandler.abort();
      cancelIfNotNull(registerFuture);
      cancelIfNotNull(processBundleResponse);
      super.abort();
    }
  }

  public Map<String, NameContext> getPCollectionIdToNameContext() {
    return this.pcollectionIdToNameContext;
  }

  public Map<String, DataflowStepContext> getPtransformIdToUserStepContext() {
    return ptransformIdToUserStepContext;
  }

  /**
   * Returns the compound metrics recorded, by issuing a request to the SDK harness.
   *
   * <p>This includes key progress indicators in {@link BeamFnApi.Metrics.PTransform.Measured} as
   * well as user-defined metrics in {@link BeamFnApi.Metrics.User}.
   *
   * <p>Use {@link #getInputElementsConsumed(BeamFnApi.Metrics)} on the future value to extract the
   * elements consumed from the upstream read operation.
   *
   * <p>May be called at any time, including before start() and after finish().
   *
   * @throws InterruptedException
   * @throws ExecutionException
   */
  public CompletionStage<BeamFnApi.ProcessBundleProgressResponse> getProcessBundleProgress()
      throws InterruptedException, ExecutionException {
    // processBundleId may be reset if this bundle finishes asynchronously.
    String processBundleId = this.processBundleId;

    if (processBundleId == null) {
      return CompletableFuture.completedFuture(
          BeamFnApi.ProcessBundleProgressResponse.getDefaultInstance());
    }

    InstructionRequest processBundleRequest =
        InstructionRequest.newBuilder()
            .setInstructionId(idGenerator.getId())
            .setProcessBundleProgress(
                ProcessBundleProgressRequest.newBuilder().setInstructionReference(processBundleId))
            .build();

    return instructionRequestHandler
        .handle(processBundleRequest)
        .thenApply(
            response -> {
              if (!response.getError().isEmpty()) {
                throw new IllegalStateException(response.getError());
              }
              return response.getProcessBundleProgress();
            });
  }

  /** Returns the final metrics returned by the SDK harness when it completes the bundle. */
  public CompletionStage<BeamFnApi.Metrics> getFinalMetrics() {
    return getProcessBundleResponse(processBundleResponse)
        .thenApply(response -> response.getMetrics());
  }

  public CompletionStage<List<MonitoringInfo>> getFinalMonitoringInfos() {
    return getProcessBundleResponse(processBundleResponse)
        .thenApply(response -> response.getMonitoringInfosList());
  }

  public boolean hasFailed() throws ExecutionException, InterruptedException {
    if (processBundleResponse != null && processBundleResponse.toCompletableFuture().isDone()) {
      return !processBundleResponse.toCompletableFuture().get().getError().isEmpty();
    } else {
      // At the very least, we don't know that this has failed yet.
      return false;
    }
  }

  /*
   * Returns a subset of monitoring infos that refer to grpc IO.
   */
  public List<MonitoringInfo> findIOPCollectionMonitoringInfos(
      Iterable<MonitoringInfo> monitoringInfos) {
    List<MonitoringInfo> result = new ArrayList<MonitoringInfo>();
    if (grpcReadTransformReadWritePCollectionNames.isEmpty()) {
      return result;
    }

    for (MonitoringInfo mi : monitoringInfos) {
      if (mi.getUrn().equals(MonitoringInfoConstants.Urns.ELEMENT_COUNT)) {
        String pcollection =
            mi.getLabelsOrDefault(MonitoringInfoConstants.Labels.PCOLLECTION, null);
        if ((pcollection != null)
            && (grpcReadTransformReadWritePCollectionNames.contains(pcollection))) {
          result.add(mi);
        }
      }
    }

    return result;
  }

  long getInputElementsConsumed(final Iterable<MonitoringInfo> monitoringInfos) {
    if (grpcReadTransformId == null) {
      return 0;
    }

    for (MonitoringInfo mi : monitoringInfos) {
      if (mi.getUrn().equals(MonitoringInfoConstants.Urns.ELEMENT_COUNT)) {
        String pcollection =
            mi.getLabelsOrDefault(MonitoringInfoConstants.Labels.PCOLLECTION, null);
        if (pcollection != null && pcollection.equals(grpcReadTransformOutputPCollectionName)) {
          return mi.getMetric().getCounterData().getInt64Value();
        }
      }
    }

    return 0;
  }

  /** Returns the number of input elements consumed by the gRPC read, if known, otherwise 0. */
  double getInputElementsConsumed(BeamFnApi.Metrics metrics) {
    return metrics
        .getPtransformsOrDefault(
            grpcReadTransformId, BeamFnApi.Metrics.PTransform.getDefaultInstance())
        .getProcessedElements()
        .getMeasured()
        .getOutputElementCountsOrDefault(grpcReadTransformOutputName, 0);
  }

  private CompletionStage<BeamFnApi.StateResponse.Builder> delegateByStateKeyType(
      StateRequest stateRequest) {
    switch (stateRequest.getStateKey().getTypeCase()) {
      case BAG_USER_STATE:
        return handleBagUserState(stateRequest);
      case MULTIMAP_SIDE_INPUT:
        return handleMultimapSideInput(stateRequest);
      default:
        throw new UnsupportedOperationException(
            String.format(
                "Dataflow does not handle StateRequests of type %s",
                stateRequest.getStateKey().getTypeCase()));
    }
  }

  private CompletionStage<BeamFnApi.StateResponse.Builder> handleMultimapSideInput(
      StateRequest stateRequest) {
    checkState(
        stateRequest.getRequestCase() == RequestCase.GET,
        String.format(
            "MultimapSideInput state requests only support '%s' requests, received '%s'",
            RequestCase.GET, stateRequest.getRequestCase()));

    StateKey.MultimapSideInput multimapSideInputStateKey =
        stateRequest.getStateKey().getMultimapSideInput();

    SideInputReader sideInputReader =
        ptransformIdToSideInputReader.get(multimapSideInputStateKey.getPtransformId());
    checkState(
        sideInputReader != null,
        String.format("Unknown PTransform '%s'", multimapSideInputStateKey.getPtransformId()));

    PCollectionView<Materializations.MultimapView<Object, Object>> view =
        (PCollectionView<Materializations.MultimapView<Object, Object>>)
            ptransformIdToSideInputIdToPCollectionView.get(
                multimapSideInputStateKey.getPtransformId(),
                multimapSideInputStateKey.getSideInputId());
    checkState(
        view != null,
        String.format(
            "Unknown side input '%s' on PTransform '%s'",
            multimapSideInputStateKey.getSideInputId(),
            multimapSideInputStateKey.getPtransformId()));
    checkState(
        Materializations.MULTIMAP_MATERIALIZATION_URN.equals(
            view.getViewFn().getMaterialization().getUrn()),
        String.format(
            "Unknown materialization for side input '%s' on PTransform '%s' with urn '%s'",
            multimapSideInputStateKey.getSideInputId(),
            multimapSideInputStateKey.getPtransformId(),
            view.getViewFn().getMaterialization().getUrn()));
    checkState(
        view.getCoderInternal() instanceof KvCoder,
        String.format(
            "Materialization of side input '%s' on PTransform '%s' expects %s but received %s.",
            multimapSideInputStateKey.getSideInputId(),
            multimapSideInputStateKey.getPtransformId(),
            KvCoder.class.getSimpleName(),
            view.getCoderInternal().getClass().getSimpleName()));
    Coder<Object> keyCoder = ((KvCoder) view.getCoderInternal()).getKeyCoder();
    Coder<Object> valueCoder = ((KvCoder) view.getCoderInternal()).getValueCoder();

    BoundedWindow window;
    try {
      // TODO: Use EncodedWindow instead of decoding the window.
      window =
          view.getWindowingStrategyInternal()
              .getWindowFn()
              .windowCoder()
              .decode(multimapSideInputStateKey.getWindow().newInput());
    } catch (IOException e) {
      throw new IllegalArgumentException(
          String.format(
              "Unable to decode window for side input '%s' on PTransform '%s'.",
              multimapSideInputStateKey.getSideInputId(),
              multimapSideInputStateKey.getPtransformId()),
          e);
    }

    Object userKey;
    try {
      // TODO: Use the encoded representation of the key.
      userKey = keyCoder.decode(multimapSideInputStateKey.getKey().newInput());
    } catch (IOException e) {
      throw new IllegalArgumentException(
          String.format(
              "Unable to decode user key for side input '%s' on PTransform '%s'.",
              multimapSideInputStateKey.getSideInputId(),
              multimapSideInputStateKey.getPtransformId()),
          e);
    }

    Materializations.MultimapView<Object, Object> sideInput = sideInputReader.get(view, window);
    Iterable<Object> values = sideInput.get(userKey);
    try {
      // TODO: Chunk the requests and use a continuation key to support side input values
      // that are larger then 2 GiBs.
      // TODO: Use the raw value so we don't go through a decode/encode cycle for no reason.
      return CompletableFuture.completedFuture(
          StateResponse.newBuilder()
              .setGet(StateGetResponse.newBuilder().setData(encodeAndConcat(values, valueCoder))));
    } catch (IOException e) {
      throw new IllegalArgumentException(
          String.format(
              "Unable to encode values for side input '%s' on PTransform '%s'.",
              multimapSideInputStateKey.getSideInputId(),
              multimapSideInputStateKey.getPtransformId()),
          e);
    }
  }

  private CompletionStage<BeamFnApi.StateResponse.Builder> handleBagUserState(
      StateRequest stateRequest) {
    StateKey.BagUserState bagUserStateKey = stateRequest.getStateKey().getBagUserState();
    DataflowStepContext userStepContext =
        ptransformIdToUserStepContext.get(bagUserStateKey.getPtransformId());
    checkState(
        userStepContext != null,
        String.format("Unknown PTransform id '%s'", bagUserStateKey.getPtransformId()));
    // TODO: We should not be required to hold onto a pointer to the bag states for the
    // user. InMemoryStateInternals assumes that the Java garbage collector does the clean-up work
    // but instead StateInternals should hold its own references and write out any data and
    // clear references when the MapTask within Dataflow completes like how WindmillStateInternals
    // works.
    BagState<ByteString> state =
        userStateData.computeIfAbsent(
            stateRequest.getStateKey(),
            unused ->
                userStepContext
                    .stateInternals()
                    .state(
                        // TODO: Once we have access to the ParDoPayload, use its windowing strategy
                        // to decode the window for the well known window types. Longer term we need
                        // to swap
                        // to use the encoded version and not rely on needing to decode the entire
                        // window.
                        StateNamespaces.window(GlobalWindow.Coder.INSTANCE, GlobalWindow.INSTANCE),
                        StateTags.bag(bagUserStateKey.getUserStateId(), ByteStringCoder.of())));
    switch (stateRequest.getRequestCase()) {
      case GET:
        return CompletableFuture.completedFuture(
            StateResponse.newBuilder()
                .setGet(StateGetResponse.newBuilder().setData(concat(state.read()))));
      case APPEND:
        state.add(stateRequest.getAppend().getData());
        return CompletableFuture.completedFuture(
            StateResponse.newBuilder().setAppend(StateAppendResponse.getDefaultInstance()));
      case CLEAR:
        state.clear();
        return CompletableFuture.completedFuture(
            StateResponse.newBuilder().setClear(StateClearResponse.getDefaultInstance()));
      default:
        throw new IllegalArgumentException(
            String.format("Unknown request type %s", stateRequest.getRequestCase()));
    }
  }

  @Override
  public boolean supportsRestart() {
    return true;
  }

  private static CompletionStage<BeamFnApi.InstructionResponse> throwIfFailure(
      CompletionStage<InstructionResponse> responseFuture) {
    return responseFuture.thenApply(
        response -> {
          if (!response.getError().isEmpty()) {
            throw new IllegalStateException(
                String.format(
                    "Client failed to process %s with error [%s].",
                    response.getInstructionId(), response.getError()));
          }
          return response;
        });
  }

  private static CompletionStage<BeamFnApi.ProcessBundleResponse> getProcessBundleResponse(
      CompletionStage<InstructionResponse> responseFuture) {
    return throwIfFailure(responseFuture)
        .thenApply(
            response -> {
              switch (response.getResponseCase()) {
                case PROCESS_BUNDLE:
                  return response.getProcessBundle();
                default:
                  throw new IllegalStateException(
                      String.format(
                          "SDK harness returned wrong kind of response to ProcessBundleRequest: %s",
                          TextFormat.printToString(response)));
              }
            });
  }

  private static CompletionStage<BeamFnApi.RegisterResponse> getRegisterResponse(
      CompletionStage<InstructionResponse> responseFuture)
      throws ExecutionException, InterruptedException {
    return throwIfFailure(responseFuture)
        .thenApply(
            response -> {
              switch (response.getResponseCase()) {
                case REGISTER:
                  return response.getRegister();
                default:
                  throw new IllegalStateException(
                      String.format(
                          "SDK harness returned wrong kind of response to RegisterRequest: %s",
                          TextFormat.printToString(response)));
              }
            });
  }

  private static void cancelIfNotNull(CompletionStage<?> future) {
    if (future != null) {
      // TODO: add cancel(boolean) to MoreFutures
      future.toCompletableFuture().cancel(true);
    }
  }

  private ByteString concat(Iterable<ByteString> values) {
    ByteString rval = ByteString.EMPTY;
    if (values != null) {
      for (ByteString value : values) {
        rval = rval.concat(value);
      }
    }
    return rval;
  }

  static ByteString encodeAndConcat(Iterable<Object> values, Coder valueCoder) throws IOException {
    ByteString.Output out = ByteString.newOutput();
    if (values != null) {
      for (Object value : values) {
        int size = out.size();
        valueCoder.encode(value, out);
        // Pad empty values by one byte as per the Beam Fn data transfer specification.
        if (size == out.size()) {
          out.write(0);
        }
      }
    }
    return out.toByteString();
  }

  @Override
  public String toString() {
    return MoreObjects.toStringHelper(this)
        .add("processBundleId", processBundleId)
        .add("processBundleDescriptors", registerRequest.getProcessBundleDescriptorList())
        .toString();
  }
}