SimpleDoFnRunner.java

/*
 * 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 com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;

import com.google.common.collect.Iterables;
import com.google.common.collect.Sets;
import java.util.Collection;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Set;
import javax.annotation.Nullable;
import org.apache.beam.runners.core.DoFnRunners.OutputManager;
import org.apache.beam.runners.core.ExecutionContext.StepContext;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.state.State;
import org.apache.beam.sdk.state.StateSpec;
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.state.Timer;
import org.apache.beam.sdk.state.TimerSpec;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.DoFn.FinishBundleContext;
import org.apache.beam.sdk.transforms.DoFn.OnTimerContext;
import org.apache.beam.sdk.transforms.DoFn.ProcessContext;
import org.apache.beam.sdk.transforms.DoFn.StartBundleContext;
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.splittabledofn.RestrictionTracker;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.transforms.windowing.GlobalWindows;
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
import org.apache.beam.sdk.transforms.windowing.WindowFn;
import org.apache.beam.sdk.util.SystemDoFnInternal;
import org.apache.beam.sdk.util.UserCodeException;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.WindowingStrategy;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.joda.time.format.PeriodFormat;

/**
 * Runs a {@link DoFn} by constructing the appropriate contexts and passing them in.
 *
 * <p>Also, if the {@link DoFn} observes the window of the element, then {@link SimpleDoFnRunner}
 * explodes windows of the input {@link WindowedValue} and calls {@link DoFn.ProcessElement} for
 * each window individually.
 *
 * @param <InputT> the type of the {@link DoFn} (main) input elements
 * @param <OutputT> the type of the {@link DoFn} (main) output elements
 */
public class SimpleDoFnRunner<InputT, OutputT> implements DoFnRunner<InputT, OutputT> {

  /** The {@link DoFn} being run. */
  private final DoFn<InputT, OutputT> fn;

  /** The {@link DoFnInvoker} being run. */
  private final DoFnInvoker<InputT, OutputT> invoker;

  /** The context used for running the {@link DoFn}. */
  private final DoFnContext<InputT, OutputT> context;

  private final OutputManager outputManager;

  private final TupleTag<OutputT> mainOutputTag;

  private final boolean observesWindow;

  private final DoFnSignature signature;

  private final Coder<BoundedWindow> windowCoder;

  private final Duration allowedLateness;

  // Because of setKey(Object), we really must refresh stateInternals() at each access
  private final StepContext stepContext;

  public SimpleDoFnRunner(
      PipelineOptions options,
      DoFn<InputT, OutputT> fn,
      SideInputReader sideInputReader,
      OutputManager outputManager,
      TupleTag<OutputT> mainOutputTag,
      List<TupleTag<?>> additionalOutputTags,
      StepContext stepContext,
      WindowingStrategy<?, ?> windowingStrategy) {
    this.fn = fn;
    this.signature = DoFnSignatures.getSignature(fn.getClass());
    this.observesWindow = signature.processElement().observesWindow() || !sideInputReader.isEmpty();
    this.invoker = DoFnInvokers.invokerFor(fn);
    this.outputManager = outputManager;
    this.mainOutputTag = mainOutputTag;
    this.stepContext = stepContext;

    // This is a cast of an _invariant_ coder. But we are assured by pipeline validation
    // that it really is the coder for whatever BoundedWindow subclass is provided
    @SuppressWarnings("unchecked")
    Coder<BoundedWindow> untypedCoder =
        (Coder<BoundedWindow>) windowingStrategy.getWindowFn().windowCoder();
    this.windowCoder = untypedCoder;
    this.allowedLateness = windowingStrategy.getAllowedLateness();

    this.context =
        new DoFnContext<>(
            options,
            fn,
            sideInputReader,
            outputManager,
            mainOutputTag,
            additionalOutputTags,
            stepContext,
            windowingStrategy.getWindowFn());
  }

  @Override
  public void startBundle() {
    DoFnStartBundleContext<InputT, OutputT> startBundleContext =
        createStartBundleContext(fn, context);
    // This can contain user code. Wrap it in case it throws an exception.
    try {
      invoker.invokeStartBundle(startBundleContext);
    } catch (Throwable t) {
      // Exception in user code.
      throw wrapUserCodeException(t);
    }
  }

