/*
 * 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 static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;

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.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.auth.VendoredCredentialsAdapter;
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.FailoverChannel;
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.sdk.values.WindowedValues;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.auth.MoreCallCredentials;
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.checkerframework.checker.initialization.qual.UnderInitialization;
import org.checkerframework.checker.initialization.qual.UnknownInitialization;
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.
 */
@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 ELEMENT_METADATA_SUPPORTED_EXPERIMENT = "element_metadata_supported";

  @SuppressWarnings("unused")
  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 ScheduledExecutorService commitFinalizerCleanupExecutor;
  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 @Nullable 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,
      ScheduledExecutorService commitFinalizerCleanupExecutor,
      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.commitFinalizerCleanupExecutor = commitFinalizerCleanupExecutor;
    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,
            commitFinalizerCleanupExecutor,
            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,
                checkNotNull(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(
      @UnderInitialization()
          StreamingDataflowWorker this, // Use receiver parameter syntax to allow annotation.
      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);
    @SuppressWarnings("methodref.receiver.bound")
    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(checkNotNull(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(
      @UnknownInitialization()
          StreamingDataflowWorker this, // Use receiver parameter syntax to allow annotation.
      long clientId,
      DataflowWorkerHarnessOptions options,
      GrpcWindmillStreamFactory windmillStreamFactory,
      StreamingWorkScheduler streamingWorkScheduler,
      ThrottlingGetDataMetricTracker getDataMetricTracker,
      MemoryMonitor memoryMonitor,
      GrpcDispatcherClient dispatcherClient) {
    WeightedSemaphore<Commit> maxCommitByteSemaphore = Commits.maxCommitByteSemaphore();
    ChannelCache channelCache =
        createChannelCache(options, checkNotNull(configFetcher), dispatcherClient);
    @SuppressWarnings("methodref.receiver.bound")
    FanOutStreamingEngineWorkerHarness fanOutStreamingEngineWorkerHarness =
        FanOutStreamingEngineWorkerHarness.create(
            createJobHeader(options, clientId),
            GetWorkBudget.builder()
                .setItems(chooseMaxBundlesOutstanding(options))
                .setBytes(MAX_GET_WORK_FETCH_BYTES)
                .build(),
            windmillStreamFactory,
            (workItem,
                serializedWorkItemSize,
                watermarks,
                processingContext,
                drainMode,
                appliedFinalizeIds,
                getWorkStreamLatencies) ->
                checkNotNull(computationStateCache)
                    .get(processingContext.computationId())
                    .ifPresent(
                        computationState -> {
                          memoryMonitor.waitForResources("GetWork");
                          streamingWorkScheduler.queueAppliedFinalizeIds(appliedFinalizeIds);
                          streamingWorkScheduler.scheduleWork(
                              computationState,
                              workItem,
                              serializedWorkItemSize,
                              watermarks,
                              processingContext,
                              drainMode,
                              getWorkStreamLatencies);
                        }),
            ChannelCachingRemoteStubFactory.create(options.getGcpCredential(), channelCache),
            GetWorkBudgetDistributors.distributeEvenly(),
            checkNotNull(dispatcherClient),
            commitWorkStream ->
                StreamingEngineWorkCommitter.builder()
                    // Share the commitByteSemaphore across all created workCommitters.
                    .setCommitByteSemaphore(maxCommitByteSemaphore)
                    .setBackendWorkerToken(commitWorkStream.backendWorkerToken())
                    .setOnCommitComplete(this::onCompleteCommit)
                    .setNumCommitSenders(
                        Math.max(options.getWindmillServiceDirectPathCommitThreads(), 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(
      @UnknownInitialization()
          StreamingDataflowWorker this, // Use receiver parameter syntax to allow annotation.
      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(windmillServer);
    @SuppressWarnings("methodref.receiver.bound")
    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(checkNotNull(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_DEFAULT) {
      return;
    }
    boolean directPath = connectivityType == ConnectivityType.CONNECTIVITY_TYPE_DIRECTPATH;
    if ((directPath
            && this.streamingWorkerHarness.get() instanceof FanOutStreamingEngineWorkerHarness)
        || (!directPath && 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;
    if (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(),
              checkNotNull(this.dispatcherClient));
      this.streamingWorkerHarness.set(newHarnessFactoryOutput.streamingWorkerHarness());
      streamingWorkerHarness.get().start();
      LOG.debug("Started FanOutStreamingEngineWorkerHarness");
    } else {
      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(
      WindmillServerStub windmillClient) {
    return StreamPoolHeartbeatSender.create(
        WindmillStreamPool.create(1, GET_DATA_STREAM_TIMEOUT, windmillClient::getDataStream));
  }

  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);
    ScheduledExecutorService commitFinalizerCleanupExecutor =
        Executors.newScheduledThreadPool(
            1,
            new ThreadFactoryBuilder()
                .setNameFormat("FinalizationCallbackCleanup-%d")
                .setDaemon(true)
                .build());
    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,
        commitFinalizerCleanupExecutor,
        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(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(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,
      GrpcDispatcherClient dispatcherClient) {
    ChannelCache channelCache =
        ChannelCache.create(
            (currentFlowControlSettings, serviceAddress) -> {
              // IsolationChannel wrapping FailoverChannel so that each active RPC gets its own
              // FailoverChannel instance. The fallback channel is created lazily, at most once,
              // only if failover is actually needed.
              return IsolationChannel.create(
                  () ->
                      FailoverChannel.create(
                          remoteChannel(
                              serviceAddress,
                              workerOptions.getWindmillServiceRpcChannelAliveTimeoutSec(),
                              currentFlowControlSettings),
                          () ->
                              remoteChannel(
                                  dispatcherClient.getDispatcherEndpoints().iterator().next(),
                                  workerOptions.getWindmillServiceRpcChannelAliveTimeoutSec(),
                                  currentFlowControlSettings),
                          MoreCallCredentials.from(
                              new VendoredCredentialsAdapter(workerOptions.getGcpCredential()))),
                  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);
    ScheduledExecutorService commitFinalizerCleanupExecutor =
        Executors.newScheduledThreadPool(
            1,
            new ThreadFactoryBuilder()
                .setNameFormat("FinalizationCallbackCleanup-%d")
                .setDaemon(true)
                .build());
    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(stubFactory);
    grpcDispatcherClient.consumeWindmillDispatcherEndpoints(
        ImmutableSet.<HostAndPort>builder()
            .add(HostAndPort.fromHost("StreamingDataflowWorkerTest"))
            .build());

    return new StreamingDataflowWorker(
        windmillServer,
        1L,
        configFetcher,
        computationStateCache,
        stateCache,
        workExecutor,
        commitFinalizerCleanupExecutor,
        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();
    if (DataflowRunner.hasExperiment(options, ELEMENT_METADATA_SUPPORTED_EXPERIMENT)) {
      WindowedValues.FullWindowedValueCoder.setMetadataSupported();
    }

    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();
      commitFinalizerCleanupExecutor.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();
    }
  }
}
