All Downloads are FREE. Search and download functionalities are using the official Maven repository.

org.apache.beam.runners.fnexecution.control.SdkHarnessClient Maven / Gradle / Ivy

Go to download

The newest version!
/*
 * 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.
 *
 * 
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 remoteInputs;
    private final Map> timerSpecs;
    private final StateDelegator stateDelegator;

    private BundleProcessor(
        ProcessBundleDescriptor processBundleDescriptor,
        List remoteInputs,
        Map> 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.
     *
     * 
The input channels for the returned {@link ActiveBundle} are derived from the instructions
     * in the {@link BeamFnApi.ProcessBundleDescriptor}.
     *
     * 
NOTE: It is important to {@link #close()} each bundle after all elements are emitted.
     *
     * 
{@code
     * try (ActiveBundle bundle = SdkHarnessClient.newBundle(...)) {
     *   FnDataReceiver inputReceiver =
     *       (FnDataReceiver) bundle.getInputReceivers().get(mainPCollectionId);
     *   // send all main input elements ...
     * }
     * }

     *
     * 
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> 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.
     *
     * 
The input channels for the returned {@link ActiveBundle} are derived from the instructions
     * in the {@link BeamFnApi.ProcessBundleDescriptor}.
     *
     * 
NOTE: It is important to {@link #close()} each bundle after all elements are emitted.
     *
     * 
{@code
     * try (ActiveBundle bundle = SdkHarnessClient.newBundle(...)) {
     *   FnDataReceiver inputReceiver =
     *       (FnDataReceiver) bundle.getInputReceivers().get(mainPCollectionId);
     *   // send all main input elements ...
     * }
     * }

     *
     * 
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> 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.
     *
     * 
The input channels for the returned {@link ActiveBundle} are derived from the instructions
     * in the {@link BeamFnApi.ProcessBundleDescriptor}.
     *
     * 
NOTE: It is important to {@link #close()} each bundle after all elements are emitted.
     *
     * 
{@code
     * try (ActiveBundle bundle = SdkHarnessClient.newBundle(...)) {
     *   FnDataReceiver inputReceiver =
     *       (FnDataReceiver) bundle.getInputReceivers().get(mainPCollectionId);
     *   // send all main input elements ...
     * }
     * }

     *
     * 
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> outputReceivers,
        Map, RemoteOutputReceiver>> 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.
     *
     * 
The input channels for the returned {@link ActiveBundle} are derived from the instructions
     * in the {@link BeamFnApi.ProcessBundleDescriptor}.
     *
     * 
NOTE: It is important to {@link #close()} each bundle after all elements are emitted.
     *
     * 
{@code
     * try (ActiveBundle bundle = SdkHarnessClient.newBundle(...)) {
     *   FnDataReceiver inputReceiver =
     *       (FnDataReceiver) bundle.getInputReceivers().get(mainPCollectionId);
     *   // send all elements ...
     * }
     * }

     */
    public ActiveBundle newBundle(
        Map> outputReceivers,
        Map, RemoteOutputReceiver>> timerReceivers,
        StateRequestHandler stateRequestHandler,
        BundleProgressHandler progressHandler,
        BundleSplitHandler splitHandler,
        BundleCheckpointHandler checkpointHandler,
        BundleFinalizationHandler finalizationHandler) {
      String bundleId = idGenerator.getId();

      final CompletionStage 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 specificResponse =
          genericResponse.thenApply(InstructionResponse::getProcessBundle);
      Optional beamFnDataInboundObserver;
      if (outputReceivers.isEmpty() && timerReceivers.isEmpty()) {
        beamFnDataInboundObserver = Optional.empty();
      } else {
        List> dataEndpoints = new ArrayList<>(outputReceivers.size());
        for (Map.Entry> receiver : outputReceivers.entrySet()) {
          dataEndpoints.add(
              DataEndpoint.create(
                  receiver.getKey(),
                  (Coder) receiver.getValue().getCoder(),
                  (FnDataReceiver) receiver.getValue().getReceiver()));
        }
        List> timerEndpoints = new ArrayList<>(timerReceivers.size());
        for (Map.Entry, RemoteOutputReceiver>> 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> 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> 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 response;
      private final BeamFnDataOutboundAggregator beamFnDataOutboundAggregator;
      private final Map> inputReceivers;
      private final Optional 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 response,
          BeamFnDataOutboundAggregator beamFnDataOutboundAggregator,
          Map> inputReceivers,
          Optional 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 getInputReceivers() {
        ImmutableMap.Builder rval = ImmutableMap.builder();
        for (Map.Entry> entry : inputReceivers.entrySet()) {
          if (!entry.getKey().isTimer()) {
            rval.put(entry.getKey().getTransformId(), entry.getValue());
          }
        }
        return rval.build();
      }

      @Override
      public Map, FnDataReceiver> getTimerReceivers() {
        ImmutableMap.Builder, FnDataReceiver> rval =
            ImmutableMap.builder();
        for (Map.Entry> entry : inputReceivers.entrySet()) {
          if (entry.getKey().isTimer()) {
            rval.put(
                KV.of(entry.getKey().getTransformId(), entry.getKey().getTimerFamilyId()),
                (FnDataReceiver) 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 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 splits = new HashMap<>();
        for (Map.Entry> 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 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:
       *
       * 
    
       *   
  • closing each {@link #getInputReceivers() input receiver}.
       *   
  • closing each {@link #getTimerReceivers() timer receiver}.
       *   
  • waiting for the SDK to say that processing the bundle is finished.
       *   
  • waiting for all inbound data clients to complete
       * 

       *
       * 
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 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:
   *
   * 
    
   *   
  • Side inputs
   *   
  • User state
   *   
  • Remote references
   * 

   *
   * 
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 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.
   *
   * 
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 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:
   *
   * 
    
   *   
  • Side inputs
   *   
  • User state
   *   
  • Remote references
   * 

   *
   * 
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 remoteInputDestinations,
      StateDelegator stateDelegator,
      Map> 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 implements FnDataReceiver {
    private final FnDataReceiver delegate;
    private long count;

    private CountingFnDataReceiver(FnDataReceiver 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 remoteInputDestinations,
      Map> 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() {}
}
    

    

    

                
 
            

    

                
 
            



    
    






    
© 2015 - 2024 Weber Informatics LLC | Privacy Policy