SimpleParDoFn implements ParDoFn { @Override @SuppressWarnings("unchecked") public void processElement(Object untypedElem) throws Exception { if (fnRunner == null) { 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(); } SimpleParDoFn(
      PipelineOptions options,
      DoFnInstanceManager doFnInstanceManager,
      SideInputReader sideInputReader,
      TupleTag<OutputT> mainOutputTag,
      Map<TupleTag<?>, Integer> outputTupleTagsToReceiverIndices,
      DataflowExecutionContext.DataflowStepContext stepContext,
      DataflowOperationContext operationContext,
      DoFnSchemaInformation doFnSchemaInformation,
      Map<String, PCollectionView<?>> sideInputMapping,
      DoFnRunnerFactory runnerFactory); @Override void startBundle(Receiver... receivers); @Override @SuppressWarnings("unchecked") void processElement(Object untypedElem); @Override void processTimers(); @Override void finishBundle(); @Override void abort();  }
@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 , null , 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); try (Closeable trackerCloser = tracker.activate()) { try (Closeable processCloser = operationContext.enterProcess()) { userParDoFn.processElement(WindowedValue.valueInGlobalWindow(5)); } } }