org.apache.flink.streaming.api.functions.async.AsyncFunction Java Examples

public AsyncWaitOperator(
		@Nonnull AsyncFunction<IN, OUT> asyncFunction,
		long timeout,
		int capacity,
		@Nonnull AsyncDataStream.OutputMode outputMode,
		@Nonnull ProcessingTimeService processingTimeService,
		@Nonnull MailboxExecutor mailboxExecutor) {
	super(asyncFunction);

	setChainingStrategy(ChainingStrategy.ALWAYS);

	Preconditions.checkArgument(capacity > 0, "The number of concurrent async operation should be greater than 0.");
	this.capacity = capacity;

	this.outputMode = Preconditions.checkNotNull(outputMode, "outputMode");

	this.timeout = timeout;

	this.processingTimeService = Preconditions.checkNotNull(processingTimeService);

	this.mailboxExecutor = mailboxExecutor;
}
 

public AsyncWaitOperator(
		AsyncFunction<IN, OUT> asyncFunction,
		long timeout,
		int capacity,
		AsyncDataStream.OutputMode outputMode) {
	super(asyncFunction);

	// TODO this is a temporary fix for the problems described under FLINK-13063 at the cost of breaking chains for
	//  AsyncOperators.
	setChainingStrategy(ChainingStrategy.HEAD);

	Preconditions.checkArgument(capacity > 0, "The number of concurrent async operation should be greater than 0.");
	this.capacity = capacity;

	this.outputMode = Preconditions.checkNotNull(outputMode, "outputMode");

	this.timeout = timeout;
}
 

private static <OUT> OneInputStreamOperatorTestHarness<Integer, OUT> createTestHarness(
		AsyncFunction<Integer, OUT> function,
		long timeout,
		int capacity,
		AsyncDataStream.OutputMode outputMode) throws Exception {

	return new OneInputStreamOperatorTestHarness<>(
		new AsyncWaitOperatorFactory<>(function, timeout, capacity, outputMode),
		IntSerializer.INSTANCE);
}
 

/**
 * Add an AsyncWaitOperator.
 *
 * @param in The {@link DataStream} where the {@link AsyncWaitOperator} will be added.
 * @param func {@link AsyncFunction} wrapped inside {@link AsyncWaitOperator}.
 * @param timeout for the asynchronous operation to complete
 * @param bufSize The max number of inputs the {@link AsyncWaitOperator} can hold inside.
 * @param mode Processing mode for {@link AsyncWaitOperator}.
 * @param <IN> Input type.
 * @param <OUT> Output type.
 * @return A new {@link SingleOutputStreamOperator}
 */
private static <IN, OUT> SingleOutputStreamOperator<OUT> addOperator(
		DataStream<IN> in,
		AsyncFunction<IN, OUT> func,
		long timeout,
		int bufSize,
		OutputMode mode) {

	TypeInformation<OUT> outTypeInfo = TypeExtractor.getUnaryOperatorReturnType(
		func,
		AsyncFunction.class,
		0,
		1,
		new int[]{1, 0},
		in.getType(),
		Utils.getCallLocationName(),
		true);

	// create transform
	AsyncWaitOperator<IN, OUT> operator = new AsyncWaitOperator<>(
		in.getExecutionEnvironment().clean(func),
		timeout,
		bufSize,
		mode);

	return in.transform("async wait operator", outTypeInfo, operator);
}
 

/**
 * This helper function is needed to check that the temporary fix for FLINK-13063 can be backwards compatible with
 * the old chaining behavior by setting the ChainingStrategy manually. TODO: remove after a proper fix for
 * FLINK-13063 is in place that allows chaining.
 */
