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

Example 1 with SnapshotResult

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

the class HeapSnapshotStrategy method asyncSnapshot.

@Override
public SnapshotResultSupplier<KeyedStateHandle> asyncSnapshot(HeapSnapshotResources<K> syncPartResource, long checkpointId, long timestamp, @Nonnull CheckpointStreamFactory streamFactory, @Nonnull CheckpointOptions checkpointOptions) {
    List<StateMetaInfoSnapshot> metaInfoSnapshots = syncPartResource.getMetaInfoSnapshots();
    if (metaInfoSnapshots.isEmpty()) {
        return snapshotCloseableRegistry -> SnapshotResult.empty();
    }
    final KeyedBackendSerializationProxy<K> serializationProxy = new KeyedBackendSerializationProxy<>(// get a serialized form already at state registration time in the future
    syncPartResource.getKeySerializer(), metaInfoSnapshots, !Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, keyGroupCompressionDecorator));
    final SupplierWithException<CheckpointStreamWithResultProvider, Exception> checkpointStreamSupplier = localRecoveryConfig.isLocalRecoveryEnabled() && !checkpointOptions.getCheckpointType().isSavepoint() ? () -> createDuplicatingStream(checkpointId, CheckpointedStateScope.EXCLUSIVE, streamFactory, localRecoveryConfig.getLocalStateDirectoryProvider().orElseThrow(LocalRecoveryConfig.localRecoveryNotEnabled())) : () -> createSimpleStream(CheckpointedStateScope.EXCLUSIVE, streamFactory);
    return (snapshotCloseableRegistry) -> {
        final Map<StateUID, Integer> stateNamesToId = syncPartResource.getStateNamesToId();
        final Map<StateUID, StateSnapshot> cowStateStableSnapshots = syncPartResource.getCowStateStableSnapshots();
        final CheckpointStreamWithResultProvider streamWithResultProvider = checkpointStreamSupplier.get();
        snapshotCloseableRegistry.registerCloseable(streamWithResultProvider);
        final CheckpointStateOutputStream localStream = streamWithResultProvider.getCheckpointOutputStream();
        final DataOutputViewStreamWrapper outView = new DataOutputViewStreamWrapper(localStream);
        serializationProxy.write(outView);
        final long[] keyGroupRangeOffsets = new long[keyGroupRange.getNumberOfKeyGroups()];
        for (int keyGroupPos = 0; keyGroupPos < keyGroupRange.getNumberOfKeyGroups(); ++keyGroupPos) {
            int keyGroupId = keyGroupRange.getKeyGroupId(keyGroupPos);
            keyGroupRangeOffsets[keyGroupPos] = localStream.getPos();
            outView.writeInt(keyGroupId);
            for (Map.Entry<StateUID, StateSnapshot> stateSnapshot : cowStateStableSnapshots.entrySet()) {
                StateSnapshot.StateKeyGroupWriter partitionedSnapshot = stateSnapshot.getValue().getKeyGroupWriter();
                try (OutputStream kgCompressionOut = keyGroupCompressionDecorator.decorateWithCompression(localStream)) {
                    DataOutputViewStreamWrapper kgCompressionView = new DataOutputViewStreamWrapper(kgCompressionOut);
                    kgCompressionView.writeShort(stateNamesToId.get(stateSnapshot.getKey()));
                    partitionedSnapshot.writeStateInKeyGroup(kgCompressionView, keyGroupId);
                }
            // this will just close the outer compression stream
            }
        }
        if (snapshotCloseableRegistry.unregisterCloseable(streamWithResultProvider)) {
            KeyGroupRangeOffsets kgOffs = new KeyGroupRangeOffsets(keyGroupRange, keyGroupRangeOffsets);
            SnapshotResult<StreamStateHandle> result = streamWithResultProvider.closeAndFinalizeCheckpointStreamResult();
            return toKeyedStateHandleSnapshotResult(result, kgOffs, KeyGroupsStateHandle::new);
        } else {
            throw new IOException("Stream already unregistered.");
        }
    };
}

Example 2 with SnapshotResult

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

the class HeapKeyedStateBackendSnapshotMigrationTest method testMapStateMigrationAfterHashMapSerRemoval.

