/*
 * 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.core;

import static org.apache.beam.sdk.transforms.DoFn.ProcessContinuation.resume;
import static org.apache.beam.sdk.transforms.DoFn.ProcessContinuation.stop;
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.contains;
import static org.hamcrest.Matchers.greaterThanOrEqualTo;
import static org.hamcrest.Matchers.hasItem;
import static org.hamcrest.Matchers.hasItems;
import static org.hamcrest.Matchers.not;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;

import java.io.Serializable;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.concurrent.Executors;
import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems.ProcessFn;
import org.apache.beam.sdk.coders.BigEndianIntegerCoder;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.InstantCoder;
import org.apache.beam.sdk.coders.SerializableCoder;
import org.apache.beam.sdk.coders.VoidCoder;
import org.apache.beam.sdk.io.range.OffsetRange;
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.testing.ResetDateTimeProvider;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.DoFnTester;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.transforms.splittabledofn.HasDefaultTracker;
import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.IsBounded;
import org.apache.beam.sdk.transforms.splittabledofn.SplitResult;
import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.FixedWindows;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
import org.apache.beam.sdk.util.WindowedValueMultiReceiver;
import org.apache.beam.sdk.values.CausedByDrain;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TimestampedValue;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.ValueInSingleWindow;
import org.apache.beam.sdk.values.WindowedValue;
import org.apache.beam.sdk.values.WindowedValues;
import org.apache.beam.sdk.values.WindowingStrategy;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;

/** Tests for {@link SplittableParDoViaKeyedWorkItems.ProcessFn}. */
@RunWith(JUnit4.class)
public class SplittableParDoProcessFnTest {
  private static final int MAX_OUTPUTS_PER_BUNDLE = 10000;
  private static final Duration MAX_BUNDLE_DURATION = Duration.standardSeconds(5);
  @Rule public final ResetDateTimeProvider dateTimeProvider = new ResetDateTimeProvider();

  // ----------------- Tests for whether the transform sets boundedness correctly --------------
  private static class SomeRestriction
      implements Serializable, HasDefaultTracker<SomeRestriction, SomeRestrictionTracker> {
    @Override
    public SomeRestrictionTracker newTracker() {
      return new SomeRestrictionTracker(this);
    }
  }

  private static class SomeRestrictionTracker extends RestrictionTracker<SomeRestriction, Void> {
    private final SomeRestriction someRestriction;

    public SomeRestrictionTracker(SomeRestriction someRestriction) {
      this.someRestriction = someRestriction;
    }

    @Override
    public boolean tryClaim(Void position) {
      return true;
    }

    @Override
    public SomeRestriction currentRestriction() {
      return someRestriction;
    }

    @Override
    public SplitResult<SomeRestriction> trySplit(double fractionOfRemainder) {
      return SplitResult.of(null, someRestriction);
    }

    @Override
    public void checkDone() {}

    @Override
    public RestrictionTracker.IsBounded isBounded() {
      return RestrictionTracker.IsBounded.BOUNDED;
    }
  }

  @Rule public TestPipeline pipeline = TestPipeline.create();

