ExecutableStageDoFnOperator extends DoFnOperator<InputT, OutputT> { @Override public void open() throws Exception { executableStage = ExecutableStage.fromPayload(payload); initializeUserState(executableStage, getKeyedStateBackend()); stageContext = contextFactory.get(jobInfo); stageBundleFactory = stageContext.getStageBundleFactory(executableStage); stateRequestHandler = getStateRequestHandler(executableStage); progressHandler = new BundleProgressHandler() { @Override public void onProgress(ProcessBundleProgressResponse progress) { if (flinkMetricContainer != null) { flinkMetricContainer.updateMetrics(stepName, progress.getMonitoringInfosList()); } } @Override public void onCompleted(ProcessBundleResponse response) { if (flinkMetricContainer != null) { flinkMetricContainer.updateMetrics(stepName, response.getMonitoringInfosList()); } } }; minEventTimeTimerTimestampInCurrentBundle = Long.MAX_VALUE; minEventTimeTimerTimestampInLastBundle = Long.MAX_VALUE; super.setPreBundleCallback(this::preBundleStartCallback); super.setBundleFinishedCallback(this::finishBundleCallback); super.open(); } ExecutableStageDoFnOperator(
      String stepName,
      Coder<WindowedValue<InputT>> windowedInputCoder,
      Map<TupleTag<?>, Coder<?>> outputCoders,
      TupleTag<OutputT> mainOutputTag,
      List<TupleTag<?>> additionalOutputTags,
      OutputManagerFactory<OutputT> outputManagerFactory,
      Map<Integer, PCollectionView<?>> sideInputTagMapping,
      Collection<PCollectionView<?>> sideInputs,
      Map<RunnerApi.ExecutableStagePayload.SideInputId, PCollectionView<?>> sideInputIds,
      PipelineOptions options,
      RunnerApi.ExecutableStagePayload payload,
      JobInfo jobInfo,
      FlinkExecutableStageContextFactory contextFactory,
      Map<String, TupleTag<?>> outputMap,
      WindowingStrategy windowingStrategy,
      Coder keyCoder,
      KeySelector<WindowedValue<InputT>, ?> keySelector); @Override void open(); @Override void setKeyContextElement1(StreamRecord record); @Override void setCurrentKey(Object key); @Override ByteBuffer getCurrentKey(); @Override void close(); @Override void dispose(); @Override long applyInputWatermarkHold(long inputWatermark); @Override long applyOutputWatermarkHold(long currentOutputWatermark, long potentialOutputWatermark);  }
@Test @SuppressWarnings("unchecked") public void testEnsureStateCleanupWithKeyedInput() throws Exception { TupleTag<Integer> mainOutput = new TupleTag<>("main-output"); DoFnOperator.MultiOutputOutputManagerFactory<Integer> outputManagerFactory = new DoFnOperator.MultiOutputOutputManagerFactory(mainOutput, VarIntCoder.of()); VarIntCoder keyCoder = VarIntCoder.of(); ExecutableStageDoFnOperator<Integer, Integer> operator = getOperator( mainOutput, Collections.emptyList(), outputManagerFactory, WindowingStrategy.globalDefault(), keyCoder, WindowedValue.getFullCoder(keyCoder, GlobalWindow.Coder.INSTANCE)); KeyedOneInputStreamOperatorTestHarness<Integer, WindowedValue<Integer>, WindowedValue<Integer>> testHarness = new KeyedOneInputStreamOperatorTestHarness( operator, val -> val, new CoderTypeInformation<>(keyCoder)); RemoteBundle bundle = Mockito.mock(RemoteBundle.class); when(bundle.getInputReceivers()) .thenReturn( ImmutableMap.<String, FnDataReceiver<WindowedValue>>builder() .put("input", Mockito.mock(FnDataReceiver.class)) .build()); when(stageBundleFactory.getBundle(any(), any(), any(), any())).thenReturn(bundle); testHarness.open(); Object doFnRunner = Whitebox.getInternalState(operator, "doFnRunner"); assertThat(doFnRunner, instanceOf(DoFnRunnerWithMetricsUpdate.class)); Object statefulDoFnRunner = Whitebox.getInternalState(doFnRunner, "delegate"); assertThat(statefulDoFnRunner, instanceOf(StatefulDoFnRunner.class)); }