content: /*
   * 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.fnexecution.control;
  
  import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
  import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
  
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  import java.util.concurrent.CompletionStage;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.Phaser;
  import java.util.concurrent.atomic.AtomicBoolean;
  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.ProcessBundleProgressResponse;
  import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleRequest;
  import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse;
  import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest;
  import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest.DesiredSplit;
  import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitResponse;
  import org.apache.beam.model.pipeline.v1.Endpoints;
  import org.apache.beam.runners.fnexecution.control.ProcessBundleDescriptors.TimerSpec;
  import org.apache.beam.runners.fnexecution.data.FnDataService;
  import org.apache.beam.runners.fnexecution.data.RemoteInputDestination;
  import org.apache.beam.runners.fnexecution.state.StateDelegator;
  import org.apache.beam.runners.fnexecution.state.StateRequestHandler;
  import org.apache.beam.sdk.coders.Coder;
  import org.apache.beam.sdk.fn.IdGenerator;
  import org.apache.beam.sdk.fn.IdGenerators;
  import org.apache.beam.sdk.fn.data.BeamFnDataInboundObserver;
  import org.apache.beam.sdk.fn.data.BeamFnDataOutboundAggregator;
  import org.apache.beam.sdk.fn.data.CloseableFnDataReceiver;
  import org.apache.beam.sdk.fn.data.DataEndpoint;
  import org.apache.beam.sdk.fn.data.FnDataReceiver;
  import org.apache.beam.sdk.fn.data.LogicalEndpoint;
  import org.apache.beam.sdk.fn.data.TimerEndpoint;
  import org.apache.beam.sdk.util.MoreFutures;
  import org.apache.beam.sdk.util.construction.Timer;
  import org.apache.beam.sdk.values.KV;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * A high-level client for an SDK harness.
   *
   * <p>This provides a Java-friendly wrapper around {@link InstructionRequestHandler} and {@link
   * CloseableFnDataReceiver}, which handle lower-level gRPC message wrangling.
   */
  @SuppressWarnings({
    "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
  })
  public class SdkHarnessClient implements AutoCloseable {
    private static final Logger LOG = LoggerFactory.getLogger(SdkHarnessClient.class);
  
    /**
     * A processor capable of creating bundles for some registered {@link ProcessBundleDescriptor}.
     */
    public class BundleProcessor {
      private final ProcessBundleDescriptor processBundleDescriptor;
      private final List<RemoteInputDestination> remoteInputs;
      private final Map<String, Map<String, TimerSpec>> timerSpecs;
      private final StateDelegator stateDelegator;
  
      private BundleProcessor(
          ProcessBundleDescriptor processBundleDescriptor,
          List<RemoteInputDestination> remoteInputs,
          Map<String, Map<String, TimerSpec>> timerSpecs,
          StateDelegator stateDelegator) {
        this.processBundleDescriptor = processBundleDescriptor;
        this.remoteInputs = remoteInputs;
        this.timerSpecs = timerSpecs;
        this.stateDelegator = stateDelegator;
      }
  
      /**
       * Start a new bundle for the given {@link BeamFnApi.ProcessBundleDescriptor} identifier.
       *
       * <p>The input channels for the returned {@link ActiveBundle} are derived from the instructions
       * in the {@link BeamFnApi.ProcessBundleDescriptor}.
       *
       * <p>NOTE: It is important to {@link #close()} each bundle after all elements are emitted.
       *
       * <pre>{@code
       * try (ActiveBundle bundle = SdkHarnessClient.newBundle(...)) {
       *   FnDataReceiver<InputT> inputReceiver =
       *       (FnDataReceiver) bundle.getInputReceivers().get(mainPCollectionId);
       *   // send all main input elements ...
       * }
       * }</pre>
       *
       * <p>An exception during {@link #close()} will be thrown if the bundle requests finalization or
       * attempts to checkpoint by returning a {@link BeamFnApi.DelayedBundleApplication}.
       */
      public ActiveBundle newBundle(
          Map<String, RemoteOutputReceiver<?>> outputReceivers,
          BundleProgressHandler progressHandler) {
        return newBundle(
            outputReceivers,
            request -> {
              throw new UnsupportedOperationException(
                  String.format(
                      "The %s does not have a registered state handler.",
                      ActiveBundle.class.getSimpleName()));
            },
            progressHandler);
      }
  
      /**
       * Start a new bundle for the given {@link BeamFnApi.ProcessBundleDescriptor} identifier.
       *
       * <p>The input channels for the returned {@link ActiveBundle} are derived from the instructions
       * in the {@link BeamFnApi.ProcessBundleDescriptor}.
       *
       * <p>NOTE: It is important to {@link #close()} each bundle after all elements are emitted.
       *
       * <pre>{@code
       * try (ActiveBundle bundle = SdkHarnessClient.newBundle(...)) {
       *   FnDataReceiver<InputT> inputReceiver =
       *       (FnDataReceiver) bundle.getInputReceivers().get(mainPCollectionId);
       *   // send all main input elements ...
       * }
       * }</pre>
       *
       * <p>An exception during {@link #close()} will be thrown if the bundle requests finalization or
       * attempts to checkpoint by returning a {@link BeamFnApi.DelayedBundleApplication}.
       */
      public ActiveBundle newBundle(
          Map<String, RemoteOutputReceiver<?>> outputReceivers,
          StateRequestHandler stateRequestHandler,
          BundleProgressHandler progressHandler) {
        return newBundle(
            outputReceivers,
            Collections.emptyMap(),
            stateRequestHandler,
            progressHandler,
            BundleSplitHandler.unsupported(),
            request -> {
              throw new UnsupportedOperationException(
                  String.format(
                      "The %s does not have a registered bundle checkpoint handler.",
                      ActiveBundle.class.getSimpleName()));
            },
            bundleId -> {
              throw new UnsupportedOperationException(
                  String.format(
                      "The %s does not have a registered bundle finalization handler.",
                      ActiveBundle.class.getSimpleName()));
            });
      }
  
      /**
       * Start a new bundle for the given {@link BeamFnApi.ProcessBundleDescriptor} identifier.
       *
       * <p>The input channels for the returned {@link ActiveBundle} are derived from the instructions
       * in the {@link BeamFnApi.ProcessBundleDescriptor}.
       *
       * <p>NOTE: It is important to {@link #close()} each bundle after all elements are emitted.
       *
       * <pre>{@code
       * try (ActiveBundle bundle = SdkHarnessClient.newBundle(...)) {
       *   FnDataReceiver<InputT> inputReceiver =
       *       (FnDataReceiver) bundle.getInputReceivers().get(mainPCollectionId);
       *   // send all main input elements ...
       * }
       * }</pre>
       *
       * <p>An exception during {@link #close()} will be thrown if the bundle requests finalization if
       * {@link BundleFinalizationHandler} is {@code null} or attempts to checkpoint by returning a
       * {@link BeamFnApi.DelayedBundleApplication} .
       */
      public ActiveBundle newBundle(
          Map<String, RemoteOutputReceiver<?>> outputReceivers,
          Map<KV<String, String>, RemoteOutputReceiver<Timer<?>>> timerReceivers,
          StateRequestHandler stateRequestHandler,
          BundleProgressHandler progressHandler,
          BundleFinalizationHandler finalizationHandler,
          BundleCheckpointHandler checkpointHandler) {
        return newBundle(
            outputReceivers,
            timerReceivers,
            stateRequestHandler,
            progressHandler,
            BundleSplitHandler.unsupported(),
            checkpointHandler == null
                ? request -> {
                  throw new UnsupportedOperationException(
                      String.format(
                          "The %s does not have a registered bundle checkpoint handler.",
                          ActiveBundle.class.getSimpleName()));
                }
                : checkpointHandler,
            finalizationHandler == null
                ? bundleId -> {
                  throw new UnsupportedOperationException(
                      String.format(
                          "The %s does not have a registered bundle finalization handler.",
                          ActiveBundle.class.getSimpleName()));
                }
                : finalizationHandler);
      }
  
      /**
       * Start a new bundle for the given {@link BeamFnApi.ProcessBundleDescriptor} identifier.
       *
       * <p>The input channels for the returned {@link ActiveBundle} are derived from the instructions
       * in the {@link BeamFnApi.ProcessBundleDescriptor}.
       *
       * <p>NOTE: It is important to {@link #close()} each bundle after all elements are emitted.
       *
       * <pre>{@code
       * try (ActiveBundle bundle = SdkHarnessClient.newBundle(...)) {
       *   FnDataReceiver<InputT> inputReceiver =
       *       (FnDataReceiver) bundle.getInputReceivers().get(mainPCollectionId);
       *   // send all elements ...
       * }
       * }</pre>
       */
      public ActiveBundle newBundle(
          Map<String, RemoteOutputReceiver<?>> outputReceivers,
          Map<KV<String, String>, RemoteOutputReceiver<Timer<?>>> timerReceivers,
          StateRequestHandler stateRequestHandler,
          BundleProgressHandler progressHandler,
          BundleSplitHandler splitHandler,
          BundleCheckpointHandler checkpointHandler,
          BundleFinalizationHandler finalizationHandler) {
        String bundleId = idGenerator.getId();
  
        final CompletionStage<BeamFnApi.InstructionResponse> genericResponse =
            fnApiControlClient.handle(
                BeamFnApi.InstructionRequest.newBuilder()
                    .setInstructionId(bundleId)
                    .setProcessBundle(
                        BeamFnApi.ProcessBundleRequest.newBuilder()
                            .setProcessBundleDescriptorId(processBundleDescriptor.getId())
                            .addAllCacheTokens(stateRequestHandler.getCacheTokens()))
                    .build());
        LOG.debug(
            "Sent {} with ID {} for {} with ID {}",
            ProcessBundleRequest.class.getSimpleName(),
            bundleId,
            ProcessBundleDescriptor.class.getSimpleName(),
            processBundleDescriptor.getId());
  
        CompletionStage<BeamFnApi.ProcessBundleResponse> specificResponse =
            genericResponse.thenApply(InstructionResponse::getProcessBundle);
        Optional<BeamFnDataInboundObserver> beamFnDataInboundObserver;
        if (outputReceivers.isEmpty() && timerReceivers.isEmpty()) {
          beamFnDataInboundObserver = Optional.empty();
        } else {
          List<DataEndpoint<?>> dataEndpoints = new ArrayList<>(outputReceivers.size());
          for (Map.Entry<String, RemoteOutputReceiver<?>> receiver : outputReceivers.entrySet()) {
            dataEndpoints.add(
                DataEndpoint.create(
                    receiver.getKey(),
                    (Coder<Object>) receiver.getValue().getCoder(),
                    (FnDataReceiver<Object>) receiver.getValue().getReceiver()));
          }
          List<TimerEndpoint<?>> timerEndpoints = new ArrayList<>(timerReceivers.size());
          for (Map.Entry<KV<String, String>, RemoteOutputReceiver<Timer<?>>> timerReceiver :
              timerReceivers.entrySet()) {
            timerEndpoints.add(
                TimerEndpoint.create(
                    timerReceiver.getKey().getKey(),
                    timerReceiver.getKey().getValue(),
                    timerReceiver.getValue().getCoder(),
                    timerReceiver.getValue().getReceiver()));
          }
          beamFnDataInboundObserver =
              Optional.of(BeamFnDataInboundObserver.forConsumers(dataEndpoints, timerEndpoints));
          fnApiDataService.registerReceiver(bundleId, beamFnDataInboundObserver.get());
        }
  
        ImmutableMap.Builder<LogicalEndpoint, FnDataReceiver<?>> receiverBuilder =
            ImmutableMap.builder();
        BeamFnDataOutboundAggregator beamFnDataOutboundAggregator =
            fnApiDataService.createOutboundAggregator(() -> bundleId, false);
        for (RemoteInputDestination remoteInput : remoteInputs) {
          LogicalEndpoint endpoint = LogicalEndpoint.data(bundleId, remoteInput.getPTransformId());
          receiverBuilder.put(
              endpoint,
              new CountingFnDataReceiver(
                  beamFnDataOutboundAggregator.registerOutputDataLocation(
                      remoteInput.getPTransformId(), remoteInput.getCoder())));
        }
  
        for (Map.Entry<String, Map<String, TimerSpec>> entry : timerSpecs.entrySet()) {
          for (TimerSpec timerSpec : entry.getValue().values()) {
            LogicalEndpoint endpoint =
                LogicalEndpoint.timer(bundleId, timerSpec.transformId(), timerSpec.timerId());
            receiverBuilder.put(
                endpoint,
                beamFnDataOutboundAggregator.registerOutputTimersLocation(
                    timerSpec.transformId(), timerSpec.timerId(), timerSpec.coder()));
          }
        }
        beamFnDataOutboundAggregator.start();
  
        return new ActiveBundle(
            bundleId,
            specificResponse,
            beamFnDataOutboundAggregator,
            receiverBuilder.build(),
            beamFnDataInboundObserver,
            stateDelegator.registerForProcessBundleInstructionId(bundleId, stateRequestHandler),
            progressHandler,
            splitHandler,
            checkpointHandler,
            finalizationHandler);
      }
  
      /** An active bundle for a particular {@link BeamFnApi.ProcessBundleDescriptor}. */
      public class ActiveBundle implements RemoteBundle {
        private final String bundleId;
        private final CompletionStage<BeamFnApi.ProcessBundleResponse> response;
        private final BeamFnDataOutboundAggregator beamFnDataOutboundAggregator;
        private final Map<LogicalEndpoint, FnDataReceiver<?>> inputReceivers;
        private final Optional<BeamFnDataInboundObserver> beamFnDataInboundObserver;
        private final StateDelegator.Registration stateRegistration;
        private final BundleProgressHandler progressHandler;
        private final BundleSplitHandler splitHandler;
        private final BundleCheckpointHandler checkpointHandler;
        private final BundleFinalizationHandler finalizationHandler;
        private final Phaser outstandingRequests;
        private final AtomicBoolean isClosed;
        private boolean bundleIsCompleted;
  
        private ActiveBundle(
            String bundleId,
            CompletionStage<ProcessBundleResponse> response,
            BeamFnDataOutboundAggregator beamFnDataOutboundAggregator,
            Map<LogicalEndpoint, FnDataReceiver<?>> inputReceivers,
            Optional<BeamFnDataInboundObserver> beamFnDataInboundObserver,
            StateDelegator.Registration stateRegistration,
            BundleProgressHandler progressHandler,
            BundleSplitHandler splitHandler,
            BundleCheckpointHandler checkpointHandler,
            BundleFinalizationHandler finalizationHandler) {
          this.bundleId = bundleId;
          this.response = response;
          this.beamFnDataOutboundAggregator = beamFnDataOutboundAggregator;
          this.inputReceivers = inputReceivers;
          this.beamFnDataInboundObserver = beamFnDataInboundObserver;
          this.stateRegistration = stateRegistration;
          this.progressHandler = progressHandler;
          this.splitHandler = splitHandler;
          this.checkpointHandler = checkpointHandler;
          this.finalizationHandler = finalizationHandler;
          this.outstandingRequests = new Phaser(1);
          this.isClosed = new AtomicBoolean(false);
  
          // Ensure that we mark when the bundle is completed
          this.response.whenComplete(
              (processBundleResponse, throwable) -> {
                synchronized (ActiveBundle.this) {
                  this.bundleIsCompleted = true;
                }
              });
        }
  
        /** Returns an id used to represent this bundle. */
        @Override
        public String getId() {
          return bundleId;
        }
  
        /**
         * Get a map of PCollection ids to {@link FnDataReceiver receiver}s which consume input
         * elements, forwarding them to the remote environment.
         */
        @Override
        public Map<String, FnDataReceiver> getInputReceivers() {
          ImmutableMap.Builder<String, FnDataReceiver> rval = ImmutableMap.builder();
          for (Map.Entry<LogicalEndpoint, FnDataReceiver<?>> entry : inputReceivers.entrySet()) {
            if (!entry.getKey().isTimer()) {
              rval.put(entry.getKey().getTransformId(), entry.getValue());
            }
          }
          return rval.build();
        }
  
        @Override
        public Map<KV<String, String>, FnDataReceiver<Timer>> getTimerReceivers() {
          ImmutableMap.Builder<KV<String, String>, FnDataReceiver<Timer>> rval =
              ImmutableMap.builder();
          for (Map.Entry<LogicalEndpoint, FnDataReceiver<?>> entry : inputReceivers.entrySet()) {
            if (entry.getKey().isTimer()) {
              rval.put(
                  KV.of(entry.getKey().getTransformId(), entry.getKey().getTimerFamilyId()),
                  (FnDataReceiver<Timer>) entry.getValue());
            }
          }
          return rval.build();
        }
  
        @Override
        public void requestProgress() {
          synchronized (this) {
            if (bundleIsCompleted) {
              return;
            }
            outstandingRequests.register();
          }
          InstructionRequest request =
              InstructionRequest.newBuilder()
                  .setInstructionId(idGenerator.getId())
                  .setProcessBundleProgress(
                      ProcessBundleProgressRequest.newBuilder().setInstructionId(bundleId).build())
                  .build();
          CompletionStage<InstructionResponse> response = fnApiControlClient.handle(request);
          response
              .whenComplete(
                  (instructionResponse, throwable) -> {
                    // Don't forward empty responses.
                    if (ProcessBundleProgressResponse.getDefaultInstance()
                        .equals(instructionResponse.getProcessBundleProgress())) {
                      return;
                    }
                    progressHandler.onProgress(instructionResponse.getProcessBundleProgress());
                  })
              .whenComplete((instructionResponse, throwable) -> outstandingRequests.arrive());
        }
  
        @Override
        public void split(double fractionOfRemainder) {
          synchronized (this) {
            if (bundleIsCompleted) {
              return;
            }
            outstandingRequests.register();
          }
          Map<String, DesiredSplit> splits = new HashMap<>();
          for (Map.Entry<LogicalEndpoint, FnDataReceiver<?>> ptransformToInput :
              inputReceivers.entrySet()) {
            if (!ptransformToInput.getKey().isTimer()) {
              splits.put(
                  ptransformToInput.getKey().getTransformId(),
                  DesiredSplit.newBuilder()
                      .setFractionOfRemainder(fractionOfRemainder)
                      .setEstimatedInputElements(
                          ((CountingFnDataReceiver) ptransformToInput.getValue()).getCount())
                      .build());
            }
          }
          InstructionRequest request =
              InstructionRequest.newBuilder()
                  .setInstructionId(idGenerator.getId())
                  .setProcessBundleSplit(
                      ProcessBundleSplitRequest.newBuilder()
                          .setInstructionId(bundleId)
                          .putAllDesiredSplits(splits)
                          .build())
                  .build();
          CompletionStage<InstructionResponse> response = fnApiControlClient.handle(request);
          response
              .whenComplete(
                  (instructionResponse, throwable) -> {
                    // Don't forward empty responses representing the failure to split.
                    if (ProcessBundleSplitResponse.getDefaultInstance()
                        .equals(instructionResponse.getProcessBundleSplit())) {
                      return;
                    }
                    splitHandler.split(instructionResponse.getProcessBundleSplit());
                  })
              .whenComplete((instructionResponse, throwable) -> outstandingRequests.arrive());
        }
  
        /**
         * Blocks until bundle processing is finished. This is comprised of:
         *
         * <ul>
         *   <li>closing each {@link #getInputReceivers() input receiver}.
         *   <li>closing each {@link #getTimerReceivers() timer receiver}.
         *   <li>waiting for the SDK to say that processing the bundle is finished.
         *   <li>waiting for all inbound data clients to complete
         * </ul>
         *
         * <p>This method will throw an exception if bundle processing has failed. {@link
         * Throwable#getSuppressed()} will return all the reasons as to why processing has failed.
         */
        @Override
        public void close() throws Exception {
          if (isClosed.getAndSet(true)) {
            return;
          }
  
          Exception exception = null;
          try {
            beamFnDataOutboundAggregator.sendOrCollectBufferedDataAndFinishOutboundStreams();
          } catch (Exception e) {
            exception = e;
          } finally {
            beamFnDataOutboundAggregator.discard();
          }
          try {
            // We don't have to worry about the completion stage.
            if (exception == null) {
              BeamFnApi.ProcessBundleResponse completedResponse = MoreFutures.get(response);
              if (completedResponse.hasElements()) {
                beamFnDataInboundObserver.get().accept(completedResponse.getElements());
              }
  
              outstandingRequests.arriveAndAwaitAdvance();
  
              progressHandler.onCompleted(completedResponse);
              if (completedResponse.getResidualRootsCount() > 0) {
                checkpointHandler.onCheckpoint(completedResponse);
              }
              if (completedResponse.getRequiresFinalization()) {
                finalizationHandler.requestsFinalization(bundleId);
              }
            } else {
              // TODO: [https://github.com/apache/beam/issues/18756] Handle aborting the bundle being
              // processed.
              throw new IllegalStateException(
                  "Processing bundle failed, TODO: [https://github.com/apache/beam/issues/18756] abort bundle.");
            }
          } catch (Exception e) {
            if (exception == null) {
              exception = e;
            } else {
              exception.addSuppressed(e);
            }
          }
          try {
            if (exception == null) {
              stateRegistration.deregister();
            } else {
              stateRegistration.abort();
            }
          } catch (Exception e) {
            if (exception == null) {
              exception = e;
            } else {
              exception.addSuppressed(e);
            }
          }
          if (beamFnDataInboundObserver.isPresent()) {
            try {
              if (exception == null) {
                beamFnDataInboundObserver.get().awaitCompletion();
                fnApiDataService.unregisterReceiver(bundleId);
              } else {
                beamFnDataInboundObserver.get().close();
              }
            } catch (Exception e) {
              if (exception == null) {
                exception = e;
              } else {
                exception.addSuppressed(e);
              }
            }
          }
          if (exception != null) {
            throw exception;
          }
        }
      }
    }
  
    private final IdGenerator idGenerator;
    private final InstructionRequestHandler fnApiControlClient;
    private final FnDataService fnApiDataService;
  
    private final ConcurrentHashMap<String, BundleProcessor> clientProcessors;
  
    private SdkHarnessClient(
        InstructionRequestHandler fnApiControlClient,
        FnDataService fnApiDataService,
        IdGenerator idGenerator) {
      this.fnApiDataService = fnApiDataService;
      this.idGenerator = idGenerator;
      this.fnApiControlClient = fnApiControlClient;
      this.clientProcessors = new ConcurrentHashMap<>();
    }
  
    public InstructionRequestHandler getInstructionRequestHandler() {
      return fnApiControlClient;
    }
  
    /**
     * Creates a client for a particular SDK harness. It is the responsibility of the caller to ensure
     * that these correspond to the same SDK harness, so control plane and data plane messages can be
     * correctly associated.
     */
    public static SdkHarnessClient usingFnApiClient(
        InstructionRequestHandler fnApiControlClient, FnDataService fnApiDataService) {
      return new SdkHarnessClient(
          fnApiControlClient, fnApiDataService, IdGenerators.incrementingLongs());
    }
  
    public SdkHarnessClient withIdGenerator(IdGenerator idGenerator) {
      return new SdkHarnessClient(fnApiControlClient, fnApiDataService, idGenerator);
    }
  
    /**
     * Provides {@link BundleProcessor} that is capable of processing bundles not containing timers or
     * state accesses such as:
     *
     * <ul>
     *   <li>Side inputs
     *   <li>User state
     *   <li>Remote references
     * </ul>
     *
     * <p>Note that bundle processors are cached based upon the {@link ProcessBundleDescriptor#getId()
     * process bundle descriptor id}. A previously created instance may be returned.
     */
    public BundleProcessor getProcessor(
        BeamFnApi.ProcessBundleDescriptor descriptor,
        List<RemoteInputDestination> remoteInputDesinations) {
      checkState(
          !descriptor.hasStateApiServiceDescriptor(),
          "The %s cannot support a %s containing a state %s.",
          BundleProcessor.class.getSimpleName(),
          BeamFnApi.ProcessBundleDescriptor.class.getSimpleName(),
          Endpoints.ApiServiceDescriptor.class.getSimpleName());
      return getProcessor(descriptor, remoteInputDesinations, NoOpStateDelegator.INSTANCE);
    }
  
    /**
     * Provides {@link BundleProcessor} that is capable of processing bundles not containing timers.
     *
     * <p>Note that bundle processors are cached based upon the {@link ProcessBundleDescriptor#getId()
     * process bundle descriptor id}. A previously created instance may be returned.
     */
    public BundleProcessor getProcessor(
        BeamFnApi.ProcessBundleDescriptor descriptor,
        List<RemoteInputDestination> remoteInputDesinations,
        StateDelegator stateDelegator) {
      checkState(
          !descriptor.hasTimerApiServiceDescriptor(),
          "The %s cannot support a %s containing a timer %s.",
          BundleProcessor.class.getSimpleName(),
          BeamFnApi.ProcessBundleDescriptor.class.getSimpleName(),
          Endpoints.ApiServiceDescriptor.class.getSimpleName());
      return getProcessor(descriptor, remoteInputDesinations, stateDelegator, Collections.EMPTY_MAP);
    }
  
    /**
     * Provides {@link BundleProcessor} that is capable of processing bundles containing timers and
     * state accesses such as:
     *
     * <ul>
     *   <li>Side inputs
     *   <li>User state
     *   <li>Remote references
     * </ul>
     *
     * <p>Note that bundle processors are cached based upon the {@link ProcessBundleDescriptor#getId()
     * process bundle descriptor id}. A previously created instance may be returned.
     */
    public BundleProcessor getProcessor(
        BeamFnApi.ProcessBundleDescriptor descriptor,
        List<RemoteInputDestination> remoteInputDestinations,
        StateDelegator stateDelegator,
        Map<String, Map<String, TimerSpec>> timerSpecs) {
      @SuppressWarnings("unchecked")
      BundleProcessor bundleProcessor =
          clientProcessors.computeIfAbsent(
              descriptor.getId(),
              s -> create(descriptor, remoteInputDestinations, timerSpecs, stateDelegator));
      checkArgument(
          bundleProcessor.processBundleDescriptor.equals(descriptor),
          "The provided %s with id %s collides with an existing %s with the same id but "
              + "containing different contents.",
          BeamFnApi.ProcessBundleDescriptor.class.getSimpleName(),
          descriptor.getId(),
          BeamFnApi.ProcessBundleDescriptor.class.getSimpleName());
      return bundleProcessor;
    }
  
    /**
     * A {@link StateDelegator} that issues zero state requests to any provided {@link
     * StateRequestHandler state handlers}.
     */
    private static class NoOpStateDelegator implements StateDelegator {
      private static final NoOpStateDelegator INSTANCE = new NoOpStateDelegator();
  
      @Override
      public Registration registerForProcessBundleInstructionId(
          String processBundleInstructionId, StateRequestHandler handler) {
        return Registration.INSTANCE;
      }
  
      /** The corresponding registration for a {@link NoOpStateDelegator} that does nothing. */
      private static class Registration implements StateDelegator.Registration {
        private static final Registration INSTANCE = new Registration();
  
        @Override
        public void deregister() {}
  
        @Override
        public void abort() {}
      }
    }
  
    /** A {@link FnDataReceiver} which counts the number of elements that have been accepted. */
    private static class CountingFnDataReceiver<T> implements FnDataReceiver<T> {
      private final FnDataReceiver<T> delegate;
      private long count;
  
      private CountingFnDataReceiver(FnDataReceiver<T> delegate) {
        this.delegate = delegate;
      }
  
      public long getCount() {
        return count;
      }
  
      @Override
      public void accept(T input) throws Exception {
        delegate.accept(input);
        count += 1;
      }
    }
  
    /** Registers a {@link BeamFnApi.ProcessBundleDescriptor} for future processing. */
    private BundleProcessor create(
        BeamFnApi.ProcessBundleDescriptor processBundleDescriptor,
        List<RemoteInputDestination> remoteInputDestinations,
        Map<String, Map<String, TimerSpec>> timerSpecs,
        StateDelegator stateDelegator) {
  
      LOG.debug("Registering {}", processBundleDescriptor);
      // TODO: validate that all the necessary data endpoints are known
      fnApiControlClient.registerProcessBundleDescriptor(processBundleDescriptor);
      BundleProcessor bundleProcessor =
          new BundleProcessor(
              processBundleDescriptor, remoteInputDestinations, timerSpecs, stateDelegator);
  
      return bundleProcessor;
    }
  
    @Override
    public void close() {}
  }
  