Examples with CheckpointListener org.apache.flink.api.common.state.CheckpointListener used on opensource projects

Example 1 with CheckpointListener

use of org.apache.flink.api.common.state.CheckpointListener in project flink by apache.

the class EmbeddedRocksDBStateBackendTest method testSharedIncrementalStateDeRegistration.

@Test
public void testSharedIncrementalStateDeRegistration() throws Exception {
    if (enableIncrementalCheckpointing) {
        CheckpointableKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
        try {
            ValueStateDescriptor<String> kvId = new ValueStateDescriptor<>("id", String.class, null);
            kvId.initializeSerializerUnlessSet(new ExecutionConfig());
            ValueState<String> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
            Queue<IncrementalRemoteKeyedStateHandle> previousStateHandles = new LinkedList<>();
            SharedStateRegistry sharedStateRegistry = spy(new SharedStateRegistryImpl());
            for (int checkpointId = 0; checkpointId < 3; ++checkpointId) {
                reset(sharedStateRegistry);
                backend.setCurrentKey(checkpointId);
                state.update("Hello-" + checkpointId);
                RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot = backend.snapshot(checkpointId, checkpointId, createStreamFactory(), CheckpointOptions.forCheckpointWithDefaultLocation());
                snapshot.run();
                SnapshotResult<KeyedStateHandle> snapshotResult = snapshot.get();
                IncrementalRemoteKeyedStateHandle stateHandle = (IncrementalRemoteKeyedStateHandle) snapshotResult.getJobManagerOwnedSnapshot();
                Map<StateHandleID, StreamStateHandle> sharedState = new HashMap<>(stateHandle.getSharedState());
                stateHandle.registerSharedStates(sharedStateRegistry, checkpointId);
                for (Map.Entry<StateHandleID, StreamStateHandle> e : sharedState.entrySet()) {
                    verify(sharedStateRegistry).registerReference(stateHandle.createSharedStateRegistryKeyFromFileName(e.getKey()), e.getValue(), checkpointId);
                }
                previousStateHandles.add(stateHandle);
                ((CheckpointListener) backend).notifyCheckpointComplete(checkpointId);
                if (previousStateHandles.size() > 1) {
                    previousStateHandles.remove().discardState();
                }
            }
            while (!previousStateHandles.isEmpty()) {
                reset(sharedStateRegistry);
                previousStateHandles.remove().discardState();
            }
        } finally {
            IOUtils.closeQuietly(backend);
            backend.dispose();
        }
    }
}