Examples with CheckpointStreamFactory org.apache.flink.runtime.state.CheckpointStreamFactory used on opensource projects

Example 1 with CheckpointStreamFactory

use of org.apache.flink.runtime.state.CheckpointStreamFactory in project flink by apache.

the class StreamTaskTest method testAsyncCheckpointingConcurrentCloseBeforeAcknowledge.

/**
	 * FLINK-5667
	 *
	 * Tests that a concurrent cancel operation discards the state handles of a not yet
	 * acknowledged checkpoint and prevents sending an acknowledge message to the
	 * CheckpointCoordinator. The situation can only happen if the cancel call is executed
	 * before Environment.acknowledgeCheckpoint().
	 */
@Test
public void testAsyncCheckpointingConcurrentCloseBeforeAcknowledge() throws Exception {
    final long checkpointId = 42L;
    final long timestamp = 1L;
    final OneShotLatch createSubtask = new OneShotLatch();
    final OneShotLatch completeSubtask = new OneShotLatch();
    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);
    whenNew(SubtaskState.class).withAnyArguments().thenAnswer(new Answer<SubtaskState>() {

        @Override
        public SubtaskState answer(InvocationOnMock invocation) throws Throwable {
            createSubtask.trigger();
            completeSubtask.await();
            return new SubtaskState((ChainedStateHandle<StreamStateHandle>) invocation.getArguments()[0], (ChainedStateHandle<OperatorStateHandle>) invocation.getArguments()[1], (ChainedStateHandle<OperatorStateHandle>) invocation.getArguments()[2], (KeyGroupsStateHandle) invocation.getArguments()[3], (KeyGroupsStateHandle) invocation.getArguments()[4]);
        }
    });
    StreamTask<?, AbstractStreamOperator<?>> streamTask = mock(StreamTask.class, Mockito.CALLS_REAL_METHODS);
    CheckpointMetaData checkpointMetaData = new CheckpointMetaData(checkpointId, timestamp);
    streamTask.setEnvironment(mockEnvironment);
    StreamOperator<?> streamOperator = mock(StreamOperator.class, withSettings().extraInterfaces(StreamCheckpointedOperator.class));
    KeyGroupsStateHandle managedKeyedStateHandle = mock(KeyGroupsStateHandle.class);
    KeyGroupsStateHandle rawKeyedStateHandle = mock(KeyGroupsStateHandle.class);
    OperatorStateHandle managedOperatorStateHandle = mock(OperatorStateHandle.class);
    OperatorStateHandle rawOperatorStateHandle = mock(OperatorStateHandle.class);
    OperatorSnapshotResult operatorSnapshotResult = new OperatorSnapshotResult(new DoneFuture<>(managedKeyedStateHandle), new DoneFuture<>(rawKeyedStateHandle), new DoneFuture<>(managedOperatorStateHandle), new DoneFuture<>(rawOperatorStateHandle));
    when(streamOperator.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class))).thenReturn(operatorSnapshotResult);
    StreamOperator<?>[] streamOperators = { streamOperator };
    OperatorChain<Void, AbstractStreamOperator<Void>> operatorChain = mock(OperatorChain.class);
    when(operatorChain.getAllOperators()).thenReturn(streamOperators);
    StreamStateHandle streamStateHandle = mock(StreamStateHandle.class);
    CheckpointStreamFactory.CheckpointStateOutputStream outStream = mock(CheckpointStreamFactory.CheckpointStateOutputStream.class);
    when(outStream.closeAndGetHandle()).thenReturn(streamStateHandle);
    CheckpointStreamFactory mockStreamFactory = mock(CheckpointStreamFactory.class);
    when(mockStreamFactory.createCheckpointStateOutputStream(anyLong(), anyLong())).thenReturn(outStream);
    AbstractStateBackend mockStateBackend = mock(AbstractStateBackend.class);
    when(mockStateBackend.createStreamFactory(any(JobID.class), anyString())).thenReturn(mockStreamFactory);
    ExecutorService executor = Executors.newFixedThreadPool(1);
    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", executor);
    Whitebox.setInternalState(streamTask, "configuration", new StreamConfig(new Configuration()));
    Whitebox.setInternalState(streamTask, "stateBackend", mockStateBackend);
    streamTask.triggerCheckpoint(checkpointMetaData, CheckpointOptions.forFullCheckpoint());
    createSubtask.await();
    streamTask.cancel();
    completeSubtask.trigger();
    // wait for the completion of the async task
    executor.shutdown();
    if (!executor.awaitTermination(10000L, TimeUnit.MILLISECONDS)) {
        fail("Executor did not shut down within the given timeout. This indicates that the " + "checkpointing did not resume.");
    }
    // check that the checkpoint has not been acknowledged
    verify(mockEnvironment, never()).acknowledgeCheckpoint(eq(checkpointId), any(CheckpointMetrics.class), any(SubtaskState.class));
    // check that the state handles have been discarded
    verify(managedKeyedStateHandle).discardState();
    verify(rawKeyedStateHandle).discardState();
    verify(managedOperatorStateHandle).discardState();
    verify(rawOperatorStateHandle).discardState();
}

Example 2 with CheckpointStreamFactory

use of org.apache.flink.runtime.state.CheckpointStreamFactory 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 3 with CheckpointStreamFactory

use of org.apache.flink.runtime.state.CheckpointStreamFactory in project flink by apache.

the class AbstractStreamOperator method snapshotLegacyOperatorState.

