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

Example 21 with SharedStateRegistry

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

the class ChangelogStateBackendTestUtils method testMaterializedRestoreForPriorityQueue.

public static void testMaterializedRestoreForPriorityQueue(StateBackend stateBackend, Environment env, CheckpointStreamFactory streamFactory) throws Exception {
    SharedStateRegistry sharedStateRegistry = new SharedStateRegistryImpl();
    String fieldName = "key-grouped-priority-queue";
    ChangelogKeyedStateBackend<Integer> keyedBackend = (ChangelogKeyedStateBackend<Integer>) createKeyedBackend(stateBackend, env);
    CompletableFuture<Void> asyncComplete = new CompletableFuture<>();
    PeriodicMaterializationManager periodicMaterializationManager = periodicMaterializationManager(keyedBackend, asyncComplete);
    try {
        KeyGroupedInternalPriorityQueue<TestType> priorityQueue = keyedBackend.create(fieldName, new TestType.V1TestTypeSerializer());
        TestType elementA100 = new TestType("a", 100);
        TestType elementA10 = new TestType("a", 10);
        TestType elementA20 = new TestType("a", 20);
        assertTrue(priorityQueue.add(elementA100));
        assertTrue(priorityQueue.add(elementA10));
        assertFalse(priorityQueue.add(elementA20));
        assertFalse(priorityQueue.add(elementA10));
        List<TestType> actualList = new ArrayList<>();
        try (CloseableIterator<TestType> iterator = priorityQueue.iterator()) {
            iterator.forEachRemaining(actualList::add);
        }
        assertThat(actualList, containsInAnyOrder(elementA100, elementA10, elementA20));
        materialize(keyedBackend, periodicMaterializationManager);
        TestType elementB9 = new TestType("b", 9);
        assertTrue(priorityQueue.add(elementB9));
        materialize(keyedBackend, periodicMaterializationManager);
        TestType elementC9 = new TestType("c", 9);
        TestType elementC8 = new TestType("c", 8);
        assertFalse(priorityQueue.add(elementC9));
        assertTrue(priorityQueue.add(elementC8));
        KeyedStateHandle snapshot = runSnapshot(keyedBackend.snapshot(682375462378L, 2, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation()), sharedStateRegistry);
        IOUtils.closeQuietly(keyedBackend);
        keyedBackend.dispose();
        // make sure the asycn phase completes successfully
        if (asyncComplete.isCompletedExceptionally()) {
            asyncComplete.get();
        }
        // ============================ restore snapshot ===============================
        keyedBackend = (ChangelogKeyedStateBackend<Integer>) restoreKeyedBackend(stateBackend, snapshot, env);
        snapshot.discardState();
        KeyGroupedInternalPriorityQueue<TestType> priorityQueueRestored = keyedBackend.create(fieldName, new TestType.V1TestTypeSerializer());
        List<TestType> actualListRestore = new ArrayList<>();
        try (CloseableIterator<TestType> iterator = priorityQueueRestored.iterator()) {
            iterator.forEachRemaining(actualListRestore::add);
        }
        assertThat(actualListRestore, containsInAnyOrder(elementA100, elementA10, elementA20, elementB9, elementC9, elementC8));
        assertFalse(priorityQueueRestored.add(new TestType("d", 11)));
    } finally {
        IOUtils.closeQuietly(keyedBackend);
        keyedBackend.dispose();
    }
}

Example 22 with SharedStateRegistry

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

the class ChangelogStateBackendTestUtils method testMaterializedRestore.

public static void testMaterializedRestore(StateBackend stateBackend, StateTtlConfig stateTtlConfig, Environment env, CheckpointStreamFactory streamFactory) throws Exception {
    SharedStateRegistry sharedStateRegistry = new SharedStateRegistryImpl();
    TypeInformation<StateBackendTestBase.TestPojo> pojoType = new GenericTypeInfo<>(StateBackendTestBase.TestPojo.class);
    ValueStateDescriptor<StateBackendTestBase.TestPojo> kvId = new ValueStateDescriptor<>("id", pojoType);
    if (stateTtlConfig.isEnabled()) {
        kvId.enableTimeToLive(stateTtlConfig);
    }
    ChangelogKeyedStateBackend<Integer> keyedBackend = (ChangelogKeyedStateBackend<Integer>) createKeyedBackend(stateBackend, env);
    CompletableFuture<Void> asyncComplete = new CompletableFuture<>();
    PeriodicMaterializationManager periodicMaterializationManager = periodicMaterializationManager(keyedBackend, asyncComplete);
    try {
        ValueState<StateBackendTestBase.TestPojo> state = keyedBackend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
        keyedBackend.setCurrentKey(1);
        state.update(new StateBackendTestBase.TestPojo("u1", 1));
        keyedBackend.setCurrentKey(2);
        state.update(new StateBackendTestBase.TestPojo("u2", 2));
        materialize(keyedBackend, periodicMaterializationManager);
        keyedBackend.setCurrentKey(2);
        state.update(new StateBackendTestBase.TestPojo("u2", 22));
        keyedBackend.setCurrentKey(3);
        state.update(new StateBackendTestBase.TestPojo("u3", 3));
        materialize(keyedBackend, periodicMaterializationManager);
        keyedBackend.setCurrentKey(4);
        state.update(new StateBackendTestBase.TestPojo("u4", 4));
        keyedBackend.setCurrentKey(2);
        state.update(new StateBackendTestBase.TestPojo("u2", 222));
        KeyedStateHandle snapshot = runSnapshot(keyedBackend.snapshot(682375462378L, 2, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation()), sharedStateRegistry);
        IOUtils.closeQuietly(keyedBackend);
        keyedBackend.dispose();
        // make sure the asycn phase completes successfully
        if (asyncComplete.isCompletedExceptionally()) {
            asyncComplete.get();
        }
        // ============================ restore snapshot ===============================
        env.getExecutionConfig().registerKryoType(StateBackendTestBase.TestPojo.class);
        keyedBackend = (ChangelogKeyedStateBackend<Integer>) restoreKeyedBackend(stateBackend, snapshot, env);
        snapshot.discardState();
        state = keyedBackend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
        keyedBackend.setCurrentKey(1);
        assertEquals(new StateBackendTestBase.TestPojo("u1", 1), state.value());
        keyedBackend.setCurrentKey(2);
        assertEquals(new StateBackendTestBase.TestPojo("u2", 222), state.value());
        keyedBackend.setCurrentKey(3);
        assertEquals(new StateBackendTestBase.TestPojo("u3", 3), state.value());
    } finally {
        IOUtils.closeQuietly(keyedBackend);
        keyedBackend.dispose();
    }
}