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

Example 16 with KeyedStateHandle

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

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

the class OperatorSnapshotFuturesTest method testCancelAndCleanup.

/**
 * Tests that all runnable futures in an OperatorSnapshotResult are properly cancelled and if
 * the StreamStateHandle result is retrievable that the state handle are discarded.
 */
@Test
public void testCancelAndCleanup() throws Exception {
    OperatorSnapshotFutures operatorSnapshotResult = new OperatorSnapshotFutures();
    operatorSnapshotResult.cancel();
    KeyedStateHandle keyedManagedStateHandle = mock(KeyedStateHandle.class);
    SnapshotResult<KeyedStateHandle> keyedStateManagedResult = SnapshotResult.of(keyedManagedStateHandle);
    RunnableFuture<SnapshotResult<KeyedStateHandle>> keyedStateManagedFuture = spy(DoneFuture.of(keyedStateManagedResult));
    KeyedStateHandle keyedRawStateHandle = mock(KeyedStateHandle.class);
    SnapshotResult<KeyedStateHandle> keyedStateRawResult = SnapshotResult.of(keyedRawStateHandle);
    RunnableFuture<SnapshotResult<KeyedStateHandle>> keyedStateRawFuture = spy(DoneFuture.of(keyedStateRawResult));
    OperatorStateHandle operatorManagedStateHandle = mock(OperatorStreamStateHandle.class);
    SnapshotResult<OperatorStateHandle> operatorStateManagedResult = SnapshotResult.of(operatorManagedStateHandle);
    RunnableFuture<SnapshotResult<OperatorStateHandle>> operatorStateManagedFuture = spy(DoneFuture.of(operatorStateManagedResult));
    OperatorStateHandle operatorRawStateHandle = mock(OperatorStreamStateHandle.class);
    SnapshotResult<OperatorStateHandle> operatorStateRawResult = SnapshotResult.of(operatorRawStateHandle);
    RunnableFuture<SnapshotResult<OperatorStateHandle>> operatorStateRawFuture = spy(DoneFuture.of(operatorStateRawResult));
    InputChannelStateHandle inputChannelRawStateHandle = mock(InputChannelStateHandle.class);
    SnapshotResult<StateObjectCollection<InputChannelStateHandle>> inputChannelStateRawResult = SnapshotResult.of(StateObjectCollection.singleton(inputChannelRawStateHandle));
    Future<SnapshotResult<StateObjectCollection<InputChannelStateHandle>>> inputChannelStateRawFuture = spy(DoneFuture.of(inputChannelStateRawResult));
    ResultSubpartitionStateHandle resultSubpartitionRawStateHandle = mock(ResultSubpartitionStateHandle.class);
    SnapshotResult<StateObjectCollection<ResultSubpartitionStateHandle>> resultSubpartitionStateRawResult = SnapshotResult.of(StateObjectCollection.singleton(resultSubpartitionRawStateHandle));
    Future<SnapshotResult<StateObjectCollection<ResultSubpartitionStateHandle>>> resultSubpartitionStateRawFuture = spy(DoneFuture.of(resultSubpartitionStateRawResult));
    operatorSnapshotResult = new OperatorSnapshotFutures(keyedStateManagedFuture, keyedStateRawFuture, operatorStateManagedFuture, operatorStateRawFuture, inputChannelStateRawFuture, resultSubpartitionStateRawFuture);
    operatorSnapshotResult.cancel();
    verify(keyedStateManagedFuture).cancel(true);
    verify(keyedStateRawFuture).cancel(true);
    verify(operatorStateManagedFuture).cancel(true);
    verify(operatorStateRawFuture).cancel(true);
    verify(inputChannelStateRawFuture).cancel(true);
    verify(resultSubpartitionStateRawFuture).cancel(true);
    verify(keyedManagedStateHandle).discardState();
    verify(keyedRawStateHandle).discardState();
    verify(operatorManagedStateHandle).discardState();
    verify(operatorRawStateHandle).discardState();
    verify(inputChannelRawStateHandle).discardState();
    verify(resultSubpartitionRawStateHandle).discardState();
}

Example 18 with KeyedStateHandle

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

the class SavepointStateBackendSwitchTestBase method takeSavepoint.