@SuppressWarnings("deprecation")
@Deprecated
@Override
public StreamStateHandle snapshotLegacyOperatorState(long checkpointId, long timestamp, CheckpointOptions checkpointOptions) throws Exception {
    if (this instanceof StreamCheckpointedOperator) {
        CheckpointStreamFactory factory = getCheckpointStreamFactory(checkpointOptions);
        final CheckpointStreamFactory.CheckpointStateOutputStream outStream = factory.createCheckpointStateOutputStream(checkpointId, timestamp);
        getContainingTask().getCancelables().registerClosable(outStream);
        try {
            ((StreamCheckpointedOperator) this).snapshotState(outStream, checkpointId, timestamp);
            return outStream.closeAndGetHandle();
        } finally {
            getContainingTask().getCancelables().unregisterClosable(outStream);
            outStream.close();
        }
    } else {
        return null;
    }
}

Example 4 with CheckpointStreamFactory

use of org.apache.flink.runtime.state.CheckpointStreamFactory in project flink by apache.

the class AbstractStreamOperatorTest method testFailingBackendSnapshotMethod.

/**
	 * Tests that a failing snapshot method call to the keyed state backend will trigger the closing
	 * of the StateSnapshotContextSynchronousImpl and the cancellation of the
	 * OperatorSnapshotResult. The latter is supposed to also cancel all assigned futures.
	 */
@Test
public void testFailingBackendSnapshotMethod() throws Exception {
    final long checkpointId = 42L;
    final long timestamp = 1L;
    final Exception failingException = new Exception("Test exception");
    final CloseableRegistry closeableRegistry = new CloseableRegistry();
    RunnableFuture<KeyGroupsStateHandle> futureKeyGroupStateHandle = mock(RunnableFuture.class);
    RunnableFuture<OperatorStateHandle> futureOperatorStateHandle = mock(RunnableFuture.class);
    StateSnapshotContextSynchronousImpl context = mock(StateSnapshotContextSynchronousImpl.class);
    when(context.getKeyedStateStreamFuture()).thenReturn(futureKeyGroupStateHandle);
    when(context.getOperatorStateStreamFuture()).thenReturn(futureOperatorStateHandle);
    OperatorSnapshotResult operatorSnapshotResult = spy(new OperatorSnapshotResult());
    whenNew(StateSnapshotContextSynchronousImpl.class).withAnyArguments().thenReturn(context);
    whenNew(OperatorSnapshotResult.class).withAnyArguments().thenReturn(operatorSnapshotResult);
    CheckpointStreamFactory streamFactory = mock(CheckpointStreamFactory.class);
    StreamTask<Void, AbstractStreamOperator<Void>> containingTask = mock(StreamTask.class);
    when(containingTask.getCancelables()).thenReturn(closeableRegistry);
    AbstractStreamOperator<Void> operator = mock(AbstractStreamOperator.class);
    when(operator.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class))).thenCallRealMethod();
    // The amount of mocking in this test makes it necessary to make the
    // getCheckpointStreamFactory method visible for the test and to
    // overwrite its behaviour.
    when(operator.getCheckpointStreamFactory(any(CheckpointOptions.class))).thenReturn(streamFactory);
    doReturn(containingTask).when(operator).getContainingTask();
    RunnableFuture<OperatorStateHandle> futureManagedOperatorStateHandle = mock(RunnableFuture.class);
    OperatorStateBackend operatorStateBackend = mock(OperatorStateBackend.class);
    when(operatorStateBackend.snapshot(eq(checkpointId), eq(timestamp), eq(streamFactory), any(CheckpointOptions.class))).thenReturn(futureManagedOperatorStateHandle);
    AbstractKeyedStateBackend<?> keyedStateBackend = mock(AbstractKeyedStateBackend.class);
    when(keyedStateBackend.snapshot(eq(checkpointId), eq(timestamp), eq(streamFactory), eq(CheckpointOptions.forFullCheckpoint()))).thenThrow(failingException);
    Whitebox.setInternalState(operator, "operatorStateBackend", operatorStateBackend);
    Whitebox.setInternalState(operator, "keyedStateBackend", keyedStateBackend);
    Whitebox.setInternalState(operator, "checkpointStreamFactory", streamFactory);
    try {
        operator.snapshotState(checkpointId, timestamp, CheckpointOptions.forFullCheckpoint());
        fail("Exception expected.");
    } catch (Exception e) {
        assertEquals(failingException, e.getCause());
    }
    // verify that the context has been closed, the operator snapshot result has been cancelled
    // and that all futures have been cancelled.
    verify(context).close();
    verify(operatorSnapshotResult).cancel();
    verify(futureKeyGroupStateHandle).cancel(anyBoolean());
    verify(futureOperatorStateHandle).cancel(anyBoolean());
    verify(futureKeyGroupStateHandle).cancel(anyBoolean());
}

Example 5 with CheckpointStreamFactory

use of org.apache.flink.runtime.state.CheckpointStreamFactory in project flink by apache.

the class StateSnapshotContextSynchronousImplTest method setUp.

@Before
public void setUp() throws Exception {
    CloseableRegistry closableRegistry = new CloseableRegistry();
    CheckpointStreamFactory streamFactory = new MemCheckpointStreamFactory(1024);
    KeyGroupRange keyGroupRange = new KeyGroupRange(0, 2);
    this.snapshotContext = new StateSnapshotContextSynchronousImpl(42, 4711, streamFactory, keyGroupRange, closableRegistry);
}