Examples with OperatorSubtaskState org.apache.flink.runtime.checkpoint.OperatorSubtaskState used on opensource projects

Example 36 with OperatorSubtaskState

use of org.apache.flink.runtime.checkpoint.OperatorSubtaskState in project flink by apache.

the class KeyedStateBootstrapOperatorTest method getState.

private OperatorSubtaskState getState(KeyedStateBootstrapOperator<Long, Long> bootstrapOperator, KeyedOneInputStreamOperatorTestHarness<Long, Long, TaggedOperatorSubtaskState> harness) throws Exception {
    OperatorSubtaskState state;
    bootstrapOperator.endInput();
    state = harness.extractOutputValues().get(0).state;
    return state;
}

Example 37 with OperatorSubtaskState

use of org.apache.flink.runtime.checkpoint.OperatorSubtaskState in project flink by apache.

the class KeyedProcessOperatorTest method testSnapshotAndRestore.

@Test
public void testSnapshotAndRestore() throws Exception {
    final int expectedKey = 5;
    KeyedProcessOperator<Integer, Integer, String> operator = new KeyedProcessOperator<>(new BothTriggeringFlatMapFunction(expectedKey));
    OneInputStreamOperatorTestHarness<Integer, String> testHarness = new KeyedOneInputStreamOperatorTestHarness<>(operator, new IdentityKeySelector<Integer>(), BasicTypeInfo.INT_TYPE_INFO);
    testHarness.setup();
    testHarness.open();
    testHarness.processElement(new StreamRecord<>(expectedKey, 12L));
    // snapshot and restore from scratch
    OperatorSubtaskState snapshot = testHarness.snapshot(0, 0);
    testHarness.close();
    operator = new KeyedProcessOperator<>(new BothTriggeringFlatMapFunction(expectedKey));
    testHarness = new KeyedOneInputStreamOperatorTestHarness<>(operator, new IdentityKeySelector<Integer>(), BasicTypeInfo.INT_TYPE_INFO);
    testHarness.setup();
    testHarness.initializeState(snapshot);
    testHarness.open();
    testHarness.setProcessingTime(5);
    testHarness.processWatermark(new Watermark(6));
    ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
    expectedOutput.add(new StreamRecord<>("PROC:1777"));
    expectedOutput.add(new StreamRecord<>("EVENT:1777", 6L));
    expectedOutput.add(new Watermark(6));
    TestHarnessUtil.assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());
    testHarness.close();
}

Example 38 with OperatorSubtaskState

use of org.apache.flink.runtime.checkpoint.OperatorSubtaskState in project flink by apache.

the class OperatorSnapshotFinalizerTest method testRunAndExtract.

/**
 * Test that the runnable futures are executed and the result is correctly extracted.
 */
@Test
public void testRunAndExtract() throws Exception {
    Random random = new Random(0x42);
    KeyedStateHandle keyedTemplate = StateHandleDummyUtil.createNewKeyedStateHandle(new KeyGroupRange(0, 0));
    OperatorStateHandle operatorTemplate = StateHandleDummyUtil.createNewOperatorStateHandle(2, random);
    InputChannelStateHandle inputChannelTemplate = StateHandleDummyUtil.createNewInputChannelStateHandle(2, random);
    ResultSubpartitionStateHandle resultSubpartitionTemplate = StateHandleDummyUtil.createNewResultSubpartitionStateHandle(2, random);
    SnapshotResult<KeyedStateHandle> manKeyed = withLocalState(deepDummyCopy(keyedTemplate), deepDummyCopy(keyedTemplate));
    SnapshotResult<KeyedStateHandle> rawKeyed = withLocalState(deepDummyCopy(keyedTemplate), deepDummyCopy(keyedTemplate));
    SnapshotResult<OperatorStateHandle> manOper = withLocalState(deepDummyCopy(operatorTemplate), deepDummyCopy(operatorTemplate));
    SnapshotResult<OperatorStateHandle> rawOper = withLocalState(deepDummyCopy(operatorTemplate), deepDummyCopy(operatorTemplate));
    SnapshotResult<StateObjectCollection<InputChannelStateHandle>> inputChannel = withLocalState(singleton(deepDummyCopy(inputChannelTemplate)), singleton(deepDummyCopy(inputChannelTemplate)));
    SnapshotResult<StateObjectCollection<ResultSubpartitionStateHandle>> resultSubpartition = withLocalState(singleton(deepDummyCopy(resultSubpartitionTemplate)), singleton(deepDummyCopy(resultSubpartitionTemplate)));
    OperatorSnapshotFutures snapshotFutures = new OperatorSnapshotFutures(new PseudoNotDoneFuture<>(manKeyed), new PseudoNotDoneFuture<>(rawKeyed), new PseudoNotDoneFuture<>(manOper), new PseudoNotDoneFuture<>(rawOper), new PseudoNotDoneFuture<>(inputChannel), new PseudoNotDoneFuture<>(resultSubpartition));
    for (Future<?> f : snapshotFutures.getAllFutures()) {
        assertFalse(f.isDone());
    }
    OperatorSnapshotFinalizer finalizer = new OperatorSnapshotFinalizer(snapshotFutures);
    for (Future<?> f : snapshotFutures.getAllFutures()) {
        assertTrue(f.isDone());
    }
    Map<SnapshotResult<?>, Function<OperatorSubtaskState, ? extends StateObject>> map = new HashMap<>();
    map.put(manKeyed, headExtractor(OperatorSubtaskState::getManagedKeyedState));
    map.put(rawKeyed, headExtractor(OperatorSubtaskState::getRawKeyedState));
    map.put(manOper, headExtractor(OperatorSubtaskState::getManagedOperatorState));
    map.put(rawOper, headExtractor(OperatorSubtaskState::getRawOperatorState));
    map.put(inputChannel, OperatorSubtaskState::getInputChannelState);
    map.put(resultSubpartition, OperatorSubtaskState::getResultSubpartitionState);
    for (Map.Entry<SnapshotResult<?>, Function<OperatorSubtaskState, ? extends StateObject>> e : map.entrySet()) {
        assertEquals(e.getKey().getJobManagerOwnedSnapshot(), e.getValue().apply(finalizer.getJobManagerOwnedState()));
    }
    for (Map.Entry<SnapshotResult<?>, Function<OperatorSubtaskState, ? extends StateObject>> e : map.entrySet()) {
        assertEquals(e.getKey().getTaskLocalSnapshot(), e.getValue().apply(finalizer.getTaskLocalState()));
    }
}