  /**
   * A helper for testing {@link ProcessFn} on 1 element (but possibly over multiple {@link
   * DoFn.ProcessElement} calls).
   */
  private static class ProcessFnTester<
          InputT, OutputT, RestrictionT, PositionT, WatermarkEstimatorStateT>
      implements AutoCloseable {
    private final DoFnTester<KeyedWorkItem<byte[], KV<InputT, RestrictionT>>, OutputT> tester;
    private Instant currentProcessingTime;

    private InMemoryTimerInternals timerInternals;
    private TestInMemoryStateInternals<String> stateInternals;
    private InMemoryBundleFinalizer bundleFinalizer;
    private final ProcessFn<InputT, OutputT, RestrictionT, PositionT, WatermarkEstimatorStateT>
        processFn;

    ProcessFnTester(
        Instant currentProcessingTime,
        final DoFn<InputT, OutputT> fn,
        Coder<InputT> inputCoder,
        Coder<RestrictionT> restrictionCoder,
        Coder<WatermarkEstimatorStateT> watermarkEstimatorStateCoder,
        int maxOutputsPerBundle,
        Duration maxBundleDuration)
        throws Exception {
      this(
          currentProcessingTime,
          fn,
          inputCoder,
          restrictionCoder,
          watermarkEstimatorStateCoder,
          maxOutputsPerBundle,
          maxBundleDuration,
          Collections.emptyMap(),
          NullSideInputReader.empty());
    }

    ProcessFnTester(
        Instant currentProcessingTime,
        final DoFn<InputT, OutputT> fn,
        Coder<InputT> inputCoder,
        Coder<RestrictionT> restrictionCoder,
        Coder<WatermarkEstimatorStateT> watermarkEstimatorStateCoder,
        int maxOutputsPerBundle,
        Duration maxBundleDuration,
        Map<String, PCollectionView<?>> sideInputMapping,
        SideInputReader sideInputReader)
        throws Exception {
      // The exact windowing strategy doesn't matter in this test, but it should be able to
      // encode IntervalWindow's because that's what all tests here use.
      WindowingStrategy<InputT, BoundedWindow> windowingStrategy =
          (WindowingStrategy) WindowingStrategy.of(FixedWindows.of(Duration.standardSeconds(1)));
      this.processFn =
          new ProcessFn<>(
              fn,
              inputCoder,
              restrictionCoder,
              watermarkEstimatorStateCoder,
              windowingStrategy,
              sideInputMapping);
      this.tester = DoFnTester.of(processFn);
      this.timerInternals = new InMemoryTimerInternals();
      this.stateInternals = new TestInMemoryStateInternals<>("dummy");
      processFn.setStateInternalsFactory(key -> stateInternals);
      processFn.setTimerInternalsFactory(key -> timerInternals);
      processFn.setSideInputReader(sideInputReader);
      processFn.setProcessElementInvoker(
          new OutputAndTimeBoundedSplittableProcessElementInvoker<>(
              fn,
              tester.getPipelineOptions(),
              new DoFnTesterWindowedValueReceiver(tester),
              tester.getMainOutputTag(),
              sideInputReader,
              Executors.newSingleThreadScheduledExecutor(Executors.defaultThreadFactory()),
              maxOutputsPerBundle,
              maxBundleDuration,
              () -> bundleFinalizer));
      // Do not clone since ProcessFn references non-serializable DoFnTester itself
      // through the state/timer/output callbacks.
      this.tester.setCloningBehavior(DoFnTester.CloningBehavior.DO_NOT_CLONE);
      this.tester.startBundle();
      timerInternals.advanceProcessingTime(currentProcessingTime);

      this.currentProcessingTime = currentProcessingTime;
    }

    @Override
    public void close() throws Exception {
      tester.close();
    }

    /** Performs a seed {@link DoFn.ProcessElement} call feeding the element and restriction. */
    void startElement(InputT element, RestrictionT restriction) throws Exception {
      startElement(
          WindowedValues.of(
              KV.of(element, restriction),
              currentProcessingTime,
              GlobalWindow.INSTANCE,
              PaneInfo.ON_TIME_AND_ONLY_FIRING));
    }

    void startElement(WindowedValue<KV<InputT, RestrictionT>> windowedValue) throws Exception {
      tester.processElement(
          KeyedWorkItems.elementsWorkItem(
              "key".getBytes(StandardCharsets.UTF_8), Collections.singletonList(windowedValue)));
    }

    boolean advanceDrain() throws Exception {

      List<TimerInternals.TimerData> timers = new ArrayList<>();
      timers.add(
          TimerInternals.TimerData.of(
              "",
              StateNamespaces.window(GlobalWindow.Coder.INSTANCE, GlobalWindow.INSTANCE),
              new Instant(Long.MAX_VALUE),
              new Instant(Long.MAX_VALUE),
              TimeDomain.EVENT_TIME,
              CausedByDrain.CAUSED_BY_DRAIN));
      tester.processElement(
          KeyedWorkItems.timersWorkItem("key".getBytes(StandardCharsets.UTF_8), timers));
      return true;
    }

    /**
     * Advances processing time by a given duration and, if any timers fired, performs a non-seed
     * {@link DoFn.ProcessElement} call, feeding it the timers.
     */
    boolean advanceProcessingTimeBy(Duration duration) throws Exception {
      currentProcessingTime = currentProcessingTime.plus(duration);
      timerInternals.advanceProcessingTime(currentProcessingTime);

      List<TimerInternals.TimerData> timers = new ArrayList<>();
      TimerInternals.TimerData nextTimer;
      while ((nextTimer = timerInternals.removeNextProcessingTimer()) != null) {
        timers.add(nextTimer);
      }
      if (timers.isEmpty()) {
        return false;
      }
      tester.processElement(
          KeyedWorkItems.timersWorkItem("key".getBytes(StandardCharsets.UTF_8), timers));
      return true;
    }

    List<TimestampedValue<OutputT>> peekOutputElementsInWindow(BoundedWindow window) {
      return tester.peekOutputElementsInWindow(window);
    }

    List<OutputT> takeOutputElements() {
      return tester.takeOutputElements();
    }

    public Instant getWatermarkHold() {
      return stateInternals.earliestWatermarkHold();
    }
  }

