org.apache.flink.streaming.util.MockStreamConfig Java Examples

private MockEnvironment setupEnvironment(boolean[] outputAvailabilities) {
	final Configuration configuration = new Configuration();
	new MockStreamConfig(configuration, outputAvailabilities.length);

	final List<ResultPartitionWriter> writers = new ArrayList<>(outputAvailabilities.length);
	for (int i = 0; i < outputAvailabilities.length; i++) {
		writers.add(new AvailabilityTestResultPartitionWriter(outputAvailabilities[i]));
	}

	final MockEnvironment environment = new MockEnvironmentBuilder()
		.setTaskConfiguration(configuration)
		.build();
	environment.addOutputs(writers);
	return environment;
}
 

/**
 * Test case for FLINK-5638: Tests that the async wait operator can be closed even if the
 * emitter is currently waiting on the checkpoint lock (e.g. in the case of two chained async
 * wait operators where the latter operator's queue is currently full).
 *
 * <p>Note that this test does not enforce the exact strict ordering because with the fix it is no
 * longer possible. However, it provokes the described situation without the fix.
 */
@Test(timeout = 10000L)
public void testClosingWithBlockedEmitter() throws Exception {
	final Object lock = new Object();

	ArgumentCaptor<Throwable> failureReason = ArgumentCaptor.forClass(Throwable.class);

	MockEnvironment environment = createMockEnvironment();

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

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

	final OneShotLatch closingLatch = new OneShotLatch();
	final OneShotLatch outputLatch = new OneShotLatch();

	Output<StreamRecord<Integer>> output = mock(Output.class);
	doAnswer(new Answer() {
		@Override
		public Object answer(InvocationOnMock invocation) throws Throwable {
			assertTrue("Output should happen under the checkpoint lock.", Thread.currentThread().holdsLock(lock));

			outputLatch.trigger();

			// wait until we're in the closing method of the operator
			while (!closingLatch.isTriggered()) {
				lock.wait();
			}

			return null;
		}
	}).when(output).collect(any(StreamRecord.class));

	AsyncWaitOperator<Integer, Integer> operator = new TestAsyncWaitOperator<>(
		new MyAsyncFunction(),
		1000L,
		1,
		AsyncDataStream.OutputMode.ORDERED,
		closingLatch);

	operator.setup(
		containingTask,
		streamConfig,
		output);

	operator.open();

	synchronized (lock) {
		operator.processElement(new StreamRecord<>(42));
	}

	outputLatch.await();

	synchronized (lock) {
		operator.close();
	}
}
 

/**
 * 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));
}
 

/**
 * Test case for FLINK-5638: Tests that the async wait operator can be closed even if the
 * emitter is currently waiting on the checkpoint lock (e.g. in the case of two chained async
 * wait operators where the latter operator's queue is currently full).
 *
 * <p>Note that this test does not enforce the exact strict ordering because with the fix it is no
 * longer possible. However, it provokes the described situation without the fix.
 */
@Test(timeout = 10000L)
public void testClosingWithBlockedEmitter() throws Exception {
	final Object lock = new Object();

	ArgumentCaptor<Throwable> failureReason = ArgumentCaptor.forClass(Throwable.class);

	MockEnvironment environment = createMockEnvironment();

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

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

	final OneShotLatch closingLatch = new OneShotLatch();
	final OneShotLatch outputLatch = new OneShotLatch();

	Output<StreamRecord<Integer>> output = mock(Output.class);
	doAnswer(new Answer() {
		@Override
		public Object answer(InvocationOnMock invocation) throws Throwable {
			assertTrue("Output should happen under the checkpoint lock.", Thread.currentThread().holdsLock(lock));

			outputLatch.trigger();

			// wait until we're in the closing method of the operator
			while (!closingLatch.isTriggered()) {
				lock.wait();
			}

			return null;
		}
	}).when(output).collect(any(StreamRecord.class));

	AsyncWaitOperator<Integer, Integer> operator = new TestAsyncWaitOperator<>(
		new MyAsyncFunction(),
		1000L,
		1,
		AsyncDataStream.OutputMode.ORDERED,
		closingLatch);

	operator.setup(
		containingTask,
		streamConfig,
		output);

	operator.open();

	synchronized (lock) {
		operator.processElement(new StreamRecord<>(42));
	}

	outputLatch.await();

	synchronized (lock) {
		operator.close();
	}
}
 

/**
 * 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));
}