private <IN, OUT> SingleOutputStreamOperator<OUT> addAsyncOperatorLegacyChained(
	DataStream<IN> in,
	AsyncFunction<IN, OUT> func,
	long timeout,
	int bufSize,
	AsyncDataStream.OutputMode mode) {

	TypeInformation<OUT> outTypeInfo = TypeExtractor.getUnaryOperatorReturnType(
		func,
		AsyncFunction.class,
		0,
		1,
		new int[]{1, 0},
		in.getType(),
		Utils.getCallLocationName(),
		true);

	// create transform
	AsyncWaitOperatorFactory<IN, OUT> factory = new AsyncWaitOperatorFactory<>(
		in.getExecutionEnvironment().clean(func),
		timeout,
		bufSize,
		mode);

	factory.setChainingStrategy(ChainingStrategy.ALWAYS);

	return in.transform("async wait operator", outTypeInfo, factory);
}
 

public AsyncWaitOperatorFactory(
		AsyncFunction<IN, OUT> asyncFunction,
		long timeout,
		int capacity,
		AsyncDataStream.OutputMode outputMode) {
	this.asyncFunction = asyncFunction;
	this.timeout = timeout;
	this.capacity = capacity;
	this.outputMode = outputMode;
	this.chainingStrategy = ChainingStrategy.ALWAYS;
}
 

/**
 * Add an AsyncWaitOperator. The order to process input records is guaranteed to be the same as
 * input ones.
 *
 * @param in Input {@link DataStream}
 * @param func {@link AsyncFunction}
 * @param timeout for the asynchronous operation to complete
 * @param timeUnit of the given timeout
 * @param <IN> Type of input record
 * @param <OUT> Type of output record
 * @return A new {@link SingleOutputStreamOperator}.
 */
public static <IN, OUT> SingleOutputStreamOperator<OUT> orderedWait(
		DataStream<IN> in,
		AsyncFunction<IN, OUT> func,
		long timeout,
		TimeUnit timeUnit) {
	return addOperator(
		in,
		func,
		timeUnit.toMillis(timeout),
		DEFAULT_QUEUE_CAPACITY,
		OutputMode.ORDERED);
}
 

/**
 * Add an AsyncWaitOperator. The order of output stream records may be reordered.
 * @param in Input {@link DataStream}
 * @param func {@link AsyncFunction}
 * @param timeout for the asynchronous operation to complete
 * @param timeUnit of the given timeout
 * @param <IN> Type of input record
 * @param <OUT> Type of output record
 * @return A new {@link SingleOutputStreamOperator}.
 */
public static <IN, OUT> SingleOutputStreamOperator<OUT> unorderedWait(
		DataStream<IN> in,
		AsyncFunction<IN, OUT> func,
		long timeout,
		TimeUnit timeUnit) {
	return addOperator(
		in,
		func,
		timeUnit.toMillis(timeout),
		DEFAULT_QUEUE_CAPACITY,
		OutputMode.UNORDERED);
}
 

/**
 * Add an AsyncWaitOperator.
 *
 * @param in The {@link DataStream} where the {@link AsyncWaitOperator} will be added.
 * @param func {@link AsyncFunction} wrapped inside {@link AsyncWaitOperator}.
 * @param timeout for the asynchronous operation to complete
 * @param bufSize The max number of inputs the {@link AsyncWaitOperator} can hold inside.
 * @param mode Processing mode for {@link AsyncWaitOperator}.
 * @param <IN> Input type.
 * @param <OUT> Output type.
 * @return A new {@link SingleOutputStreamOperator}
 */
private static <IN, OUT> SingleOutputStreamOperator<OUT> addOperator(
		DataStream<IN> in,
		AsyncFunction<IN, OUT> func,
		long timeout,
		int bufSize,
		OutputMode mode) {

	TypeInformation<OUT> outTypeInfo = TypeExtractor.getUnaryOperatorReturnType(
		func,
		AsyncFunction.class,
		0,
		1,
		new int[]{1, 0},
		in.getType(),
		Utils.getCallLocationName(),
		true);

	// create transform
	AsyncWaitOperatorFactory<IN, OUT> operatorFactory = new AsyncWaitOperatorFactory<>(
		in.getExecutionEnvironment().clean(func),
		timeout,
		bufSize,
		mode);

	return in.transform("async wait operator", outTypeInfo, operatorFactory);
}
 