private void takeSavepoint(CheckpointableKeyedStateBackend<String> keyedBackend, File pathToWrite, MapStateDescriptor<Long, Long> stateDescr, ValueStateDescriptor<Long> valueStateDescriptor, ListStateDescriptor<Long> listStateDescriptor, Integer namespace1, Integer namespace2, Integer namespace3, Integer namespace4) throws Exception {
    InternalMapState<String, Integer, Long, Long> mapState = keyedBackend.createInternalState(IntSerializer.INSTANCE, stateDescr);
    InternalValueState<String, Integer, Long> valueState = keyedBackend.createInternalState(IntSerializer.INSTANCE, valueStateDescriptor);
    InternalListState<String, Integer, Long> listState = keyedBackend.createInternalState(IntSerializer.INSTANCE, listStateDescriptor);
    keyedBackend.setCurrentKey("abc");
    mapState.setCurrentNamespace(namespace1);
    mapState.put(33L, 33L);
    mapState.put(55L, 55L);
    mapState.setCurrentNamespace(namespace2);
    mapState.put(22L, 22L);
    mapState.put(11L, 11L);
    listState.setCurrentNamespace(namespace2);
    listState.add(4L);
    listState.add(5L);
    listState.add(6L);
    mapState.setCurrentNamespace(namespace3);
    mapState.put(44L, 44L);
    keyedBackend.setCurrentKey("mno");
    mapState.setCurrentNamespace(namespace3);
    mapState.put(11L, 11L);
    mapState.put(22L, 22L);
    mapState.put(33L, 33L);
    mapState.put(44L, 44L);
    mapState.put(55L, 55L);
    valueState.setCurrentNamespace(namespace3);
    valueState.update(1239L);
    listState.setCurrentNamespace(namespace3);
    listState.add(1L);
    listState.add(2L);
    listState.add(3L);
    mapState.setCurrentNamespace(namespace4);
    mapState.put(1L, 1L);
    // HEAP state backend will keep an empty map as an entry in the underlying State Table
    // we should skip such entries when serializing
    Iterator<Map.Entry<Long, Long>> iterator = mapState.iterator();
    while (iterator.hasNext()) {
        iterator.next();
        iterator.remove();
    }
    KeyGroupedInternalPriorityQueue<TimerHeapInternalTimer<String, Integer>> priorityQueue = keyedBackend.create("event-time", new TimerSerializer<>(keyedBackend.getKeySerializer(), IntSerializer.INSTANCE));
    priorityQueue.add(new TimerHeapInternalTimer<>(1234L, "mno", namespace3));
    priorityQueue.add(new TimerHeapInternalTimer<>(2345L, "mno", namespace2));
    priorityQueue.add(new TimerHeapInternalTimer<>(3456L, "mno", namespace3));
    SnapshotStrategyRunner<KeyedStateHandle, ? extends FullSnapshotResources<?>> savepointRunner = StreamOperatorStateHandler.prepareCanonicalSavepoint(keyedBackend, new CloseableRegistry());
    RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot = savepointRunner.snapshot(0L, 0L, new MemCheckpointStreamFactory(4 * 1024 * 1024), new CheckpointOptions(SavepointType.savepoint(SavepointFormatType.CANONICAL), CheckpointStorageLocationReference.getDefault()));
    snapshot.run();
    try (BufferedOutputStream bis = new BufferedOutputStream(new FileOutputStream(pathToWrite))) {
        InstantiationUtil.serializeObject(bis, snapshot.get());
    }
}

Example 19 with KeyedStateHandle

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

the class SavepointStateBackendSwitchTestBase method switchStateBackend.