  @Override
  public void processElement(WindowedValue<InputT> compressedElem) {
    if (observesWindow) {
      for (WindowedValue<InputT> elem : compressedElem.explodeWindows()) {
        invokeProcessElement(elem);
      }
    } else {
      invokeProcessElement(compressedElem);
    }
  }

  @Override
  public void onTimer(
      String timerId, BoundedWindow window, Instant timestamp, TimeDomain timeDomain) {

    // The effective timestamp is when derived elements will have their timestamp set, if not
    // otherwise specified. If this is an event time timer, then they have the timestamp of the
    // timer itself. Otherwise, they are set to the input timestamp, which is by definition
    // non-late.
    Instant effectiveTimestamp;
    switch (timeDomain) {
      case EVENT_TIME:
        effectiveTimestamp = timestamp;
        break;

      case PROCESSING_TIME:
      case SYNCHRONIZED_PROCESSING_TIME:
        effectiveTimestamp = context.stepContext.timerInternals().currentInputWatermarkTime();
        break;

      default:
        throw new IllegalArgumentException(
            String.format("Unknown time domain: %s", timeDomain));
    }

    OnTimerArgumentProvider<InputT, OutputT> argumentProvider =
        new OnTimerArgumentProvider<>(
            fn, context, window, allowedLateness, effectiveTimestamp, timeDomain);
    invoker.invokeOnTimer(timerId, argumentProvider);
  }

  private void invokeProcessElement(WindowedValue<InputT> elem) {
    final DoFnProcessContext<InputT, OutputT> processContext = createProcessContext(elem);

    // This can contain user code. Wrap it in case it throws an exception.
    try {
      invoker.invokeProcessElement(processContext);
    } catch (Exception ex) {
      throw wrapUserCodeException(ex);
    }
  }

  @Override
  public void finishBundle() {
    DoFnFinishBundleContext<InputT, OutputT> finishBundleContext =
        createFinishBundleContext(fn, context);
    // This can contain user code. Wrap it in case it throws an exception.
    try {
      invoker.invokeFinishBundle(finishBundleContext);
    } catch (Throwable t) {
      // Exception in user code.
      throw wrapUserCodeException(t);
    }
  }

  private DoFnStartBundleContext<InputT, OutputT> createStartBundleContext(
      DoFn<InputT, OutputT> fn, DoFnContext<InputT, OutputT> context) {
    return new DoFnStartBundleContext<>(this.fn, this.context);
  }

  private DoFnFinishBundleContext<InputT, OutputT> createFinishBundleContext(
      DoFn<InputT, OutputT> fn, DoFnContext<InputT, OutputT> context) {
    return new DoFnFinishBundleContext<>(fn, context);
  }

  /** Returns a new {@link DoFn.ProcessContext} for the given element. */
  private DoFnProcessContext<InputT, OutputT> createProcessContext(WindowedValue<InputT> elem) {
    return new DoFnProcessContext<InputT, OutputT>(fn, context, elem, allowedLateness);
  }

  private RuntimeException wrapUserCodeException(Throwable t) {
    throw UserCodeException.wrapIf(!isSystemDoFn(), t);
  }

  private boolean isSystemDoFn() {
    return invoker.getClass().isAnnotationPresent(SystemDoFnInternal.class);
  }

  /**
   * A concrete implementation of {@code DoFn.Context} used for running a {@link DoFn}.
   *
   * @param <InputT> the type of the {@link DoFn} (main) input elements
   * @param <OutputT> the type of the {@link DoFn} (main) output elements
   */
  private static class DoFnContext<InputT, OutputT> {
    private static final int MAX_SIDE_OUTPUTS = 1000;

    final PipelineOptions options;
    final DoFn<InputT, OutputT> fn;
    final SideInputReader sideInputReader;
    final OutputManager outputManager;
    final TupleTag<OutputT> mainOutputTag;
    final StepContext stepContext;
    final WindowFn<?, ?> windowFn;