public AsyncLookupJoinWithCalcRunner(
		GeneratedFunction<AsyncFunction<RowData, Object>> generatedFetcher,
		GeneratedFunction<FlatMapFunction<RowData, RowData>> generatedCalc,
		GeneratedResultFuture<TableFunctionResultFuture<RowData>> generatedResultFuture,
		TypeInformation<?> fetcherReturnType,
		RowDataTypeInfo rightRowTypeInfo,
		boolean isLeftOuterJoin,
		int asyncBufferCapacity) {
	super(generatedFetcher, generatedResultFuture, fetcherReturnType,
		rightRowTypeInfo, isLeftOuterJoin, asyncBufferCapacity);
	this.rightRowTypeInfo = rightRowTypeInfo;
	this.generatedCalc = generatedCalc;
}
 

public AsyncLookupJoinRunner(
		GeneratedFunction<AsyncFunction<RowData, Object>> generatedFetcher,
		GeneratedResultFuture<TableFunctionResultFuture<RowData>> generatedResultFuture,
		TypeInformation<?> fetcherReturnType,
		RowDataTypeInfo rightRowTypeInfo,
		boolean isLeftOuterJoin,
		int asyncBufferCapacity) {
	this.generatedFetcher = generatedFetcher;
	this.generatedResultFuture = generatedResultFuture;
	this.isLeftOuterJoin = isLeftOuterJoin;
	this.asyncBufferCapacity = asyncBufferCapacity;
	this.fetcherReturnType = fetcherReturnType;
	this.rightRowTypeInfo = rightRowTypeInfo;
}
 

public static DataStream<Row> translate(
        DataStream<Row> inputStream,
        RealTimeTransForm transForm)
{
    RowTypeInfo streamRowType = (RowTypeInfo) inputStream.getType();
    AsyncFunction<Row, Row> asyncFunction = new RichAsyncFunctionImpl(transForm, streamRowType);

    DataStream<Row> joinResultStream = AsyncDataStream.orderedWait(
            inputStream, asyncFunction,
            1000, TimeUnit.MILLISECONDS, // 超时时间
            100);  // 进行中的异步请求的最大数量

    return joinResultStream;
}
 

/**
 * This helper function is needed to check that the temporary fix for FLINK-13063 can be backwards compatible with
 * the old chaining behavior by setting the ChainingStrategy manually. TODO: remove after a proper fix for
 * FLINK-13063 is in place that allows chaining.
 */
private <IN, OUT> SingleOutputStreamOperator<OUT> addAsyncOperatorLegacyChained(
	DataStream<IN> in,
	AsyncFunction<IN, OUT> func,
	long timeout,
	int bufSize,
	AsyncDataStream.OutputMode mode) {

	TypeInformation<OUT> outTypeInfo = TypeExtractor.getUnaryOperatorReturnType(
		func,
		AsyncFunction.class,
		0,
		1,
		new int[]{1, 0},
		in.getType(),
		Utils.getCallLocationName(),
		true);

	// create transform
	AsyncWaitOperator<IN, OUT> operator = new AsyncWaitOperator<>(
		in.getExecutionEnvironment().clean(func),
		timeout,
		bufSize,
		mode);

	operator.setChainingStrategy(ChainingStrategy.ALWAYS);

	return in.transform("async wait operator", outTypeInfo, operator);
}
 

private void testTimeoutExceptionHandling(AsyncDataStream.OutputMode outputMode) throws Exception {
	AsyncFunction<Integer, Integer> asyncFunction = new NoOpAsyncFunction<>();
	long timeout = 10L; // 1 milli second

	AsyncWaitOperator<Integer, Integer> asyncWaitOperator = new AsyncWaitOperator<>(
		asyncFunction,
		timeout,
		2,
		outputMode);

	final MockEnvironment mockEnvironment = createMockEnvironment();
	mockEnvironment.setExpectedExternalFailureCause(Throwable.class);

	OneInputStreamOperatorTestHarness<Integer, Integer> harness = new OneInputStreamOperatorTestHarness<>(
		asyncWaitOperator,
		IntSerializer.INSTANCE,
		mockEnvironment);

	harness.open();

	synchronized (harness.getCheckpointLock()) {
		harness.processElement(1, 1L);
	}

	harness.setProcessingTime(10L);

	synchronized (harness.getCheckpointLock()) {
		harness.close();
	}
}
 

