Java Code Examples for org.apache.flink.runtime.operators.testutils.MockEnvironment#setExpectedExternalFailureCause()

/**
 * This test checks the async exceptions handling wraps the message and cause as an AsynchronousException
 * and propagates this to the environment.
 */
@Test
public void streamTaskAsyncExceptionHandler_handleException_forwardsMessageProperly() {
	MockEnvironment mockEnvironment = MockEnvironment.builder().build();
	RuntimeException expectedException = new RuntimeException("RUNTIME EXCEPTION");

	final StreamTask.StreamTaskAsyncExceptionHandler asyncExceptionHandler = new StreamTask.StreamTaskAsyncExceptionHandler(mockEnvironment);

	mockEnvironment.setExpectedExternalFailureCause(AsynchronousException.class);
	final String expectedErrorMessage = "EXPECTED_ERROR MESSAGE";

	asyncExceptionHandler.handleAsyncException(expectedErrorMessage, expectedException);

	// expect an AsynchronousException containing the supplied error details
	Optional<? extends Throwable> actualExternalFailureCause = mockEnvironment.getActualExternalFailureCause();
	final Throwable actualException = actualExternalFailureCause
		.orElseThrow(() -> new AssertionError("Expected exceptional completion"));

	assertThat(actualException, instanceOf(AsynchronousException.class));
	assertThat(actualException.getMessage(), is("EXPECTED_ERROR MESSAGE"));
	assertThat(actualException.getCause(), is(expectedException));
}
 

private void testUserExceptionHandling(AsyncDataStream.OutputMode outputMode) throws Exception {
	UserExceptionAsyncFunction asyncWaitFunction = new UserExceptionAsyncFunction();
	long timeout = 2000L;

	AsyncWaitOperator<Integer, Integer> asyncWaitOperator = new AsyncWaitOperator<>(
		asyncWaitFunction,
		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);
	}

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

	assertTrue(harness.getEnvironment().getActualExternalFailureCause().isPresent());
}
 

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();
	}
}
 

private void testUserExceptionHandling(AsyncDataStream.OutputMode outputMode) throws Exception {
	UserExceptionAsyncFunction asyncWaitFunction = new UserExceptionAsyncFunction();
	long timeout = 2000L;

	AsyncWaitOperator<Integer, Integer> asyncWaitOperator = new AsyncWaitOperator<>(
		asyncWaitFunction,
		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);
	}

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

	assertTrue(harness.getEnvironment().getActualExternalFailureCause().isPresent());
}
 

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();
	}
}
 

private void testAsyncTimeout(LazyAsyncFunction lazyAsyncFunction,
		Optional<Class<? extends Throwable>> expectedException,
		StreamRecord<Integer>... expectedRecords) throws Exception {
	final long timeout = 10L;

	final AsyncWaitOperator<Integer, Integer> operator = new AsyncWaitOperator<>(
		lazyAsyncFunction,
		timeout,
		2,
		AsyncDataStream.OutputMode.ORDERED);

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

	final OneInputStreamOperatorTestHarness<Integer, Integer> testHarness =
		new OneInputStreamOperatorTestHarness<>(operator, IntSerializer.INSTANCE, mockEnvironment);

	final long initialTime = 0L;
	final ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();

	testHarness.open();

	testHarness.setProcessingTime(initialTime);

	synchronized (testHarness.getCheckpointLock()) {
		testHarness.processElement(new StreamRecord<>(1, initialTime));
		testHarness.setProcessingTime(initialTime + 5L);
		testHarness.processElement(new StreamRecord<>(2, initialTime + 5L));
	}

	// trigger the timeout of the first stream record
	testHarness.setProcessingTime(initialTime + timeout + 1L);

	// allow the second async stream record to be processed
	lazyAsyncFunction.countDown();

	// wait until all async collectors in the buffer have been emitted out.
	synchronized (testHarness.getCheckpointLock()) {
		testHarness.close();
	}

	expectedOutput.addAll(Arrays.asList(expectedRecords));

	TestHarnessUtil.assertOutputEquals("Output with watermark was not correct.", expectedOutput, testHarness.getOutput());

	if (expectedException.isPresent()) {
		assertTrue(mockEnvironment.getActualExternalFailureCause().isPresent());
		assertTrue(ExceptionUtils.findThrowable(
			mockEnvironment.getActualExternalFailureCause().get(),
			expectedException.get()).isPresent());
	}
}
 