    /**
     * The set of known output tags, some of which may be undeclared, so we can throw an exception
     * when it exceeds {@link #MAX_SIDE_OUTPUTS}.
     */
    private Set<TupleTag<?>> outputTags;

    public DoFnContext(
        PipelineOptions options,
        DoFn<InputT, OutputT> fn,
        SideInputReader sideInputReader,
        OutputManager outputManager,
        TupleTag<OutputT> mainOutputTag,
        List<TupleTag<?>> additionalOutputTags,
        StepContext stepContext,
        WindowFn<?, ?> windowFn) {
      this.options = options;
      this.fn = fn;
      this.sideInputReader = sideInputReader;
      this.outputManager = outputManager;
      this.mainOutputTag = mainOutputTag;
      this.outputTags = Sets.newHashSet();

      outputTags.add(mainOutputTag);
      for (TupleTag<?> additionalOutputTag : additionalOutputTags) {
        outputTags.add(additionalOutputTag);
      }

      this.stepContext = stepContext;
      this.windowFn = windowFn;
    }

    //////////////////////////////////////////////////////////////////////////////

    public PipelineOptions getPipelineOptions() {
      return options;
    }

    <T, W extends BoundedWindow> WindowedValue<T> makeWindowedValue(
        T output, Instant timestamp, Collection<W> windows, PaneInfo pane) {
      final Instant inputTimestamp = timestamp;

      if (timestamp == null) {
        timestamp = BoundedWindow.TIMESTAMP_MIN_VALUE;
      }

      if (windows == null) {
        try {
          // The windowFn can never succeed at accessing the element, so its type does not
          // matter here
          @SuppressWarnings("unchecked")
          WindowFn<Object, W> objectWindowFn = (WindowFn<Object, W>) windowFn;
          windows =
              objectWindowFn.assignWindows(
                  objectWindowFn.new AssignContext() {
                    @Override
                    public Object element() {
                      throw new UnsupportedOperationException(
                          "WindowFn attempted to access input element when none was available");
                    }

                    @Override
                    public Instant timestamp() {
                      if (inputTimestamp == null) {
                        throw new UnsupportedOperationException(
                            "WindowFn attempted to access input timestamp when none was available");
                      }
                      return inputTimestamp;
                    }

                    @Override
                    public W window() {
                      throw new UnsupportedOperationException(
                          "WindowFn attempted to access input windows when none were available");
                    }
                  });
        } catch (Exception e) {
          throw UserCodeException.wrap(e);
        }
      }

      return WindowedValue.of(output, timestamp, windows, pane);
    }

    public <T> T sideInput(PCollectionView<T> view, BoundedWindow sideInputWindow) {
      if (!sideInputReader.contains(view)) {
        throw new IllegalArgumentException("calling sideInput() with unknown view");
      }
      return sideInputReader.get(view, sideInputWindow);
    }

    void outputWindowedValue(
        OutputT output,
        Instant timestamp,
        Collection<? extends BoundedWindow> windows,
        PaneInfo pane) {
      outputWindowedValue(makeWindowedValue(output, timestamp, windows, pane));
    }

    void outputWindowedValue(WindowedValue<OutputT> windowedElem) {
      outputManager.output(mainOutputTag, windowedElem);
      if (stepContext != null) {
        stepContext.noteOutput(windowedElem);
      }
    }

    private <T> void outputWindowedValue(
        TupleTag<T> tag,
        T output,
        Instant timestamp,
        Collection<? extends BoundedWindow> windows,
        PaneInfo pane) {
      outputWindowedValue(tag, makeWindowedValue(output, timestamp, windows, pane));
    }

    private <T> void outputWindowedValue(TupleTag<T> tag, WindowedValue<T> windowedElem) {
      if (!outputTags.contains(tag)) {
        // This tag wasn't declared nor was it seen before during this execution.
        // Thus, this must be a new, undeclared and unconsumed output.
        // To prevent likely user errors, enforce the limit on the number of side
        // outputs.
        if (outputTags.size() >= MAX_SIDE_OUTPUTS) {
          throw new IllegalArgumentException(
              "the number of outputs has exceeded a limit of " + MAX_SIDE_OUTPUTS);
        }
        outputTags.add(tag);
      }

      outputManager.output(tag, windowedElem);
      if (stepContext != null) {
        stepContext.noteOutput(tag, windowedElem);
      }
    }
  }