public TestAsyncWaitOperator(
		AsyncFunction<IN, OUT> asyncFunction,
		long timeout,
		int capacity,
		AsyncDataStream.OutputMode outputMode,
		OneShotLatch closingLatch) {
	super(asyncFunction, timeout, capacity, outputMode);

	this.closingLatch = Preconditions.checkNotNull(closingLatch);
}
 

/**
 * Add an AsyncWaitOperator. The order to process input records is guaranteed to be the same as
 * input ones.
 *
 * @param in Input {@link DataStream}
 * @param func {@link AsyncFunction}
 * @param timeout for the asynchronous operation to complete
 * @param timeUnit of the given timeout
 * @param <IN> Type of input record
 * @param <OUT> Type of output record
 * @return A new {@link SingleOutputStreamOperator}.
 */
public static <IN, OUT> SingleOutputStreamOperator<OUT> orderedWait(
		DataStream<IN> in,
		AsyncFunction<IN, OUT> func,
		long timeout,
		TimeUnit timeUnit) {
	return addOperator(
		in,
		func,
		timeUnit.toMillis(timeout),
		DEFAULT_QUEUE_CAPACITY,
		OutputMode.ORDERED);
}
 

/**
 * Add an AsyncWaitOperator. The order of output stream records may be reordered.
 * @param in Input {@link DataStream}
 * @param func {@link AsyncFunction}
 * @param timeout for the asynchronous operation to complete
 * @param timeUnit of the given timeout
 * @param <IN> Type of input record
 * @param <OUT> Type of output record
 * @return A new {@link SingleOutputStreamOperator}.
 */
public static <IN, OUT> SingleOutputStreamOperator<OUT> unorderedWait(
		DataStream<IN> in,
		AsyncFunction<IN, OUT> func,
		long timeout,
		TimeUnit timeUnit) {
	return addOperator(
		in,
		func,
		timeUnit.toMillis(timeout),
		DEFAULT_QUEUE_CAPACITY,
		OutputMode.UNORDERED);
}
 

/**
 * Add an AsyncWaitOperator. The order of output stream records may be reordered.
 * @param in Input {@link DataStream}
 * @param func {@link AsyncFunction}
 * @param timeout for the asynchronous operation to complete
 * @param timeUnit of the given timeout
 * @param <IN> Type of input record
 * @param <OUT> Type of output record
 * @return A new {@link SingleOutputStreamOperator}.
 */
public static <IN, OUT> SingleOutputStreamOperator<OUT> unorderedWait(
		DataStream<IN> in,
		AsyncFunction<IN, OUT> func,
		long timeout,
		TimeUnit timeUnit) {
	return addOperator(
		in,
		func,
		timeUnit.toMillis(timeout),
		DEFAULT_QUEUE_CAPACITY,
		OutputMode.UNORDERED);
}
 

/**
 * Add an AsyncWaitOperator. The order to process input records is guaranteed to be the same as
 * input ones.
 *
 * @param in Input {@link DataStream}
 * @param func {@link AsyncFunction}
 * @param timeout for the asynchronous operation to complete
 * @param timeUnit of the given timeout
 * @param <IN> Type of input record
 * @param <OUT> Type of output record
 * @return A new {@link SingleOutputStreamOperator}.
 */
public static <IN, OUT> SingleOutputStreamOperator<OUT> orderedWait(
		DataStream<IN> in,
		AsyncFunction<IN, OUT> func,
		long timeout,
		TimeUnit timeUnit) {
	return addOperator(
		in,
		func,
		timeUnit.toMillis(timeout),
		DEFAULT_QUEUE_CAPACITY,
		OutputMode.ORDERED);
}
 

public AsyncWaitOperator(
		AsyncFunction<IN, OUT> asyncFunction,
		long timeout,
		int capacity,
		AsyncDataStream.OutputMode outputMode) {
	super(asyncFunction);
	chainingStrategy = ChainingStrategy.ALWAYS;

	Preconditions.checkArgument(capacity > 0, "The number of concurrent async operation should be greater than 0.");
	this.capacity = capacity;

	this.outputMode = Preconditions.checkNotNull(outputMode, "outputMode");

	this.timeout = timeout;
}
 

