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

Example 46 with KeyedStateHandle

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

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

Example 48 with KeyedStateHandle

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

the class ChangelogBackendRestoreOperation method restore.

public static <K, T extends ChangelogStateHandle> ChangelogKeyedStateBackend<K> restore(StateChangelogHandleReader<T> changelogHandleReader, ClassLoader classLoader, Collection<ChangelogStateBackendHandle> stateHandles, BaseBackendBuilder<K> baseBackendBuilder, DeltaBackendBuilder<K> changelogBackendBuilder) throws Exception {
    Collection<KeyedStateHandle> baseState = extractBaseState(stateHandles);
    AbstractKeyedStateBackend<K> baseBackend = baseBackendBuilder.apply(baseState);
    ChangelogKeyedStateBackend<K> changelogBackend = changelogBackendBuilder.apply(baseBackend, stateHandles);
    for (ChangelogStateBackendHandle handle : stateHandles) {
        if (handle != null) {
            // null is empty state (no change)
            readBackendHandle(changelogBackend, handle, changelogHandleReader, classLoader);
        }
    }
    return changelogBackend;
}

Example 49 with KeyedStateHandle

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

the class StreamTaskStateInitializerImplTest method testWithRestore.

@SuppressWarnings("unchecked")
@Test
public void testWithRestore() throws Exception {
    StateBackend mockingBackend = spy(new StateBackend() {

        @Override
        public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(Environment env, JobID jobID, String operatorIdentifier, TypeSerializer<K> keySerializer, int numberOfKeyGroups, KeyGroupRange keyGroupRange, TaskKvStateRegistry kvStateRegistry, TtlTimeProvider ttlTimeProvider, MetricGroup metricGroup, @Nonnull Collection<KeyedStateHandle> stateHandles, CloseableRegistry cancelStreamRegistry) throws Exception {
            return mock(AbstractKeyedStateBackend.class);
        }

        @Override
        public OperatorStateBackend createOperatorStateBackend(Environment env, String operatorIdentifier, @Nonnull Collection<OperatorStateHandle> stateHandles, CloseableRegistry cancelStreamRegistry) throws Exception {
            return mock(OperatorStateBackend.class);
        }
    });
    OperatorID operatorID = new OperatorID(47L, 11L);
    TaskStateSnapshot taskStateSnapshot = new TaskStateSnapshot();
    Random random = new Random(0x42);
    OperatorSubtaskState operatorSubtaskState = OperatorSubtaskState.builder().setManagedOperatorState(new OperatorStreamStateHandle(Collections.singletonMap("a", new OperatorStateHandle.StateMetaInfo(new long[] { 0, 10 }, SPLIT_DISTRIBUTE)), CheckpointTestUtils.createDummyStreamStateHandle(random, null))).setRawOperatorState(new OperatorStreamStateHandle(Collections.singletonMap("_default_", new OperatorStateHandle.StateMetaInfo(new long[] { 0, 20, 30 }, SPLIT_DISTRIBUTE)), CheckpointTestUtils.createDummyStreamStateHandle(random, null))).setManagedKeyedState(CheckpointTestUtils.createDummyKeyGroupStateHandle(random, null)).setRawKeyedState(CheckpointTestUtils.createDummyKeyGroupStateHandle(random, null)).setInputChannelState(singleton(createNewInputChannelStateHandle(10, random))).setResultSubpartitionState(singleton(createNewResultSubpartitionStateHandle(10, random))).build();
    taskStateSnapshot.putSubtaskStateByOperatorID(operatorID, operatorSubtaskState);
    JobManagerTaskRestore jobManagerTaskRestore = new JobManagerTaskRestore(42L, taskStateSnapshot);
    StreamTaskStateInitializer streamTaskStateManager = streamTaskStateManager(mockingBackend, jobManagerTaskRestore, false);
    AbstractStreamOperator<?> streamOperator = mock(AbstractStreamOperator.class);
    when(streamOperator.getOperatorID()).thenReturn(operatorID);
    TypeSerializer<?> typeSerializer = new IntSerializer();
    CloseableRegistry closeableRegistry = new CloseableRegistry();
    StreamOperatorStateContext stateContext = streamTaskStateManager.streamOperatorStateContext(streamOperator.getOperatorID(), streamOperator.getClass().getSimpleName(), new TestProcessingTimeService(), streamOperator, typeSerializer, closeableRegistry, new UnregisteredMetricsGroup(), 1.0, false);
    OperatorStateBackend operatorStateBackend = stateContext.operatorStateBackend();
    CheckpointableKeyedStateBackend<?> keyedStateBackend = stateContext.keyedStateBackend();
    InternalTimeServiceManager<?> timeServiceManager = stateContext.internalTimerServiceManager();
    CloseableIterable<KeyGroupStatePartitionStreamProvider> keyedStateInputs = stateContext.rawKeyedStateInputs();
    CloseableIterable<StatePartitionStreamProvider> operatorStateInputs = stateContext.rawOperatorStateInputs();
    Assert.assertTrue("Expected the context to be restored", stateContext.isRestored());
    Assert.assertEquals(OptionalLong.of(42L), stateContext.getRestoredCheckpointId());
    Assert.assertNotNull(operatorStateBackend);
    Assert.assertNotNull(keyedStateBackend);
    // this is deactivated on purpose so that it does not attempt to consume the raw keyed
    // state.
    Assert.assertNull(timeServiceManager);
    Assert.assertNotNull(keyedStateInputs);
    Assert.assertNotNull(operatorStateInputs);
    int count = 0;
    for (KeyGroupStatePartitionStreamProvider keyedStateInput : keyedStateInputs) {
        ++count;
    }
    Assert.assertEquals(1, count);
    count = 0;
    for (StatePartitionStreamProvider operatorStateInput : operatorStateInputs) {
        ++count;
    }
    Assert.assertEquals(3, count);
    checkCloseablesRegistered(closeableRegistry, operatorStateBackend, keyedStateBackend, keyedStateInputs, operatorStateInputs);
}