  /**
   * A concrete implementation of {@link DoFn.StartBundleContext}.
   */
  private class DoFnStartBundleContext<InputT, OutputT>
      extends DoFn<InputT, OutputT>.StartBundleContext
      implements DoFnInvoker.ArgumentProvider<InputT, OutputT> {
    private final DoFn<InputT, OutputT> fn;
    private final DoFnContext<InputT, OutputT> context;

    private DoFnStartBundleContext(DoFn<InputT, OutputT> fn, DoFnContext<InputT, OutputT> context) {
      fn.super();
      this.fn = fn;
      this.context = context;
    }

    @Override
    public PipelineOptions getPipelineOptions() {
      return context.getPipelineOptions();
    }

    @Override
    public BoundedWindow window() {
      throw new UnsupportedOperationException(
          "Cannot access window outside of @ProcessElement and @OnTimer methods.");
    }

    @Override
    public StartBundleContext startBundleContext(DoFn<InputT, OutputT> doFn) {
      return this;
    }

    @Override
    public FinishBundleContext finishBundleContext(DoFn<InputT, OutputT> doFn) {
      throw new UnsupportedOperationException(
          "Cannot access FinishBundleContext outside of @FinishBundle method.");
    }

    @Override
    public ProcessContext processContext(DoFn<InputT, OutputT> doFn) {
      throw new UnsupportedOperationException(
          "Cannot access ProcessContext outside of @ProcessElement method.");
    }

    @Override
    public OnTimerContext onTimerContext(DoFn<InputT, OutputT> doFn) {
      throw new UnsupportedOperationException(
          "Cannot access OnTimerContext outside of @OnTimer methods.");
    }

    @Override
    public RestrictionTracker<?> restrictionTracker() {
      throw new UnsupportedOperationException(
          "Cannot access RestrictionTracker outside of @ProcessElement method.");
    }

    @Override
    public State state(String stateId) {
      throw new UnsupportedOperationException(
          "Cannot access state outside of @ProcessElement and @OnTimer methods.");
    }

    @Override
    public Timer timer(String timerId) {
      throw new UnsupportedOperationException(
          "Cannot access timers outside of @ProcessElement and @OnTimer methods.");
    }
  }

  /**
   * B
   * A concrete implementation of {@link DoFn.FinishBundleContext}.
   */
  private class DoFnFinishBundleContext<InputT, OutputT>
      extends DoFn<InputT, OutputT>.FinishBundleContext
      implements DoFnInvoker.ArgumentProvider<InputT, OutputT> {
    private final DoFnContext<InputT, OutputT> context;

    private DoFnFinishBundleContext(
        DoFn<InputT, OutputT> fn, DoFnContext<InputT, OutputT> context) {
      fn.super();
      this.context = context;
    }

    @Override
    public PipelineOptions getPipelineOptions() {
      return context.getPipelineOptions();
    }

    @Override
    public BoundedWindow window() {
      throw new UnsupportedOperationException(
          "Cannot access window outside of @ProcessElement and @OnTimer methods.");
    }

    @Override
    public StartBundleContext startBundleContext(DoFn<InputT, OutputT> doFn) {
      throw new UnsupportedOperationException(
          "Cannot access StartBundleContext outside of @StartBundle method.");
    }

    @Override
    public FinishBundleContext finishBundleContext(DoFn<InputT, OutputT> doFn) {
      return this;
    }

    @Override
    public ProcessContext processContext(DoFn<InputT, OutputT> doFn) {
      throw new UnsupportedOperationException(
          "Cannot access ProcessContext outside of @ProcessElement method.");
    }

    @Override
    public OnTimerContext onTimerContext(DoFn<InputT, OutputT> doFn) {
      throw new UnsupportedOperationException(
          "Cannot access OnTimerContext outside of @OnTimer methods.");
    }

    @Override
    public RestrictionTracker<?> restrictionTracker() {
      throw new UnsupportedOperationException(
          "Cannot access RestrictionTracker outside of @ProcessElement method.");
    }

    @Override
    public State state(String stateId) {
      throw new UnsupportedOperationException(
          "Cannot access state outside of @ProcessElement and @OnTimer methods.");
    }

    @Override
    public Timer timer(String timerId) {
      throw new UnsupportedOperationException(
          "Cannot access timers outside of @ProcessElement and @OnTimer methods.");
    }

    @Override
    public void output(OutputT output, Instant timestamp, BoundedWindow window) {
      context.outputWindowedValue(WindowedValue.of(output, timestamp, window, PaneInfo.NO_FIRING));
    }

    @Override
    public <T> void output(TupleTag<T> tag, T output, Instant timestamp, BoundedWindow window) {
      context.outputWindowedValue(
          tag, WindowedValue.of(output, timestamp, window, PaneInfo.NO_FIRING));
    }
  }