Example 39 with OperatorSubtaskState

use of org.apache.flink.runtime.checkpoint.OperatorSubtaskState in project flink by apache.

the class StateInitializationContextImplTest method setUp.

@Before
public void setUp() throws Exception {
    this.writtenKeyGroups = 0;
    this.writtenOperatorStates = new HashSet<>();
    this.closableRegistry = new CloseableRegistry();
    ByteArrayOutputStreamWithPos out = new ByteArrayOutputStreamWithPos(64);
    List<KeyedStateHandle> keyedStateHandles = new ArrayList<>(NUM_HANDLES);
    int prev = 0;
    for (int i = 0; i < NUM_HANDLES; ++i) {
        out.reset();
        int size = i % 4;
        int end = prev + size;
        DataOutputView dov = new DataOutputViewStreamWrapper(out);
        KeyGroupRangeOffsets offsets = new KeyGroupRangeOffsets(i == 9 ? KeyGroupRange.EMPTY_KEY_GROUP_RANGE : new KeyGroupRange(prev, end));
        prev = end + 1;
        for (int kg : offsets.getKeyGroupRange()) {
            offsets.setKeyGroupOffset(kg, out.getPosition());
            dov.writeInt(kg);
            ++writtenKeyGroups;
        }
        KeyedStateHandle handle = new KeyGroupsStateHandle(offsets, new ByteStateHandleCloseChecking("kg-" + i, out.toByteArray()));
        keyedStateHandles.add(handle);
    }
    List<OperatorStateHandle> operatorStateHandles = new ArrayList<>(NUM_HANDLES);
    for (int i = 0; i < NUM_HANDLES; ++i) {
        int size = i % 4;
        out.reset();
        DataOutputView dov = new DataOutputViewStreamWrapper(out);
        LongArrayList offsets = new LongArrayList(size);
        for (int s = 0; s < size; ++s) {
            offsets.add(out.getPosition());
            int val = i * NUM_HANDLES + s;
            dov.writeInt(val);
            writtenOperatorStates.add(val);
        }
        Map<String, OperatorStateHandle.StateMetaInfo> offsetsMap = new HashMap<>();
        offsetsMap.put(DefaultOperatorStateBackend.DEFAULT_OPERATOR_STATE_NAME, new OperatorStateHandle.StateMetaInfo(offsets.toArray(), OperatorStateHandle.Mode.SPLIT_DISTRIBUTE));
        OperatorStateHandle operatorStateHandle = new OperatorStreamStateHandle(offsetsMap, new ByteStateHandleCloseChecking("os-" + i, out.toByteArray()));
        operatorStateHandles.add(operatorStateHandle);
    }
    OperatorSubtaskState operatorSubtaskState = OperatorSubtaskState.builder().setRawOperatorState(new StateObjectCollection<>(operatorStateHandles)).setRawKeyedState(new StateObjectCollection<>(keyedStateHandles)).build();
    OperatorID operatorID = new OperatorID();
    TaskStateSnapshot taskStateSnapshot = new TaskStateSnapshot();
    taskStateSnapshot.putSubtaskStateByOperatorID(operatorID, operatorSubtaskState);
    JobManagerTaskRestore jobManagerTaskRestore = new JobManagerTaskRestore(0L, taskStateSnapshot);
    TaskStateManager manager = new TaskStateManagerImpl(new JobID(), new ExecutionAttemptID(), new TestTaskLocalStateStore(), new InMemoryStateChangelogStorage(), jobManagerTaskRestore, mock(CheckpointResponder.class));
    DummyEnvironment environment = new DummyEnvironment("test", 1, 0, prev);
    environment.setTaskStateManager(manager);
    StateBackend stateBackend = new MemoryStateBackend(1024);
    StreamTaskStateInitializer streamTaskStateManager = new StreamTaskStateInitializerImpl(environment, stateBackend, TtlTimeProvider.DEFAULT, new InternalTimeServiceManager.Provider() {

        @Override
        public <K> InternalTimeServiceManager<K> create(CheckpointableKeyedStateBackend<K> keyedStatedBackend, ClassLoader userClassloader, KeyContext keyContext, ProcessingTimeService processingTimeService, Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates) throws Exception {
            // stream.
            return null;
        }
    });
    AbstractStreamOperator<?> mockOperator = mock(AbstractStreamOperator.class);
    when(mockOperator.getOperatorID()).thenReturn(operatorID);
    StreamOperatorStateContext stateContext = streamTaskStateManager.streamOperatorStateContext(operatorID, "TestOperatorClass", mock(ProcessingTimeService.class), mockOperator, // consumed by the timer service.
    IntSerializer.INSTANCE, closableRegistry, new UnregisteredMetricsGroup(), 1.0, false);
    OptionalLong restoredCheckpointId = stateContext.getRestoredCheckpointId();
    this.initializationContext = new StateInitializationContextImpl(restoredCheckpointId.isPresent() ? restoredCheckpointId.getAsLong() : null, stateContext.operatorStateBackend(), mock(KeyedStateStore.class), stateContext.rawKeyedStateInputs(), stateContext.rawOperatorStateInputs());
}