  private static class DoFnTesterWindowedValueReceiver implements WindowedValueMultiReceiver {
    private final DoFnTester<?, ?> tester;

    private DoFnTesterWindowedValueReceiver(DoFnTester<?, ?> tester) {
      this.tester = tester;
    }

    @Override
    public <OutputT> void output(TupleTag<OutputT> tag, WindowedValue<OutputT> output) {
      for (BoundedWindow window : output.getWindows()) {
        tester
            .getMutableOutput(tag)
            .add(
                ValueInSingleWindow.of(
                    output.getValue(), output.getTimestamp(), window, output.getPaneInfo()));
      }
    }
  }

  /** A simple splittable {@link DoFn} that's actually monolithic. */
  private static class ToStringFn extends DoFn<Integer, String> {
    @ProcessElement
    public void process(ProcessContext c, RestrictionTracker<SomeRestriction, Void> tracker) {
      checkState(tracker.tryClaim(null));
      c.output(c.element().toString() + "a");
      c.output(c.element().toString() + "b");
      c.output(c.element().toString() + "c");
    }

    @GetInitialRestriction
    public SomeRestriction getInitialRestriction(
        @SuppressWarnings("unused") @Element Integer elem) {
      return new SomeRestriction();
    }
  }

  @Test
  public void testTrivialProcessFnPropagatesOutputWindowAndTimestamp() throws Exception {
    // Tests that ProcessFn correctly propagates the window and timestamp of the element
    // inside the KeyedWorkItem.
    // The underlying DoFn is actually monolithic, so this doesn't test splitting.
    DoFn<Integer, String> fn = new ToStringFn();

    Instant base = Instant.now();

    IntervalWindow w =
        new IntervalWindow(
            base.minus(Duration.standardMinutes(1)), base.plus(Duration.standardMinutes(1)));

    ProcessFnTester<Integer, String, SomeRestriction, Void, Void> tester =
        new ProcessFnTester<>(
            base,
            fn,
            BigEndianIntegerCoder.of(),
            SerializableCoder.of(SomeRestriction.class),
            VoidCoder.of(),
            MAX_OUTPUTS_PER_BUNDLE,
            MAX_BUNDLE_DURATION);
    tester.startElement(
        WindowedValues.of(
            KV.of(42, new SomeRestriction()),
            base,
            Collections.singletonList(w),
            PaneInfo.ON_TIME_AND_ONLY_FIRING));

    assertEquals(
        Arrays.asList(
            TimestampedValue.of("42a", base),
            TimestampedValue.of("42b", base),
            TimestampedValue.of("42c", base)),
        tester.peekOutputElementsInWindow(w));
  }

  private static class WatermarkUpdateFn extends DoFn<Instant, String> {
    @ProcessElement
    public void process(
        ProcessContext c,
        RestrictionTracker<OffsetRange, Long> tracker,
        ManualWatermarkEstimator<Instant> watermarkEstimator,
        CausedByDrain causedByDrain) {
      for (long i = tracker.currentRestriction().getFrom(); tracker.tryClaim(i); ++i) {
        watermarkEstimator.setWatermark(c.element().plus(Duration.standardSeconds(i)));
        c.output(String.valueOf(i));
        if (causedByDrain == CausedByDrain.CAUSED_BY_DRAIN) {
          c.output("drains");
        }
      }
    }

    @GetInitialRestriction
    public OffsetRange getInitialRestriction(@SuppressWarnings("unused") @Element Instant elem) {
      throw new IllegalStateException("Expected to be supplied explicitly in this test");
    }

    @NewTracker
    public OffsetRangeTracker newTracker(@Restriction OffsetRange range) {
      return new OffsetRangeTracker(range);
    }

    @GetInitialWatermarkEstimatorState
    public Instant getInitialWatermarkEstimatorState() {
      return GlobalWindow.TIMESTAMP_MIN_VALUE;
    }

    @NewWatermarkEstimator
    public WatermarkEstimators.Manual newWatermarkEstimator(
        @WatermarkEstimatorState Instant watermarkEstimatorState) {
      return new WatermarkEstimators.Manual(watermarkEstimatorState);
    }
  }

