@Override
  public ParDoFn create(
      PipelineOptions options,
      CloudObject cloudUserFn,
      @Nullable List<SideInputInfo> sideInputInfos,
      TupleTag<?> mainOutputTag,
      Map<TupleTag<?>, Integer> outputTupleTagsToReceiverIndices,
      DataflowExecutionContext<?> executionContext,
      DataflowOperationContext operationContext)
      throws Exception {

    DoFnInstanceManager instanceManager =
        fnCache.get(
            operationContext.nameContext().systemName(),
            () ->
                DoFnInstanceManagers.cloningPool(doFnExtractor.getDoFnInfo(cloudUserFn), options));

    DoFnInfo<?, ?> doFnInfo = instanceManager.peek();

    DataflowExecutionContext.DataflowStepContext stepContext =
        executionContext.getStepContext(operationContext);

    Iterable<PCollectionView<?>> sideInputViews = doFnInfo.getSideInputViews();
    SideInputReader sideInputReader =
        executionContext.getSideInputReader(sideInputInfos, sideInputViews, operationContext);

    if (doFnInfo.getDoFn() instanceof BatchStatefulParDoOverrides.BatchStatefulDoFn) {
      // HACK: BatchStatefulDoFn is a class from DataflowRunner's overrides
      // that just instructs the worker to execute it differently. This will
      // be replaced by metadata in the Runner API payload
      BatchStatefulParDoOverrides.BatchStatefulDoFn fn =
          (BatchStatefulParDoOverrides.BatchStatefulDoFn) doFnInfo.getDoFn();
      DoFn underlyingFn = fn.getUnderlyingDoFn();

      return new BatchModeUngroupingParDoFn(
          (BatchModeExecutionContext.StepContext) stepContext,
          new SimpleParDoFn(
              options,
              DoFnInstanceManagers.singleInstance(doFnInfo.withFn(underlyingFn)),
              sideInputReader,
              doFnInfo.getMainOutput(),
              outputTupleTagsToReceiverIndices,
              stepContext,
              operationContext,
              doFnInfo.getDoFnSchemaInformation(),
              doFnInfo.getSideInputMapping(),
              runnerFactory));

    } else if (doFnInfo.getDoFn() instanceof StreamingPCollectionViewWriterFn) {
      // HACK: StreamingPCollectionViewWriterFn is a class from
      // DataflowPipelineTranslator. Using the class as an indicator is a migration path
      // to simply having an indicator string.

      checkArgument(
          stepContext instanceof StreamingModeExecutionContext.StreamingModeStepContext,
          "stepContext must be a StreamingModeStepContext to use StreamingPCollectionViewWriterFn");
      DataflowRunner.StreamingPCollectionViewWriterFn<Object> writerFn =
          (StreamingPCollectionViewWriterFn<Object>) doFnInfo.getDoFn();
      return new StreamingPCollectionViewWriterParDoFn(
          (StreamingModeExecutionContext.StreamingModeStepContext) stepContext,
          writerFn.getView().getTagInternal(),
          writerFn.getDataCoder(),
          (Coder<BoundedWindow>) doFnInfo.getWindowingStrategy().getWindowFn().windowCoder());
    } else {
      return new SimpleParDoFn(
          options,
          instanceManager,
          sideInputReader,
          doFnInfo.getMainOutput(),
          outputTupleTagsToReceiverIndices,
          stepContext,
          operationContext,
          doFnInfo.getDoFnSchemaInformation(),
          doFnInfo.getSideInputMapping(),
          runnerFactory);
    }
  }
@Test
  public void testCleanupRegistered() throws Exception {
    PipelineOptions options = PipelineOptionsFactory.create();
    CounterSet counters = new CounterSet();
    DoFn<?, ?> initialFn = new TestStatefulDoFn();
    CloudObject cloudObject =
        getCloudObject(
            initialFn,
            WindowingStrategy.globalDefault().withWindowFn(FixedWindows.of(Duration.millis(10))));

    TimerInternals timerInternals = mock(TimerInternals.class);

    DataflowStepContext stepContext = mock(DataflowStepContext.class);
    when(stepContext.timerInternals()).thenReturn(timerInternals);

    DataflowExecutionContext<DataflowStepContext> executionContext =
        mock(DataflowExecutionContext.class);
    TestOperationContext operationContext = TestOperationContext.create(counters);
    when(executionContext.getStepContext(operationContext)).thenReturn(stepContext);
    when(executionContext.getSideInputReader(any(), any(), any()))
        .thenReturn(NullSideInputReader.empty());

    ParDoFn parDoFn =
        factory.create(
            options,
            cloudObject,
            Collections.emptyList(),
            MAIN_OUTPUT,
            ImmutableMap.of(MAIN_OUTPUT, 0),
            executionContext,
            operationContext);

    Receiver rcvr = new OutputReceiver();
    parDoFn.startBundle(rcvr);

    IntervalWindow firstWindow = new IntervalWindow(new Instant(0), new Instant(10));
    parDoFn.processElement(
        WindowedValue.of("foo", new Instant(1), firstWindow, PaneInfo.NO_FIRING));

    verify(stepContext)
        .setStateCleanupTimer(
            SimpleParDoFn.CLEANUP_TIMER_ID,
            firstWindow,
            IntervalWindow.getCoder(),
            firstWindow.maxTimestamp().plus(Duration.millis(1L)),
            firstWindow.maxTimestamp().plus(Duration.millis(1L)));
  }