/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.beam.runners.dataflow.worker;

import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;

import com.google.api.services.dataflow.model.CounterUpdate;
import com.google.api.services.dataflow.model.SideInputInfo;
import java.io.Closeable;
import java.io.IOException;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.AtomicLong;
import javax.annotation.Nullable;
import org.apache.beam.runners.core.SideInputReader;
import org.apache.beam.runners.core.StateInternals;
import org.apache.beam.runners.core.StateNamespaces;
import org.apache.beam.runners.core.TimerInternals;
import org.apache.beam.runners.core.TimerInternals.TimerData;
import org.apache.beam.runners.core.metrics.ExecutionStateTracker;
import org.apache.beam.runners.core.metrics.ExecutionStateTracker.ExecutionState;
import org.apache.beam.runners.dataflow.worker.StreamingModeExecutionContext.StepContext;
import org.apache.beam.runners.dataflow.worker.counters.CounterFactory;
import org.apache.beam.runners.dataflow.worker.counters.NameContext;
import org.apache.beam.runners.dataflow.worker.profiler.ScopedProfiler.ProfileScope;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalDataId;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalDataRequest;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.Timer;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.io.UnboundedSource;
import org.apache.beam.sdk.metrics.MetricsContainer;
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Optional;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.FluentIterable;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** {@link DataflowExecutionContext} for use in streaming mode. */
public class StreamingModeExecutionContext extends DataflowExecutionContext<StepContext> {

  private static final Logger LOG = LoggerFactory.getLogger(StreamingModeExecutionContext.class);

  /**
   * The current user-facing key for this execution context.
   *
   * <p>This field is set to a new key upon each call to {@link #start}, which corresponds to one
   * input bundle. For non-keyed bundles, the key will be {@code null}.
   *
   * <p>This key should not be mistaken for the sharding key of the computation, which is always
   * present.
   */
  @Nullable private Object key = null;

  private final String computationId;
  private final Map<TupleTag<?>, Map<BoundedWindow, Object>> sideInputCache;

  // Per-key cache of active Reader objects in use by this process.
  private final ImmutableMap<String, String> stateNameMap;
  private final WindmillStateCache.ForComputation stateCache;

  private Windmill.WorkItem work;
  private StateFetcher stateFetcher;
  private Windmill.WorkItemCommitRequest.Builder outputBuilder;
  private UnboundedSource.UnboundedReader<?> activeReader;
  private volatile long backlogBytes;
  private final ReaderCache readerCache;

  public StreamingModeExecutionContext(
      CounterFactory counterFactory,
      String computationId,
      ReaderCache readerCache,
      Map<String, String> stateNameMap,
      WindmillStateCache.ForComputation stateCache,
      MetricsContainerRegistry<StreamingStepMetricsContainer> metricsContainerRegistry,
      DataflowExecutionStateTracker executionStateTracker,
      StreamingModeExecutionStateRegistry executionStateRegistry,
      long sinkByteLimit) {
    super(
        counterFactory,
        metricsContainerRegistry,
        executionStateTracker,
        executionStateRegistry,
        sinkByteLimit);
    this.computationId = computationId;
    this.readerCache = readerCache;
    this.sideInputCache = new HashMap<>();
    this.stateNameMap = ImmutableMap.copyOf(stateNameMap);
    this.stateCache = stateCache;
    this.backlogBytes = UnboundedSource.UnboundedReader.BACKLOG_UNKNOWN;
  }