public TestAsyncWaitOperator(
		AsyncFunction<IN, OUT> asyncFunction,
		long timeout,
		int capacity,
		AsyncDataStream.OutputMode outputMode,
		OneShotLatch closingLatch) {
	super(asyncFunction, timeout, capacity, outputMode);

	this.closingLatch = Preconditions.checkNotNull(closingLatch);
}
 

private void testTimeoutExceptionHandling(AsyncDataStream.OutputMode outputMode) throws Exception {
	AsyncFunction<Integer, Integer> asyncFunction = new NoOpAsyncFunction<>();
	long timeout = 10L; // 1 milli second

	AsyncWaitOperator<Integer, Integer> asyncWaitOperator = new AsyncWaitOperator<>(
		asyncFunction,
		timeout,
		2,
		outputMode);

	final MockEnvironment mockEnvironment = createMockEnvironment();
	mockEnvironment.setExpectedExternalFailureCause(Throwable.class);

	OneInputStreamOperatorTestHarness<Integer, Integer> harness = new OneInputStreamOperatorTestHarness<>(
		asyncWaitOperator,
		IntSerializer.INSTANCE,
		mockEnvironment);

	harness.open();

	synchronized (harness.getCheckpointLock()) {
		harness.processElement(1, 1L);
	}

	harness.setProcessingTime(10L);

	synchronized (harness.getCheckpointLock()) {
		harness.close();
	}
}
 

public AsyncLookupJoinRunner(
		GeneratedFunction<AsyncFunction<BaseRow, Object>> generatedFetcher,
		GeneratedResultFuture<TableFunctionResultFuture<BaseRow>> generatedResultFuture,
		TypeInformation<?> fetcherReturnType,
		BaseRowTypeInfo rightRowTypeInfo,
		boolean isLeftOuterJoin,
		int asyncBufferCapacity) {
	this.generatedFetcher = generatedFetcher;
	this.generatedResultFuture = generatedResultFuture;
	this.isLeftOuterJoin = isLeftOuterJoin;
	this.asyncBufferCapacity = asyncBufferCapacity;
	this.fetcherReturnType = fetcherReturnType;
	this.rightRowTypeInfo = rightRowTypeInfo;
}
 

public AsyncLookupJoinWithCalcRunner(
		GeneratedFunction<AsyncFunction<BaseRow, Object>> generatedFetcher,
		GeneratedFunction<FlatMapFunction<BaseRow, BaseRow>> generatedCalc,
		GeneratedResultFuture<TableFunctionResultFuture<BaseRow>> generatedResultFuture,
		TypeInformation<?> fetcherReturnType,
		BaseRowTypeInfo rightRowTypeInfo,
		boolean isLeftOuterJoin,
		int asyncBufferCapacity) {
	super(generatedFetcher, generatedResultFuture, fetcherReturnType,
		rightRowTypeInfo, isLeftOuterJoin, asyncBufferCapacity);
	this.rightRowTypeInfo = rightRowTypeInfo;
	this.generatedCalc = generatedCalc;
}
 

/**
 * Add an AsyncWaitOperator.
 *
 * @param in The {@link DataStream} where the {@link AsyncWaitOperator} will be added.
 * @param func {@link AsyncFunction} wrapped inside {@link AsyncWaitOperator}.
 * @param timeout for the asynchronous operation to complete
 * @param bufSize The max number of inputs the {@link AsyncWaitOperator} can hold inside.
 * @param mode Processing mode for {@link AsyncWaitOperator}.
 * @param <IN> Input type.
 * @param <OUT> Output type.
 * @return A new {@link SingleOutputStreamOperator}
 */