  /**
   * A concrete implementation of {@link DoFn.ProcessContext} used for running a {@link DoFn} over a
   * single element.
   *
   * @param <InputT> the type of the {@link DoFn} (main) input elements
   * @param <OutputT> the type of the {@link DoFn} (main) output elements
   */
  private class DoFnProcessContext<InputT, OutputT> extends DoFn<InputT, OutputT>.ProcessContext
      implements DoFnInvoker.ArgumentProvider<InputT, OutputT> {

    final DoFn<InputT, OutputT> fn;
    final DoFnContext<InputT, OutputT> context;
    final WindowedValue<InputT> windowedValue;
    private final Duration allowedLateness;

    /** Lazily initialized; should only be accessed via {@link #getNamespace()}. */
    @Nullable private StateNamespace namespace;

    /**
     * The state namespace for this context.
     *
     * <p>Any call to {@link #getNamespace()} when more than one window is present will crash; this
     * represents a bug in the runner or the {@link DoFnSignature}, since values must be in exactly
     * one window when state or timers are relevant.
     */
    private StateNamespace getNamespace() {
      if (namespace == null) {
        namespace = StateNamespaces.window(windowCoder, window());
      }
      return namespace;
    }

    private DoFnProcessContext(
        DoFn<InputT, OutputT> fn,
        DoFnContext<InputT, OutputT> context,
        WindowedValue<InputT> windowedValue,
        Duration allowedLateness) {
      fn.super();
      this.fn = fn;
      this.context = context;
      this.windowedValue = windowedValue;
      this.allowedLateness = allowedLateness;
    }

    @Override
    public PipelineOptions getPipelineOptions() {
      return context.getPipelineOptions();
    }

    @Override
    public InputT element() {
      return windowedValue.getValue();
    }

    @Override
    public <T> T sideInput(PCollectionView<T> view) {
      checkNotNull(view, "View passed to sideInput cannot be null");
      Iterator<? extends BoundedWindow> windowIter = windows().iterator();
      BoundedWindow window;
      if (!windowIter.hasNext()) {
        if (context.windowFn instanceof GlobalWindows) {
          // TODO: Remove this once GroupByKeyOnly no longer outputs elements
          // without windows
          window = GlobalWindow.INSTANCE;
        } else {
          throw new IllegalStateException(
              "sideInput called when main input element is not in any windows");
        }
      } else {
        window = windowIter.next();
        if (windowIter.hasNext()) {
          throw new IllegalStateException(
              "sideInput called when main input element is in multiple windows");
        }
      }
      return context.sideInput(
          view, view.getWindowMappingFn().getSideInputWindow(window));
    }

    @Override
    public PaneInfo pane() {
      return windowedValue.getPane();
    }

    @Override
    public void updateWatermark(Instant watermark) {
      throw new UnsupportedOperationException("Only splittable DoFn's can use updateWatermark()");
    }

    @Override
    public void output(OutputT output) {
      context.outputWindowedValue(windowedValue.withValue(output));
    }

    @Override
    public void outputWithTimestamp(OutputT output, Instant timestamp) {
      checkTimestamp(timestamp);
      context.outputWindowedValue(
          output, timestamp, windowedValue.getWindows(), windowedValue.getPane());
    }

    @Override
    public <T> void output(TupleTag<T> tag, T output) {
      checkNotNull(tag, "Tag passed to output cannot be null");
      context.outputWindowedValue(tag, windowedValue.withValue(output));
    }

    @Override
    public <T> void outputWithTimestamp(TupleTag<T> tag, T output, Instant timestamp) {
      checkNotNull(tag, "Tag passed to outputWithTimestamp cannot be null");
      checkTimestamp(timestamp);
      context.outputWindowedValue(
          tag, output, timestamp, windowedValue.getWindows(), windowedValue.getPane());
    }

    @Override
    public Instant timestamp() {
      return windowedValue.getTimestamp();
    }

    public Collection<? extends BoundedWindow> windows() {
      return windowedValue.getWindows();
    }

    @SuppressWarnings("deprecation") // Allowed Skew is deprecated for users, but must be respected
    private void checkTimestamp(Instant timestamp) {
      // The documentation of getAllowedTimestampSkew explicitly permits Long.MAX_VALUE to be used
      // for infinite skew. Defend against underflow in that case for timestamps before the epoch
      if (fn.getAllowedTimestampSkew().getMillis() != Long.MAX_VALUE
          && timestamp.isBefore(windowedValue.getTimestamp().minus(fn.getAllowedTimestampSkew()))) {
        throw new IllegalArgumentException(
            String.format(
                "Cannot output with timestamp %s. Output timestamps must be no earlier than the "
                    + "timestamp of the current input (%s) minus the allowed skew (%s). See the "
                    + "DoFn#getAllowedTimestampSkew() Javadoc for details on changing the allowed "
                    + "skew.",
                timestamp,
                windowedValue.getTimestamp(),
                PeriodFormat.getDefault().print(fn.getAllowedTimestampSkew().toPeriod())));
      }
    }

    @Override
    public BoundedWindow window() {
      return Iterables.getOnlyElement(windowedValue.getWindows());
    }

    @Override
    public StartBundleContext startBundleContext(DoFn<InputT, OutputT> doFn) {
      throw new UnsupportedOperationException("StartBundleContext parameters are not supported.");
    }

    @Override
    public FinishBundleContext finishBundleContext(DoFn<InputT, OutputT> doFn) {
      throw new UnsupportedOperationException("FinishBundleContext parameters are not supported.");
    }

    @Override
    public DoFn<InputT, OutputT>.ProcessContext processContext(DoFn<InputT, OutputT> doFn) {
      return this;
    }

    @Override
    public OnTimerContext onTimerContext(DoFn<InputT, OutputT> doFn) {
      throw new UnsupportedOperationException(
          "Cannot access OnTimerContext outside of @OnTimer methods.");
    }

    @Override
    public RestrictionTracker<?> restrictionTracker() {
      throw new UnsupportedOperationException("RestrictionTracker parameters are not supported.");
    }

    @Override
    public State state(String stateId) {
      try {
        StateSpec<?> spec =
            (StateSpec<?>) signature.stateDeclarations().get(stateId).field().get(fn);
        return stepContext
            .stateInternals()
            .state(getNamespace(), StateTags.tagForSpec(stateId, (StateSpec) spec));
      } catch (IllegalAccessException e) {
        throw new RuntimeException(e);
      }
    }

    @Override
    public Timer timer(String timerId) {
      try {
        TimerSpec spec =
            (TimerSpec) signature.timerDeclarations().get(timerId).field().get(fn);
        return new TimerInternalsTimer(
            window(), getNamespace(), allowedLateness, timerId, spec, stepContext.timerInternals());
      } catch (IllegalAccessException e) {
        throw new RuntimeException(e);
      }
    }
  }