/**
 * Tests that in case of a failing AsyncCheckpointRunnable all operator snapshot results are
 * cancelled and all non partitioned state handles are discarded.
 */
@Test
public void testFailingAsyncCheckpointRunnable() throws Exception {
	final long checkpointId = 42L;
	final long timestamp = 1L;

	MockEnvironment mockEnvironment = new MockEnvironmentBuilder().build();
	StreamTask<?, ?> streamTask = spy(new EmptyStreamTask(mockEnvironment));
	CheckpointMetaData checkpointMetaData = new CheckpointMetaData(checkpointId, timestamp);

	// mock the operators
	StreamOperator<?> streamOperator1 = mock(StreamOperator.class);
	StreamOperator<?> streamOperator2 = mock(StreamOperator.class);
	StreamOperator<?> streamOperator3 = mock(StreamOperator.class);

	// mock the new state operator snapshots
	OperatorSnapshotFutures operatorSnapshotResult1 = mock(OperatorSnapshotFutures.class);
	OperatorSnapshotFutures operatorSnapshotResult2 = mock(OperatorSnapshotFutures.class);
	OperatorSnapshotFutures operatorSnapshotResult3 = mock(OperatorSnapshotFutures.class);

	RunnableFuture<SnapshotResult<OperatorStateHandle>> failingFuture = mock(RunnableFuture.class);
	when(failingFuture.get()).thenThrow(new ExecutionException(new Exception("Test exception")));

	when(operatorSnapshotResult3.getOperatorStateRawFuture()).thenReturn(failingFuture);

	when(streamOperator1.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class), any(CheckpointStreamFactory.class))).thenReturn(operatorSnapshotResult1);
	when(streamOperator2.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class), any(CheckpointStreamFactory.class))).thenReturn(operatorSnapshotResult2);
	when(streamOperator3.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class), any(CheckpointStreamFactory.class))).thenReturn(operatorSnapshotResult3);

	OperatorID operatorID1 = new OperatorID();
	OperatorID operatorID2 = new OperatorID();
	OperatorID operatorID3 = new OperatorID();
	when(streamOperator1.getOperatorID()).thenReturn(operatorID1);
	when(streamOperator2.getOperatorID()).thenReturn(operatorID2);
	when(streamOperator3.getOperatorID()).thenReturn(operatorID3);

	StreamOperator<?>[] streamOperators = {streamOperator1, streamOperator2, streamOperator3};

	OperatorChain<Void, AbstractStreamOperator<Void>> operatorChain = mock(OperatorChain.class);
	when(operatorChain.getAllOperators()).thenReturn(streamOperators);

	Whitebox.setInternalState(streamTask, "isRunning", true);
	Whitebox.setInternalState(streamTask, "lock", new Object());
	Whitebox.setInternalState(streamTask, "operatorChain", operatorChain);
	Whitebox.setInternalState(streamTask, "cancelables", new CloseableRegistry());
	Whitebox.setInternalState(streamTask, "asyncOperationsThreadPool", newDirectExecutorService());
	Whitebox.setInternalState(streamTask, "configuration", new StreamConfig(new Configuration()));
	Whitebox.setInternalState(streamTask, "checkpointStorage", new MemoryBackendCheckpointStorage(new JobID(), null, null, Integer.MAX_VALUE));

	CheckpointExceptionHandlerFactory checkpointExceptionHandlerFactory = new CheckpointExceptionHandlerFactory();
	CheckpointExceptionHandler checkpointExceptionHandler =
		checkpointExceptionHandlerFactory.createCheckpointExceptionHandler(true, mockEnvironment);
	Whitebox.setInternalState(streamTask, "synchronousCheckpointExceptionHandler", checkpointExceptionHandler);

	StreamTask.AsyncCheckpointExceptionHandler asyncCheckpointExceptionHandler =
		new StreamTask.AsyncCheckpointExceptionHandler(streamTask);
	Whitebox.setInternalState(streamTask, "asynchronousCheckpointExceptionHandler", asyncCheckpointExceptionHandler);

	mockEnvironment.setExpectedExternalFailureCause(Throwable.class);
	streamTask.triggerCheckpoint(checkpointMetaData, CheckpointOptions.forCheckpointWithDefaultLocation());

	verify(streamTask).handleAsyncException(anyString(), any(Throwable.class));

	verify(operatorSnapshotResult1).cancel();
	verify(operatorSnapshotResult2).cancel();
	verify(operatorSnapshotResult3).cancel();
}
 