private static <IN, OUT> SingleOutputStreamOperator<OUT> addOperator(
		DataStream<IN> in,
		AsyncFunction<IN, OUT> func,
		long timeout,
		int bufSize,
		OutputMode mode) {

	TypeInformation<OUT> outTypeInfo = TypeExtractor.getUnaryOperatorReturnType(
		func,
		AsyncFunction.class,
		0,
		1,
		new int[]{1, 0},
		in.getType(),
		Utils.getCallLocationName(),
		true);

	// create transform
	AsyncWaitOperator<IN, OUT> operator = new AsyncWaitOperator<>(
		in.getExecutionEnvironment().clean(func),
		timeout,
		bufSize,
		mode);

	return in.transform("async wait operator", outTypeInfo, operator);
}
 

public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        final ParameterTool params = ParameterTool.fromArgs(args);

        final String statePath;
        final String cpMode;
        final int maxCount;
        final long sleepFactor;
        final float failRatio;
        final String mode;
        final int taskNum;
        final String timeType;
        final long shutdownWaitTS;
        final long timeout;

        try {
            // check the configuration for the job
            statePath = params.get("fsStatePath", null);
            cpMode = params.get("checkpointMode", "exactly_once");
            maxCount = params.getInt("maxCount", 100000);
            sleepFactor = params.getLong("sleepFactor", 100);
            failRatio = params.getFloat("failRatio", 0.001f);
            mode = params.get("waitMode", "ordered");
            taskNum = params.getInt("waitOperatorParallelism", 1);
            timeType = params.get("eventType", "EventTime");
            shutdownWaitTS = params.getLong("shutdownWaitTS", 20000);
            timeout = params.getLong("timeout", 10000L);
        } catch (Exception e) {
            printUsage();
            throw e;
        }

        if (statePath != null) {
            env.setStateBackend(new FsStateBackend(statePath));
        }

        if (EXACTLY_ONCE_MODE.equals(cpMode)) {
            env.enableCheckpointing(1000L, CheckpointingMode.EXACTLY_ONCE);
        } else {
            env.enableCheckpointing(1000L, CheckpointingMode.AT_LEAST_ONCE);
        }

        if (EVENT_TIME.equals(timeType)) {
            env.setStreamTimeCharacteristic(TimeCharacteristic.EventTime);
        } else if (INGESTION_TIME.equals(timeType)) {
            env.setStreamTimeCharacteristic(TimeCharacteristic.IngestionTime);
        }

        DataStream<Integer> inputStream = env.addSource(new SimpleSource(maxCount));

        // create async function, which will *wait* for a while to simulate the process of async i/o
        AsyncFunction<Integer, String> function =
                new SampleAsyncFunction(sleepFactor, failRatio, shutdownWaitTS);

        // add async operator to streaming job
        DataStream<String> result;
        if (ORDERED.equals(mode)) {
            result = AsyncDataStream.orderedWait(
                    inputStream,
                    function,
                    timeout,
                    TimeUnit.MILLISECONDS,
                    20).setParallelism(taskNum);
        } else {
            result = AsyncDataStream.unorderedWait(
                    inputStream,
                    function,
                    timeout,
                    TimeUnit.MILLISECONDS,
                    20).setParallelism(taskNum);
        }

        result.flatMap(new FlatMapFunction<String, Tuple2<String, Integer>>() {
            private static final long serialVersionUID = -938116068682344455L;

            @Override
            public void flatMap(String value, Collector<Tuple2<String, Integer>> out) throws Exception {
                out.collect(new Tuple2<>(value, 1));
            }
        })
                .keyBy(0)
                .sum(1)
                .print();

        env.execute("Async IO Example");
    }
 

