Examples with OperatorSnapshotResult org.apache.flink.streaming.api.operators.OperatorSnapshotResult used on opensource projects

Example 1 with OperatorSnapshotResult

use of org.apache.flink.streaming.api.operators.OperatorSnapshotResult in project flink by apache.

the class AbstractStreamOperatorTestHarness method snapshot.

/**
	 * Calls {@link StreamOperator#snapshotState(long, long, CheckpointOptions)}.
	 */
public OperatorStateHandles snapshot(long checkpointId, long timestamp) throws Exception {
    CheckpointStreamFactory streamFactory = stateBackend.createStreamFactory(new JobID(), "test_op");
    OperatorSnapshotResult operatorStateResult = operator.snapshotState(checkpointId, timestamp, CheckpointOptions.forFullCheckpoint());
    KeyGroupsStateHandle keyedManaged = FutureUtil.runIfNotDoneAndGet(operatorStateResult.getKeyedStateManagedFuture());
    KeyGroupsStateHandle keyedRaw = FutureUtil.runIfNotDoneAndGet(operatorStateResult.getKeyedStateRawFuture());
    OperatorStateHandle opManaged = FutureUtil.runIfNotDoneAndGet(operatorStateResult.getOperatorStateManagedFuture());
    OperatorStateHandle opRaw = FutureUtil.runIfNotDoneAndGet(operatorStateResult.getOperatorStateRawFuture());
    return new OperatorStateHandles(0, null, keyedManaged != null ? Collections.singletonList(keyedManaged) : null, keyedRaw != null ? Collections.singletonList(keyedRaw) : null, opManaged != null ? Collections.singletonList(opManaged) : null, opRaw != null ? Collections.singletonList(opRaw) : null);
}

Example 2 with OperatorSnapshotResult

use of org.apache.flink.streaming.api.operators.OperatorSnapshotResult in project flink by apache.

the class StreamTaskTest method testFailingCheckpointStreamOperator.

@Test
public void testFailingCheckpointStreamOperator() throws Exception {
    final long checkpointId = 42L;
    final long timestamp = 1L;
    TaskInfo mockTaskInfo = mock(TaskInfo.class);
    when(mockTaskInfo.getTaskNameWithSubtasks()).thenReturn("foobar");
    when(mockTaskInfo.getIndexOfThisSubtask()).thenReturn(0);
    Environment mockEnvironment = mock(Environment.class);
    when(mockEnvironment.getTaskInfo()).thenReturn(mockTaskInfo);
    StreamTask<?, AbstractStreamOperator<?>> streamTask = mock(StreamTask.class, Mockito.CALLS_REAL_METHODS);
    CheckpointMetaData checkpointMetaData = new CheckpointMetaData(checkpointId, timestamp);
    streamTask.setEnvironment(mockEnvironment);
    // mock the operators
    StreamOperator<?> streamOperator1 = mock(StreamOperator.class, withSettings().extraInterfaces(StreamCheckpointedOperator.class));
    StreamOperator<?> streamOperator2 = mock(StreamOperator.class, withSettings().extraInterfaces(StreamCheckpointedOperator.class));
    StreamOperator<?> streamOperator3 = mock(StreamOperator.class, withSettings().extraInterfaces(StreamCheckpointedOperator.class));
    // mock the returned snapshots
    OperatorSnapshotResult operatorSnapshotResult1 = mock(OperatorSnapshotResult.class);
    OperatorSnapshotResult operatorSnapshotResult2 = mock(OperatorSnapshotResult.class);
    final Exception testException = new Exception("Test exception");
    when(streamOperator1.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class))).thenReturn(operatorSnapshotResult1);
    when(streamOperator2.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class))).thenReturn(operatorSnapshotResult2);
    when(streamOperator3.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class))).thenThrow(testException);
    // mock the returned legacy snapshots
    StreamStateHandle streamStateHandle1 = mock(StreamStateHandle.class);
    StreamStateHandle streamStateHandle2 = mock(StreamStateHandle.class);
    StreamStateHandle streamStateHandle3 = mock(StreamStateHandle.class);
    when(streamOperator1.snapshotLegacyOperatorState(anyLong(), anyLong(), any(CheckpointOptions.class))).thenReturn(streamStateHandle1);
    when(streamOperator2.snapshotLegacyOperatorState(anyLong(), anyLong(), any(CheckpointOptions.class))).thenReturn(streamStateHandle2);
    when(streamOperator3.snapshotLegacyOperatorState(anyLong(), anyLong(), any(CheckpointOptions.class))).thenReturn(streamStateHandle3);
    // set up the task
    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, "configuration", new StreamConfig(new Configuration()));
    try {
        streamTask.triggerCheckpoint(checkpointMetaData, CheckpointOptions.forFullCheckpoint());
        fail("Expected test exception here.");
    } catch (Exception e) {
        assertEquals(testException, e.getCause());
    }
    verify(operatorSnapshotResult1).cancel();
    verify(operatorSnapshotResult2).cancel();
    verify(streamStateHandle1).discardState();
    verify(streamStateHandle2).discardState();
    verify(streamStateHandle3).discardState();
}