@Test
public void testMapStateMigrationAfterHashMapSerRemoval() throws Exception {
    ClassLoader cl = getClass().getClassLoader();
    URL resource = cl.getResource("heap_keyed_statebackend_1_5_map.snapshot");
    Preconditions.checkNotNull(resource, "Binary snapshot resource not found!");
    final SnapshotResult<KeyedStateHandle> stateHandles;
    try (BufferedInputStream bis = new BufferedInputStream((new FileInputStream(resource.getFile())))) {
        stateHandles = InstantiationUtil.deserializeObject(bis, Thread.currentThread().getContextClassLoader());
    }
    final KeyedStateHandle stateHandle = stateHandles.getJobManagerOwnedSnapshot();
    try (final HeapKeyedStateBackend<String> keyedBackend = createKeyedBackend(StateObjectCollection.singleton(stateHandle))) {
        final Integer namespace1 = 1;
        final Integer namespace2 = 2;
        final Integer namespace3 = 3;
        final MapStateDescriptor<Long, Long> stateDescr = new MapStateDescriptor<>("my-map-state", Long.class, Long.class);
        stateDescr.initializeSerializerUnlessSet(new ExecutionConfig());
        InternalMapState<String, Integer, Long, Long> state = keyedBackend.createInternalState(IntSerializer.INSTANCE, stateDescr);
        keyedBackend.setCurrentKey("abc");
        state.setCurrentNamespace(namespace1);
        assertEquals(33L, (long) state.get(33L));
        assertEquals(55L, (long) state.get(55L));
        assertEquals(2, getStateSize(state));
        state.setCurrentNamespace(namespace2);
        assertEquals(22L, (long) state.get(22L));
        assertEquals(11L, (long) state.get(11L));
        assertEquals(2, getStateSize(state));
        state.setCurrentNamespace(namespace3);
        assertEquals(44L, (long) state.get(44L));
        assertEquals(1, getStateSize(state));
        keyedBackend.setCurrentKey("def");
        state.setCurrentNamespace(namespace1);
        assertEquals(11L, (long) state.get(11L));
        assertEquals(44L, (long) state.get(44L));
        assertEquals(2, getStateSize(state));
        state.setCurrentNamespace(namespace3);
        assertEquals(22L, (long) state.get(22L));
        assertEquals(55L, (long) state.get(55L));
        assertEquals(33L, (long) state.get(33L));
        assertEquals(3, getStateSize(state));
        keyedBackend.setCurrentKey("jkl");
        state.setCurrentNamespace(namespace1);
        assertEquals(11L, (long) state.get(11L));
        assertEquals(22L, (long) state.get(22L));
        assertEquals(33L, (long) state.get(33L));
        assertEquals(44L, (long) state.get(44L));
        assertEquals(55L, (long) state.get(55L));
        assertEquals(5, getStateSize(state));
        keyedBackend.setCurrentKey("mno");
        state.setCurrentNamespace(namespace3);
        assertEquals(11L, (long) state.get(11L));
        assertEquals(22L, (long) state.get(22L));
        assertEquals(33L, (long) state.get(33L));
        assertEquals(44L, (long) state.get(44L));
        assertEquals(55L, (long) state.get(55L));
        assertEquals(5, getStateSize(state));
        RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot = keyedBackend.snapshot(1L, 1L, new MemCheckpointStreamFactory(4 * 1024 * 1024), CheckpointOptions.forCheckpointWithDefaultLocation());
        snapshot.run();
    }
}

Example 3 with SnapshotResult

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

the class EmbeddedRocksDBStateBackendTest method testReleasingSnapshotAfterBackendClosed.

@Test
public void testReleasingSnapshotAfterBackendClosed() throws Exception {
    setupRocksKeyedStateBackend();
    try {
        RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot = keyedStateBackend.snapshot(0L, 0L, testStreamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());
        RocksDB spyDB = keyedStateBackend.db;
        if (!enableIncrementalCheckpointing) {
            verify(spyDB, times(1)).getSnapshot();
            verify(spyDB, times(0)).releaseSnapshot(any(Snapshot.class));
        }
        // Ensure every RocksObjects not closed yet
        for (RocksObject rocksCloseable : allCreatedCloseables) {
            verify(rocksCloseable, times(0)).close();
        }
        snapshot.cancel(true);
        this.keyedStateBackend.dispose();
        verify(spyDB, times(1)).close();
        assertEquals(true, keyedStateBackend.isDisposed());
        // Ensure every RocksObjects was closed exactly once
        for (RocksObject rocksCloseable : allCreatedCloseables) {
            verify(rocksCloseable, times(1)).close();
        }
    } finally {
        keyedStateBackend.dispose();
        keyedStateBackend = null;
    }
    verifyRocksDBStateUploaderClosed();
}

Example 4 with SnapshotResult

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

use of org.apache.flink.runtime.state.SnapshotResult 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();
}