  /**
   * A concrete implementation of {@link DoFnInvoker.ArgumentProvider} used for running a {@link
   * DoFn} on a timer.
   *
   * @param <InputT> the type of the {@link DoFn} (main) input elements
   * @param <OutputT> the type of the {@link DoFn} (main) output elements
   */
  private class OnTimerArgumentProvider<InputT, OutputT>
      extends DoFn<InputT, OutputT>.OnTimerContext
      implements DoFnInvoker.ArgumentProvider<InputT, OutputT> {

    final DoFn<InputT, OutputT> fn;
    final DoFnContext<InputT, OutputT> context;
    private final BoundedWindow window;
    private final Instant timestamp;
    private final TimeDomain timeDomain;
    private final Duration allowedLateness;

    /** Lazily initialized; should only be accessed via {@link #getNamespace()}. */
    private StateNamespace namespace;

    /**
     * The state namespace for this context.
     *
     * <p>Any call to {@link #getNamespace()} when more than one window is present will crash; this
     * represents a bug in the runner or the {@link DoFnSignature}, since values must be in exactly
     * one window when state or timers are relevant.
     */
    private StateNamespace getNamespace() {
      if (namespace == null) {
        namespace = StateNamespaces.window(windowCoder, window);
      }
      return namespace;
    }

    private OnTimerArgumentProvider(
        DoFn<InputT, OutputT> fn,
        DoFnContext<InputT, OutputT> context,
        BoundedWindow window,
        Duration allowedLateness,
        Instant timestamp,
        TimeDomain timeDomain) {
      fn.super();
      this.fn = fn;
      this.context = context;
      this.window = window;
      this.allowedLateness = allowedLateness;
      this.timestamp = timestamp;
      this.timeDomain = timeDomain;
    }

    @Override
    public Instant timestamp() {
      return timestamp;
    }

    @Override
    public BoundedWindow window() {
      return window;
    }

    @Override
    public StartBundleContext startBundleContext(DoFn<InputT, OutputT> doFn) {
      throw new UnsupportedOperationException("StartBundleContext parameters are not supported.");
    }

    @Override
    public FinishBundleContext finishBundleContext(DoFn<InputT, OutputT> doFn) {
      throw new UnsupportedOperationException("FinishBundleContext parameters are not supported.");
    }

    @Override
    public TimeDomain timeDomain() {
      return timeDomain;
    }


    @Override
    public ProcessContext processContext(DoFn<InputT, OutputT> doFn) {
      throw new UnsupportedOperationException("ProcessContext parameters are not supported.");
    }

    @Override
    public OnTimerContext onTimerContext(DoFn<InputT, OutputT> doFn) {
      return this;
    }

    @Override
    public RestrictionTracker<?> restrictionTracker() {
      throw new UnsupportedOperationException("RestrictionTracker parameters are not supported.");
    }

    @Override
    public State state(String stateId) {
      try {
        StateSpec<?> spec =
            (StateSpec<?>) signature.stateDeclarations().get(stateId).field().get(fn);
        return stepContext
            .stateInternals()
            .state(getNamespace(), StateTags.tagForSpec(stateId, (StateSpec) spec));
      } catch (IllegalAccessException e) {
        throw new RuntimeException(e);
      }
    }

    @Override
    public Timer timer(String timerId) {
      try {
        TimerSpec spec =
            (TimerSpec) signature.timerDeclarations().get(timerId).field().get(fn);
        return new TimerInternalsTimer(
            window, getNamespace(), allowedLateness, timerId, spec, stepContext.timerInternals());
      } catch (IllegalAccessException e) {
        throw new RuntimeException(e);
      }
    }

    @Override
    public PipelineOptions getPipelineOptions() {
      return context.getPipelineOptions();
    }

    @Override
    public void output(OutputT output) {
      context.outputWindowedValue(
          output, timestamp(), Collections.singleton(window()), PaneInfo.NO_FIRING);
    }

    @Override
    public void outputWithTimestamp(OutputT output, Instant timestamp) {
      context.outputWindowedValue(
          output, timestamp, Collections.singleton(window()), PaneInfo.NO_FIRING);
    }

    @Override
    public <T> void output(TupleTag<T> tag, T output) {
      context.outputWindowedValue(
          tag, output, timestamp, Collections.singleton(window()), PaneInfo.NO_FIRING);
    }

    @Override
    public <T> void outputWithTimestamp(TupleTag<T> tag, T output, Instant timestamp) {
      context.outputWindowedValue(
          tag, output, timestamp, Collections.singleton(window()), PaneInfo.NO_FIRING);
    }
  }