  /**
   * Execution states in Streaming are shared between multiple map-task executors. Thus this class
   * needs to be thread safe for multiple writers. A single stage could have have multiple executors
   * running concurrently.
   */
  public static class StreamingModeExecutionState
      extends DataflowOperationContext.DataflowExecutionState {

    // AtomicLong is used because this value is written in two places:
    // 1. The sampling thread calls takeSample to increment the time spent in this state
    // 2. The reporting thread calls extractUpdate which reads the current sum *AND* sets it to 0.
    private final AtomicLong totalMillisInState = new AtomicLong();

    // The worker that created this state.  Used to report lulls back to the worker.
    private final StreamingDataflowWorker worker;

    public StreamingModeExecutionState(
        NameContext nameContext,
        String stateName,
        MetricsContainer metricsContainer,
        ProfileScope profileScope,
        StreamingDataflowWorker worker) {
      // TODO: Take in the requesting step name and side input index for streaming.
      super(nameContext, stateName, null, null, metricsContainer, profileScope);
      this.worker = worker;
    }

    /*
     * Report the lull to the StreamingDataflowWorker that is stuck in addition to logging the
     * lull.
     */
    @Override
    public void reportLull(Thread trackedThread, long millis) {
      super.reportLull(trackedThread, millis);
      // Also report the failure to the list of pending failures to report on the worker thread
      // so that the failure gets communicated to the StreamingDataflowWorker.
      String errorMessage = getLullMessage(trackedThread, Duration.millis(millis));
      worker.addFailure(errorMessage);
    }

    /**
     * Take sample is only called by the ExecutionStateSampler thread. It is the only place that
     * increments totalMillisInState, however the reporting thread periodically calls extractUpdate
     * which will read the sum and reset it to 0, so totalMillisInState does have multiple writers.
     */
    @Override
    public void takeSample(long millisSinceLastSample) {
      totalMillisInState.addAndGet(millisSinceLastSample);
    }

    /**
     * Extract updates in the form of a {@link CounterUpdate}.
     *
     * <p>Non-final updates are extracted periodically and report the physical value as a delta.
     * This requires setting the totalMillisInState back to 0.
     *
     * <p>Final updates should never be requested from a Streaming job since the work unit never
     * completes.
     */
    @Nullable
    @Override
    public CounterUpdate extractUpdate(boolean isFinalUpdate) {
      // Streaming reports deltas, so isFinalUpdate doesn't matter, and should never be true.
      long sum = totalMillisInState.getAndSet(0);
      return sum == 0 ? null : createUpdate(false, sum);
    }
  }

  /**
   * Implementation of DataflowExecutionStateRegistry that creates Streaming versions of
   * ExecutionState.
   */
  public static class StreamingModeExecutionStateRegistry extends DataflowExecutionStateRegistry {
    private final StreamingDataflowWorker worker;

    public StreamingModeExecutionStateRegistry(StreamingDataflowWorker worker) {
      this.worker = worker;
    }

    @Override
    protected DataflowOperationContext.DataflowExecutionState createState(
        NameContext nameContext,
        String stateName,
        String requestingStepName,
        Integer inputIndex,
        MetricsContainer container,
        ProfileScope profileScope) {
      return new StreamingModeExecutionState(
          nameContext, stateName, container, profileScope, worker);
    }
  }

  @VisibleForTesting
  public long getBacklogBytes() {
    return backlogBytes;
  }

  public void start(
      @Nullable Object key,
      Windmill.WorkItem work,
      Instant inputDataWatermark,
      @Nullable Instant outputDataWatermark,
      @Nullable Instant synchronizedProcessingTime,
      WindmillStateReader stateReader,
      StateFetcher stateFetcher,
      Windmill.WorkItemCommitRequest.Builder outputBuilder) {
    this.key = key;
    this.work = work;
    this.stateFetcher = stateFetcher;
    this.outputBuilder = outputBuilder;
    this.sideInputCache.clear();
    clearSinkFullHint();

    Instant processingTime = Instant.now();
    // Ensure that the processing time is greater than any fired processing time
    // timers.  Otherwise a trigger could ignore the timer and orphan the window.
    for (Windmill.Timer timer : work.getTimers().getTimersList()) {
      if (timer.getType() == Windmill.Timer.Type.REALTIME) {
        Instant inferredFiringTime =
            WindmillTimeUtils.windmillToHarnessTimestamp(timer.getTimestamp())
                .plus(Duration.millis(1));
        if (inferredFiringTime.isAfter(processingTime)) {
          processingTime = inferredFiringTime;
        }
      }
    }

    for (StepContext stepContext : getAllStepContexts()) {
      stepContext.start(
          stateReader,
          inputDataWatermark,
          processingTime,
          outputDataWatermark,
          synchronizedProcessingTime);
    }
  }