private void testAsyncTimeout(LazyAsyncFunction lazyAsyncFunction,
		Optional<Class<? extends Throwable>> expectedException,
		StreamRecord<Integer>... expectedRecords) throws Exception {
	final long timeout = 10L;

	final AsyncWaitOperator<Integer, Integer> operator = new AsyncWaitOperator<>(
		lazyAsyncFunction,
		timeout,
		2,
		AsyncDataStream.OutputMode.ORDERED);

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

	final OneInputStreamOperatorTestHarness<Integer, Integer> testHarness =
		new OneInputStreamOperatorTestHarness<>(operator, IntSerializer.INSTANCE, mockEnvironment);

	final long initialTime = 0L;
	final ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();

	testHarness.open();

	testHarness.setProcessingTime(initialTime);

	synchronized (testHarness.getCheckpointLock()) {
		testHarness.processElement(new StreamRecord<>(1, initialTime));
		testHarness.setProcessingTime(initialTime + 5L);
		testHarness.processElement(new StreamRecord<>(2, initialTime + 5L));
	}

	// trigger the timeout of the first stream record
	testHarness.setProcessingTime(initialTime + timeout + 1L);

	// allow the second async stream record to be processed
	lazyAsyncFunction.countDown();

	// wait until all async collectors in the buffer have been emitted out.
	synchronized (testHarness.getCheckpointLock()) {
		testHarness.close();
	}

	expectedOutput.addAll(Arrays.asList(expectedRecords));

	TestHarnessUtil.assertOutputEquals("Output with watermark was not correct.", expectedOutput, testHarness.getOutput());

	if (expectedException.isPresent()) {
		assertTrue(mockEnvironment.getActualExternalFailureCause().isPresent());
		assertTrue(ExceptionUtils.findThrowable(
			mockEnvironment.getActualExternalFailureCause().get(),
			expectedException.get()).isPresent());
	}
}
 

private void testAsyncTimeout(
		LazyAsyncFunction lazyAsyncFunction,
		Optional<Class<? extends Throwable>> expectedException,
		StreamRecord<Integer>... expectedRecords) throws Exception {
	final long timeout = 10L;

	final OneInputStreamOperatorTestHarness<Integer, Integer> testHarness =
		createTestHarness(lazyAsyncFunction, timeout, 2, AsyncDataStream.OutputMode.ORDERED);

	final MockEnvironment mockEnvironment = testHarness.getEnvironment();
	mockEnvironment.setExpectedExternalFailureCause(Throwable.class);

	final long initialTime = 0L;
	final ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();

	testHarness.open();

	testHarness.setProcessingTime(initialTime);

	synchronized (testHarness.getCheckpointLock()) {
		testHarness.processElement(new StreamRecord<>(1, initialTime));
		testHarness.setProcessingTime(initialTime + 5L);
		testHarness.processElement(new StreamRecord<>(2, initialTime + 5L));
	}

	// trigger the timeout of the first stream record
	testHarness.setProcessingTime(initialTime + timeout + 1L);

	// allow the second async stream record to be processed
	lazyAsyncFunction.countDown();

	// wait until all async collectors in the buffer have been emitted out.
	synchronized (testHarness.getCheckpointLock()) {
		testHarness.endInput();
		testHarness.close();
	}

	expectedOutput.addAll(Arrays.asList(expectedRecords));

	TestHarnessUtil.assertOutputEquals("Output with watermark was not correct.", expectedOutput, testHarness.getOutput());

	if (expectedException.isPresent()) {
		assertTrue(mockEnvironment.getActualExternalFailureCause().isPresent());
		assertTrue(ExceptionUtils.findThrowable(
			mockEnvironment.getActualExternalFailureCause().get(),
			expectedException.get()).isPresent());
	}
}