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.dataflow.worker;
  
  import static org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillChannels.remoteChannel;
  
  import com.google.api.services.dataflow.model.MapTask;
  import com.google.auto.value.AutoValue;
  import java.io.PrintWriter;
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  import java.util.Random;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.ConcurrentMap;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  import java.util.concurrent.ScheduledExecutorService;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicBoolean;
  import java.util.concurrent.atomic.AtomicReference;
  import java.util.function.Consumer;
  import java.util.function.Function;
  import java.util.function.Supplier;
  import javax.annotation.Nullable;
  import org.apache.beam.runners.core.metrics.MetricsLogger;
  import org.apache.beam.runners.dataflow.DataflowRunner;
  import org.apache.beam.runners.dataflow.options.DataflowWorkerHarnessOptions;
  import org.apache.beam.runners.dataflow.worker.status.DebugCapture;
  import org.apache.beam.runners.dataflow.worker.status.WorkerStatusPages;
  import org.apache.beam.runners.dataflow.worker.streaming.ComputationState;
  import org.apache.beam.runners.dataflow.worker.streaming.ComputationStateCache;
  import org.apache.beam.runners.dataflow.worker.streaming.StageInfo;
  import org.apache.beam.runners.dataflow.worker.streaming.WeightedSemaphore;
  import org.apache.beam.runners.dataflow.worker.streaming.WorkHeartbeatResponseProcessor;
  import org.apache.beam.runners.dataflow.worker.streaming.config.ComputationConfig;
  import org.apache.beam.runners.dataflow.worker.streaming.config.FixedGlobalConfigHandle;
  import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingApplianceComputationConfigFetcher;
  import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingEngineComputationConfigFetcher;
  import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfig;
  import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfigHandle;
  import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfigHandleImpl;
  import org.apache.beam.runners.dataflow.worker.streaming.harness.FanOutStreamingEngineWorkerHarness;
  import org.apache.beam.runners.dataflow.worker.streaming.harness.SingleSourceWorkerHarness;
  import org.apache.beam.runners.dataflow.worker.streaming.harness.SingleSourceWorkerHarness.GetWorkSender;
  import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingCounters;
  import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingWorkerHarness;
  import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingWorkerStatusPages;
  import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingWorkerStatusReporter;
  import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor;
  import org.apache.beam.runners.dataflow.worker.util.MemoryMonitor;
  import org.apache.beam.runners.dataflow.worker.windmill.ApplianceWindmillClient;
  import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
  import org.apache.beam.runners.dataflow.worker.windmill.Windmill.ConnectivityType;
  import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
  import org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub;
  import org.apache.beam.runners.dataflow.worker.windmill.appliance.JniWindmillApplianceServer;
  import org.apache.beam.runners.dataflow.worker.windmill.client.CloseableStream;
  import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.GetDataStream;
  import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStreamPool;
  import org.apache.beam.runners.dataflow.worker.windmill.client.commits.Commit;
  import org.apache.beam.runners.dataflow.worker.windmill.client.commits.Commits;
  import org.apache.beam.runners.dataflow.worker.windmill.client.commits.CompleteCommit;
  import org.apache.beam.runners.dataflow.worker.windmill.client.commits.StreamingApplianceWorkCommitter;
  import org.apache.beam.runners.dataflow.worker.windmill.client.commits.StreamingEngineWorkCommitter;
  import org.apache.beam.runners.dataflow.worker.windmill.client.commits.WorkCommitter;
  import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.ApplianceGetDataClient;
  import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.GetDataClient;
  import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.StreamPoolGetDataClient;
  import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.ThrottlingGetDataMetricTracker;
  import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.ChannelzServlet;
  import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.GrpcDispatcherClient;
  import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.GrpcWindmillServer;
  import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.GrpcWindmillStreamFactory;
  import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.ChannelCache;
  import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.ChannelCachingRemoteStubFactory;
  import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.IsolationChannel;
  import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillStubFactoryFactory;
  import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillStubFactoryFactoryImpl;
  import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateCache;
  import org.apache.beam.runners.dataflow.worker.windmill.work.budget.GetWorkBudget;
  import org.apache.beam.runners.dataflow.worker.windmill.work.budget.GetWorkBudgetDistributors;
  import org.apache.beam.runners.dataflow.worker.windmill.work.processing.StreamingWorkScheduler;
  import org.apache.beam.runners.dataflow.worker.windmill.work.processing.failures.FailureTracker;
  import org.apache.beam.runners.dataflow.worker.windmill.work.processing.failures.StreamingApplianceFailureTracker;
  import org.apache.beam.runners.dataflow.worker.windmill.work.processing.failures.StreamingEngineFailureTracker;
  import org.apache.beam.runners.dataflow.worker.windmill.work.processing.failures.WorkFailureProcessor;
  import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.ActiveWorkRefresher;
  import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.ApplianceHeartbeatSender;
  import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.HeartbeatSender;
  import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.StreamPoolHeartbeatSender;
  import org.apache.beam.sdk.annotations.Internal;
  import org.apache.beam.sdk.fn.IdGenerator;
  import org.apache.beam.sdk.fn.IdGenerators;
  import org.apache.beam.sdk.fn.JvmInitializers;
  import org.apache.beam.sdk.io.FileSystems;
  import org.apache.beam.sdk.io.gcp.bigquery.BigQuerySinkMetrics;
  import org.apache.beam.sdk.metrics.MetricsEnvironment;
  import org.apache.beam.sdk.util.construction.CoderTranslation;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheStats;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.net.HostAndPort;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder;
  import org.joda.time.Duration;
  import org.joda.time.Instant;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * <b>For internal use only.</b>
   *
   * <p>Implements a Streaming Dataflow worker.
   */
  @SuppressWarnings({
    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
  })
  @Internal
  public final class StreamingDataflowWorker {
  
    /**
     * Sinks are marked 'full' in {@link StreamingModeExecutionContext} once the amount of data sinked
     * (across all the sinks, if there are more than one) reaches this limit. This serves as hint for
     * readers to stop producing more. This can be disabled with 'disable_limiting_bundle_sink_bytes'
     * experiment.
     */
    public static final int MAX_SINK_BYTES = 10_000_000;
  
    private static final Logger LOG = LoggerFactory.getLogger(StreamingDataflowWorker.class);
  
    /**
     * Maximum number of threads for processing. Currently, each thread processes one key at a time.
     */
    private static final int MAX_PROCESSING_THREADS = 300;
  
    /** The idGenerator to generate unique id globally. */
    private static final IdGenerator ID_GENERATOR = IdGenerators.decrementingLongs();
  
    /** Maximum size of the result of a GetWork request. */
    private static final long MAX_GET_WORK_FETCH_BYTES = 64L << 20; // 64m
  
    /** Maximum number of failure stacktraces to report in each update sent to backend. */
    private static final int MAX_FAILURES_TO_REPORT_IN_UPDATE = 1000;
  
    private static final long THREAD_EXPIRATION_TIME_SEC = 60;
    private static final Duration COMMIT_STREAM_TIMEOUT = Duration.standardMinutes(1);
    private static final Duration GET_DATA_STREAM_TIMEOUT = Duration.standardSeconds(30);
    private static final int DEFAULT_STATUS_PORT = 8081;
    private static final Random CLIENT_ID_GENERATOR = new Random();
    private static final String CHANNELZ_PATH = "/channelz";
    private static final String BEAM_FN_API_EXPERIMENT = "beam_fn_api";
    private static final String STREAMING_ENGINE_USE_JOB_SETTINGS_FOR_HEARTBEAT_POOL_EXPERIMENT =
        "streaming_engine_use_job_settings_for_heartbeat_pool";
    // Experiment make the monitor within BoundedQueueExecutor fair
    public static final String BOUNDED_QUEUE_EXECUTOR_USE_FAIR_MONITOR_EXPERIMENT =
        "windmill_bounded_queue_executor_use_fair_monitor";
  
    private final WindmillStateCache stateCache;
    private AtomicReference<StreamingWorkerStatusPages> statusPages = new AtomicReference<>();
    private final ComputationConfig.Fetcher configFetcher;
    private final ComputationStateCache computationStateCache;
    private final BoundedQueueExecutor workUnitExecutor;
    private final AtomicReference<StreamingWorkerHarness> streamingWorkerHarness =
        new AtomicReference<>();
    private final AtomicBoolean running = new AtomicBoolean();
    private final DataflowWorkerHarnessOptions options;
    private final BackgroundMemoryMonitor memoryMonitor;
    private final ReaderCache readerCache;
    private final DataflowExecutionStateSampler sampler = DataflowExecutionStateSampler.instance();
    private final ActiveWorkRefresher activeWorkRefresher;
    private final StreamingWorkerStatusReporter workerStatusReporter;
    private final int numCommitThreads;
    private final Supplier<Instant> clock;
    private final GrpcDispatcherClient dispatcherClient;
    private final ExecutorService harnessSwitchExecutor;
    private final long clientId;
    private final WindmillServerStub windmillServer;
    private final GrpcWindmillStreamFactory windmillStreamFactory;
    private final StreamingWorkScheduler streamingWorkScheduler;
    private final ThrottlingGetDataMetricTracker getDataMetricTracker;
  
    private StreamingDataflowWorker(
        WindmillServerStub windmillServer,
        long clientId,
        ComputationConfig.Fetcher configFetcher,
        ComputationStateCache computationStateCache,
        WindmillStateCache windmillStateCache,
        BoundedQueueExecutor workUnitExecutor,
        DataflowMapTaskExecutorFactory mapTaskExecutorFactory,
        DataflowWorkerHarnessOptions options,
        HotKeyLogger hotKeyLogger,
        Supplier<Instant> clock,
        StreamingWorkerStatusReporter streamingWorkerStatusReporter,
        FailureTracker failureTracker,
        WorkFailureProcessor workFailureProcessor,
        StreamingCounters streamingCounters,
        MemoryMonitor memoryMonitor,
        GrpcWindmillStreamFactory windmillStreamFactory,
        ScheduledExecutorService activeWorkRefreshExecutorFn,
        ConcurrentMap<String, StageInfo> stageInfoMap,
        @Nullable GrpcDispatcherClient dispatcherClient) {
      // Register standard file systems.
      FileSystems.setDefaultPipelineOptions(options);
      this.configFetcher = configFetcher;
      this.computationStateCache = computationStateCache;
      this.stateCache = windmillStateCache;
      this.readerCache =
          new ReaderCache(
              Duration.standardSeconds(options.getReaderCacheTimeoutSec()),
              Executors.newCachedThreadPool());
      this.options = options;
      this.workUnitExecutor = workUnitExecutor;
      this.harnessSwitchExecutor =
          Executors.newSingleThreadExecutor(
              new ThreadFactoryBuilder().setNameFormat("HarnessSwitchExecutor").build());
      this.clock = clock;
      this.memoryMonitor = BackgroundMemoryMonitor.create(memoryMonitor);
      this.numCommitThreads =
          options.isEnableStreamingEngine()
              ? Math.max(options.getWindmillServiceCommitThreads(), 1)
              : 1;
      this.dispatcherClient = dispatcherClient;
      this.clientId = clientId;
      this.windmillServer = windmillServer;
      this.windmillStreamFactory = windmillStreamFactory;
      this.streamingWorkScheduler =
          StreamingWorkScheduler.create(
              options,
              clock,
              readerCache,
              mapTaskExecutorFactory,
              workUnitExecutor,
              this.stateCache::forComputation,
              failureTracker,
              workFailureProcessor,
              streamingCounters,
              hotKeyLogger,
              sampler,
              ID_GENERATOR,
              this.configFetcher.getGlobalConfigHandle(),
              stageInfoMap);
      this.getDataMetricTracker = new ThrottlingGetDataMetricTracker(memoryMonitor);
      StreamingWorkerHarnessFactoryOutput harnessFactoryOutput;
      // Different implementations on whether the harness is streaming engine
      // direct path, streaming engine cloud path, or streaming appliance.
      if (options.isEnableStreamingEngine()) {
        if (options.getIsWindmillServiceDirectPathEnabled()) {
          harnessFactoryOutput =
              createFanOutStreamingEngineWorkerHarness(
                  clientId,
                  options,
                  windmillStreamFactory,
                  streamingWorkScheduler,
                  getDataMetricTracker,
                  memoryMonitor,
                  this.dispatcherClient);
        } else {
          harnessFactoryOutput =
              createSingleSourceWorkerHarness(
                  clientId,
                  options,
                  windmillServer,
                  streamingWorkScheduler,
                  getDataMetricTracker,
                  memoryMonitor);
        }
      } else { // Appliance
        harnessFactoryOutput =
            createApplianceWorkerHarness(
                clientId,
                options,
                windmillServer,
                streamingWorkScheduler,
                getDataMetricTracker,
                memoryMonitor);
      }
      this.streamingWorkerHarness.set(harnessFactoryOutput.streamingWorkerHarness());
      this.workerStatusReporter = streamingWorkerStatusReporter;
      this.activeWorkRefresher =
          new ActiveWorkRefresher(
              clock,
              options.getActiveWorkRefreshPeriodMillis(),
              options.isEnableStreamingEngine()
                  ? Math.max(options.getStuckCommitDurationMillis(), 0)
                  : 0,
              computationStateCache::getAllPresentComputations,
              sampler,
              activeWorkRefreshExecutorFn,
              getDataMetricTracker::trackHeartbeats);
  
      this.statusPages.set(
          createStatusPageBuilder(
                  this.options, this.windmillStreamFactory, this.memoryMonitor.memoryMonitor())
              .setClock(this.clock)
              .setClientId(this.clientId)
              .setIsRunning(this.running)
              .setStateCache(this.stateCache)
              .setComputationStateCache(this.computationStateCache)
              .setWorkUnitExecutor(this.workUnitExecutor)
              .setGlobalConfigHandle(this.configFetcher.getGlobalConfigHandle())
              .setChannelzServlet(harnessFactoryOutput.channelzServlet())
              .setGetDataStatusProvider(harnessFactoryOutput.getDataStatusProvider())
              .setCurrentActiveCommitBytes(harnessFactoryOutput.currentActiveCommitBytesProvider())
              .setChannelCache(harnessFactoryOutput.channelCache())
              .build());
  
      LOG.debug("isDirectPathEnabled: {}", options.getIsWindmillServiceDirectPathEnabled());
      LOG.debug("windmillServiceEnabled: {}", options.isEnableStreamingEngine());
      LOG.debug("WindmillServiceEndpoint: {}", options.getWindmillServiceEndpoint());
      LOG.debug("WindmillServicePort: {}", options.getWindmillServicePort());
      LOG.debug("LocalWindmillHostport: {}", options.getLocalWindmillHostport());
    }
  
    private StreamingWorkerHarnessFactoryOutput createApplianceWorkerHarness(
        long clientId,
        DataflowWorkerHarnessOptions options,
        WindmillServerStub windmillServer,
        StreamingWorkScheduler streamingWorkScheduler,
        ThrottlingGetDataMetricTracker getDataMetricTracker,
        MemoryMonitor memoryMonitor) {
      Windmill.GetWorkRequest request =
          Windmill.GetWorkRequest.newBuilder()
              .setClientId(clientId)
              .setMaxItems(chooseMaxBundlesOutstanding(options))
              .setMaxBytes(MAX_GET_WORK_FETCH_BYTES)
              .build();
  
      GetDataClient getDataClient = new ApplianceGetDataClient(windmillServer, getDataMetricTracker);
      HeartbeatSender heartbeatSender = new ApplianceHeartbeatSender(windmillServer::getData);
      WorkCommitter workCommitter =
          StreamingApplianceWorkCommitter.create(windmillServer::commitWork, this::onCompleteCommit);
      GetWorkSender getWorkSender = GetWorkSender.forAppliance(() -> windmillServer.getWork(request));
  
      return StreamingWorkerHarnessFactoryOutput.builder()
          .setStreamingWorkerHarness(
              SingleSourceWorkerHarness.builder()
                  .setStreamingWorkScheduler(streamingWorkScheduler)
                  .setWorkCommitter(workCommitter)
                  .setGetDataClient(getDataClient)
                  .setComputationStateFetcher(this.computationStateCache::get)
                  .setWaitForResources(() -> memoryMonitor.waitForResources("GetWork"))
                  .setHeartbeatSender(heartbeatSender)
                  .setGetWorkSender(getWorkSender)
                  .build())
          .setGetDataStatusProvider(getDataClient::printHtml)
          .setCurrentActiveCommitBytesProvider(workCommitter::currentActiveCommitBytes)
          .setChannelzServlet(null) // Appliance doesn't use ChannelzServlet
          .setChannelCache(null) // Appliance doesn't use ChannelCache
          .build();
    }
  
    private StreamingWorkerHarnessFactoryOutput createFanOutStreamingEngineWorkerHarness(
        long clientId,
        DataflowWorkerHarnessOptions options,
        GrpcWindmillStreamFactory windmillStreamFactory,
        StreamingWorkScheduler streamingWorkScheduler,
        ThrottlingGetDataMetricTracker getDataMetricTracker,
        MemoryMonitor memoryMonitor,
        GrpcDispatcherClient dispatcherClient) {
      WeightedSemaphore<Commit> maxCommitByteSemaphore = Commits.maxCommitByteSemaphore();
      ChannelCache channelCache = createChannelCache(options, configFetcher);
      FanOutStreamingEngineWorkerHarness fanOutStreamingEngineWorkerHarness =
          FanOutStreamingEngineWorkerHarness.create(
              createJobHeader(options, clientId),
              GetWorkBudget.builder()
                  .setItems(chooseMaxBundlesOutstanding(options))
                  .setBytes(MAX_GET_WORK_FETCH_BYTES)
                  .build(),
              windmillStreamFactory,
              (workItem,
                  serializedWorkItemSize,
                  watermarks,
                  processingContext,
                  getWorkStreamLatencies) ->
                  computationStateCache
                      .get(processingContext.computationId())
                      .ifPresent(
                          computationState -> {
                            memoryMonitor.waitForResources("GetWork");
                            streamingWorkScheduler.scheduleWork(
                                computationState,
                                workItem,
                                serializedWorkItemSize,
                                watermarks,
                                processingContext,
                                getWorkStreamLatencies);
                          }),
              ChannelCachingRemoteStubFactory.create(options.getGcpCredential(), channelCache),
              GetWorkBudgetDistributors.distributeEvenly(),
              Preconditions.checkNotNull(dispatcherClient),
              commitWorkStream ->
                  StreamingEngineWorkCommitter.builder()
                      // Share the commitByteSemaphore across all created workCommitters.
                      .setCommitByteSemaphore(maxCommitByteSemaphore)
                      .setBackendWorkerToken(commitWorkStream.backendWorkerToken())
                      .setOnCommitComplete(this::onCompleteCommit)
                      .setNumCommitSenders(Math.max(options.getWindmillServiceCommitThreads(), 1))
                      .setCommitWorkStreamFactory(
                          () -> CloseableStream.create(commitWorkStream, () -> {}))
                      .build(),
              getDataMetricTracker);
      ChannelzServlet channelzServlet =
          createChannelzServlet(
              options, fanOutStreamingEngineWorkerHarness::currentWindmillEndpoints);
      return StreamingWorkerHarnessFactoryOutput.builder()
          .setStreamingWorkerHarness(fanOutStreamingEngineWorkerHarness)
          .setGetDataStatusProvider(getDataMetricTracker::printHtml)
          .setCurrentActiveCommitBytesProvider(
              fanOutStreamingEngineWorkerHarness::currentActiveCommitBytes)
          .setChannelzServlet(channelzServlet)
          .setChannelCache(channelCache)
          .build();
    }
  
    private StreamingWorkerHarnessFactoryOutput createSingleSourceWorkerHarness(
        long clientId,
        DataflowWorkerHarnessOptions options,
        WindmillServerStub windmillServer,
        StreamingWorkScheduler streamingWorkScheduler,
        ThrottlingGetDataMetricTracker getDataMetricTracker,
        MemoryMonitor memoryMonitor) {
      Windmill.GetWorkRequest request =
          Windmill.GetWorkRequest.newBuilder()
              .setClientId(clientId)
              .setMaxItems(chooseMaxBundlesOutstanding(options))
              .setMaxBytes(MAX_GET_WORK_FETCH_BYTES)
              .build();
      WindmillStreamPool<GetDataStream> getDataStreamPool =
          WindmillStreamPool.create(
              Math.max(1, options.getWindmillGetDataStreamCount()),
              GET_DATA_STREAM_TIMEOUT,
              windmillServer::getDataStream);
      GetDataClient getDataClient =
          new StreamPoolGetDataClient(getDataMetricTracker, getDataStreamPool);
      HeartbeatSender heartbeatSender =
          createStreamingEngineHeartbeatSender(
              options, windmillServer, getDataStreamPool, configFetcher.getGlobalConfigHandle());
      WorkCommitter workCommitter =
          StreamingEngineWorkCommitter.builder()
              .setCommitWorkStreamFactory(
                  WindmillStreamPool.create(
                          numCommitThreads, COMMIT_STREAM_TIMEOUT, windmillServer::commitWorkStream)
                      ::getCloseableStream)
              .setCommitByteSemaphore(Commits.maxCommitByteSemaphore())
              .setNumCommitSenders(numCommitThreads)
              .setOnCommitComplete(this::onCompleteCommit)
              .build();
      GetWorkSender getWorkSender =
          GetWorkSender.forStreamingEngine(
              receiver -> windmillServer.getWorkStream(request, receiver));
      ChannelzServlet channelzServlet =
          createChannelzServlet(options, windmillServer::getWindmillServiceEndpoints);
      return StreamingWorkerHarnessFactoryOutput.builder()
          .setStreamingWorkerHarness(
              SingleSourceWorkerHarness.builder()
                  .setStreamingWorkScheduler(streamingWorkScheduler)
                  .setWorkCommitter(workCommitter)
                  .setGetDataClient(getDataClient)
                  .setComputationStateFetcher(this.computationStateCache::get)
                  .setWaitForResources(() -> memoryMonitor.waitForResources("GetWork"))
                  .setHeartbeatSender(heartbeatSender)
                  .setGetWorkSender(getWorkSender)
                  .build())
          .setGetDataStatusProvider(getDataClient::printHtml)
          .setCurrentActiveCommitBytesProvider(workCommitter::currentActiveCommitBytes)
          .setChannelzServlet(channelzServlet)
          .setChannelCache(null) // SingleSourceWorkerHarness doesn't use ChannelCache
          .build();
    }
  
    private void switchStreamingWorkerHarness(ConnectivityType connectivityType) {
      if ((connectivityType == ConnectivityType.CONNECTIVITY_TYPE_DIRECTPATH
              && this.streamingWorkerHarness.get() instanceof FanOutStreamingEngineWorkerHarness)
          || (connectivityType == ConnectivityType.CONNECTIVITY_TYPE_CLOUDPATH
              && streamingWorkerHarness.get() instanceof SingleSourceWorkerHarness)) {
        return;
      }
      // Stop the current status pages before switching the harness.
      this.statusPages.get().stop();
      LOG.debug("Stopped StreamingWorkerStatusPages before switching connectivity type.");
      StreamingWorkerHarnessFactoryOutput newHarnessFactoryOutput = null;
      if (connectivityType == ConnectivityType.CONNECTIVITY_TYPE_DIRECTPATH) {
        // If dataflow experiment `enable_windmill_service_direct_path` is not set for
        // the job, do not switch to FanOutStreamingEngineWorkerHarness. This is because
        // `enable_windmill_service_direct_path` is tied to SDK version and is only
        // enabled for job running with SDK above the cut off version,
        // and we do not want jobs below the cutoff to switch to
        // FanOutStreamingEngineWorkerHarness
        if (!options.getIsWindmillServiceDirectPathEnabled()) {
          LOG.info(
              "Dataflow experiment `enable_windmill_service_direct_path` is not set for the job. Job"
                  + " cannot switch to connectivity type DIRECTPATH. Job will continue running on"
                  + " CLOUDPATH");
          return;
        }
        LOG.info("Switching connectivity type from CLOUDPATH to DIRECTPATH");
        LOG.debug("Shutting down to SingleSourceWorkerHarness");
        this.streamingWorkerHarness.get().shutdown();
        newHarnessFactoryOutput =
            createFanOutStreamingEngineWorkerHarness(
                this.clientId,
                this.options,
                this.windmillStreamFactory,
                this.streamingWorkScheduler,
                this.getDataMetricTracker,
                this.memoryMonitor.memoryMonitor(),
                this.dispatcherClient);
        this.streamingWorkerHarness.set(newHarnessFactoryOutput.streamingWorkerHarness());
        streamingWorkerHarness.get().start();
        LOG.debug("Started FanOutStreamingEngineWorkerHarness");
      } else if (connectivityType == ConnectivityType.CONNECTIVITY_TYPE_CLOUDPATH) {
        LOG.info("Switching connectivity type from DIRECTPATH to CLOUDPATH");
        LOG.debug("Shutting down FanOutStreamingEngineWorkerHarness");
        streamingWorkerHarness.get().shutdown();
        newHarnessFactoryOutput =
            createSingleSourceWorkerHarness(
                this.clientId,
                this.options,
                this.windmillServer,
                this.streamingWorkScheduler,
                this.getDataMetricTracker,
                this.memoryMonitor.memoryMonitor());
        this.streamingWorkerHarness.set(newHarnessFactoryOutput.streamingWorkerHarness());
        streamingWorkerHarness.get().start();
        LOG.debug("Started SingleSourceWorkerHarness");
      }
      this.statusPages.set(
          createStatusPageBuilder(
                  this.options, this.windmillStreamFactory, this.memoryMonitor.memoryMonitor())
              .setClock(this.clock)
              .setClientId(this.clientId)
              .setIsRunning(this.running)
              .setStateCache(this.stateCache)
              .setComputationStateCache(this.computationStateCache)
              .setWorkUnitExecutor(this.workUnitExecutor)
              .setGlobalConfigHandle(this.configFetcher.getGlobalConfigHandle())
              .setChannelzServlet(newHarnessFactoryOutput.channelzServlet())
              .setGetDataStatusProvider(newHarnessFactoryOutput.getDataStatusProvider())
              .setCurrentActiveCommitBytes(newHarnessFactoryOutput.currentActiveCommitBytesProvider())
              .setChannelCache(newHarnessFactoryOutput.channelCache())
              .build());
      this.statusPages.get().start(this.options);
      LOG.info("Started new StreamingWorkerStatusPages instance.");
    }
  
    private static StreamingWorkerStatusPages.Builder createStatusPageBuilder(
        DataflowWorkerHarnessOptions options,
        GrpcWindmillStreamFactory windmillStreamFactory,
        MemoryMonitor memoryMonitor) {
      WorkerStatusPages workerStatusPages =
          WorkerStatusPages.create(DEFAULT_STATUS_PORT, memoryMonitor);
  
      StreamingWorkerStatusPages.Builder streamingStatusPages =
          StreamingWorkerStatusPages.builder().setStatusPages(workerStatusPages);
  
      return options.isEnableStreamingEngine()
          ? streamingStatusPages
              .setDebugCapture(
                  new DebugCapture.Manager(options, workerStatusPages.getDebugCapturePages()))
              .setWindmillStreamFactory(windmillStreamFactory)
          : streamingStatusPages;
    }
  
    private static ChannelzServlet createChannelzServlet(
        DataflowWorkerHarnessOptions options,
        Supplier<ImmutableSet<HostAndPort>> windmillEndpointProvider) {
      return new ChannelzServlet(CHANNELZ_PATH, options, windmillEndpointProvider);
    }
  
    private static HeartbeatSender createStreamingEngineHeartbeatSender(
        DataflowWorkerHarnessOptions options,
        WindmillServerStub windmillClient,
        WindmillStreamPool<GetDataStream> getDataStreamPool,
        StreamingGlobalConfigHandle globalConfigHandle) {
      // Experiment gates the logic till backend changes are rollback safe
      if (!DataflowRunner.hasExperiment(
              options, STREAMING_ENGINE_USE_JOB_SETTINGS_FOR_HEARTBEAT_POOL_EXPERIMENT)
          || options.getUseSeparateWindmillHeartbeatStreams() != null) {
        return StreamPoolHeartbeatSender.create(
            Boolean.TRUE.equals(options.getUseSeparateWindmillHeartbeatStreams())
                ? WindmillStreamPool.create(1, GET_DATA_STREAM_TIMEOUT, windmillClient::getDataStream)
                : getDataStreamPool);
  
      } else {
        return StreamPoolHeartbeatSender.create(
            WindmillStreamPool.create(1, GET_DATA_STREAM_TIMEOUT, windmillClient::getDataStream),
            getDataStreamPool,
            globalConfigHandle);
      }
    }
  
    public static StreamingDataflowWorker fromOptions(DataflowWorkerHarnessOptions options) {
      long clientId = CLIENT_ID_GENERATOR.nextLong();
      MemoryMonitor memoryMonitor = MemoryMonitor.fromOptions(options);
      ConcurrentMap<String, StageInfo> stageInfo = new ConcurrentHashMap<>();
      StreamingCounters streamingCounters = StreamingCounters.create();
      WorkUnitClient dataflowServiceClient = new DataflowWorkUnitClient(options, LOG);
      BoundedQueueExecutor workExecutor = createWorkUnitExecutor(options);
      WindmillStateCache windmillStateCache =
          WindmillStateCache.builder()
              .setSizeMb(options.getWorkerCacheMb())
              .setSupportMapViaMultimap(options.isEnableStreamingEngine())
              .build();
  
      GrpcWindmillStreamFactory.Builder windmillStreamFactoryBuilder =
          createGrpcwindmillStreamFactoryBuilder(options, clientId);
  
      ConfigFetcherComputationStateCacheAndWindmillClient
          configFetcherComputationStateCacheAndWindmillClient =
              createConfigFetcherComputationStateCacheAndWindmillClient(
                  options,
                  dataflowServiceClient,
                  windmillStreamFactoryBuilder,
                  configFetcher ->
                      ComputationStateCache.create(
                          configFetcher,
                          workExecutor,
                          windmillStateCache::forComputation,
                          ID_GENERATOR));
  
      ComputationStateCache computationStateCache =
          configFetcherComputationStateCacheAndWindmillClient.computationStateCache();
      WindmillServerStub windmillServer =
          configFetcherComputationStateCacheAndWindmillClient.windmillServer();
  
      FailureTracker failureTracker =
          options.isEnableStreamingEngine()
              ? StreamingEngineFailureTracker.create(
                  MAX_FAILURES_TO_REPORT_IN_UPDATE, options.getMaxStackTraceDepthToReport())
              : StreamingApplianceFailureTracker.create(
                  MAX_FAILURES_TO_REPORT_IN_UPDATE,
                  options.getMaxStackTraceDepthToReport(),
                  windmillServer::reportStats);
  
      Supplier<Instant> clock = Instant::now;
      WorkFailureProcessor workFailureProcessor =
          WorkFailureProcessor.create(
              workExecutor,
              failureTracker,
              () -> Optional.ofNullable(memoryMonitor.tryToDumpHeap()),
              clock);
      StreamingWorkerStatusReporter workerStatusReporter =
          StreamingWorkerStatusReporter.builder()
              .setDataflowServiceClient(dataflowServiceClient)
              .setAllStageInfo(stageInfo::values)
              .setFailureTracker(failureTracker)
              .setStreamingCounters(streamingCounters)
              .setMemoryMonitor(memoryMonitor)
              .setWorkExecutor(workExecutor)
              .setWindmillHarnessUpdateReportingPeriodMillis(
                  options.getWindmillHarnessUpdateReportingPeriod().getMillis())
              .setPerWorkerMetricsUpdateReportingPeriodMillis(
                  options.getPerWorkerMetricsUpdateReportingPeriodMillis())
              .build();
  
      return new StreamingDataflowWorker(
          windmillServer,
          clientId,
          configFetcherComputationStateCacheAndWindmillClient.configFetcher(),
          computationStateCache,
          windmillStateCache,
          workExecutor,
          IntrinsicMapTaskExecutorFactory.defaultFactory(),
          options,
          new HotKeyLogger(),
          clock,
          workerStatusReporter,
          failureTracker,
          workFailureProcessor,
          streamingCounters,
          memoryMonitor,
          configFetcherComputationStateCacheAndWindmillClient.windmillStreamFactory(),
          Executors.newSingleThreadScheduledExecutor(
              new ThreadFactoryBuilder().setNameFormat("RefreshWork").build()),
          stageInfo,
          configFetcherComputationStateCacheAndWindmillClient.windmillDispatcherClient());
    }
  
    /**
     * {@link ComputationConfig.Fetcher}, {@link ComputationStateCache}, and {@link
     * WindmillServerStub} are constructed in different orders due to cyclic dependencies depending on
     * the underlying implementation. This method simplifies creating them and returns an object with
     * all of these dependencies initialized.
     */
    private static ConfigFetcherComputationStateCacheAndWindmillClient
        createConfigFetcherComputationStateCacheAndWindmillClient(
            DataflowWorkerHarnessOptions options,
            WorkUnitClient dataflowServiceClient,
            GrpcWindmillStreamFactory.Builder windmillStreamFactoryBuilder,
            Function<ComputationConfig.Fetcher, ComputationStateCache> computationStateCacheFactory) {
      if (options.isEnableStreamingEngine()) {
        GrpcDispatcherClient dispatcherClient =
            GrpcDispatcherClient.create(options, new WindmillStubFactoryFactoryImpl(options));
        ComputationConfig.Fetcher configFetcher =
            StreamingEngineComputationConfigFetcher.create(
                options.getGlobalConfigRefreshPeriod().getMillis(), dataflowServiceClient);
        configFetcher.getGlobalConfigHandle().registerConfigObserver(dispatcherClient::onJobConfig);
        ComputationStateCache computationStateCache =
            computationStateCacheFactory.apply(configFetcher);
        GrpcWindmillStreamFactory windmillStreamFactory =
            windmillStreamFactoryBuilder
                .setProcessHeartbeatResponses(
                    new WorkHeartbeatResponseProcessor(computationStateCache::get))
                .setHealthCheckIntervalMillis(
                    options.getWindmillServiceStreamingRpcHealthCheckPeriodMs())
                .build();
        return ConfigFetcherComputationStateCacheAndWindmillClient.builder()
            .setWindmillDispatcherClient(dispatcherClient)
            .setConfigFetcher(configFetcher)
            .setComputationStateCache(computationStateCache)
            .setWindmillStreamFactory(windmillStreamFactory)
            .setWindmillServer(
                GrpcWindmillServer.create(options, windmillStreamFactory, dispatcherClient))
            .build();
      }
  
      // Build with local Windmill client.
      if (options.getWindmillServiceEndpoint() != null
          || options.getLocalWindmillHostport().startsWith("grpc:")) {
        GrpcDispatcherClient dispatcherClient =
            GrpcDispatcherClient.create(options, new WindmillStubFactoryFactoryImpl(options));
        GrpcWindmillStreamFactory windmillStreamFactory =
            windmillStreamFactoryBuilder
                .setHealthCheckIntervalMillis(
                    options.getWindmillServiceStreamingRpcHealthCheckPeriodMs())
                .build();
        GrpcWindmillServer windmillServer =
            GrpcWindmillServer.create(options, windmillStreamFactory, dispatcherClient);
        ComputationConfig.Fetcher configFetcher =
            createApplianceComputationConfigFetcher(windmillServer);
        return ConfigFetcherComputationStateCacheAndWindmillClient.builder()
            .setWindmillDispatcherClient(dispatcherClient)
            .setWindmillServer(windmillServer)
            .setWindmillStreamFactory(windmillStreamFactory)
            .setConfigFetcher(configFetcher)
            .setComputationStateCache(computationStateCacheFactory.apply(configFetcher))
            .build();
      }
  
      WindmillServerStub windmillServer =
          new JniWindmillApplianceServer(options.getLocalWindmillHostport());
      ComputationConfig.Fetcher configFetcher =
          createApplianceComputationConfigFetcher(windmillServer);
      return ConfigFetcherComputationStateCacheAndWindmillClient.builder()
          .setWindmillStreamFactory(windmillStreamFactoryBuilder.build())
          .setWindmillServer(windmillServer)
          .setConfigFetcher(configFetcher)
          .setComputationStateCache(computationStateCacheFactory.apply(configFetcher))
          .build();
    }
  
    private static StreamingApplianceComputationConfigFetcher createApplianceComputationConfigFetcher(
        ApplianceWindmillClient windmillClient) {
      return new StreamingApplianceComputationConfigFetcher(
          windmillClient::getConfig,
          new FixedGlobalConfigHandle(StreamingGlobalConfig.builder().build()));
    }
  
    private static void validateWorkerOptions(DataflowWorkerHarnessOptions options) {
      Preconditions.checkArgument(
          options.isStreaming(),
          "%s instantiated with options indicating batch use",
          StreamingDataflowWorker.class.getName());
  
      Preconditions.checkArgument(
          !DataflowRunner.hasExperiment(options, BEAM_FN_API_EXPERIMENT),
          "%s cannot be main() class with beam_fn_api enabled",
          StreamingDataflowWorker.class.getSimpleName());
    }
  
    private static ChannelCache createChannelCache(
        DataflowWorkerHarnessOptions workerOptions, ComputationConfig.Fetcher configFetcher) {
      ChannelCache channelCache =
          ChannelCache.create(
              (currentFlowControlSettings, serviceAddress) -> {
                // IsolationChannel will create and manage separate RPC channels to the same
                // serviceAddress.
                return IsolationChannel.create(
                    () ->
                        remoteChannel(
                            serviceAddress,
                            workerOptions.getWindmillServiceRpcChannelAliveTimeoutSec(),
                            currentFlowControlSettings),
                    currentFlowControlSettings.getOnReadyThresholdBytes());
              });
      configFetcher
          .getGlobalConfigHandle()
          .registerConfigObserver(
              config ->
                  channelCache.consumeFlowControlSettings(
                      config.userWorkerJobSettings().getFlowControlSettings()));
      return channelCache;
    }
  
    @VisibleForTesting
    static StreamingDataflowWorker forTesting(
        Map<String, String> prePopulatedStateNameMappings,
        WindmillServerStub windmillServer,
        List<MapTask> mapTasks,
        DataflowMapTaskExecutorFactory mapTaskExecutorFactory,
        WorkUnitClient workUnitClient,
        DataflowWorkerHarnessOptions options,
        boolean publishCounters,
        HotKeyLogger hotKeyLogger,
        Supplier<Instant> clock,
        Function<String, ScheduledExecutorService> executorSupplier,
        StreamingGlobalConfigHandleImpl globalConfigHandle,
        int localRetryTimeoutMs,
        StreamingCounters streamingCounters,
        WindmillStubFactoryFactory stubFactory) {
      ConcurrentMap<String, StageInfo> stageInfo = new ConcurrentHashMap<>();
      BoundedQueueExecutor workExecutor = createWorkUnitExecutor(options);
      WindmillStateCache stateCache =
          WindmillStateCache.builder()
              .setSizeMb(options.getWorkerCacheMb())
              .setSupportMapViaMultimap(options.isEnableStreamingEngine())
              .build();
      ComputationConfig.Fetcher configFetcher =
          options.isEnableStreamingEngine()
              ? StreamingEngineComputationConfigFetcher.forTesting(
                  /* hasReceivedGlobalConfig= */ true,
                  options.getGlobalConfigRefreshPeriod().getMillis(),
                  workUnitClient,
                  globalConfigHandle,
                  executorSupplier)
              : new StreamingApplianceComputationConfigFetcher(
                  windmillServer::getConfig, globalConfigHandle);
      configFetcher
          .getGlobalConfigHandle()
          .registerConfigObserver(
              config -> {
                if (config.windmillServiceEndpoints().isEmpty()) {
                  LOG.warn("Received empty windmill service endpoints");
                  return;
                }
                windmillServer.setWindmillServiceEndpoints(config.windmillServiceEndpoints());
              });
      ConcurrentMap<String, String> stateNameMap =
          new ConcurrentHashMap<>(prePopulatedStateNameMappings);
      ComputationStateCache computationStateCache =
          ComputationStateCache.forTesting(
              configFetcher, workExecutor, stateCache::forComputation, ID_GENERATOR, stateNameMap);
      computationStateCache.loadCacheForTesting(
          mapTasks,
          mapTask ->
              new ComputationState(
                  mapTask.getStageName(),
                  mapTask,
                  workExecutor,
                  stateNameMap,
                  stateCache.forComputation(mapTask.getStageName())));
      MemoryMonitor memoryMonitor = MemoryMonitor.fromOptions(options);
      FailureTracker failureTracker =
          options.isEnableStreamingEngine()
              ? StreamingEngineFailureTracker.create(
                  MAX_FAILURES_TO_REPORT_IN_UPDATE, options.getMaxStackTraceDepthToReport())
              : StreamingApplianceFailureTracker.create(
                  MAX_FAILURES_TO_REPORT_IN_UPDATE,
                  options.getMaxStackTraceDepthToReport(),
                  windmillServer::reportStats);
      WorkFailureProcessor workFailureProcessor =
          WorkFailureProcessor.forTesting(
              workExecutor,
              failureTracker,
              () -> Optional.ofNullable(memoryMonitor.tryToDumpHeap()),
              clock,
              localRetryTimeoutMs);
      StreamingWorkerStatusReporter workerStatusReporter =
          StreamingWorkerStatusReporter.builder()
              .setPublishCounters(publishCounters)
              .setDataflowServiceClient(workUnitClient)
              .setAllStageInfo(stageInfo::values)
              .setFailureTracker(failureTracker)
              .setStreamingCounters(streamingCounters)
              .setMemoryMonitor(memoryMonitor)
              .setWorkExecutor(workExecutor)
              .setExecutorFactory(executorSupplier)
              .setWindmillHarnessUpdateReportingPeriodMillis(
                  options.getWindmillHarnessUpdateReportingPeriod().getMillis())
              .setPerWorkerMetricsUpdateReportingPeriodMillis(
                  options.getPerWorkerMetricsUpdateReportingPeriodMillis())
              .build();
  
      GrpcWindmillStreamFactory.Builder windmillStreamFactory =
          createGrpcwindmillStreamFactoryBuilder(options, 1)
              .setProcessHeartbeatResponses(
                  new WorkHeartbeatResponseProcessor(computationStateCache::get));
      GrpcDispatcherClient grpcDispatcherClient = GrpcDispatcherClient.create(options, stubFactory);
      grpcDispatcherClient.consumeWindmillDispatcherEndpoints(
          ImmutableSet.<HostAndPort>builder()
              .add(HostAndPort.fromHost("StreamingDataflowWorkerTest"))
              .build());
  
      return new StreamingDataflowWorker(
          windmillServer,
          1L,
          configFetcher,
          computationStateCache,
          stateCache,
          workExecutor,
          mapTaskExecutorFactory,
          options,
          hotKeyLogger,
          clock,
          workerStatusReporter,
          failureTracker,
          workFailureProcessor,
          streamingCounters,
          memoryMonitor,
          options.isEnableStreamingEngine()
              ? windmillStreamFactory
                  .setHealthCheckIntervalMillis(
                      options.getWindmillServiceStreamingRpcHealthCheckPeriodMs())
                  .build()
              : windmillStreamFactory.build(),
          executorSupplier.apply("RefreshWork"),
          stageInfo,
          grpcDispatcherClient);
    }
  
    private static GrpcWindmillStreamFactory.Builder createGrpcwindmillStreamFactoryBuilder(
        DataflowWorkerHarnessOptions options, long clientId) {
      Duration maxBackoff =
          !options.isEnableStreamingEngine() && options.getLocalWindmillHostport() != null
              ? GrpcWindmillServer.LOCALHOST_MAX_BACKOFF
              : Duration.millis(options.getWindmillServiceStreamMaxBackoffMillis());
      return GrpcWindmillStreamFactory.of(createJobHeader(options, clientId))
          .setWindmillMessagesBetweenIsReadyChecks(options.getWindmillMessagesBetweenIsReadyChecks())
          .setMaxBackOffSupplier(() -> maxBackoff)
          .setLogEveryNStreamFailures(options.getWindmillServiceStreamingLogEveryNStreamFailures())
          .setStreamingRpcBatchLimit(options.getWindmillServiceStreamingRpcBatchLimit())
          .setSendKeyedGetDataRequests(
              !options.isEnableStreamingEngine()
                  || DataflowRunner.hasExperiment(
                      options, "streaming_engine_disable_new_heartbeat_requests"))
          .setRequestBatchedGetWorkResponse(options.getWindmillRequestBatchedGetWorkResponse());
    }
  
    private static JobHeader createJobHeader(DataflowWorkerHarnessOptions options, long clientId) {
      return JobHeader.newBuilder()
          .setJobId(options.getJobId())
          .setProjectId(options.getProject())
          .setWorkerId(options.getWorkerId())
          .setClientId(clientId)
          .build();
    }
  
    private static BoundedQueueExecutor createWorkUnitExecutor(DataflowWorkerHarnessOptions options) {
      boolean useFairMonitor =
          DataflowRunner.hasExperiment(options, BOUNDED_QUEUE_EXECUTOR_USE_FAIR_MONITOR_EXPERIMENT);
      return new BoundedQueueExecutor(
          chooseMaxThreads(options),
          THREAD_EXPIRATION_TIME_SEC,
          TimeUnit.SECONDS,
          chooseMaxBundlesOutstanding(options),
          chooseMaxBytesOutstanding(options),
          new ThreadFactoryBuilder().setNameFormat("DataflowWorkUnits-%d").setDaemon(true).build(),
          useFairMonitor);
    }
  
    public static void main(String[] args) throws Exception {
      JvmInitializers.runOnStartup();
  
      DataflowWorkerHarnessHelper.initializeLogging(StreamingDataflowWorker.class);
      DataflowWorkerHarnessOptions options =
          DataflowWorkerHarnessHelper.initializeGlobalStateAndPipelineOptions(
              StreamingDataflowWorker.class, DataflowWorkerHarnessOptions.class);
      DataflowWorkerHarnessHelper.configureLogging(options);
      validateWorkerOptions(options);
  
      CoderTranslation.verifyModelCodersRegistered();
  
      LOG.debug("Creating StreamingDataflowWorker from options: {}", options);
      StreamingDataflowWorker worker = StreamingDataflowWorker.fromOptions(options);
  
      // Use the MetricsLogger container which is used by BigQueryIO to periodically log process-wide
      // metrics.
      MetricsEnvironment.setProcessWideContainer(new MetricsLogger(null));
  
      if (options.isEnableStreamingEngine()
          && !DataflowRunner.hasExperiment(options, "disable_per_worker_metrics")) {
        enableBigQueryMetrics();
      }
  
      JvmInitializers.runBeforeProcessing(options);
      worker.startStatusPages();
      worker.start();
    }
  
    private static int chooseMaxThreads(DataflowWorkerHarnessOptions options) {
      if (options.getNumberOfWorkerHarnessThreads() != 0) {
        return options.getNumberOfWorkerHarnessThreads();
      }
      return MAX_PROCESSING_THREADS;
    }
  
    private static int chooseMaxBundlesOutstanding(DataflowWorkerHarnessOptions options) {
      int maxBundles = options.getMaxBundlesFromWindmillOutstanding();
      return maxBundles > 0 ? maxBundles : chooseMaxThreads(options) + 100;
    }
  
    private static long chooseMaxBytesOutstanding(DataflowWorkerHarnessOptions options) {
      long maxMem = options.getMaxBytesFromWindmillOutstanding();
      return maxMem > 0 ? maxMem : (Runtime.getRuntime().maxMemory() / 2);
    }
  
    private static void enableBigQueryMetrics() {
      // When enabled, the Pipeline will record Per-Worker metrics that will be piped to DFE.
      StreamingStepMetricsContainer.setEnablePerWorkerMetrics(true);
      // StreamingStepMetricsContainer automatically deletes perWorkerCounters if they are zero-valued
      // for longer than 5 minutes.
      BigQuerySinkMetrics.setSupportMetricsDeletion(true);
      // Support metrics for BigQuery's Streaming Inserts write method.
      BigQuerySinkMetrics.setSupportStreamingInsertsMetrics(true);
    }
  
    @VisibleForTesting
    void reportPeriodicWorkerUpdatesForTest() {
      workerStatusReporter.reportPeriodicWorkerUpdates();
    }
  
    @VisibleForTesting
    public boolean workExecutorIsEmpty() {
      return workUnitExecutor.executorQueueIsEmpty();
    }
  
    @VisibleForTesting
    int numCommitThreads() {
      return numCommitThreads;
    }
  
    @VisibleForTesting
    CacheStats getStateCacheStats() {
      return stateCache.getCacheStats();
    }
  
    @VisibleForTesting
    ComputationStateCache getComputationStateCache() {
      return computationStateCache;
    }
  
    public void start() {
      running.set(true);
      configFetcher.start();
      memoryMonitor.start();
      streamingWorkerHarness.get().start();
      sampler.start();
      workerStatusReporter.start();
      activeWorkRefresher.start();
      configFetcher
          .getGlobalConfigHandle()
          .registerConfigObserver(
              streamingGlobalConfig -> {
                ConnectivityType connectivityType =
                    streamingGlobalConfig.userWorkerJobSettings().getConnectivityType();
                if (connectivityType != ConnectivityType.CONNECTIVITY_TYPE_DEFAULT) {
                  LOG.debug("Switching to connectivityType: {}.", connectivityType);
                  harnessSwitchExecutor.execute(() -> switchStreamingWorkerHarness(connectivityType));
                }
              });
    }
  
    /** Starts the status page server for debugging. May be omitted for lighter weight testing. */
    private void startStatusPages() {
      statusPages.get().start(options);
    }
  
    @VisibleForTesting
    StreamingWorkerHarness getStreamingWorkerHarness() {
      return streamingWorkerHarness.get();
    }
  
    @VisibleForTesting
    ExecutorService getHarnessSwitchExecutor() {
      return harnessSwitchExecutor;
    }
  
    @VisibleForTesting
    void stop() {
      try {
        configFetcher.stop();
        activeWorkRefresher.stop();
        statusPages.get().stop();
        running.set(false);
        harnessSwitchExecutor.shutdown();
        streamingWorkerHarness.get().shutdown();
        memoryMonitor.shutdown();
        workUnitExecutor.shutdown();
        computationStateCache.closeAndInvalidateAll();
        workerStatusReporter.stop();
      } catch (Exception e) {
        LOG.warn("Exception while shutting down: ", e);
      }
    }
  
    private void onCompleteCommit(CompleteCommit completeCommit) {
      if (completeCommit.status() != Windmill.CommitStatus.OK) {
        readerCache.invalidateReader(
            WindmillComputationKey.create(
                completeCommit.computationId(), completeCommit.shardedKey()));
        stateCache
            .forComputation(completeCommit.computationId())
            .invalidate(completeCommit.shardedKey());
      }
  
      computationStateCache
          .getIfPresent(completeCommit.computationId())
          .ifPresent(
              state ->
                  state.completeWorkAndScheduleNextWorkForKey(
                      completeCommit.shardedKey(), completeCommit.workId()));
    }
  
    @AutoValue
    abstract static class ConfigFetcherComputationStateCacheAndWindmillClient {
  
      private static Builder builder() {
        return new AutoValue_StreamingDataflowWorker_ConfigFetcherComputationStateCacheAndWindmillClient
            .Builder();
      }
  
      abstract ComputationConfig.Fetcher configFetcher();
  
      abstract ComputationStateCache computationStateCache();
  
      abstract WindmillServerStub windmillServer();
  
      abstract GrpcWindmillStreamFactory windmillStreamFactory();
  
      abstract @Nullable GrpcDispatcherClient windmillDispatcherClient();
  
      @AutoValue.Builder
      abstract static class Builder {
  
        abstract Builder setConfigFetcher(ComputationConfig.Fetcher value);
  
        abstract Builder setComputationStateCache(ComputationStateCache value);
  
        abstract Builder setWindmillServer(WindmillServerStub value);
  
        abstract Builder setWindmillStreamFactory(GrpcWindmillStreamFactory value);
  
        abstract Builder setWindmillDispatcherClient(GrpcDispatcherClient value);
  
        abstract ConfigFetcherComputationStateCacheAndWindmillClient build();
      }
    }
  
    /**
     * Monitors memory pressure on a background executor. May be used to throttle calls, blocking if
     * there is memory pressure.
     */
    @AutoValue
    abstract static class BackgroundMemoryMonitor {
  
      private static BackgroundMemoryMonitor create(MemoryMonitor memoryMonitor) {
        return new AutoValue_StreamingDataflowWorker_BackgroundMemoryMonitor(
            memoryMonitor,
            Executors.newSingleThreadScheduledExecutor(
                new ThreadFactoryBuilder()
                    .setNameFormat("MemoryMonitor")
                    .setPriority(Thread.MIN_PRIORITY)
                    .build()));
      }
  
      abstract MemoryMonitor memoryMonitor();
  
      abstract ExecutorService executor();
  
      private void start() {
        executor().execute(memoryMonitor());
      }
  
      private void shutdown() {
        memoryMonitor().stop();
        executor().shutdown();
      }
    }
  
    /**
     * Holds the {@link StreamingWorkerHarness} and its associated dependencies that are created
     * together.
     */
    @AutoValue
    abstract static class StreamingWorkerHarnessFactoryOutput {
      static Builder builder() {
        return new AutoValue_StreamingDataflowWorker_StreamingWorkerHarnessFactoryOutput.Builder();
      }
  
      abstract StreamingWorkerHarness streamingWorkerHarness();
  
      abstract Consumer<PrintWriter> getDataStatusProvider();
  
      abstract Supplier<Long> currentActiveCommitBytesProvider();
  
      abstract @Nullable ChannelzServlet channelzServlet();
  
      abstract @Nullable ChannelCache channelCache();
  
      @AutoValue.Builder
      abstract static class Builder {
        abstract Builder setStreamingWorkerHarness(StreamingWorkerHarness value);
  
        abstract Builder setGetDataStatusProvider(Consumer<PrintWriter> value);
  
        abstract Builder setCurrentActiveCommitBytesProvider(Supplier<Long> value);
  
        abstract Builder setChannelzServlet(@Nullable ChannelzServlet value);
  
        abstract Builder setChannelCache(@Nullable ChannelCache value);
  
        abstract StreamingWorkerHarnessFactoryOutput build();
      }
    }
  }
  