  private static class GetSizeFn extends DoFn<Integer, String> {
    @ProcessElement
    public ProcessContinuation process(
        ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker) {
      for (long i = tracker.currentRestriction().getFrom(); tracker.tryClaim(i); ++i) {
        c.output(String.valueOf(i));
        if (i == 2) {
          return resume();
        }
      }
      return stop();
    }

    @GetInitialRestriction
    public OffsetRange getInitialRestriction() {
      return new OffsetRange(0, 10);
    }

    @NewTracker
    public OffsetRangeTracker newTracker(@Restriction OffsetRange range) {
      return new OffsetRangeTracker(range);
    }

    @GetSize
    public double getSize(@Restriction OffsetRange range) {
      return range.getTo() - range.getFrom();
    }
  }

  // Used to check that backlog can be computed from the restriction tracker if GetSize is not
  // defined.
  private static class SdfWithoutGetSize extends DoFn<Integer, String> {
    @ProcessElement
    public ProcessContinuation process(
        ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker) {
      for (long i = tracker.currentRestriction().getFrom(); tracker.tryClaim(i); ++i) {
        c.output(String.valueOf(i));
        if (i == 2) {
          return resume();
        }
      }
      return stop();
    }

    @GetInitialRestriction
    public OffsetRange getInitialRestriction() {
      return new OffsetRange(0, 10);
    }

    @NewTracker
    public OffsetRangeTracker newTracker(@Restriction OffsetRange range) {
      return new OffsetRangeTracker(range);
    }
  }

  @Test
  public void testDrains() throws Exception {
    DoFn<Instant, String> fn = new WatermarkUpdateFn();
    Instant base = Instant.now();

    ProcessFnTester<Instant, String, OffsetRange, Long, Instant> tester =
        new ProcessFnTester<>(
            base,
            fn,
            InstantCoder.of(),
            SerializableCoder.of(OffsetRange.class),
            InstantCoder.of(),
            3,
            MAX_BUNDLE_DURATION);

    tester.startElement(base, new OffsetRange(0, 8));
    assertThat(tester.takeOutputElements(), hasItems("0", "1", "2"));
    assertEquals(base.plus(Duration.standardSeconds(2)), tester.getWatermarkHold());
    assertTrue(tester.advanceDrain());
    assertThat(tester.takeOutputElements(), hasItems("3", "4", "drains", "drains"));
  }

  @Test
  public void testUpdatesWatermark() throws Exception {
    DoFn<Instant, String> fn = new WatermarkUpdateFn();
    Instant base = Instant.now();

    ProcessFnTester<Instant, String, OffsetRange, Long, Instant> tester =
        new ProcessFnTester<>(
            base,
            fn,
            InstantCoder.of(),
            SerializableCoder.of(OffsetRange.class),
            InstantCoder.of(),
            3,
            MAX_BUNDLE_DURATION);

    tester.startElement(base, new OffsetRange(0, 8));
    assertThat(tester.takeOutputElements(), hasItems("0", "1", "2"));
    assertEquals(base.plus(Duration.standardSeconds(2)), tester.getWatermarkHold());

    assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
    assertThat(tester.takeOutputElements(), hasItems("3", "4", "5"));
    assertEquals(base.plus(Duration.standardSeconds(5)), tester.getWatermarkHold());

    assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
    assertThat(tester.takeOutputElements(), hasItems("6", "7"));
    assertEquals(null, tester.getWatermarkHold());
  }

  /** A simple splittable {@link DoFn} that outputs the given element every 5 seconds forever. */
  private static class SelfInitiatedResumeFn extends DoFn<Integer, String> {
    @ProcessElement
    public ProcessContinuation process(
        ProcessContext c, RestrictionTracker<SomeRestriction, Void> tracker) {
      checkState(tracker.tryClaim(null));
      c.output(c.element().toString());
      return resume().withResumeDelay(Duration.standardSeconds(5));
    }

    @GetInitialRestriction
    public SomeRestriction getInitialRestriction(
        @SuppressWarnings("unused") @Element Integer elem) {
      return new SomeRestriction();
    }
  }

