/*
 * 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.flink.translation.wrappers.streaming;

import static org.apache.flink.util.Preconditions.checkArgument;

import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.Serializable;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.locks.Lock;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.apache.beam.runners.core.DoFnRunner;
import org.apache.beam.runners.core.DoFnRunners;
import org.apache.beam.runners.core.InMemoryBundleFinalizer;
import org.apache.beam.runners.core.NullSideInputReader;
import org.apache.beam.runners.core.ProcessFnRunner;
import org.apache.beam.runners.core.PushbackSideInputDoFnRunner;
import org.apache.beam.runners.core.SideInputHandler;
import org.apache.beam.runners.core.SideInputReader;
import org.apache.beam.runners.core.SimplePushbackSideInputDoFnRunner;
import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems;
import org.apache.beam.runners.core.StateInternals;
import org.apache.beam.runners.core.StateNamespace;
import org.apache.beam.runners.core.StateNamespaces.WindowNamespace;
import org.apache.beam.runners.core.StatefulDoFnRunner;
import org.apache.beam.runners.core.StepContext;
import org.apache.beam.runners.core.TimerInternals;
import org.apache.beam.runners.core.TimerInternals.TimerData;
import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
import org.apache.beam.runners.flink.FlinkPipelineOptions;
import org.apache.beam.runners.flink.adapter.FlinkKey;
import org.apache.beam.runners.flink.metrics.DoFnRunnerWithMetricsUpdate;
import org.apache.beam.runners.flink.metrics.FlinkMetricContainer;
import org.apache.beam.runners.flink.translation.types.CoderTypeSerializer;
import org.apache.beam.runners.flink.translation.utils.CheckpointStats;
import org.apache.beam.runners.flink.translation.utils.Workarounds;
import org.apache.beam.runners.flink.translation.wrappers.streaming.stableinput.BufferingDoFnRunner;
import org.apache.beam.runners.flink.translation.wrappers.streaming.state.FlinkBroadcastStateInternals;
import org.apache.beam.runners.flink.translation.wrappers.streaming.state.FlinkStateInternals;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.StructuredCoder;
import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.metrics.MetricName;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.state.StateSpec;
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.DoFn.BundleFinalizer;
import org.apache.beam.sdk.transforms.DoFnSchemaInformation;
import org.apache.beam.sdk.transforms.join.RawUnionValue;
import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature;
import org.apache.beam.sdk.transforms.reflect.DoFnSignatures;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.util.NoopLock;
import org.apache.beam.sdk.util.WindowedValueMultiReceiver;
import org.apache.beam.sdk.util.WindowedValueReceiver;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.WindowedValue;
import org.apache.beam.sdk.values.WindowingStrategy;
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.Joiner;
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.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
import org.apache.flink.api.common.operators.ProcessingTimeService.ProcessingTimeCallback;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.state.MapState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.typeutils.base.StringSerializer;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.core.execution.CheckpointingMode;
import org.apache.flink.runtime.state.InternalPriorityQueue;
import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.runtime.state.OperatorStateBackend;
import org.apache.flink.runtime.state.StateInitializationContext;
import org.apache.flink.runtime.state.StateSnapshotContext;
import org.apache.flink.streaming.api.graph.StreamConfig;
import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
import org.apache.flink.streaming.api.operators.InternalTimeServiceManagerImpl;
import org.apache.flink.streaming.api.operators.InternalTimer;
import org.apache.flink.streaming.api.operators.InternalTimerService;
import org.apache.flink.streaming.api.operators.InternalTimerServiceImpl;
import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
import org.apache.flink.streaming.api.operators.Output;
import org.apache.flink.streaming.api.operators.Triggerable;
import org.apache.flink.streaming.api.operators.TwoInputStreamOperator;
import org.apache.flink.streaming.api.operators.sorted.state.BatchExecutionInternalTimeService;
import org.apache.flink.streaming.api.operators.sorted.state.BatchExecutionInternalTimeServiceManager;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.runtime.tasks.ProcessingTimeService;
import org.apache.flink.streaming.runtime.tasks.StreamTask;
import org.apache.flink.util.OutputTag;
import org.apache.flink.util.function.BiConsumerWithException;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.joda.time.Instant;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Flink operator for executing {@link DoFn DoFns}.
 *
 * @param <InputT> the input type of the {@link DoFn}
 * @param <OutputT> the output type of the {@link DoFn}
 */
