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

Example 1 with OperatorStateBackend

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

use of org.apache.flink.runtime.state.OperatorStateBackend in project beam by apache.

the class FlinkBroadcastStateInternalsTest method initStateInternals.

@Before
public void initStateInternals() {
    MemoryStateBackend backend = new MemoryStateBackend();
    try {
        OperatorStateBackend operatorStateBackend = backend.createOperatorStateBackend(new DummyEnvironment("test", 1, 0), "");
        underTest = new FlinkBroadcastStateInternals<>(1, operatorStateBackend);
    } catch (Exception e) {
        throw new RuntimeException(e);
    }
}

Example 3 with OperatorStateBackend

use of org.apache.flink.runtime.state.OperatorStateBackend in project beam by apache.

the class FlinkSplitStateInternalsTest method initStateInternals.

@Before
public void initStateInternals() {
    MemoryStateBackend backend = new MemoryStateBackend();
    try {
        OperatorStateBackend operatorStateBackend = backend.createOperatorStateBackend(new DummyEnvironment("test", 1, 0), "");
        underTest = new FlinkSplitStateInternals<>(operatorStateBackend);
    } catch (Exception e) {
        throw new RuntimeException(e);
    }
}

Example 4 with OperatorStateBackend

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

the class StreamTask method createOperatorStateBackend.

public OperatorStateBackend createOperatorStateBackend(StreamOperator<?> op, Collection<OperatorStateHandle> restoreStateHandles) throws Exception {
    Environment env = getEnvironment();
    String opId = createOperatorIdentifier(op, getConfiguration().getVertexID());
    OperatorStateBackend operatorStateBackend = stateBackend.createOperatorStateBackend(env, opId);
    // let operator state backend participate in the operator lifecycle, i.e. make it responsive to cancelation
    cancelables.registerClosable(operatorStateBackend);
    // restore if we have some old state
    if (null != restoreStateHandles) {
        operatorStateBackend.restore(restoreStateHandles);
    }
    return operatorStateBackend;
}