public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        final ParameterTool params = ParameterTool.fromArgs(args);

        final String statePath;
        final String cpMode;
        final int maxCount;
        final long sleepFactor;
        final float failRatio;
        final String mode;
        final int taskNum;
        final String timeType;
        final long shutdownWaitTS;
        final long timeout;

        try {
            // check the configuration for the job
            statePath = params.get("fsStatePath", null);
            cpMode = params.get("checkpointMode", "exactly_once");
            maxCount = params.getInt("maxCount", 100000);
            sleepFactor = params.getLong("sleepFactor", 100);
            failRatio = params.getFloat("failRatio", 0.001f);
            mode = params.get("waitMode", "ordered");
            taskNum = params.getInt("waitOperatorParallelism", 1);
            timeType = params.get("eventType", "EventTime");
            shutdownWaitTS = params.getLong("shutdownWaitTS", 20000);
            timeout = params.getLong("timeout", 10000L);
        } catch (Exception e) {
            printUsage();
            throw e;
        }

        if (statePath != null) {
            env.setStateBackend(new FsStateBackend(statePath));
        }

        if (EXACTLY_ONCE_MODE.equals(cpMode)) {
            env.enableCheckpointing(1000L, CheckpointingMode.EXACTLY_ONCE);
        } else {
            env.enableCheckpointing(1000L, CheckpointingMode.AT_LEAST_ONCE);
        }

        if (EVENT_TIME.equals(timeType)) {
            env.setStreamTimeCharacteristic(TimeCharacteristic.EventTime);
        } else if (INGESTION_TIME.equals(timeType)) {
            env.setStreamTimeCharacteristic(TimeCharacteristic.IngestionTime);
        }

        DataStream<Integer> inputStream = env.addSource(new SimpleSource(maxCount));

        // create async function, which will *wait* for a while to simulate the process of async i/o
        AsyncFunction<Integer, String> function =
                new SampleAsyncFunction(sleepFactor, failRatio, shutdownWaitTS);

        // add async operator to streaming job
        DataStream<String> result;
        if (ORDERED.equals(mode)) {
            result = AsyncDataStream.orderedWait(
                    inputStream,
                    function,
                    timeout,
                    TimeUnit.MILLISECONDS,
                    20).setParallelism(taskNum);
        } else {
            result = AsyncDataStream.unorderedWait(
                    inputStream,
                    function,
                    timeout,
                    TimeUnit.MILLISECONDS,
                    20).setParallelism(taskNum);
        }

        result.flatMap(new FlatMapFunction<String, Tuple2<String, Integer>>() {
            private static final long serialVersionUID = -938116068682344455L;

            @Override
            public void flatMap(String value, Collector<Tuple2<String, Integer>> out) throws Exception {
                out.collect(new Tuple2<>(value, 1));
            }
        })
                .keyBy(0)
                .sum(1)
                .print();

        env.execute("Async IO Example");
    }
 

/**
 * FLINK-5652
 * Tests that registered timers are properly canceled upon completion of a
 * {@link StreamRecordQueueEntry} in order to avoid resource leaks because TriggerTasks hold
 * a reference on the StreamRecordQueueEntry.
 */
@Test
public void testTimeoutCleanup() throws Exception {
	final Object lock = new Object();

	final long timeout = 100000L;
	final long timestamp = 1L;

	Environment environment = createMockEnvironment();

	ScheduledFuture<?> scheduledFuture = mock(ScheduledFuture.class);

	ProcessingTimeService processingTimeService = mock(ProcessingTimeService.class);
	when(processingTimeService.getCurrentProcessingTime()).thenReturn(timestamp);
	doReturn(scheduledFuture).when(processingTimeService).registerTimer(anyLong(), any(ProcessingTimeCallback.class));

	StreamTask<?, ?> containingTask = mock(StreamTask.class);
	when(containingTask.getEnvironment()).thenReturn(environment);
	when(containingTask.getCheckpointLock()).thenReturn(lock);
	when(containingTask.getProcessingTimeService()).thenReturn(processingTimeService);

	StreamConfig streamConfig = new MockStreamConfig();
	streamConfig.setTypeSerializerIn1(IntSerializer.INSTANCE);

	Output<StreamRecord<Integer>> output = mock(Output.class);

	AsyncWaitOperator<Integer, Integer> operator = new AsyncWaitOperator<>(
		new AsyncFunction<Integer, Integer>() {
			private static final long serialVersionUID = -3718276118074877073L;

			@Override
			public void asyncInvoke(Integer input, ResultFuture<Integer> resultFuture) throws Exception {
				resultFuture.complete(Collections.singletonList(input));
			}
		},
		timeout,
		1,
		AsyncDataStream.OutputMode.UNORDERED);

	operator.setup(
		containingTask,
		streamConfig,
		output);

	operator.open();

	final StreamRecord<Integer> streamRecord = new StreamRecord<>(42, timestamp);

	synchronized (lock) {
		// processing an element will register a timeout
		operator.processElement(streamRecord);
	}

	synchronized (lock) {
		// closing the operator waits until all inputs have been processed
		operator.close();
	}

	// check that we actually outputted the result of the single input
	verify(output).collect(eq(streamRecord));
	verify(processingTimeService).registerTimer(eq(processingTimeService.getCurrentProcessingTime() + timeout), any(ProcessingTimeCallback.class));

	// check that we have cancelled our registered timeout
	verify(scheduledFuture).cancel(eq(true));
}
 