  @Override
  public StepContext createStepContext(DataflowOperationContext operationContext) {
    return new StepContext(operationContext);
  }

  @Override
  protected SideInputReader getSideInputReader(
      Iterable<? extends SideInputInfo> sideInputInfos, DataflowOperationContext operationContext) {
    throw new UnsupportedOperationException(
        "Cannot call getSideInputReader for StreamingDataflowWorker: "
            + "the MapTask specification should not have had any SideInputInfo descriptors "
            + "since the streaming runner does not yet support them.");
  }

  @Override
  protected SideInputReader getSideInputReaderForViews(
      Iterable<? extends PCollectionView<?>> views) {
    return StreamingModeSideInputReader.of(views, this);
  }

  /**
   * Fetches the requested sideInput, and maintains a view of the cache that doesn't remove items
   * until the active work item is finished.
   *
   * <p>If the side input was not ready, throws {@code IllegalStateException} if the state is
   * {@literal CACHED_IN_WORKITEM} or returns null otherwise.
   *
   * <p>If the side input was ready and null, returns {@literal Optional.absent()}. If the side
   * input was ready and non-null returns {@literal Optional.present(...)}.
   */
  @Nullable
  private <T> Optional<T> fetchSideInput(
      PCollectionView<T> view,
      BoundedWindow sideInputWindow,
      String stateFamily,
      StateFetcher.SideInputState state,
      Supplier<Closeable> scopedReadStateSupplier) {
    Map<BoundedWindow, Object> tagCache = sideInputCache.get(view.getTagInternal());
    if (tagCache == null) {
      tagCache = new HashMap<>();
      sideInputCache.put(view.getTagInternal(), tagCache);
    }

    if (tagCache.containsKey(sideInputWindow)) {
      @SuppressWarnings("unchecked")
      T typed = (T) tagCache.get(sideInputWindow);
      return Optional.fromNullable(typed);
    } else {
      if (state == StateFetcher.SideInputState.CACHED_IN_WORKITEM) {
        throw new IllegalStateException(
            "Expected side input to be cached. Tag: " + view.getTagInternal().getId());
      }
      Optional<T> fetched =
          stateFetcher.fetchSideInput(
              view, sideInputWindow, stateFamily, state, scopedReadStateSupplier);
      if (fetched != null) {
        tagCache.put(sideInputWindow, fetched.orNull());
      }
      return fetched;
    }
  }

  public Iterable<Windmill.GlobalDataId> getSideInputNotifications() {
    return work.getGlobalDataIdNotificationsList();
  }

  private List<Timer> getFiredTimers() {
    return work.getTimers().getTimersList();
  }

  @Nullable
  public ByteString getSerializedKey() {
    return work == null ? null : work.getKey();
  }

  public long getWorkToken() {
    return work.getWorkToken();
  }

  public Windmill.WorkItem getWork() {
    return work;
  }

  public Windmill.WorkItemCommitRequest.Builder getOutputBuilder() {
    return outputBuilder;
  }

  /**
   * Returns cached reader for this key if one exists. The reader is removed from the cache. NOTE:
   * The caller is responsible for the reader and should appropriately close it as required.
   */
  public UnboundedSource.UnboundedReader<?> getCachedReader() {
    return readerCache.acquireReader(computationId, getSerializedKey(), getWork().getCacheToken());
  }

  public void setActiveReader(UnboundedSource.UnboundedReader<?> reader) {
    checkState(activeReader == null, "not expected to be overwritten");
    activeReader = reader;
  }