@Test
public void switchStateBackend() throws Exception {
    final File pathToWrite = tempFolder.newFile();
    final MapStateDescriptor<Long, Long> mapStateDescriptor = new MapStateDescriptor<>("my-map-state", Long.class, Long.class);
    mapStateDescriptor.initializeSerializerUnlessSet(new ExecutionConfig());
    final ValueStateDescriptor<Long> valueStateDescriptor = new ValueStateDescriptor<>("my-value-state", Long.class);
    valueStateDescriptor.initializeSerializerUnlessSet(new ExecutionConfig());
    final ListStateDescriptor<Long> listStateDescriptor = new ListStateDescriptor<>("my-list-state", Long.class);
    listStateDescriptor.initializeSerializerUnlessSet(new ExecutionConfig());
    final Integer namespace1 = 1;
    final Integer namespace2 = 2;
    final Integer namespace3 = 3;
    final Integer namespace4 = 4;
    try (final CheckpointableKeyedStateBackend<String> keyedBackend = fromBackend.createBackend(KEY_GROUP_RANGE, NUM_KEY_GROUPS, Collections.emptyList())) {
        takeSavepoint(keyedBackend, pathToWrite, mapStateDescriptor, valueStateDescriptor, listStateDescriptor, namespace1, namespace2, namespace3, namespace4);
    }
    final SnapshotResult<KeyedStateHandle> stateHandles;
    try (BufferedInputStream bis = new BufferedInputStream((new FileInputStream(pathToWrite)))) {
        stateHandles = InstantiationUtil.deserializeObject(bis, Thread.currentThread().getContextClassLoader());
    }
    final KeyedStateHandle stateHandle = stateHandles.getJobManagerOwnedSnapshot();
    try (final CheckpointableKeyedStateBackend<String> keyedBackend = toBackend.createBackend(KEY_GROUP_RANGE, NUM_KEY_GROUPS, StateObjectCollection.singleton(stateHandle))) {
        verifyRestoredState(mapStateDescriptor, valueStateDescriptor, listStateDescriptor, namespace1, namespace2, namespace3, namespace4, keyedBackend);
    }
}

Example 20 with KeyedStateHandle

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

the class StreamOperatorStateHandler method snapshotState.

@VisibleForTesting
void snapshotState(CheckpointedStreamOperator streamOperator, Optional<InternalTimeServiceManager<?>> timeServiceManager, String operatorName, long checkpointId, long timestamp, CheckpointOptions checkpointOptions, CheckpointStreamFactory factory, OperatorSnapshotFutures snapshotInProgress, StateSnapshotContextSynchronousImpl snapshotContext, boolean isUsingCustomRawKeyedState) throws CheckpointException {
    try {
        if (timeServiceManager.isPresent()) {
            checkState(keyedStateBackend != null, "keyedStateBackend should be available with timeServiceManager");
            final InternalTimeServiceManager<?> manager = timeServiceManager.get();
            boolean requiresLegacyRawKeyedStateSnapshots = keyedStateBackend instanceof AbstractKeyedStateBackend && ((AbstractKeyedStateBackend<?>) keyedStateBackend).requiresLegacySynchronousTimerSnapshots(checkpointOptions.getCheckpointType());
            if (requiresLegacyRawKeyedStateSnapshots) {
                checkState(!isUsingCustomRawKeyedState, "Attempting to snapshot timers to raw keyed state, but this operator has custom raw keyed state to write.");
                manager.snapshotToRawKeyedState(snapshotContext.getRawKeyedOperatorStateOutput(), operatorName);
            }
        }
        streamOperator.snapshotState(snapshotContext);
        snapshotInProgress.setKeyedStateRawFuture(snapshotContext.getKeyedStateStreamFuture());
        snapshotInProgress.setOperatorStateRawFuture(snapshotContext.getOperatorStateStreamFuture());
        if (null != operatorStateBackend) {
            snapshotInProgress.setOperatorStateManagedFuture(operatorStateBackend.snapshot(checkpointId, timestamp, factory, checkpointOptions));
        }
        if (null != keyedStateBackend) {
            if (isCanonicalSavepoint(checkpointOptions.getCheckpointType())) {
                SnapshotStrategyRunner<KeyedStateHandle, ? extends FullSnapshotResources<?>> snapshotRunner = prepareCanonicalSavepoint(keyedStateBackend, closeableRegistry);
                snapshotInProgress.setKeyedStateManagedFuture(snapshotRunner.snapshot(checkpointId, timestamp, factory, checkpointOptions));
            } else {
                snapshotInProgress.setKeyedStateManagedFuture(keyedStateBackend.snapshot(checkpointId, timestamp, factory, checkpointOptions));
            }
        }
    } catch (Exception snapshotException) {
        try {
            snapshotInProgress.cancel();
        } catch (Exception e) {
            snapshotException.addSuppressed(e);
        }
        String snapshotFailMessage = "Could not complete snapshot " + checkpointId + " for operator " + operatorName + ".";
        try {
            snapshotContext.closeExceptionally();
        } catch (IOException e) {
            snapshotException.addSuppressed(e);
        }
        throw new CheckpointException(snapshotFailMessage, CheckpointFailureReason.CHECKPOINT_DECLINED, snapshotException);
    }
}