/**
 * FLINK-5652
 * Tests that registered timers are properly canceled upon completion of a
 * {@link StreamRecordQueueEntry} in order to avoid resource leaks because TriggerTasks hold
 * a reference on the StreamRecordQueueEntry.
 */
@Test
public void testTimeoutCleanup() throws Exception {
	final Object lock = new Object();

	final long timeout = 100000L;
	final long timestamp = 1L;

	Environment environment = createMockEnvironment();

	ScheduledFuture<?> scheduledFuture = mock(ScheduledFuture.class);

	ProcessingTimeService processingTimeService = mock(ProcessingTimeService.class);
	when(processingTimeService.getCurrentProcessingTime()).thenReturn(timestamp);
	doReturn(scheduledFuture).when(processingTimeService).registerTimer(anyLong(), any(ProcessingTimeCallback.class));

	StreamTask<?, ?> containingTask = mock(StreamTask.class);
	when(containingTask.getEnvironment()).thenReturn(environment);
	when(containingTask.getCheckpointLock()).thenReturn(lock);
	when(containingTask.getProcessingTimeService()).thenReturn(processingTimeService);

	StreamConfig streamConfig = new MockStreamConfig();
	streamConfig.setTypeSerializerIn1(IntSerializer.INSTANCE);

	Output<StreamRecord<Integer>> output = mock(Output.class);

	AsyncWaitOperator<Integer, Integer> operator = new AsyncWaitOperator<>(
		new AsyncFunction<Integer, Integer>() {
			private static final long serialVersionUID = -3718276118074877073L;

			@Override
			public void asyncInvoke(Integer input, ResultFuture<Integer> resultFuture) throws Exception {
				resultFuture.complete(Collections.singletonList(input));
			}
		},
		timeout,
		1,
		AsyncDataStream.OutputMode.UNORDERED);

	operator.setup(
		containingTask,
		streamConfig,
		output);

	operator.open();

	final StreamRecord<Integer> streamRecord = new StreamRecord<>(42, timestamp);

	synchronized (lock) {
		// processing an element will register a timeout
		operator.processElement(streamRecord);
	}

	synchronized (lock) {
		// closing the operator waits until all inputs have been processed
		operator.close();
	}

	// check that we actually outputted the result of the single input
	verify(output).collect(eq(streamRecord));
	verify(processingTimeService).registerTimer(eq(processingTimeService.getCurrentProcessingTime() + timeout), any(ProcessingTimeCallback.class));

	// check that we have cancelled our registered timeout
	verify(scheduledFuture).cancel(eq(true));
}
 

/**
 * Add an AsyncWaitOperator. The order of output stream records may be reordered.
 *
 * @param in Input {@link DataStream}
 * @param func {@link AsyncFunction}
 * @param timeout for the asynchronous operation to complete
 * @param timeUnit of the given timeout
 * @param capacity The max number of async i/o operation that can be triggered
 * @param <IN> Type of input record
 * @param <OUT> Type of output record
 * @return A new {@link SingleOutputStreamOperator}.
 */
public static <IN, OUT> SingleOutputStreamOperator<OUT> unorderedWait(
		DataStream<IN> in,
		AsyncFunction<IN, OUT> func,
		long timeout,
		TimeUnit timeUnit,
		int capacity) {
	return addOperator(in, func, timeUnit.toMillis(timeout), capacity, OutputMode.UNORDERED);
}