@Override
  @SuppressWarnings("unchecked")
  public void processElement(Object untypedElem) throws Exception {
    if (fnRunner == null) {
      // If we need to run reallyStartBundle in here, we need to make sure to switch the state
      // sampler into the start state.
      try (Closeable start = operationContext.enterStart()) {
        reallyStartBundle();
      }
    }

    WindowedValue<InputT> elem = (WindowedValue<InputT>) untypedElem;

    if (fnSignature != null && fnSignature.stateDeclarations().size() > 0) {
      registerStateCleanup(
          (WindowingStrategy<?, BoundedWindow>) getDoFnInfo().getWindowingStrategy(),
          (Collection<BoundedWindow>) elem.getWindows());
    }

    outputsPerElementTracker.onProcessElement();
    fnRunner.processElement(elem);
    outputsPerElementTracker.onProcessElementSuccess();
  }
@Test
  public void testStateTracking() throws Exception {
    ExecutionStateTracker tracker = ExecutionStateTracker.newForTest();

    TestOperationContext operationContext =
        TestOperationContext.create(
            new CounterSet(),
            NameContextsForTests.nameContextForTest(),
            new MetricsContainerImpl(NameContextsForTests.ORIGINAL_NAME),
            tracker);

    class StateTestingDoFn extends DoFn<Integer, String> {

      private boolean startCalled = false;

      @StartBundle
      public void startBundle() throws Exception {
        startCalled = true;
        assertThat(tracker.getCurrentState(), equalTo(operationContext.getStartState()));
      }

      @ProcessElement
      public void processElement(ProcessContext c) throws Exception {
        assertThat(startCalled, equalTo(true));
        assertThat(tracker.getCurrentState(), equalTo(operationContext.getProcessState()));
      }
    }

    StateTestingDoFn fn = new StateTestingDoFn();
    DoFnInfo<?, ?> fnInfo =
        DoFnInfo.forFn(
            fn,
            WindowingStrategy.globalDefault(),
            null /* side input views */,
            null /* input coder */,
            MAIN_OUTPUT,
            DoFnSchemaInformation.create(),
            Collections.emptyMap());

    ParDoFn userParDoFn =
        new SimpleParDoFn<>(
            options,
            DoFnInstanceManagers.singleInstance(fnInfo),
            NullSideInputReader.empty(),
            MAIN_OUTPUT,
            ImmutableMap.of(MAIN_OUTPUT, 0, new TupleTag<>("declared"), 1),
            BatchModeExecutionContext.forTesting(
                    options, operationContext.counterFactory(), "testStage")
                .getStepContext(operationContext),
            operationContext,
            DoFnSchemaInformation.create(),
            Collections.emptyMap(),
            SimpleDoFnRunnerFactory.INSTANCE);

    // This test ensures proper behavior of the state sampling even with lazy initialization.
    try (Closeable trackerCloser = tracker.activate()) {
      try (Closeable processCloser = operationContext.enterProcess()) {
        userParDoFn.processElement(WindowedValue.valueInGlobalWindow(5));
      }
    }
  }