  /** Invalidate the state and reader caches for this computation and key. */
  public void invalidateCache() {
    ByteString key = getSerializedKey();
    if (key != null) {
      readerCache.invalidateReader(computationId, key);
      if (activeReader != null) {
        try {
          activeReader.close();
        } catch (IOException e) {
          LOG.warn("Failed to close reader for {}-{}", computationId, key.toStringUtf8(), e);
        }
      }
      activeReader = null;
      stateCache.invalidate(key);
    }
  }

  public UnboundedSource.CheckpointMark getReaderCheckpoint(
      Coder<? extends UnboundedSource.CheckpointMark> coder) {
    try {
      ByteString state = work.getSourceState().getState();
      if (state.isEmpty()) {
        return null;
      }
      return coder.decode(state.newInput(), Coder.Context.OUTER);
    } catch (IOException e) {
      throw new RuntimeException("Exception while decoding checkpoint", e);
    }
  }

  public Map<Long, Runnable> flushState() {
    Map<Long, Runnable> callbacks = new HashMap<>();

    for (StepContext stepContext : getAllStepContexts()) {
      stepContext.flushState();
    }

    if (activeReader != null) {
      Windmill.SourceState.Builder sourceStateBuilder =
          outputBuilder.getSourceStateUpdatesBuilder();
      final UnboundedSource.CheckpointMark checkpointMark = activeReader.getCheckpointMark();
      final Instant watermark = activeReader.getWatermark();
      long id = ThreadLocalRandom.current().nextLong();
      sourceStateBuilder.addFinalizeIds(id);
      callbacks.put(
          id,
          () -> {
            try {
              checkpointMark.finalizeCheckpoint();
            } catch (IOException e) {
              throw new RuntimeException("Exception while finalizing checkpoint", e);
            }
          });

      @SuppressWarnings("unchecked")
      Coder<UnboundedSource.CheckpointMark> checkpointCoder =
          ((UnboundedSource<?, UnboundedSource.CheckpointMark>) activeReader.getCurrentSource())
              .getCheckpointMarkCoder();
      if (checkpointCoder != null) {
        ByteString.Output stream = ByteString.newOutput();
        try {
          checkpointCoder.encode(checkpointMark, stream, Coder.Context.OUTER);
        } catch (IOException e) {
          throw new RuntimeException("Exception while encoding checkpoint", e);
        }
        sourceStateBuilder.setState(stream.toByteString());
      }
      outputBuilder.setSourceWatermark(WindmillTimeUtils.harnessToWindmillTimestamp(watermark));

      backlogBytes = activeReader.getSplitBacklogBytes();
      if (backlogBytes == UnboundedSource.UnboundedReader.BACKLOG_UNKNOWN
          && WorkerCustomSources.isFirstUnboundedSourceSplit(getSerializedKey())) {
        // Only call getTotalBacklogBytes() on the first split.
        backlogBytes = activeReader.getTotalBacklogBytes();
      }
      outputBuilder.setSourceBacklogBytes(backlogBytes);

      readerCache.cacheReader(
          computationId, getSerializedKey(), getWork().getCacheToken(), activeReader);
      activeReader = null;
    }
    return callbacks;
  }

  interface StreamingModeStepContext {
    boolean issueSideInputFetch(
        PCollectionView<?> view, BoundedWindow w, StateFetcher.SideInputState s);

    void addBlockingSideInput(Windmill.GlobalDataRequest blocked);

    void addBlockingSideInputs(Iterable<Windmill.GlobalDataRequest> blocked);

    StateInternals stateInternals();

    Iterable<Windmill.GlobalDataId> getSideInputNotifications();

    /** Writes the given {@code PCollectionView} data to a globally accessible location. */
    <T, W extends BoundedWindow> void writePCollectionViewData(
        TupleTag<?> tag,
        Iterable<T> data,
        Coder<Iterable<T>> dataCoder,
        W window,
        Coder<W> windowCoder)
        throws IOException;
  }

  String getStateFamily(NameContext nameContext) {
    return nameContext.userName() == null ? null : stateNameMap.get(nameContext.userName());
  }