  @Test
  public void testResumeSetsTimer() throws Exception {
    DoFn<Integer, String> fn = new SelfInitiatedResumeFn();
    Instant base = Instant.now();
    dateTimeProvider.setDateTimeFixed(base.getMillis());
    ProcessFnTester<Integer, String, SomeRestriction, Void, Void> tester =
        new ProcessFnTester<>(
            base,
            fn,
            BigEndianIntegerCoder.of(),
            SerializableCoder.of(SomeRestriction.class),
            VoidCoder.of(),
            MAX_OUTPUTS_PER_BUNDLE,
            MAX_BUNDLE_DURATION);

    tester.startElement(42, new SomeRestriction());
    assertThat(tester.takeOutputElements(), contains("42"));

    // Should resume after 5 seconds: advancing by 3 seconds should have no effect.
    assertFalse(tester.advanceProcessingTimeBy(Duration.standardSeconds(3)));
    assertTrue(tester.takeOutputElements().isEmpty());

    // 6 seconds should be enough  should invoke the fn again.
    assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(3)));
    assertThat(tester.takeOutputElements(), contains("42"));

    // Should again resume after 5 seconds: advancing by 3 seconds should again have no effect.
    assertFalse(tester.advanceProcessingTimeBy(Duration.standardSeconds(3)));
    assertTrue(tester.takeOutputElements().isEmpty());

    // 6 seconds should again be enough.
    assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(3)));
    assertThat(tester.takeOutputElements(), contains("42"));
  }

  /** A splittable {@link DoFn} that generates the sequence [init, init + total). */
  private static class CounterFn extends DoFn<Integer, String> {
    private final int numOutputsPerCall;

    public CounterFn(int numOutputsPerCall) {
      this.numOutputsPerCall = numOutputsPerCall;
    }

    @ProcessElement
    public ProcessContinuation process(
        ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker) {
      for (long i = tracker.currentRestriction().getFrom(), numIterations = 0;
          tracker.tryClaim(i);
          ++i, ++numIterations) {
        c.output(String.valueOf(c.element() + i));
        if (numIterations == numOutputsPerCall - 1) {
          return resume();
        }
      }
      return stop();
    }

    @GetInitialRestriction
    public OffsetRange getInitialRestriction(@SuppressWarnings("unused") @Element Integer elem) {
      throw new UnsupportedOperationException("Expected to be supplied explicitly in this test");
    }
  }

  @Test
  public void testResumeCarriesOverState() throws Exception {
    DoFn<Integer, String> fn = new CounterFn(1);
    Instant base = Instant.now();
    dateTimeProvider.setDateTimeFixed(base.getMillis());
    ProcessFnTester<Integer, String, OffsetRange, Long, Void> tester =
        new ProcessFnTester<>(
            base,
            fn,
            BigEndianIntegerCoder.of(),
            SerializableCoder.of(OffsetRange.class),
            VoidCoder.of(),
            MAX_OUTPUTS_PER_BUNDLE,
            MAX_BUNDLE_DURATION);

    tester.startElement(42, new OffsetRange(0, 3));
    assertThat(tester.takeOutputElements(), contains("42"));
    assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
    assertThat(tester.takeOutputElements(), contains("43"));
    assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
    assertThat(tester.takeOutputElements(), contains("44"));
    // Should not resume the null residual.
    assertFalse(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
    // After outputting all 3 items, should not output anything more.
    assertEquals(0, tester.takeOutputElements().size());
    // Should also not ask to resume.
    assertFalse(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
  }

  @Test
  public void testCheckpointsAfterNumOutputs() throws Exception {
    int max = 100;
    DoFn<Integer, String> fn = new CounterFn(Integer.MAX_VALUE);
    Instant base = Instant.now();
    int baseIndex = 42;

    ProcessFnTester<Integer, String, OffsetRange, Long, Void> tester =
        new ProcessFnTester<>(
            base,
            fn,
            BigEndianIntegerCoder.of(),
            SerializableCoder.of(OffsetRange.class),
            VoidCoder.of(),
            max,
            MAX_BUNDLE_DURATION);

    List<String> elements;

    // Create an fn that attempts to 2x output more than checkpointing allows.
    tester.startElement(baseIndex, new OffsetRange(0, 2 * max + max / 2));
    elements = tester.takeOutputElements();
    assertEquals(max, elements.size());
    // Should output the range [0, max)
    assertThat(elements, hasItem(String.valueOf(baseIndex)));
    assertThat(elements, hasItem(String.valueOf(baseIndex + max - 1)));

    assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
    elements = tester.takeOutputElements();
    assertEquals(max, elements.size());
    // Should output the range [max, 2*max)
    assertThat(elements, hasItem(String.valueOf(baseIndex + max)));
    assertThat(elements, hasItem(String.valueOf(baseIndex + 2 * max - 1)));

    assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
    elements = tester.takeOutputElements();
    assertEquals(max / 2, elements.size());
    // Should output the range [2*max, 2*max + max/2)
    assertThat(elements, hasItem(String.valueOf(baseIndex + 2 * max)));
    assertThat(elements, hasItem(String.valueOf(baseIndex + 2 * max + max / 2 - 1)));
    assertThat(elements, not(hasItem(String.valueOf(baseIndex + 2 * max + max / 2))));
  }

  @Test
  public void testCheckpointsAfterDuration() throws Exception {
    // Don't bound number of outputs.
    int max = Integer.MAX_VALUE;
    // But bound bundle duration - the bundle should terminate.
    Duration maxBundleDuration = Duration.standardSeconds(1);
    // Create an fn that attempts to 2x output more than checkpointing allows.
    DoFn<Integer, String> fn = new CounterFn(Integer.MAX_VALUE);
    Instant base = Instant.now();
    int baseIndex = 42;

    ProcessFnTester<Integer, String, OffsetRange, Long, Void> tester =
        new ProcessFnTester<>(
            base,
            fn,
            BigEndianIntegerCoder.of(),
            SerializableCoder.of(OffsetRange.class),
            VoidCoder.of(),
            max,
            maxBundleDuration);

    List<String> elements;

    tester.startElement(baseIndex, new OffsetRange(0, Long.MAX_VALUE));
    // Bundle should terminate, and should do at least some processing.
    elements = tester.takeOutputElements();
    assertFalse(elements.isEmpty());
    // Bundle should have run for at least the requested duration.
    assertThat(
        Instant.now().getMillis() - base.getMillis(),
        greaterThanOrEqualTo(maxBundleDuration.getMillis()));
  }

  private static class LifecycleVerifyingFn extends DoFn<Integer, String> {
    private enum State {
      BEFORE_SETUP,
      OUTSIDE_BUNDLE,
      INSIDE_BUNDLE,
      TORN_DOWN
    }

    private State state = State.BEFORE_SETUP;

    @ProcessElement
    public void process(
        @SuppressWarnings("unused") ProcessContext c,
        @SuppressWarnings("unused") RestrictionTracker<SomeRestriction, Void> tracker) {
      assertEquals(State.INSIDE_BUNDLE, state);
    }

    @GetInitialRestriction
    public SomeRestriction getInitialRestriction(
        @SuppressWarnings("unused") @Element Integer element) {
      return new SomeRestriction();
    }

    @Setup
    public void setup() {
      assertEquals(State.BEFORE_SETUP, state);
      state = State.OUTSIDE_BUNDLE;
    }

    @Teardown
    public void tearDown() {
      assertEquals(State.OUTSIDE_BUNDLE, state);
      state = State.TORN_DOWN;
    }

    @StartBundle
    public void startBundle() {
      assertEquals(State.OUTSIDE_BUNDLE, state);
      state = State.INSIDE_BUNDLE;
    }

    @FinishBundle
    public void finishBundle() {
      assertEquals(State.INSIDE_BUNDLE, state);
      state = State.OUTSIDE_BUNDLE;
    }
  }

  @Test
  public void testInvokesLifecycleMethods() throws Exception {
    DoFn<Integer, String> fn = new LifecycleVerifyingFn();
    try (ProcessFnTester<Integer, String, SomeRestriction, Void, Void> tester =
        new ProcessFnTester<>(
            Instant.now(),
            fn,
            BigEndianIntegerCoder.of(),
            SerializableCoder.of(SomeRestriction.class),
            VoidCoder.of(),
            MAX_OUTPUTS_PER_BUNDLE,
            MAX_BUNDLE_DURATION)) {
      tester.startElement(42, new SomeRestriction());
    }
  }

  @Test
  public void testReportsBacklog() throws Exception {
    DoFn<Integer, String> fn = new GetSizeFn();
    Instant base = Instant.now();
    final List<Double> backlogs = new ArrayList<>();

    try (ProcessFnTester<Integer, String, OffsetRange, Long, Void> tester =
        new ProcessFnTester<>(
            base,
            fn,
            BigEndianIntegerCoder.of(),
            SerializableCoder.of(OffsetRange.class),
            VoidCoder.of(),
            MAX_OUTPUTS_PER_BUNDLE,
            MAX_BUNDLE_DURATION)) {
      tester.processFn.setBacklogBytesCallback(backlogs::add);

      tester.startElement(42, new OffsetRange(0, 10));
      // First call outputs 0, 1, and 2, and then resumes.
      // The residual range should be [3, 10), so size is 7.
      assertEquals(1, backlogs.size());
      assertEquals(7.0, backlogs.get(0), 0.001);
    }
  }

  @Test
  public void testReportsBacklogWithoutGetSize() throws Exception {
    DoFn<Integer, String> fn = new SdfWithoutGetSize();
    Instant base = Instant.now();
    final List<Double> backlogs = new ArrayList<>();

    try (ProcessFnTester<Integer, String, OffsetRange, Long, Void> tester =
        new ProcessFnTester<>(
            base,
            fn,
            BigEndianIntegerCoder.of(),
            SerializableCoder.of(OffsetRange.class),
            VoidCoder.of(),
            MAX_OUTPUTS_PER_BUNDLE,
            MAX_BUNDLE_DURATION)) {
      tester.processFn.setBacklogBytesCallback(backlogs::add);

      tester.startElement(42, new OffsetRange(0, 10));
      // First call outputs 0, 1, and 2, and then resumes.
      // The residual range should be [3, 10), so size is 7.
      assertEquals(1, backlogs.size());
      assertEquals(7.0, backlogs.get(0), 0.001);
    }
  }

  private static class TruncateFn extends DoFn<Integer, String> {
    private final boolean truncateToNull;
    private final List<String> calls = new ArrayList<>();

    public TruncateFn(boolean truncateToNull) {
      this.truncateToNull = truncateToNull;
    }

    @ProcessElement
    public ProcessContinuation process(
        ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker) {
      for (long i = tracker.currentRestriction().getFrom(); tracker.tryClaim(i); ++i) {
        c.output(c.element() + ":" + i);
        if (i == 2) {
          return resume();
        }
      }
      return stop();
    }

    @GetInitialRestriction
    public OffsetRange getInitialRestriction() {
      return new OffsetRange(0, 10);
    }

    @NewTracker
    public OffsetRangeTracker newTracker(@Restriction OffsetRange range) {
      return new OffsetRangeTracker(range);
    }

    @TruncateRestriction
    public RestrictionTracker.TruncateResult<OffsetRange> truncate(
        @Restriction OffsetRange restriction, @Element Integer element) {
      calls.add("truncate:" + element + ":" + restriction);
      if (truncateToNull) {
        return null;
      }
      // Truncate so that we only process one more element.
      return RestrictionTracker.TruncateResult.of(
          new OffsetRange(restriction.getFrom(), restriction.getFrom() + 1));
    }
  }

  @Test
  public void testTruncateRestrictionOnDrain() throws Exception {
    TruncateFn fn = new TruncateFn(false);
    Instant base = Instant.now();

    try (ProcessFnTester<Integer, String, OffsetRange, Long, Void> tester =
        new ProcessFnTester<>(
            base,
            fn,
            BigEndianIntegerCoder.of(),
            SerializableCoder.of(OffsetRange.class),
            VoidCoder.of(),
            MAX_OUTPUTS_PER_BUNDLE,
            MAX_BUNDLE_DURATION)) {
      tester.startElement(42, new OffsetRange(0, 10));
      assertThat(tester.takeOutputElements(), contains("42:0", "42:1", "42:2"));

      assertTrue(tester.advanceDrain());
      assertThat(tester.takeOutputElements(), contains("42:3"));
      assertEquals(Collections.singletonList("truncate:42:[3, 10)"), fn.calls);
      assertEquals(null, tester.getWatermarkHold());
    }
  }

  @Test
  public void testTruncateRestrictionReturnsNullOnDrain() throws Exception {
    TruncateFn fn = new TruncateFn(true);
    Instant base = Instant.now();

    try (ProcessFnTester<Integer, String, OffsetRange, Long, Void> tester =
        new ProcessFnTester<>(
            base,
            fn,
            BigEndianIntegerCoder.of(),
            SerializableCoder.of(OffsetRange.class),
            VoidCoder.of(),
            MAX_OUTPUTS_PER_BUNDLE,
            MAX_BUNDLE_DURATION)) {
      tester.startElement(42, new OffsetRange(0, 10));
      assertThat(tester.takeOutputElements(), contains("42:0", "42:1", "42:2"));

      assertTrue(tester.advanceDrain());
      assertTrue(tester.takeOutputElements().isEmpty());
      assertEquals(Collections.singletonList("truncate:42:[3, 10)"), fn.calls);
      assertEquals(null, tester.getWatermarkHold());
    }
  }

  private static class TruncateWithSideInputFn extends DoFn<Integer, String> {
    private final List<String> calls = new ArrayList<>();

    @ProcessElement
    public ProcessContinuation process(
        ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker) {
      for (long i = tracker.currentRestriction().getFrom(); tracker.tryClaim(i); ++i) {
        c.output(c.element() + ":" + i);
        if (i == 2) {
          return resume();
        }
      }
      return stop();
    }

    @GetInitialRestriction
    public OffsetRange getInitialRestriction() {
      return new OffsetRange(0, 10);
    }

    @NewTracker
    public OffsetRangeTracker newTracker(@Restriction OffsetRange range) {
      return new OffsetRangeTracker(range);
    }

    @TruncateRestriction
    public RestrictionTracker.TruncateResult<OffsetRange> truncate(
        @Restriction OffsetRange restriction,
        @Element Integer element,
        @SideInput("sideInput") String sideInput) {
      calls.add("truncate:" + element + ":" + sideInput + ":" + restriction);
      return RestrictionTracker.TruncateResult.of(
          new OffsetRange(restriction.getFrom(), restriction.getFrom() + 1));
    }
  }

  @Test
  public void testTruncateRestrictionWithSideInputOnDrain() throws Exception {
    TruncateWithSideInputFn fn = new TruncateWithSideInputFn();
    Instant base = Instant.now();
    PCollectionView<String> view =
        TestPipeline.create().apply(Create.of("sideValue")).apply(View.asSingleton());
    Map<String, PCollectionView<?>> sideInputMapping = Collections.singletonMap("sideInput", view);
    SideInputReader sideInputReader =
        new SideInputReader() {
          @Override
          public <T> T get(PCollectionView<T> v, BoundedWindow window) {
            if (v.equals(view)) {
              return (T) "sideValue";
            }
            throw new NoSuchElementException();
          }

          @Override
          public <T> boolean contains(PCollectionView<T> v) {
            return v.equals(view);
          }

          @Override
          public boolean isEmpty() {
            return false;
          }
        };

    try (ProcessFnTester<Integer, String, OffsetRange, Long, Void> tester =
        new ProcessFnTester<>(
            base,
            fn,
            BigEndianIntegerCoder.of(),
            SerializableCoder.of(OffsetRange.class),
            VoidCoder.of(),
            MAX_OUTPUTS_PER_BUNDLE,
            MAX_BUNDLE_DURATION,
            sideInputMapping,
            sideInputReader)) {
      tester.startElement(42, new OffsetRange(0, 10));
      assertThat(tester.takeOutputElements(), contains("42:0", "42:1", "42:2"));

      assertTrue(tester.advanceDrain());
      assertThat(tester.takeOutputElements(), contains("42:3"));
      assertEquals(Collections.singletonList("truncate:42:sideValue:[3, 10)"), fn.calls);
    }
  }

  private static class UnboundedOffsetRangeTracker extends OffsetRangeTracker {
    public UnboundedOffsetRangeTracker(OffsetRange range) {
      super(range);
    }

    @Override
    public IsBounded isBounded() {
      return IsBounded.UNBOUNDED;
    }
  }

  // Tests that if we don't override TruncateRestriction, the default TruncateRestriction
  // implementation is used (which for unbounded restrictions stops processing immediately).
  private static class DefaultTruncateUnboundedFn extends DoFn<Integer, String> {
    @ProcessElement
    public ProcessContinuation process(
        ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker) {
      for (long i = tracker.currentRestriction().getFrom(); tracker.tryClaim(i); ++i) {
        c.output(c.element() + ":" + i);
        if (i == 2) {
          return resume();
        }
      }
      return stop();
    }

    @GetInitialRestriction
    public OffsetRange getInitialRestriction() {
      return new OffsetRange(0, 10);
    }

    @NewTracker
    public RestrictionTracker<OffsetRange, Long> newTracker(@Restriction OffsetRange range) {
      return new UnboundedOffsetRangeTracker(range);
    }
  }

  @Test
  public void testDefaultTruncateRestrictionUnboundedStopsOnDrain() throws Exception {
    DefaultTruncateUnboundedFn fn = new DefaultTruncateUnboundedFn();
    Instant base = Instant.now();

    try (ProcessFnTester<Integer, String, OffsetRange, Long, Void> tester =
        new ProcessFnTester<>(
            base,
            fn,
            BigEndianIntegerCoder.of(),
            SerializableCoder.of(OffsetRange.class),
            VoidCoder.of(),
            MAX_OUTPUTS_PER_BUNDLE,
            MAX_BUNDLE_DURATION)) {
      tester.startElement(42, new OffsetRange(0, 10));
      assertThat(tester.takeOutputElements(), contains("42:0", "42:1", "42:2"));

      assertTrue(tester.advanceDrain());
      assertTrue(tester.takeOutputElements().isEmpty());
      assertEquals(null, tester.getWatermarkHold());
    }
  }
}