  private static class TimerInternalsTimer implements Timer {
    private final TimerInternals timerInternals;

    // The window and the namespace represent the same thing, but the namespace is a cached
    // and specially encoded form. Since the namespace can be cached across timers, it is
    // passed in whole rather than being computed here.
    private final BoundedWindow window;
    private final Duration allowedLateness;
    private final StateNamespace namespace;
    private final String timerId;
    private final TimerSpec spec;
    private Duration period = Duration.ZERO;
    private Duration offset = Duration.ZERO;

    public TimerInternalsTimer(
        BoundedWindow window,
        StateNamespace namespace,
        Duration allowedLateness,
        String timerId,
        TimerSpec spec,
        TimerInternals timerInternals) {
      this.window = window;
      this.allowedLateness = allowedLateness;
      this.namespace = namespace;
      this.timerId = timerId;
      this.spec = spec;
      this.timerInternals = timerInternals;
    }

    @Override
    public void set(Instant target) {
      verifyAbsoluteTimeDomain();
      verifyTargetTime(target);
      setUnderlyingTimer(target);
    }

    @Override
    public void setRelative() {
      Instant target;
      Instant now = getCurrentTime();
      if (period.equals(Duration.ZERO)) {
        target = now.plus(offset);
      } else {
        long millisSinceStart = now.plus(offset).getMillis() % period.getMillis();
        target = millisSinceStart == 0 ? now : now.plus(period).minus(millisSinceStart);
      }
      target = minTargetAndGcTime(target);
      setUnderlyingTimer(target);
    }