  class StepContext extends DataflowExecutionContext.DataflowStepContext
      implements StreamingModeStepContext {

    private WindmillStateInternals<Object> stateInternals;

    private WindmillTimerInternals systemTimerInternals;
    private WindmillTimerInternals userTimerInternals;
    private final String stateFamily;
    private final Supplier<Closeable> scopedReadStateSupplier;

    public StepContext(DataflowOperationContext operationContext) {
      super(operationContext.nameContext());
      this.stateFamily =
          StreamingModeExecutionContext.this.getStateFamily(operationContext.nameContext());

      this.scopedReadStateSupplier =
          new Supplier<Closeable>() {
            private ExecutionState readState = operationContext.newExecutionState("windmill-read");

            @Override
            public Closeable get() {
              ExecutionStateTracker tracker =
                  StreamingModeExecutionContext.this.getExecutionStateTracker();
              if (tracker == null) {
                return null;
              }
              return tracker.enterState(readState);
            }
          };
    }

    /** Update the {@code stateReader} used by this {@code StepContext}. */
    public void start(
        WindmillStateReader stateReader,
        Instant inputDataWatermark,
        Instant processingTime,
        @Nullable Instant outputDataWatermark,
        @Nullable Instant synchronizedProcessingTime) {
      this.stateInternals =
          new WindmillStateInternals<>(
              key,
              stateFamily,
              stateReader,
              work.getIsNewKey(),
              stateCache.forKey(
                  getSerializedKey(), stateFamily, getWork().getCacheToken(), getWorkToken()),
              scopedReadStateSupplier);

      this.systemTimerInternals =
          new WindmillTimerInternals(
              stateFamily,
              WindmillNamespacePrefix.SYSTEM_NAMESPACE_PREFIX,
              inputDataWatermark,
              processingTime,
              outputDataWatermark,
              synchronizedProcessingTime);

      this.userTimerInternals =
          new WindmillTimerInternals(
              stateFamily,
              WindmillNamespacePrefix.USER_NAMESPACE_PREFIX,
              inputDataWatermark,
              processingTime,
              outputDataWatermark,
              synchronizedProcessingTime);

      this.cachedFiredTimers = null;
      this.cachedFiredUserTimers = null;
    }

    public void flushState() {
      stateInternals.persist(outputBuilder);
      systemTimerInternals.persistTo(outputBuilder);
      userTimerInternals.persistTo(outputBuilder);
    }

    // Lazily initialized
    private Iterator<TimerData> cachedFiredTimers = null;

    @Override
    public <W extends BoundedWindow> TimerData getNextFiredTimer(Coder<W> windowCoder) {
      if (cachedFiredTimers == null) {
        cachedFiredTimers =
            FluentIterable.<Timer>from(StreamingModeExecutionContext.this.getFiredTimers())
                .filter(
                    timer ->
                        WindmillTimerInternals.isSystemTimer(timer)
                            && timer.getStateFamily().equals(stateFamily))
                .transform(
                    timer ->
                        WindmillTimerInternals.windmillTimerToTimerData(
                            WindmillNamespacePrefix.SYSTEM_NAMESPACE_PREFIX, timer, windowCoder))
                .iterator();
      }

      if (!cachedFiredTimers.hasNext()) {
        return null;
      }
      return cachedFiredTimers.next();
    }

    // Lazily initialized
    private Iterator<TimerData> cachedFiredUserTimers = null;

    public <W extends BoundedWindow> TimerData getNextFiredUserTimer(Coder<W> windowCoder) {
      if (cachedFiredUserTimers == null) {
        cachedFiredUserTimers =
            FluentIterable.<Timer>from(StreamingModeExecutionContext.this.getFiredTimers())
                .filter(
                    timer ->
                        WindmillTimerInternals.isUserTimer(timer)
                            && timer.getStateFamily().equals(stateFamily))
                .transform(
                    timer ->
                        WindmillTimerInternals.windmillTimerToTimerData(
                            WindmillNamespacePrefix.USER_NAMESPACE_PREFIX, timer, windowCoder))
                .iterator();
      }

      if (!cachedFiredUserTimers.hasNext()) {
        return null;
      }
      TimerData nextTimer = cachedFiredUserTimers.next();
      // User timers must be explicitly deleted when delivered, to release the implied hold
      userTimerInternals.deleteTimer(nextTimer);
      return nextTimer;
    }

    @Override
    public <W extends BoundedWindow> void setStateCleanupTimer(
        String timerId, W window, Coder<W> windowCoder, Instant cleanupTime) {
      timerInternals()
          .setTimer(
              StateNamespaces.window(windowCoder, window),
              timerId,
              cleanupTime,
              TimeDomain.EVENT_TIME);
    }

    @Override
    public DataflowStepContext namespacedToUser() {
      return new UserStepContext(this);
    }

    @Override
    public Iterable<Windmill.GlobalDataId> getSideInputNotifications() {
      return StreamingModeExecutionContext.this.getSideInputNotifications();
    }

    @Override
    public <T, W extends BoundedWindow> void writePCollectionViewData(
        TupleTag<?> tag,
        Iterable<T> data,
        Coder<Iterable<T>> dataCoder,
        W window,
        Coder<W> windowCoder)
        throws IOException {
      if (getSerializedKey().size() != 0) {
        throw new IllegalStateException("writePCollectionViewData must follow a Combine.globally");
      }

      ByteString.Output dataStream = ByteString.newOutput();
      dataCoder.encode(data, dataStream, Coder.Context.OUTER);

      ByteString.Output windowStream = ByteString.newOutput();
      windowCoder.encode(window, windowStream, Coder.Context.OUTER);

      if (stateFamily == null) {
        throw new IllegalStateException(
            "Tried to write view data for stateless step: " + getNameContext());
      }

      Windmill.GlobalData.Builder builder =
          Windmill.GlobalData.newBuilder()
              .setDataId(
                  Windmill.GlobalDataId.newBuilder()
                      .setTag(tag.getId())
                      .setVersion(windowStream.toByteString())
                      .build())
              .setData(dataStream.toByteString())
              .setStateFamily(stateFamily);

      outputBuilder.addGlobalDataUpdates(builder.build());
    }

    /** Fetch the given side input asynchronously and return true if it is present. */
    @Override
    public boolean issueSideInputFetch(
        PCollectionView<?> view, BoundedWindow mainInputWindow, StateFetcher.SideInputState state) {
      BoundedWindow sideInputWindow = view.getWindowMappingFn().getSideInputWindow(mainInputWindow);
      return fetchSideInput(view, sideInputWindow, stateFamily, state, scopedReadStateSupplier)
          != null;
    }

    /** Note that there is data on the current key that is blocked on the given side input. */
    @Override
    public void addBlockingSideInput(Windmill.GlobalDataRequest sideInput) {
      checkState(
          stateFamily != null,
          "Tried to set global data request for stateless step: " + getNameContext());
      sideInput =
          Windmill.GlobalDataRequest.newBuilder(sideInput).setStateFamily(stateFamily).build();
      outputBuilder.addGlobalDataRequests(sideInput);
      outputBuilder.addGlobalDataIdRequests(sideInput.getDataId());
    }

    /** Note that there is data on the current key that is blocked on the given side inputs. */
    @Override
    public void addBlockingSideInputs(Iterable<Windmill.GlobalDataRequest> sideInputs) {
      for (Windmill.GlobalDataRequest sideInput : sideInputs) {
        addBlockingSideInput(sideInput);
      }
    }

    @Override
    public StateInternals stateInternals() {
      checkState(
          stateFamily != null, "Tried to access state for stateless step: " + getNameContext());
      return checkNotNull(stateInternals);
    }

    @Override
    public TimerInternals timerInternals() {
      checkState(
          stateFamily != null, "Tried to access timers for stateless step: " + getNameContext());
      return checkNotNull(systemTimerInternals);
    }

    public TimerInternals userTimerInternals() {
      checkState(
          stateFamily != null,
          "Tried to access user timers for stateless step: " + getNameContext());
      return checkNotNull(userTimerInternals);
    }
  }

