AsyncWaitOperator extends AbstractUdfStreamOperator<OUT, AsyncFunction<IN, OUT>> implements OneInputStreamOperator<IN, OUT>, OperatorActions { @Override public void processElement(StreamRecord<IN> element) throws Exception { final StreamRecordQueueEntry<OUT> streamRecordBufferEntry = new StreamRecordQueueEntry<>(element); if (timeout > 0L) { long timeoutTimestamp = timeout + getProcessingTimeService().getCurrentProcessingTime(); final ScheduledFuture<?> timerFuture = getProcessingTimeService().registerTimer( timeoutTimestamp, new ProcessingTimeCallback() { @Override public void onProcessingTime(long timestamp) throws Exception { streamRecordBufferEntry.completeExceptionally( new TimeoutException("Async function call has timed out.")); } }); streamRecordBufferEntry.onComplete( (StreamElementQueueEntry<Collection<OUT>> value) -> { timerFuture.cancel(true); }, executor); } addAsyncBufferEntry(streamRecordBufferEntry); userFunction.asyncInvoke(element.getValue(), streamRecordBufferEntry); } AsyncWaitOperator(
			AsyncFunction<IN, OUT> asyncFunction,
			long timeout,
			int capacity,
			AsyncDataStream.OutputMode outputMode); @Override void setup(StreamTask<?, ?> containingTask, StreamConfig config, Output<StreamRecord<OUT>> output); @Override void open(); @Override void processElement(StreamRecord<IN> element); @Override void processWatermark(Watermark mark); @Override void snapshotState(StateSnapshotContext context); @Override void initializeState(StateInitializationContext context); @Override void close(); @Override void dispose(); @Override void failOperator(Throwable throwable);  }