    @Override
    public Timer offset(Duration offset) {
      this.offset = offset;
      return this;
    }

    @Override
    public Timer align(Duration period) {
      this.period = period;
      return this;
    }

    /**
     * For event time timers the target time should be prior to window GC time. So it return
     * min(time to set, GC Time of window).
     */
    private Instant minTargetAndGcTime(Instant target) {
      if (TimeDomain.EVENT_TIME.equals(spec.getTimeDomain())) {
        Instant windowExpiry = LateDataUtils.garbageCollectionTime(window, allowedLateness);
        if (target.isAfter(windowExpiry)) {
          return windowExpiry;
        }
      }
      return target;
    }

    /**
     * Ensures that the target time is reasonable. For event time timers this means that the
     * time should be prior to window GC time.
     */
    private void verifyTargetTime(Instant target) {
      if (TimeDomain.EVENT_TIME.equals(spec.getTimeDomain())) {
        Instant windowExpiry = window.maxTimestamp().plus(allowedLateness);
        checkArgument(!target.isAfter(windowExpiry),
            "Attempted to set event time timer for %s but that is after"
            + " the expiration of window %s", target, windowExpiry);
      }
    }

    /** Verifies that the time domain of this timer is acceptable for absolute timers. */
    private void verifyAbsoluteTimeDomain() {
      if (!TimeDomain.EVENT_TIME.equals(spec.getTimeDomain())) {
        throw new IllegalStateException(
            "Cannot only set relative timers in processing time domain."
                + " Use #setRelative()");
      }
    }

    /**
     * Sets the timer for the target time without checking anything about whether it is
     * a reasonable thing to do. For example, absolute processing time timers are not
     * really sensible since the user has no way to compute a good choice of time.
     */
    private void setUnderlyingTimer(Instant target) {
      timerInternals.setTimer(namespace, timerId, target, spec.getTimeDomain());
    }

    private Instant getCurrentTime() {
      switch(spec.getTimeDomain()) {
        case EVENT_TIME:
          return timerInternals.currentInputWatermarkTime();
        case PROCESSING_TIME:
          return timerInternals.currentProcessingTime();
        case SYNCHRONIZED_PROCESSING_TIME:
          return timerInternals.currentSynchronizedProcessingTime();
        default:
          throw new IllegalStateException(
              String.format("Timer created for unknown time domain %s", spec.getTimeDomain()));
      }
    }

  }
}