  /**
   * A specialized {@link StepContext} that uses provided {@link StateInternals} and {@link
   * TimerInternals} for user state and timers.
   */
  private static class UserStepContext extends DataflowStepContext
      implements StreamingModeStepContext {

    private final StreamingModeExecutionContext.StepContext wrapped;

    public UserStepContext(StreamingModeExecutionContext.StepContext wrapped) {
      super(wrapped.getNameContext());
      this.wrapped = wrapped;
    }

    @Override
    public boolean issueSideInputFetch(
        PCollectionView<?> view, BoundedWindow w, StateFetcher.SideInputState s) {
      return wrapped.issueSideInputFetch(view, w, s);
    }

    @Override
    public void addBlockingSideInput(GlobalDataRequest blocked) {
      wrapped.addBlockingSideInput(blocked);
    }

    @Override
    public void addBlockingSideInputs(Iterable<GlobalDataRequest> blocked) {
      wrapped.addBlockingSideInputs(blocked);
    }

    @Override
    public StateInternals stateInternals() {
      return wrapped.stateInternals();
    }

    @Override
    public Iterable<GlobalDataId> getSideInputNotifications() {
      return wrapped.getSideInputNotifications();
    }

    @Override
    public <T, W extends BoundedWindow> void writePCollectionViewData(
        TupleTag<?> tag,
        Iterable<T> data,
        Coder<Iterable<T>> dataCoder,
        W window,
        Coder<W> windowCoder)
        throws IOException {
      throw new IllegalStateException("User DoFns cannot write PCollectionView data");
    }

    @Override
    public TimerInternals timerInternals() {
      return wrapped.userTimerInternals();
    }

    @Override
    public <W extends BoundedWindow> TimerData getNextFiredTimer(Coder<W> windowCoder) {
      return wrapped.getNextFiredUserTimer(windowCoder);
    }

    @Override
    public <W extends BoundedWindow> void setStateCleanupTimer(
        String timerId, W window, Coder<W> windowCoder, Instant cleanupTime) {
      throw new UnsupportedOperationException(
          String.format(
              "setStateCleanupTimer should not be called on %s, only on a system %s",
              getClass().getSimpleName(),
              StreamingModeExecutionContext.StepContext.class.getSimpleName()));
    }

    @Override
    public DataflowStepContext namespacedToUser() {
      return this;
    }
  }

  /** A {@link SideInputReader} that fetches side inputs from the streaming worker's cache. */
  public static class StreamingModeSideInputReader implements SideInputReader {
    private StreamingModeExecutionContext context;
    private Set<PCollectionView<?>> viewSet;

    private StreamingModeSideInputReader(
        Iterable<? extends PCollectionView<?>> views, StreamingModeExecutionContext context) {
      this.context = context;
      this.viewSet = ImmutableSet.copyOf(views);
    }

    public static StreamingModeSideInputReader of(
        Iterable<? extends PCollectionView<?>> views, StreamingModeExecutionContext context) {
      return new StreamingModeSideInputReader(views, context);
    }

    @Override
    public <T> T get(PCollectionView<T> view, BoundedWindow window) {
      if (!contains(view)) {
        throw new RuntimeException("get() called with unknown view");
      }

      // We are only fetching the cached value here, so we don't need stateFamily or
      // readStateSupplier.
      return context
          .fetchSideInput(
              view,
              window,
              null /* unused stateFamily */,
              StateFetcher.SideInputState.CACHED_IN_WORKITEM,
              null /* unused readStateSupplier */)
          .orNull();
    }

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

    @Override
    public boolean isEmpty() {
      return viewSet.isEmpty();
    }
  }
}