@Test public void testOperatorChainWithProcessingTime() throws Exception { JobVertex chainedVertex = createChainedVertex(false); final OneInputStreamTask<Integer, Integer> task = new OneInputStreamTask<>(); final OneInputStreamTaskTestHarness<Integer, Integer> testHarness = new OneInputStreamTaskTestHarness<>(task, 1, 1, BasicTypeInfo.INT_TYPE_INFO, BasicTypeInfo.INT_TYPE_INFO); testHarness.setupOutputForSingletonOperatorChain(); testHarness.taskConfig = chainedVertex.getConfiguration(); final StreamConfig streamConfig = testHarness.getStreamConfig(); final StreamConfig operatorChainStreamConfig = new StreamConfig(chainedVertex.getConfiguration()); final AsyncWaitOperator<Integer, Integer> headOperator = operatorChainStreamConfig.getStreamOperator(AsyncWaitOperatorTest.class.getClassLoader()); streamConfig.setStreamOperator(headOperator); testHarness.invoke(); testHarness.waitForTaskRunning(); long initialTimestamp = 0L; testHarness.processElement(new StreamRecord<>(5, initialTimestamp)); testHarness.processElement(new StreamRecord<>(6, initialTimestamp + 1L)); testHarness.processElement(new StreamRecord<>(7, initialTimestamp + 2L)); testHarness.processElement(new StreamRecord<>(8, initialTimestamp + 3L)); testHarness.processElement(new StreamRecord<>(9, initialTimestamp + 4L)); testHarness.endInput(); testHarness.waitForTaskCompletion(); ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>(); expectedOutput.add(new StreamRecord<>(22, initialTimestamp)); expectedOutput.add(new StreamRecord<>(26, initialTimestamp + 1L)); expectedOutput.add(new StreamRecord<>(30, initialTimestamp + 2L)); expectedOutput.add(new StreamRecord<>(34, initialTimestamp + 3L)); expectedOutput.add(new StreamRecord<>(38, initialTimestamp + 4L)); TestHarnessUtil.assertOutputEqualsSorted( "Test for chained operator with AsyncWaitOperator failed", expectedOutput, testHarness.getOutput(), new StreamRecordComparator()); }
@Test public void testStateSnapshotAndRestore() throws Exception { final OneInputStreamTask<Integer, Integer> task = new OneInputStreamTask<>(); final OneInputStreamTaskTestHarness<Integer, Integer> testHarness = new OneInputStreamTaskTestHarness<>(task, 1, 1, BasicTypeInfo.INT_TYPE_INFO, BasicTypeInfo.INT_TYPE_INFO); testHarness.setupOutputForSingletonOperatorChain(); AsyncWaitOperator<Integer, Integer> operator = new AsyncWaitOperator<>( new LazyAsyncFunction(), TIMEOUT, 3, AsyncDataStream.OutputMode.ORDERED); final StreamConfig streamConfig = testHarness.getStreamConfig(); OperatorID operatorID = new OperatorID(42L, 4711L); streamConfig.setStreamOperator(operator); streamConfig.setOperatorID(operatorID); final AcknowledgeStreamMockEnvironment env = new AcknowledgeStreamMockEnvironment( testHarness.jobConfig, testHarness.taskConfig, testHarness.getExecutionConfig(), testHarness.memorySize, new MockInputSplitProvider(), testHarness.bufferSize); testHarness.invoke(env); testHarness.waitForTaskRunning(); final long initialTime = 0L; testHarness.processElement(new StreamRecord<>(1, initialTime + 1)); testHarness.processElement(new StreamRecord<>(2, initialTime + 2)); testHarness.processElement(new StreamRecord<>(3, initialTime + 3)); testHarness.processElement(new StreamRecord<>(4, initialTime + 4)); testHarness.waitForInputProcessing(); final long checkpointId = 1L; final long checkpointTimestamp = 1L; final CheckpointMetaData checkpointMetaData = new CheckpointMetaData(checkpointId, checkpointTimestamp); task.triggerCheckpoint(checkpointMetaData, CheckpointOptions.forCheckpoint()); env.getCheckpointLatch().await(); assertEquals(checkpointId, env.getCheckpointId()); LazyAsyncFunction.countDown(); testHarness.endInput(); testHarness.waitForTaskCompletion(); final OneInputStreamTask<Integer, Integer> restoredTask = new OneInputStreamTask<>(); TaskStateSnapshot subtaskStates = env.getCheckpointStateHandles(); restoredTask.setInitialState(subtaskStates); final OneInputStreamTaskTestHarness<Integer, Integer> restoredTaskHarness = new OneInputStreamTaskTestHarness<>(restoredTask, BasicTypeInfo.INT_TYPE_INFO, BasicTypeInfo.INT_TYPE_INFO); restoredTaskHarness.setupOutputForSingletonOperatorChain(); AsyncWaitOperator<Integer, Integer> restoredOperator = new AsyncWaitOperator<>( new MyAsyncFunction(), TIMEOUT, 6, AsyncDataStream.OutputMode.ORDERED); restoredTaskHarness.getStreamConfig().setStreamOperator(restoredOperator); restoredTaskHarness.getStreamConfig().setOperatorID(operatorID); restoredTaskHarness.invoke(); restoredTaskHarness.waitForTaskRunning(); restoredTaskHarness.processElement(new StreamRecord<>(5, initialTime + 5)); restoredTaskHarness.processElement(new StreamRecord<>(6, initialTime + 6)); restoredTaskHarness.processElement(new StreamRecord<>(7, initialTime + 7)); restoredTask.triggerCheckpoint(new CheckpointMetaData(checkpointId, checkpointTimestamp), CheckpointOptions.forCheckpoint()); restoredTaskHarness.processElement(new StreamRecord<>(8, initialTime + 8)); restoredTaskHarness.endInput(); restoredTaskHarness.waitForTaskCompletion(); ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>(); expectedOutput.add(new StreamRecord<>(2, initialTime + 1)); expectedOutput.add(new StreamRecord<>(4, initialTime + 2)); expectedOutput.add(new StreamRecord<>(6, initialTime + 3)); expectedOutput.add(new StreamRecord<>(8, initialTime + 4)); expectedOutput.add(new StreamRecord<>(10, initialTime + 5)); expectedOutput.add(new StreamRecord<>(12, initialTime + 6)); expectedOutput.add(new StreamRecord<>(14, initialTime + 7)); expectedOutput.add(new StreamRecord<>(16, initialTime + 8)); Iterator<Object> iterator = restoredTaskHarness.getOutput().iterator(); while (iterator.hasNext()) { if (iterator.next() instanceof CheckpointBarrier) { iterator.remove(); } } TestHarnessUtil.assertOutputEquals( "StateAndRestored Test Output was not correct.", expectedOutput, restoredTaskHarness.getOutput()); }