// We use Flink's lifecycle methods to initialize transient fields
@SuppressFBWarnings("SE_TRANSIENT_FIELD_NOT_RESTORED")
@SuppressWarnings({
  "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
  "keyfor",
  "nullness"
}) // TODO(https://github.com/apache/beam/issues/20497)
public class DoFnOperator<PreInputT, InputT, OutputT>
    extends AbstractStreamOperator<WindowedValue<OutputT>>
    implements OneInputStreamOperator<WindowedValue<PreInputT>, WindowedValue<OutputT>>,
        TwoInputStreamOperator<WindowedValue<PreInputT>, RawUnionValue, WindowedValue<OutputT>>,
        Triggerable<FlinkKey, TimerData> {

  private static final Logger LOG = LoggerFactory.getLogger(DoFnOperator.class);
  private final boolean isStreaming;

  protected DoFn<InputT, OutputT> doFn;

  protected final SerializablePipelineOptions serializedOptions;

  protected final TupleTag<OutputT> mainOutputTag;
  protected final List<TupleTag<?>> additionalOutputTags;

  protected final Collection<PCollectionView<?>> sideInputs;
  protected final Map<Integer, PCollectionView<?>> sideInputTagMapping;

  protected final WindowingStrategy<?, ?> windowingStrategy;

  protected final OutputManagerFactory<OutputT> outputManagerFactory;

  protected transient DoFnRunner<InputT, OutputT> doFnRunner;
  protected transient PushbackSideInputDoFnRunner<InputT, OutputT> pushbackDoFnRunner;
  protected transient BufferingDoFnRunner<InputT, OutputT> bufferingDoFnRunner;

  protected transient SideInputHandler sideInputHandler;

  protected transient SideInputReader sideInputReader;

  protected transient BufferedOutputManager<OutputT> outputManager;

  private transient DoFnInvoker<InputT, OutputT> doFnInvoker;

  protected transient FlinkStateInternals<?> keyedStateInternals;
  protected transient FlinkTimerInternals timerInternals;

  protected final String stepName;

  final Coder<WindowedValue<InputT>> windowedInputCoder;

  final Map<TupleTag<?>, Coder<?>> outputCoders;

  final Coder<?> keyCoder;

  final KeySelector<WindowedValue<InputT>, ?> keySelector;

  final TimerInternals.TimerDataCoderV2 timerCoder;

  /** Max number of elements to include in a bundle. */
  private final long maxBundleSize;
  /** Max duration of a bundle. */
  private final long maxBundleTimeMills;

  private final DoFnSchemaInformation doFnSchemaInformation;

  private final Map<String, PCollectionView<?>> sideInputMapping;

  /** If true, we must process elements only after a checkpoint is finished. */
  final boolean requiresStableInput;

  /**
   * If both requiresStableInput and this parameter are true, we must flush the buffer during drain
   * operation.
   */
  final boolean enableStableInputDrain;

  final int numConcurrentCheckpoints;

  private final boolean usesOnWindowExpiration;

  private final boolean finishBundleBeforeCheckpointing;

  /** Stores new finalizations being gathered. */
  private transient InMemoryBundleFinalizer bundleFinalizer;
  /** Pending bundle finalizations which have not been acknowledged yet. */
  private transient LinkedHashMap<Long, List<InMemoryBundleFinalizer.Finalization>>
      pendingFinalizations;
  /**
   * Keep a maximum of 32 bundle finalizations for {@link
   * BundleFinalizer.Callback#onBundleSuccess()}.
   */
  private static final int MAX_NUMBER_PENDING_BUNDLE_FINALIZATIONS = 32;

  protected transient InternalTimerService<TimerData> timerService;
  // Flink 1.20 moved timeServiceManager to protected scope. No longer need delegate
  // private transient InternalTimeServiceManager<?> timeServiceManager;

  private transient PushedBackElementsHandler<WindowedValue<InputT>> pushedBackElementsHandler;

  /** Metrics container for reporting Beam metrics to Flink (null if metrics are disabled). */
  transient @Nullable FlinkMetricContainer flinkMetricContainer;

  /** Helper class to report the checkpoint duration. */
  private transient @Nullable CheckpointStats checkpointStats;

  /** A timer that finishes the current bundle after a fixed amount of time. */
  private transient ScheduledFuture<?> checkFinishBundleTimer;

  /**
   * This and the below fields need to be volatile because we use multiple threads to access these.
   * (a) the main processing thread (b) a timer thread to finish bundles by a timeout instead of the
   * number of element However, we do not need a lock because Flink makes sure to acquire the
   * "checkpointing" lock for the main processing but also for timer set via its {@code
   * timerService}.
   *
   * <p>The volatile flag can be removed once https://issues.apache.org/jira/browse/FLINK-12481 has
   * been addressed.
   */
  private transient volatile boolean bundleStarted;
  /** Number of processed elements in the current bundle. */
  private transient volatile long elementCount;
  /** Time that the last bundle was finished (to set the timer). */
  private transient volatile long lastFinishBundleTime;
  /** Callback to be executed before the current bundle is started. */
  private transient volatile Runnable preBundleCallback;
  /** Callback to be executed after the current bundle was finished. */
  private transient volatile Runnable bundleFinishedCallback;

  // Watermark state.
  // Volatile because these can be set in two mutually exclusive threads (see above).
  private transient volatile long currentInputWatermark;
  private transient volatile long currentSideInputWatermark;
  private transient volatile long currentOutputWatermark;
  private transient volatile long pushedBackWatermark;

  /** Constructor for DoFnOperator. */
  public DoFnOperator(
      @Nullable DoFn<InputT, OutputT> doFn,
      String stepName,
      Coder<WindowedValue<InputT>> inputWindowedCoder,
      Map<TupleTag<?>, Coder<?>> outputCoders,
      TupleTag<OutputT> mainOutputTag,
      List<TupleTag<?>> additionalOutputTags,
      OutputManagerFactory<OutputT> outputManagerFactory,
      WindowingStrategy<?, ?> windowingStrategy,
      Map<Integer, PCollectionView<?>> sideInputTagMapping,
      Collection<PCollectionView<?>> sideInputs,
      PipelineOptions options,
      @Nullable Coder<?> keyCoder,
      @Nullable KeySelector<WindowedValue<InputT>, ?> keySelector,
      DoFnSchemaInformation doFnSchemaInformation,
      Map<String, PCollectionView<?>> sideInputMapping) {
    this.doFn = doFn;
    this.stepName = stepName;
    this.windowedInputCoder = inputWindowedCoder;
    this.outputCoders = outputCoders;
    this.mainOutputTag = mainOutputTag;
    this.additionalOutputTags = additionalOutputTags;
    this.sideInputTagMapping = sideInputTagMapping;
    this.sideInputs = sideInputs;
    this.serializedOptions = new SerializablePipelineOptions(options);
    this.isStreaming = serializedOptions.get().as(FlinkPipelineOptions.class).isStreaming();
    this.windowingStrategy = windowingStrategy;
    this.outputManagerFactory = outputManagerFactory;

    // API removed in Flink 2.0. setChainingStrategy is now set internally.
    // setChainingStrategy(ChainingStrategy.ALWAYS);

    this.keyCoder = keyCoder;
    this.keySelector = keySelector;

    this.timerCoder =
        TimerInternals.TimerDataCoderV2.of(windowingStrategy.getWindowFn().windowCoder());

    FlinkPipelineOptions flinkOptions = options.as(FlinkPipelineOptions.class);

    this.maxBundleSize = flinkOptions.getMaxBundleSize();
    Preconditions.checkArgument(maxBundleSize > 0, "Bundle size must be at least 1");
    this.maxBundleTimeMills = flinkOptions.getMaxBundleTimeMills();
    Preconditions.checkArgument(maxBundleTimeMills > 0, "Bundle time must be at least 1");
    this.doFnSchemaInformation = doFnSchemaInformation;
    this.sideInputMapping = sideInputMapping;

    this.requiresStableInput = isRequiresStableInput(doFn);

    this.usesOnWindowExpiration =
        doFn != null && DoFnSignatures.getSignature(doFn.getClass()).onWindowExpiration() != null;

    if (requiresStableInput) {
      Preconditions.checkState(
          CheckpointingMode.valueOf(flinkOptions.getCheckpointingMode())
              == CheckpointingMode.EXACTLY_ONCE,
          "Checkpointing mode is not set to exactly once but @RequiresStableInput is used.");
      Preconditions.checkState(
          flinkOptions.getCheckpointingInterval() > 0,
          "No checkpointing configured but pipeline uses @RequiresStableInput");
      LOG.warn(
          "Enabling stable input for transform {}. Will only process elements at most every {} milliseconds.",
          stepName,
          flinkOptions.getCheckpointingInterval()
              + Math.max(0, flinkOptions.getMinPauseBetweenCheckpoints()));
    }

    this.enableStableInputDrain = flinkOptions.getEnableStableInputDrain();

    this.numConcurrentCheckpoints = flinkOptions.getNumConcurrentCheckpoints();

    this.finishBundleBeforeCheckpointing = flinkOptions.getFinishBundleBeforeCheckpointing();
  }

  private boolean isRequiresStableInput(DoFn<InputT, OutputT> doFn) {
    // WindowDoFnOperator does not use a DoFn
    return doFn != null
        && DoFnSignatures.getSignature(doFn.getClass()).processElement().requiresStableInput();
  }

  @VisibleForTesting
  boolean getRequiresStableInput() {
    return requiresStableInput;
  }

  // allow overriding this in WindowDoFnOperator because this one dynamically creates
  // the DoFn
  protected DoFn<InputT, OutputT> getDoFn() {
    return doFn;
  }

  protected Iterable<WindowedValue<InputT>> preProcess(WindowedValue<PreInputT> input) {
    // Assume Input is PreInputT
    return Collections.singletonList((WindowedValue<InputT>) input);
  }

  // allow overriding this, for example SplittableDoFnOperator will not create a
  // stateful DoFn runner because ProcessFn, which is used for executing a Splittable DoFn
  // doesn't play by the normal DoFn rules and WindowDoFnOperator uses LateDataDroppingDoFnRunner
  protected DoFnRunner<InputT, OutputT> createWrappingDoFnRunner(
      DoFnRunner<InputT, OutputT> wrappedRunner, StepContext stepContext) {

    if (keyCoder != null) {
      StatefulDoFnRunner.CleanupTimer<InputT> cleanupTimer =
          new StatefulDoFnRunner.TimeInternalsCleanupTimer<InputT>(
              timerInternals, windowingStrategy) {
            @Override
            public void setForWindow(InputT input, BoundedWindow window) {
              if (!window.equals(GlobalWindow.INSTANCE) || usesOnWindowExpiration) {
                // Skip setting a cleanup timer for the global window as these timers
                // lead to potentially unbounded state growth in the runner, depending on key
                // cardinality. Cleanup for global window will be performed upon arrival of the
                // final watermark.
                // In the case of OnWindowExpiration, we still set the timer.
                super.setForWindow(input, window);
              }
            }
          };

      // we don't know the window type
      //      @SuppressWarnings({"unchecked", "rawtypes"})
      Coder windowCoder = windowingStrategy.getWindowFn().windowCoder();

      @SuppressWarnings({"unchecked"})
      StatefulDoFnRunner.StateCleaner<?> stateCleaner =
          new StatefulDoFnRunner.StateInternalsStateCleaner<>(
              doFn, keyedStateInternals, windowCoder);

      return DoFnRunners.defaultStatefulDoFnRunner(
          doFn,
          getInputCoder(),
          wrappedRunner,
          stepContext,
          windowingStrategy,
          cleanupTimer,
          stateCleaner,
          true /* requiresTimeSortedInput is supported */);

    } else {
      return doFnRunner;
    }
  }

  @Override
  public void setup(
      StreamTask<?, ?> containingTask,
      StreamConfig config,
      Output<StreamRecord<WindowedValue<OutputT>>> output) {

    // make sure that FileSystems is initialized correctly
    FileSystems.setDefaultPipelineOptions(serializedOptions.get());

    super.setup(containingTask, config, output);
  }

  protected boolean shoudBundleElements() {
    return isStreaming;
  }

  @Override
  public void initializeState(StateInitializationContext context) throws Exception {
    super.initializeState(context);

    ListStateDescriptor<WindowedValue<InputT>> pushedBackStateDescriptor =
        new ListStateDescriptor<>(
            "pushed-back-elements",
            new CoderTypeSerializer<>(windowedInputCoder, serializedOptions));

    if (keySelector != null) {
      pushedBackElementsHandler =
          KeyedPushedBackElementsHandler.create(
              keySelector, getKeyedStateBackend(), pushedBackStateDescriptor);
    } else {
      ListState<WindowedValue<InputT>> listState =
          getOperatorStateBackend().getListState(pushedBackStateDescriptor);
      pushedBackElementsHandler = NonKeyedPushedBackElementsHandler.create(listState);
    }

    currentInputWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis();
    currentSideInputWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis();
    currentOutputWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis();

    sideInputReader = NullSideInputReader.of(sideInputs);

    if (!sideInputs.isEmpty()) {

      FlinkBroadcastStateInternals sideInputStateInternals =
          new FlinkBroadcastStateInternals<>(
              getContainingTask().getIndexInSubtaskGroup(),
              getOperatorStateBackend(),
              serializedOptions);

      sideInputHandler = new SideInputHandler(sideInputs, sideInputStateInternals);
      sideInputReader = sideInputHandler;

      Stream<WindowedValue<InputT>> pushedBack = pushedBackElementsHandler.getElements();
      long min =
          pushedBack.map(v -> v.getTimestamp().getMillis()).reduce(Long.MAX_VALUE, Math::min);
      pushedBackWatermark = min;
    } else {
      pushedBackWatermark = Long.MAX_VALUE;
    }

    // StatefulPardo or WindowDoFn
    if (keyCoder != null) {
      keyedStateInternals =
          new FlinkStateInternals<>(
              (KeyedStateBackend) getKeyedStateBackend(),
              keyCoder,
              windowingStrategy.getWindowFn().windowCoder(),
              serializedOptions);

      if (timerService == null) {
        timerService =
            getInternalTimerService(
                "beam-timer", new CoderTypeSerializer<>(timerCoder, serializedOptions), this);
      }

      timerInternals = new FlinkTimerInternals(timerService);
      Preconditions.checkNotNull(getTimeServiceManager(), "Time service manager is not set.");
    }

    outputManager =
        outputManagerFactory.create(
            output, getLockToAcquireForStateAccessDuringBundles(), getOperatorStateBackend());
  }

  /**
   * Subclasses may provide a lock to ensure that the state backend is not accessed concurrently
   * during bundle execution.
   */
  protected Lock getLockToAcquireForStateAccessDuringBundles() {
    return NoopLock.get();
  }

  @Override
  public void open() throws Exception {
    // WindowDoFnOperator need use state and timer to get DoFn.
    // So must wait StateInternals and TimerInternals ready.
    // This will be called after initializeState()
    this.doFn = getDoFn();

    FlinkPipelineOptions options = serializedOptions.get().as(FlinkPipelineOptions.class);
    doFnInvoker = DoFnInvokers.tryInvokeSetupFor(doFn, options);

    StepContext stepContext = new FlinkStepContext();
    doFnRunner =
        DoFnRunners.simpleRunner(
            options,
            doFn,
            sideInputReader,
            outputManager,
            mainOutputTag,
            additionalOutputTags,
            stepContext,
            getInputCoder(),
            outputCoders,
            windowingStrategy,
            doFnSchemaInformation,
            sideInputMapping);

    doFnRunner =
        createBufferingDoFnRunnerIfNeeded(createWrappingDoFnRunner(doFnRunner, stepContext));
    earlyBindStateIfNeeded();

    if (!options.getDisableMetrics()) {
      flinkMetricContainer = new FlinkMetricContainer(getRuntimeContext());
      doFnRunner = new DoFnRunnerWithMetricsUpdate<>(stepName, doFnRunner, flinkMetricContainer);
      String checkpointMetricNamespace = options.getReportCheckpointDuration();
      if (checkpointMetricNamespace != null) {
        MetricName checkpointMetric =
            MetricName.named(checkpointMetricNamespace, "checkpoint_duration");
        checkpointStats =
            new CheckpointStats(
                () ->
                    flinkMetricContainer
                        .getMetricsContainer(stepName)
                        .getDistribution(checkpointMetric));
      }
    }

    elementCount = 0L;
    lastFinishBundleTime = getProcessingTimeService().getCurrentProcessingTime();

    // Schedule timer to check timeout of finish bundle.
    long bundleCheckPeriod = Math.max(maxBundleTimeMills / 2, 1);
    checkFinishBundleTimer =
        getProcessingTimeService()
            .scheduleAtFixedRate(
                timestamp -> checkInvokeFinishBundleByTime(), bundleCheckPeriod, bundleCheckPeriod);

    if (doFn instanceof SplittableParDoViaKeyedWorkItems.ProcessFn) {
      pushbackDoFnRunner =
          new ProcessFnRunner<>((DoFnRunner) doFnRunner, sideInputs, sideInputHandler);
    } else {
      pushbackDoFnRunner =
          SimplePushbackSideInputDoFnRunner.create(doFnRunner, sideInputs, sideInputHandler);
    }

    bundleFinalizer = new InMemoryBundleFinalizer();
    pendingFinalizations = new LinkedHashMap<>();
  }

  DoFnRunner<InputT, OutputT> createBufferingDoFnRunnerIfNeeded(
      DoFnRunner<InputT, OutputT> wrappedRunner) throws Exception {

    if (requiresStableInput) {
      // put this in front of the root FnRunner before any additional wrappers
      return this.bufferingDoFnRunner =
          BufferingDoFnRunner.create(
              wrappedRunner,
              "stable-input-buffer",
              windowedInputCoder,
              windowingStrategy.getWindowFn().windowCoder(),
              getOperatorStateBackend(),
              getBufferingKeyedStateBackend(),
              numConcurrentCheckpoints,
              serializedOptions);
    }
    return wrappedRunner;
  }

  /**
   * Retrieve a keyed state backend that should be used to buffer elements for
   * {@code @RequiresStableInput} functionality. By default this is the default keyed backend, but
   * can be override in {@link ExecutableStageDoFnOperator}.
   *
   * @return the keyed backend to use for element buffering
   */
  <K> @Nullable KeyedStateBackend<K> getBufferingKeyedStateBackend() {
    return getKeyedStateBackend();
  }

  private void earlyBindStateIfNeeded() throws IllegalArgumentException, IllegalAccessException {
    if (keyCoder != null) {
      if (doFn != null) {
        DoFnSignature signature = DoFnSignatures.getSignature(doFn.getClass());
        FlinkStateInternals.EarlyBinder earlyBinder =
            new FlinkStateInternals.EarlyBinder(
                getKeyedStateBackend(),
                serializedOptions,
                windowingStrategy.getWindowFn().windowCoder());
        for (DoFnSignature.StateDeclaration value : signature.stateDeclarations().values()) {
          StateSpec<?> spec =
              (StateSpec<?>) signature.stateDeclarations().get(value.id()).field().get(doFn);
          spec.bind(value.id(), earlyBinder);
        }
        if (doFnRunner instanceof StatefulDoFnRunner) {
          ((StatefulDoFnRunner<InputT, OutputT, BoundedWindow>) doFnRunner)
              .getSystemStateTags()
              .forEach(tag -> tag.getSpec().bind(tag.getId(), earlyBinder));
        }
      }
    }
  }

  void cleanUp() throws Exception {
    Optional.ofNullable(flinkMetricContainer)
        .ifPresent(FlinkMetricContainer::registerMetricsForPipelineResult);
    Optional.ofNullable(checkFinishBundleTimer).ifPresent(timer -> timer.cancel(true));
    Workarounds.deleteStaticCaches();
    Optional.ofNullable(doFnInvoker).ifPresent(DoFnInvoker::invokeTeardown);
  }

  void flushData() throws Exception {
    // This is our last change to block shutdown of this operator while
    // there are still remaining processing-time timers. Flink will ignore pending
    // processing-time timers when upstream operators have shut down and will also
    // shut down this operator with pending processing-time timers.
    if (numProcessingTimeTimers() > 0) {
      timerInternals.processPendingProcessingTimeTimers();
    }
    if (numProcessingTimeTimers() > 0) {
      throw new RuntimeException(
          "There are still "
              + numProcessingTimeTimers()
              + " processing-time timers left, this indicates a bug");
    }
    // make sure we send a +Inf watermark downstream. It can happen that we receive +Inf
    // in processWatermark*() but have holds, so we have to re-evaluate here.
    processWatermark(new Watermark(Long.MAX_VALUE));
    // Make sure to finish the current bundle
    while (bundleStarted) {
      invokeFinishBundle();
    }
    if (requiresStableInput && enableStableInputDrain) {
      // Flush any buffered events here before draining the pipeline. Note that this is best-effort
      // and requiresStableInput contract might be violated in cases where buffer processing fails.
      bufferingDoFnRunner.checkpointCompleted(Long.MAX_VALUE);
      updateOutputWatermark();
    }
    if (currentOutputWatermark < Long.MAX_VALUE) {
      throw new RuntimeException(
          String.format(
              "There are still watermark holds left when terminating operator %s Watermark held %d",
              getOperatorName(), currentOutputWatermark));
    }

    // sanity check: these should have been flushed out by +Inf watermarks
    if (!sideInputs.isEmpty()) {

      List<WindowedValue<InputT>> pushedBackElements =
          pushedBackElementsHandler.getElements().collect(Collectors.toList());

      if (pushedBackElements.size() > 0) {
        String pushedBackString = Joiner.on(",").join(pushedBackElements);
        throw new RuntimeException(
            "Leftover pushed-back data: " + pushedBackString + ". This indicates a bug.");
      }
    }
  }

  @Override
  public void finish() throws Exception {
    try {
      flushData();
    } finally {
      super.finish();
    }
  }

  @Override
  public void close() throws Exception {
    try {
      cleanUp();
    } finally {
      super.close();
    }
  }

  protected int numProcessingTimeTimers() {
    return getTimeServiceManager()
        .map(
            manager -> {
              if (timeServiceManager instanceof InternalTimeServiceManagerImpl) {
                final InternalTimeServiceManagerImpl<?> cast =
                    (InternalTimeServiceManagerImpl<?>) timeServiceManager;
                return cast.numProcessingTimeTimers();
              } else if (timeServiceManager instanceof BatchExecutionInternalTimeServiceManager) {
                return 0;
              } else {
                throw new IllegalStateException(
                    String.format(
                        "Unknown implementation of InternalTimerServiceManager. %s",
                        timeServiceManager));
              }
            })
        .orElse(0);
  }

  public long getEffectiveInputWatermark() {
    // hold back by the pushed back values waiting for side inputs
    long combinedPushedBackWatermark = pushedBackWatermark;
    if (requiresStableInput) {
      combinedPushedBackWatermark =
          Math.min(combinedPushedBackWatermark, bufferingDoFnRunner.getOutputWatermarkHold());
    }
    return Math.min(combinedPushedBackWatermark, currentInputWatermark);
  }

  public long getCurrentOutputWatermark() {
    return currentOutputWatermark;
  }

  protected final void setPreBundleCallback(Runnable callback) {
    this.preBundleCallback = callback;
  }

  protected final void setBundleFinishedCallback(Runnable callback) {
    this.bundleFinishedCallback = callback;
  }

  @Override
  public final void processElement(StreamRecord<WindowedValue<PreInputT>> streamRecord) {
    for (WindowedValue<InputT> e : preProcess(streamRecord.getValue())) {
      checkInvokeStartBundle();
      LOG.trace("Processing element {} in {}", streamRecord.getValue().getValue(), doFn.getClass());
      long oldHold = keyCoder != null ? keyedStateInternals.minWatermarkHoldMs() : -1L;
      doFnRunner.processElement(e);
      checkInvokeFinishBundleByCount();
      emitWatermarkIfHoldChanged(oldHold);
    }
  }

  @Override
  public final void processElement1(StreamRecord<WindowedValue<PreInputT>> streamRecord)
      throws Exception {
    for (WindowedValue<InputT> e : preProcess(streamRecord.getValue())) {
      checkInvokeStartBundle();
      Iterable<WindowedValue<InputT>> justPushedBack =
          pushbackDoFnRunner.processElementInReadyWindows(e);

      long min = pushedBackWatermark;
      for (WindowedValue<InputT> pushedBackValue : justPushedBack) {
        min = Math.min(min, pushedBackValue.getTimestamp().getMillis());
        pushedBackElementsHandler.pushBack(pushedBackValue);
      }
      pushedBackWatermark = min;

      checkInvokeFinishBundleByCount();
    }
  }

  /**
   * Add the side input value. Here we are assuming that views have already been materialized and
   * are sent over the wire as {@link Iterable}. Subclasses may elect to perform materialization in
   * state and receive side input incrementally instead.
   *
   * @param streamRecord
   */
  protected void addSideInputValue(StreamRecord<RawUnionValue> streamRecord) {
    @SuppressWarnings("unchecked")
    WindowedValue<Iterable<?>> value =
        (WindowedValue<Iterable<?>>) streamRecord.getValue().getValue();

    PCollectionView<?> sideInput = sideInputTagMapping.get(streamRecord.getValue().getUnionTag());
    sideInputHandler.addSideInputValue(sideInput, value);
  }

  @Override
  public final void processElement2(StreamRecord<RawUnionValue> streamRecord) throws Exception {
    // we finish the bundle because the newly arrived side-input might
    // make a view available that was previously not ready.
    // The PushbackSideInputRunner will only reset its cache of non-ready windows when
    // finishing a bundle.
    invokeFinishBundle();
    checkInvokeStartBundle();

    // add the side input, which may cause pushed back elements become eligible for processing
    addSideInputValue(streamRecord);

    List<WindowedValue<InputT>> newPushedBack = new ArrayList<>();

    Iterator<WindowedValue<InputT>> it = pushedBackElementsHandler.getElements().iterator();

    while (it.hasNext()) {
      WindowedValue<InputT> element = it.next();
      // we need to set the correct key in case the operator is
      // a (keyed) window operator
      if (keySelector != null) {
        setCurrentKey(keySelector.getKey(element));
      }

      Iterable<WindowedValue<InputT>> justPushedBack =
          pushbackDoFnRunner.processElementInReadyWindows(element);
      Iterables.addAll(newPushedBack, justPushedBack);
    }

    pushedBackElementsHandler.clear();
    long min = Long.MAX_VALUE;
    for (WindowedValue<InputT> pushedBackValue : newPushedBack) {
      min = Math.min(min, pushedBackValue.getTimestamp().getMillis());
      pushedBackElementsHandler.pushBack(pushedBackValue);
    }
    pushedBackWatermark = min;

    checkInvokeFinishBundleByCount();

    // maybe output a new watermark
    processWatermark1(new Watermark(currentInputWatermark));
  }

  @Override
  public final void processWatermark(Watermark mark) throws Exception {
    LOG.trace("Processing watermark {} in {}", mark.getTimestamp(), doFn.getClass());
    processWatermark1(mark);
  }

  @Override
  public final void processWatermark1(Watermark mark) throws Exception {
    // Flush any data buffered during snapshotState().
    outputManager.flushBuffer();

    // We do the check here because we are guaranteed to at least get the +Inf watermark on the
    // main input when the job finishes.
    if (currentSideInputWatermark >= BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()) {
      // this means we will never see any more side input
      // we also do the check here because we might have received the side-input MAX watermark
      // before receiving any main-input data
      emitAllPushedBackData();
    }

    currentInputWatermark = mark.getTimestamp();
    processInputWatermark(true);
  }

  private void processInputWatermark(boolean advanceInputWatermark) throws Exception {
    long inputWatermarkHold = applyInputWatermarkHold(getEffectiveInputWatermark());
    if (keyCoder != null && advanceInputWatermark) {
      timeServiceManager.advanceWatermark(new Watermark(inputWatermarkHold));
    }

    long potentialOutputWatermark =
        applyOutputWatermarkHold(
            currentOutputWatermark, computeOutputWatermark(inputWatermarkHold));

    maybeEmitWatermark(potentialOutputWatermark);
  }

  /**
   * Allows to apply a hold to the input watermark. By default, just passes the input watermark
   * through.
   */
  public long applyInputWatermarkHold(long inputWatermark) {
    return inputWatermark;
  }

  /**
   * Allows to apply a hold to the output watermark before it is sent out. Used to apply hold on
   * output watermark for delayed (asynchronous or buffered) processing.
   *
   * @param currentOutputWatermark the current output watermark
   * @param potentialOutputWatermark The potential new output watermark which can be adjusted, if
   *     needed. The input watermark hold has already been applied.
   * @return The new output watermark which will be emitted.
   */
  public long applyOutputWatermarkHold(long currentOutputWatermark, long potentialOutputWatermark) {
    return potentialOutputWatermark;
  }

  private long computeOutputWatermark(long inputWatermarkHold) {
    final long potentialOutputWatermark;
    if (keyCoder == null) {
      potentialOutputWatermark = inputWatermarkHold;
    } else {
      potentialOutputWatermark =
          Math.min(keyedStateInternals.minWatermarkHoldMs(), inputWatermarkHold);
    }
    return potentialOutputWatermark;
  }

  private void maybeEmitWatermark(long watermark) {
    if (watermark > currentOutputWatermark) {
      // Must invoke finishBatch before emit the +Inf watermark otherwise there are some late
      // events.
      if (watermark >= BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()) {
        invokeFinishBundle();
      }

      if (bundleStarted) {
        // do not update watermark in the middle of bundle, because it might cause
        // user-buffered data to be emitted past watermark
        return;
      }

      LOG.debug("Emitting watermark {} from {}", watermark, getOperatorName());
      currentOutputWatermark = watermark;
      output.emitWatermark(new Watermark(watermark));

      // Check if the final watermark was triggered to perform state cleanup for global window
      // TODO: Do we need to do this when OnWindowExpiration is set, since in that case we have a
      // cleanup timer?
      if (keyedStateInternals != null
          && currentOutputWatermark
              > adjustTimestampForFlink(GlobalWindow.INSTANCE.maxTimestamp().getMillis())) {
        keyedStateInternals.clearGlobalState();
      }
    }
  }

  @Override
  public final void processWatermark2(Watermark mark) throws Exception {
    currentSideInputWatermark = mark.getTimestamp();
    if (mark.getTimestamp() >= BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()) {
      // this means we will never see any more side input
      emitAllPushedBackData();

      // maybe output a new watermark
      processWatermark1(new Watermark(currentInputWatermark));
    }
  }

  /**
   * Emits all pushed-back data. This should be used once we know that there will not be any future
   * side input, i.e. that there is no point in waiting.
   */
  private void emitAllPushedBackData() throws Exception {

    Iterator<WindowedValue<InputT>> it = pushedBackElementsHandler.getElements().iterator();

    while (it.hasNext()) {
      checkInvokeStartBundle();
      WindowedValue<InputT> element = it.next();
      // we need to set the correct key in case the operator is
      // a (keyed) window operator
      setKeyContextElement1(new StreamRecord<>(element));

      doFnRunner.processElement(element);
    }

    pushedBackElementsHandler.clear();
    pushedBackWatermark = Long.MAX_VALUE;
  }

  /**
   * Check whether invoke startBundle, if it is, need to output elements that were buffered as part
   * of finishing a bundle in snapshot() first.
   *
   * <p>In order to avoid having {@link DoFnRunner#processElement(WindowedValue)} or {@link
   * DoFnRunner#onTimer(String, String, Object, BoundedWindow, Instant, Instant, TimeDomain)} not
   * between StartBundle and FinishBundle, this method needs to be called in each processElement and
   * each processWatermark and onProcessingTime. Do not need to call in onEventTime, because it has
   * been guaranteed in the processWatermark.
   */
  private void checkInvokeStartBundle() {
    if (!bundleStarted) {
      // Flush any data buffered during snapshotState().
      outputManager.flushBuffer();
      LOG.debug("Starting bundle.");
      if (preBundleCallback != null) {
        preBundleCallback.run();
      }
      pushbackDoFnRunner.startBundle();
      bundleStarted = true;
    }
  }

  /** Check whether invoke finishBundle by elements count. Called in processElement. */
  @SuppressWarnings("NonAtomicVolatileUpdate")
  @SuppressFBWarnings("VO_VOLATILE_INCREMENT")
  private void checkInvokeFinishBundleByCount() {
    if (!shoudBundleElements()) {
      return;
    }
    // We do not access this statement concurrently, but we want to make sure that each thread
    // sees the latest value, which is why we use volatile. See the class field section above
    // for more information.
    //noinspection NonAtomicOperationOnVolatileField
    elementCount++;
    if (elementCount >= maxBundleSize) {
      invokeFinishBundle();
      updateOutputWatermark();
    }
  }

  /** Check whether invoke finishBundle by timeout. */
  private void checkInvokeFinishBundleByTime() {
    if (!shoudBundleElements()) {
      return;
    }
    long now = getProcessingTimeService().getCurrentProcessingTime();
    if (now - lastFinishBundleTime >= maxBundleTimeMills) {
      invokeFinishBundle();
      scheduleForCurrentProcessingTime(ts -> updateOutputWatermark());
    }
  }

  @SuppressWarnings("FutureReturnValueIgnored")
  protected void scheduleForCurrentProcessingTime(ProcessingTimeCallback callback) {
    // We are scheduling a timer for advancing the watermark, to not delay finishing the bundle
    // and temporarily release the checkpoint lock. Otherwise, we could potentially loop when a
    // timer keeps scheduling a timer for the same timestamp.
    ProcessingTimeService timeService = getProcessingTimeService();
    timeService.registerTimer(timeService.getCurrentProcessingTime(), callback);
  }

  void updateOutputWatermark() {
    try {
      processInputWatermark(false);
    } catch (Exception ex) {
      failBundleFinalization(ex);
    }
  }

  protected final void invokeFinishBundle() {
    long previousBundleFinishTime = lastFinishBundleTime;
    if (bundleStarted) {
      LOG.debug("Finishing bundle.");
      pushbackDoFnRunner.finishBundle();
      LOG.debug("Finished bundle. Element count: {}", elementCount);
      elementCount = 0L;
      lastFinishBundleTime = getProcessingTimeService().getCurrentProcessingTime();
      bundleStarted = false;
      // callback only after current bundle was fully finalized
      // it could start a new bundle, for example resulting from timer processing
      if (bundleFinishedCallback != null) {
        LOG.debug("Invoking bundle finish callback.");
        bundleFinishedCallback.run();
      }
    }
    try {
      if (previousBundleFinishTime - getProcessingTimeService().getCurrentProcessingTime()
          > maxBundleTimeMills) {
        processInputWatermark(false);
      }
    } catch (Exception ex) {
      LOG.warn("Failed to update downstream watermark", ex);
    }
  }

  @Override
  public void prepareSnapshotPreBarrier(long checkpointId) {
    if (finishBundleBeforeCheckpointing) {
      // We finish the bundle and flush any pending data.
      // This avoids buffering any data as part of snapshotState() below.
      while (bundleStarted) {
        invokeFinishBundle();
      }
      updateOutputWatermark();
    }
  }

  @Override
  public void snapshotState(StateSnapshotContext context) throws Exception {
    if (checkpointStats != null) {
      checkpointStats.snapshotStart(context.getCheckpointId());
    }

    if (requiresStableInput) {
      // We notify the BufferingDoFnRunner to associate buffered state with this
      // snapshot id and start a new buffer for elements arriving after this snapshot.
      bufferingDoFnRunner.checkpoint(context.getCheckpointId());
    }

    int diff = pendingFinalizations.size() - MAX_NUMBER_PENDING_BUNDLE_FINALIZATIONS;
    if (diff >= 0) {
      for (Iterator<Long> iterator = pendingFinalizations.keySet().iterator(); diff >= 0; diff--) {
        iterator.next();
        iterator.remove();
      }
    }
    pendingFinalizations.put(context.getCheckpointId(), bundleFinalizer.getAndClearFinalizations());

    try {
      outputManager.openBuffer();
      // Ensure that no new bundle gets started as part of finishing a bundle
      while (bundleStarted) {
        invokeFinishBundle();
      }
      outputManager.closeBuffer();
    } catch (Exception e) {
      failBundleFinalization(e);
    }

    super.snapshotState(context);
  }

  private void failBundleFinalization(Exception e) {
    // https://jira.apache.org/jira/browse/FLINK-14653
    // Any regular exception during checkpointing will be tolerated by Flink because those
    // typically do not affect the execution flow. We need to fail hard here because errors
    // in bundle execution are application errors which are not related to checkpointing.
    throw new Error("Checkpointing failed because bundle failed to finalize.", e);
  }

  public BundleFinalizer getBundleFinalizer() {
    return bundleFinalizer;
  }

  @Override
  public void notifyCheckpointComplete(long checkpointId) throws Exception {
    if (checkpointStats != null) {
      checkpointStats.reportCheckpointDuration(checkpointId);
    }

    if (requiresStableInput) {
      // We can now release all buffered data which was held back for
      // @RequiresStableInput guarantees.
      bufferingDoFnRunner.checkpointCompleted(checkpointId);
      updateOutputWatermark();
    }

    List<InMemoryBundleFinalizer.Finalization> finalizations =
        pendingFinalizations.remove(checkpointId);
    if (finalizations != null) {
      // confirm all finalizations that were associated with the checkpoint
      for (InMemoryBundleFinalizer.Finalization finalization : finalizations) {
        finalization.getCallback().onBundleSuccess();
      }
    }

    super.notifyCheckpointComplete(checkpointId);
  }

  @Override
  public void onEventTime(InternalTimer<FlinkKey, TimerData> timer) {
    checkInvokeStartBundle();
    fireTimerInternal(timer.getKey(), timer.getNamespace());
  }

  @Override
  public void onProcessingTime(InternalTimer<FlinkKey, TimerData> timer) {
    checkInvokeStartBundle();
    fireTimerInternal(timer.getKey(), timer.getNamespace());
  }

  // allow overriding this in ExecutableStageDoFnOperator to set the key context
  protected void fireTimerInternal(FlinkKey key, TimerData timerData) {
    long oldHold = keyCoder != null ? keyedStateInternals.minWatermarkHoldMs() : -1L;
    fireTimer(timerData);
    emitWatermarkIfHoldChanged(oldHold);
  }

  void emitWatermarkIfHoldChanged(long currentWatermarkHold) {
    if (keyCoder != null) {
      long newWatermarkHold = keyedStateInternals.minWatermarkHoldMs();
      if (newWatermarkHold > currentWatermarkHold) {
        try {
          processInputWatermark(false);
        } catch (Exception ex) {
          // should not happen
          throw new IllegalStateException(ex);
        }
      }
    }
  }

  // allow overriding this in WindowDoFnOperator
  protected void fireTimer(TimerData timerData) {
    LOG.debug(
        "Firing timer: {} at {} with output time {}",
        timerData.getTimerId(),
        timerData.getTimestamp().getMillis(),
        timerData.getOutputTimestamp().getMillis());
    StateNamespace namespace = timerData.getNamespace();
    // This is a user timer, so namespace must be WindowNamespace
    checkArgument(namespace instanceof WindowNamespace);
    BoundedWindow window = ((WindowNamespace) namespace).getWindow();
    timerInternals.onFiredOrDeletedTimer(timerData);

    pushbackDoFnRunner.onTimer(
        timerData.getTimerId(),
        timerData.getTimerFamilyId(),
        keyedStateInternals.getKey(),
        window,
        timerData.getTimestamp(),
        timerData.getOutputTimestamp(),
        timerData.getDomain(),
        timerData.causedByDrain());
  }

  @SuppressWarnings("unchecked")
  Coder<InputT> getInputCoder() {
    return (Coder<InputT>) Iterables.getOnlyElement(windowedInputCoder.getCoderArguments());
  }

  /** Factory for creating an {@link BufferedOutputManager} from a Flink {@link Output}. */
  interface OutputManagerFactory<OutputT> extends Serializable {
    BufferedOutputManager<OutputT> create(
        Output<StreamRecord<WindowedValue<OutputT>>> output,
        Lock bufferLock,
        OperatorStateBackend operatorStateBackend)
        throws Exception;
  }

  /**
   * A {@link WindowedValueReceiver} that can buffer its outputs. Uses {@link
   * PushedBackElementsHandler} to buffer the data. Buffering data is necessary because no elements
   * can be emitted during {@code snapshotState} which is called when the checkpoint barrier already
   * has been sent downstream. Emitting elements would break the flow of checkpoint barrier and
   * violate exactly-once semantics.
   *
   * <p>This buffering can be deactived using {@code
   * FlinkPipelineOptions#setFinishBundleBeforeCheckpointing(true)}. If activated, we flush out
   * bundle data before the barrier is sent downstream. This is done via {@code
   * prepareSnapshotPreBarrier}. When Flink supports unaligned checkpoints, this should become the
   * default and this class should be removed as in https://github.com/apache/beam/pull/9652.
   */
  public static class BufferedOutputManager<OutputT> implements WindowedValueMultiReceiver {

    private final TupleTag<OutputT> mainTag;
    private final Map<TupleTag<?>, OutputTag<WindowedValue<?>>> tagsToOutputTags;
    private final Map<TupleTag<?>, Integer> tagsToIds;
    /**
     * A lock to be acquired before writing to the buffer. This lock will only be acquired during
     * buffering. It will not be acquired during flushing the buffer.
     */
    private final Lock bufferLock;

    private final boolean isStreaming;

    private Map<Integer, TupleTag<?>> idsToTags;
    /** Elements buffered during a snapshot, by output id. */
    @VisibleForTesting
    final PushedBackElementsHandler<KV<Integer, WindowedValue<?>>> pushedBackElementsHandler;

    protected final Output<StreamRecord<WindowedValue<OutputT>>> output;

    /** Indicates whether we are buffering data as part of snapshotState(). */
    private boolean openBuffer = false;
    /** For performance, to avoid having to access the state backend when the buffer is empty. */
    private boolean bufferIsEmpty = false;

    BufferedOutputManager(
        Output<StreamRecord<WindowedValue<OutputT>>> output,
        TupleTag<OutputT> mainTag,
        Map<TupleTag<?>, OutputTag<WindowedValue<?>>> tagsToOutputTags,
        Map<TupleTag<?>, Integer> tagsToIds,
        Lock bufferLock,
        PushedBackElementsHandler<KV<Integer, WindowedValue<?>>> pushedBackElementsHandler,
        boolean isStreaming) {
      this.output = output;
      this.mainTag = mainTag;
      this.tagsToOutputTags = tagsToOutputTags;
      this.tagsToIds = tagsToIds;
      this.bufferLock = bufferLock;
      this.idsToTags = new HashMap<>();
      for (Map.Entry<TupleTag<?>, Integer> entry : tagsToIds.entrySet()) {
        idsToTags.put(entry.getValue(), entry.getKey());
      }
      this.pushedBackElementsHandler = pushedBackElementsHandler;
      this.isStreaming = isStreaming;
    }

    void openBuffer() {
      this.openBuffer = true;
    }

    void closeBuffer() {
      this.openBuffer = false;
    }

    @Override
    public <T> void output(TupleTag<T> tag, WindowedValue<T> value) {
      // Don't buffer elements in Batch mode
      if (!openBuffer || !isStreaming) {
        emit(tag, value);
      } else {
        buffer(KV.of(tagsToIds.get(tag), value));
      }
    }

    private void buffer(KV<Integer, WindowedValue<?>> taggedValue) {
      bufferLock.lock();
      try {
        pushedBackElementsHandler.pushBack(taggedValue);
      } catch (Exception e) {
        throw new RuntimeException("Couldn't pushback element.", e);
      } finally {
        bufferLock.unlock();
        bufferIsEmpty = false;
      }
    }

    /**
     * Flush elements of bufferState to Flink Output. This method should not be invoked in {@link
     * #snapshotState(StateSnapshotContext)} because the checkpoint barrier has already been sent
     * downstream; emitting elements at this point would violate the checkpoint barrier alignment.
     *
     * <p>The buffer should be flushed before starting a new bundle when the buffer cannot be
     * concurrently accessed and thus does not need to be guarded by a lock.
     */
    void flushBuffer() {
      if (openBuffer || bufferIsEmpty) {
        // Checkpoint currently in progress or nothing buffered, do not proceed
        return;
      }
      try {
        pushedBackElementsHandler
            .getElements()
            .forEach(
                element ->
                    emit(idsToTags.get(element.getKey()), (WindowedValue) element.getValue()));
        pushedBackElementsHandler.clear();
        bufferIsEmpty = true;
      } catch (Exception e) {
        throw new RuntimeException("Couldn't flush pushed back elements.", e);
      }
    }

    private <T> void emit(TupleTag<T> tag, WindowedValue<T> value) {
      if (tag.equals(mainTag)) {
        // with tagged outputs we can't get around this because we don't
        // know our own output type...
        @SuppressWarnings("unchecked")
        WindowedValue<OutputT> castValue = (WindowedValue<OutputT>) value;
        output.collect(new StreamRecord<>(castValue));
      } else {
        @SuppressWarnings("unchecked")
        OutputTag<WindowedValue<T>> outputTag = (OutputTag) tagsToOutputTags.get(tag);
        output.collect(outputTag, new StreamRecord<>(value));
      }
    }
  }

  /** Coder for KV of id and value. It will be serialized in Flink checkpoint. */
  private static class TaggedKvCoder extends StructuredCoder<KV<Integer, WindowedValue<?>>> {

    private final Map<Integer, Coder<WindowedValue<?>>> idsToCoders;

    TaggedKvCoder(Map<Integer, Coder<WindowedValue<?>>> idsToCoders) {
      this.idsToCoders = idsToCoders;
    }

    @Override
    public void encode(KV<Integer, WindowedValue<?>> kv, OutputStream out) throws IOException {
      Coder<WindowedValue<?>> coder = idsToCoders.get(kv.getKey());
      VarIntCoder.of().encode(kv.getKey(), out);
      coder.encode(kv.getValue(), out);
    }

    @Override
    public KV<Integer, WindowedValue<?>> decode(InputStream in) throws IOException {
      Integer id = VarIntCoder.of().decode(in);
      Coder<WindowedValue<?>> coder = idsToCoders.get(id);
      WindowedValue<?> value = coder.decode(in);
      return KV.of(id, value);
    }

    @Override
    public List<? extends Coder<?>> getCoderArguments() {
      return new ArrayList<>(idsToCoders.values());
    }

    @Override
    public void verifyDeterministic() throws NonDeterministicException {
      for (Coder<?> coder : idsToCoders.values()) {
        verifyDeterministic(this, "Coder must be deterministic", coder);
      }
    }
  }

  /**
   * Implementation of {@link OutputManagerFactory} that creates an {@link BufferedOutputManager}
   * that can write to multiple logical outputs by Flink side output.
   */
  public static class MultiOutputOutputManagerFactory<OutputT>
      implements OutputManagerFactory<OutputT> {

    private final TupleTag<OutputT> mainTag;
    private final Map<TupleTag<?>, Integer> tagsToIds;
    private final Map<TupleTag<?>, OutputTag<WindowedValue<?>>> tagsToOutputTags;
    private final Map<TupleTag<?>, Coder<WindowedValue<?>>> tagsToCoders;
    private final SerializablePipelineOptions pipelineOptions;
    private final boolean isStreaming;

    // There is no side output.
    @SuppressWarnings("unchecked")
    public MultiOutputOutputManagerFactory(
        TupleTag<OutputT> mainTag,
        Coder<WindowedValue<OutputT>> mainCoder,
        SerializablePipelineOptions pipelineOptions) {
      this(
          mainTag,
          new HashMap<>(),
          ImmutableMap.<TupleTag<?>, Coder<WindowedValue<?>>>builder()
              .put(mainTag, (Coder) mainCoder)
              .build(),
          ImmutableMap.<TupleTag<?>, Integer>builder().put(mainTag, 0).build(),
          pipelineOptions);
    }

    public MultiOutputOutputManagerFactory(
        TupleTag<OutputT> mainTag,
        Map<TupleTag<?>, OutputTag<WindowedValue<?>>> tagsToOutputTags,
        Map<TupleTag<?>, Coder<WindowedValue<?>>> tagsToCoders,
        Map<TupleTag<?>, Integer> tagsToIds,
        SerializablePipelineOptions pipelineOptions) {
      this.mainTag = mainTag;
      this.tagsToOutputTags = tagsToOutputTags;
      this.tagsToCoders = tagsToCoders;
      this.tagsToIds = tagsToIds;
      this.pipelineOptions = pipelineOptions;
      this.isStreaming = pipelineOptions.get().as(FlinkPipelineOptions.class).isStreaming();
    }

    @Override
    public BufferedOutputManager<OutputT> create(
        Output<StreamRecord<WindowedValue<OutputT>>> output,
        Lock bufferLock,
        OperatorStateBackend operatorStateBackend)
        throws Exception {
      Preconditions.checkNotNull(output);
      Preconditions.checkNotNull(bufferLock);
      Preconditions.checkNotNull(operatorStateBackend);

      TaggedKvCoder taggedKvCoder = buildTaggedKvCoder();
      ListStateDescriptor<KV<Integer, WindowedValue<?>>> taggedOutputPushbackStateDescriptor =
          new ListStateDescriptor<>(
              "bundle-buffer-tag", new CoderTypeSerializer<>(taggedKvCoder, pipelineOptions));
      ListState<KV<Integer, WindowedValue<?>>> listStateBuffer =
          operatorStateBackend.getListState(taggedOutputPushbackStateDescriptor);
      PushedBackElementsHandler<KV<Integer, WindowedValue<?>>> pushedBackElementsHandler =
          NonKeyedPushedBackElementsHandler.create(listStateBuffer);

      return new BufferedOutputManager<>(
          output,
          mainTag,
          tagsToOutputTags,
          tagsToIds,
          bufferLock,
          pushedBackElementsHandler,
          isStreaming);
    }

    private TaggedKvCoder buildTaggedKvCoder() {
      ImmutableMap.Builder<Integer, Coder<WindowedValue<?>>> idsToCodersBuilder =
          ImmutableMap.builder();
      for (Map.Entry<TupleTag<?>, Integer> entry : tagsToIds.entrySet()) {
        idsToCodersBuilder.put(entry.getValue(), tagsToCoders.get(entry.getKey()));
      }
      return new TaggedKvCoder(idsToCodersBuilder.build());
    }
  }

  /**
   * {@link StepContext} for running {@link DoFn DoFns} on Flink. This does not allow accessing
   * state or timer internals.
   */
  protected class FlinkStepContext implements StepContext {

    @Override
    public StateInternals stateInternals() {
      return keyedStateInternals;
    }

    @Override
    public TimerInternals timerInternals() {
      return timerInternals;
    }

    @Override
    public BundleFinalizer bundleFinalizer() {
      return bundleFinalizer;
    }
  }

  class FlinkTimerInternals implements TimerInternals {

    private static final String PENDING_TIMERS_STATE_NAME = "pending-timers";

    /**
     * Pending Timers (=not been fired yet) by context id. The id is generated from the state
     * namespace of the timer and the timer's id. Necessary for supporting removal of existing
     * timers. In Flink removal of timers can only be done by providing id and time of the timer.
     *
     * <p>CAUTION: This map is scoped by the current active key. Do not attempt to perform any
     * calculations which span across keys.
     */
    @VisibleForTesting final MapState<String, TimerData> pendingTimersById;

    private final InternalTimerService<TimerData> timerService;

    private FlinkTimerInternals(InternalTimerService<TimerData> timerService) throws Exception {
      MapStateDescriptor<String, TimerData> pendingTimersByIdStateDescriptor =
          new MapStateDescriptor<>(
              PENDING_TIMERS_STATE_NAME,
              new StringSerializer(),
              new CoderTypeSerializer<>(timerCoder, serializedOptions));

      this.pendingTimersById = getKeyedStateStore().getMapState(pendingTimersByIdStateDescriptor);
      this.timerService = timerService;
      populateOutputTimestampQueue(timerService);
    }

    /**
     * Processes all pending processing timers. This is intended for use during shutdown. From Flink
     * 1.10 on, processing timer execution is stopped when the operator is closed. This leads to
     * problems for applications which assume all pending timers will be completed. Although Flink
     * does drain the remaining timers after close(), this is not sufficient because no new timers
     * are allowed to be scheduled anymore. This breaks Beam pipelines which rely on all processing
     * timers to be scheduled and executed.
     */
    void processPendingProcessingTimeTimers() {
      final KeyedStateBackend<Object> keyedStateBackend = getKeyedStateBackend();
      final InternalPriorityQueue<InternalTimer<Object, TimerData>> processingTimeTimersQueue =
          Workarounds.retrieveInternalProcessingTimerQueue(timerService);

      InternalTimer<Object, TimerData> internalTimer;
      while ((internalTimer = processingTimeTimersQueue.poll()) != null) {
        keyedStateBackend.setCurrentKey(internalTimer.getKey());
        TimerData timer = internalTimer.getNamespace();
        checkInvokeStartBundle();
        fireTimerInternal((FlinkKey) internalTimer.getKey(), timer);
      }
    }

    private void populateOutputTimestampQueue(InternalTimerService<TimerData> timerService)
        throws Exception {

      BiConsumerWithException<TimerData, Long, Exception> consumer =
          (timerData, stamp) ->
              keyedStateInternals.addWatermarkHoldUsage(timerData.getOutputTimestamp());
      if (timerService instanceof InternalTimerServiceImpl) {
        timerService.forEachEventTimeTimer(consumer);
        timerService.forEachProcessingTimeTimer(consumer);
      }
    }

    private String constructTimerId(String timerFamilyId, String timerId) {
      return timerFamilyId + "+" + timerId;
    }

    @Override
    public void setTimer(
        StateNamespace namespace,
        String timerId,
        String timerFamilyId,
        Instant target,
        Instant outputTimestamp,
        TimeDomain timeDomain) {
      setTimer(
          TimerData.of(timerId, timerFamilyId, namespace, target, outputTimestamp, timeDomain));
    }

    /**
     * @deprecated use {@link #setTimer(StateNamespace, String, String, Instant, Instant,
     *     TimeDomain)}.
     */
    @Deprecated
    @Override
    public void setTimer(TimerData timer) {
      try {
        LOG.debug(
            "Setting timer: {} at {} with output time {}",
            timer.getTimerId(),
            timer.getTimestamp().getMillis(),
            timer.getOutputTimestamp().getMillis());
        String contextTimerId =
            getContextTimerId(
                constructTimerId(timer.getTimerFamilyId(), timer.getTimerId()),
                timer.getNamespace());
        @Nullable final TimerData oldTimer = pendingTimersById.get(contextTimerId);
        if (!timer.equals(oldTimer)) {
          // Only one timer can exist at a time for a given timer id and context.
          // If a timer gets set twice in the same context, the second must
          // override the first. Thus, we must cancel any pending timers
          // before we set the new one.
          cancelPendingTimer(oldTimer);
          registerTimer(timer, contextTimerId);
        }
      } catch (Exception e) {
        throw new RuntimeException("Failed to set timer", e);
      }
    }

    private void registerTimer(TimerData timer, String contextTimerId) throws Exception {
      LOG.debug("Registering timer {}", timer);
      pendingTimersById.put(contextTimerId, timer);
      long time = timer.getTimestamp().getMillis();
      switch (timer.getDomain()) {
        case EVENT_TIME:
          timerService.registerEventTimeTimer(timer, adjustTimestampForFlink(time));
          break;
        case PROCESSING_TIME:
        case SYNCHRONIZED_PROCESSING_TIME:
          timerService.registerProcessingTimeTimer(timer, adjustTimestampForFlink(time));
          break;
        default:
          throw new UnsupportedOperationException("Unsupported time domain: " + timer.getDomain());
      }
      keyedStateInternals.addWatermarkHoldUsage(timer.getOutputTimestamp());
    }

    /**
     * Looks up a timer by its id. This is necessary to support canceling existing timers with the
     * same id. Flink does not provide this functionality.
     *
     * @param contextTimerId Timer ID o cancel.
     */
    private void cancelPendingTimerById(String contextTimerId) throws Exception {
      cancelPendingTimer(pendingTimersById.get(contextTimerId));
    }

    /**
     * Cancels a pending timer.
     *
     * @param timer Timer to cancel.
     */
    private void cancelPendingTimer(@Nullable TimerData timer) {
      if (timer != null) {
        deleteTimerInternal(timer);
      }
    }

    /**
     * Hook which must be called when a timer is fired or deleted to perform cleanup. Note: Make
     * sure that the state backend key is set correctly. It is best to run this in the fireTimer()
     * method.
     */
    void onFiredOrDeletedTimer(TimerData timer) {
      try {
        pendingTimersById.remove(
            getContextTimerId(
                constructTimerId(timer.getTimerFamilyId(), timer.getTimerId()),
                timer.getNamespace()));
        keyedStateInternals.removeWatermarkHoldUsage(timer.getOutputTimestamp());
      } catch (Exception e) {
        throw new RuntimeException("Failed to cleanup pending timers state.", e);
      }
    }

    /** @deprecated use {@link #deleteTimer(StateNamespace, String, String, TimeDomain)}. */
    @Deprecated
    @Override
    public void deleteTimer(StateNamespace namespace, String timerId, String timerFamilyId) {
      throw new UnsupportedOperationException("Canceling of a timer by ID is not yet supported.");
    }

    @Override
    public void deleteTimer(
        StateNamespace namespace, String timerId, String timerFamilyId, TimeDomain timeDomain) {
      try {
        cancelPendingTimerById(getContextTimerId(timerId, namespace));
      } catch (Exception e) {
        throw new RuntimeException("Failed to cancel timer", e);
      }
    }

    /** @deprecated use {@link #deleteTimer(StateNamespace, String, String, TimeDomain)}. */
    @Override
    @Deprecated
    public void deleteTimer(TimerData timer) {
      deleteTimer(
          timer.getNamespace(),
          constructTimerId(timer.getTimerFamilyId(), timer.getTimerId()),
          timer.getTimerFamilyId(),
          timer.getDomain());
    }

    void deleteTimerInternal(TimerData timer) {
      long time = timer.getTimestamp().getMillis();
      switch (timer.getDomain()) {
        case EVENT_TIME:
          timerService.deleteEventTimeTimer(timer, adjustTimestampForFlink(time));
          break;
        case PROCESSING_TIME:
        case SYNCHRONIZED_PROCESSING_TIME:
          timerService.deleteProcessingTimeTimer(timer, adjustTimestampForFlink(time));
          break;
        default:
          throw new UnsupportedOperationException("Unsupported time domain: " + timer.getDomain());
      }
      onFiredOrDeletedTimer(timer);
    }

    @Override
    public Instant currentProcessingTime() {
      return new Instant(timerService.currentProcessingTime());
    }

    @Override
    public @Nullable Instant currentSynchronizedProcessingTime() {
      return new Instant(timerService.currentProcessingTime());
    }

    @Override
    public Instant currentInputWatermarkTime() {
      if (timerService instanceof BatchExecutionInternalTimeService) {
        // In batch mode, this method will only either return BoundedWindow.TIMESTAMP_MIN_VALUE,
        // or BoundedWindow.TIMESTAMP_MAX_VALUE.
        //
        // For batch execution mode, the currentInputWatermark variable will never be updated
        // until all the records are processed. However, every time when a record with a new
        // key arrives, the Flink timer service watermark will be set to
        // MAX_WATERMARK(LONG.MAX_VALUE) so that all the timers associated with the current
        // key can fire. After that the Flink timer service watermark will be reset to
        // LONG.MIN_VALUE, so the next key will start from a fresh env as if the previous
        // records of a different key never existed. So the watermark is either Long.MIN_VALUE
        // or long MAX_VALUE. So we should just use the Flink time service watermark in batch mode.
        //
        // In Flink the watermark ranges from
        // [LONG.MIN_VALUE (-9223372036854775808), LONG.MAX_VALUE (9223372036854775807)] while the
        // beam
        // watermark range is [BoundedWindow.TIMESTAMP_MIN_VALUE (-9223372036854775),
        // BoundedWindow.TIMESTAMP_MAX_VALUE (9223372036854775)]. To ensure the timestamp visible to
        // the users follow the Beam convention, we just use the Beam range instead.
        return timerService.currentWatermark() == Long.MAX_VALUE
            ? new Instant(Long.MAX_VALUE)
            : BoundedWindow.TIMESTAMP_MIN_VALUE;
      } else {
        return new Instant(getEffectiveInputWatermark());
      }
    }

    @Override
    public @Nullable Instant currentOutputWatermarkTime() {
      return new Instant(currentOutputWatermark);
    }

    /**
     * Check whether event time timers lower or equal to the given timestamp exist. Caution: This is
     * scoped by the current key.
     */
    public boolean hasPendingEventTimeTimers(long maxTimestamp) throws Exception {
      for (TimerData timer : pendingTimersById.values()) {
        if (timer.getDomain() == TimeDomain.EVENT_TIME
            && timer.getTimestamp().getMillis() <= maxTimestamp) {
          return true;
        }
      }
      return false;
    }

    /** Unique contextual id of a timer. Used to look up any existing timers in a context. */
    private String getContextTimerId(String timerId, StateNamespace namespace) {
      return timerId + namespace.stringKey();
    }
  }

  /**
   * In Beam, a timer with timestamp {@code T} is only illegible for firing when the time has moved
   * past this time stamp, i.e. {@code T < current_time}. In the case of event time, current_time is
   * the watermark, in the case of processing time it is the system time.
   *
   * <p>Flink's TimerService has different semantics because it only ensures {@code T <=
   * current_time}.
   *
   * <p>To make up for this, we need to add one millisecond to Flink's internal timer timestamp.
   * Note that we do not modify Beam's timestamp and we are not exposing Flink's timestamp.
   *
   * <p>See also https://jira.apache.org/jira/browse/BEAM-3863
   */
  static long adjustTimestampForFlink(long beamTimerTimestamp) {
    if (beamTimerTimestamp == Long.MAX_VALUE) {
      // We would overflow, do not adjust timestamp
      return Long.MAX_VALUE;
    }
    return beamTimerTimestamp + 1;
  }
}