Example 40 with OperatorSubtaskState

use of org.apache.flink.runtime.checkpoint.OperatorSubtaskState in project flink by apache.

the class AsyncWaitOperatorTest method testRestartWithFullQueue.

/**
 * Tests that the AsyncWaitOperator can restart if checkpointed queue was full.
 *
 * <p>See FLINK-7949
 */
@Test(timeout = 10000)
public void testRestartWithFullQueue() throws Exception {
    final int capacity = 10;
    // 1. create the snapshot which contains capacity + 1 elements
    final CompletableFuture<Void> trigger = new CompletableFuture<>();
    final OneInputStreamOperatorTestHarness<Integer, Integer> snapshotHarness = createTestHarness(new ControllableAsyncFunction<>(// the NoOpAsyncFunction is like a blocking function
    trigger), 1000L, capacity, AsyncDataStream.OutputMode.ORDERED);
    snapshotHarness.open();
    final OperatorSubtaskState snapshot;
    final ArrayList<Integer> expectedOutput = new ArrayList<>(capacity);
    try {
        synchronized (snapshotHarness.getCheckpointLock()) {
            for (int i = 0; i < capacity; i++) {
                snapshotHarness.processElement(i, 0L);
                expectedOutput.add(i);
            }
        }
        synchronized (snapshotHarness.getCheckpointLock()) {
            // execute the snapshot within the checkpoint lock, because then it is guaranteed
            // that the lastElementWriter has written the exceeding element
            snapshot = snapshotHarness.snapshot(0L, 0L);
        }
        // trigger the computation to make the close call finish
        trigger.complete(null);
    } finally {
        synchronized (snapshotHarness.getCheckpointLock()) {
            snapshotHarness.close();
        }
    }
    // 2. restore the snapshot and check that we complete
    final OneInputStreamOperatorTestHarness<Integer, Integer> recoverHarness = createTestHarness(new ControllableAsyncFunction<>(CompletableFuture.completedFuture(null)), 1000L, capacity, AsyncDataStream.OutputMode.ORDERED);
    recoverHarness.initializeState(snapshot);
    synchronized (recoverHarness.getCheckpointLock()) {
        recoverHarness.open();
    }
    synchronized (recoverHarness.getCheckpointLock()) {
        recoverHarness.endInput();
        recoverHarness.close();
    }
    final ConcurrentLinkedQueue<Object> output = recoverHarness.getOutput();
    final List<Integer> outputElements = output.stream().map(r -> ((StreamRecord<Integer>) r).getValue()).collect(Collectors.toList());
    assertThat(outputElements, Matchers.equalTo(expectedOutput));
}