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

Example 91 with KeyGroupRange

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

the class StreamTaskStateInitializerImpl method keyedStatedBackend.

protected <K> CheckpointableKeyedStateBackend<K> keyedStatedBackend(TypeSerializer<K> keySerializer, String operatorIdentifierText, PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskStates, CloseableRegistry backendCloseableRegistry, MetricGroup metricGroup, double managedMemoryFraction) throws Exception {
    if (keySerializer == null) {
        return null;
    }
    String logDescription = "keyed state backend for " + operatorIdentifierText;
    TaskInfo taskInfo = environment.getTaskInfo();
    final KeyGroupRange keyGroupRange = KeyGroupRangeAssignment.computeKeyGroupRangeForOperatorIndex(taskInfo.getMaxNumberOfParallelSubtasks(), taskInfo.getNumberOfParallelSubtasks(), taskInfo.getIndexOfThisSubtask());
    // Now restore processing is included in backend building/constructing process, so we need
    // to make sure
    // each stream constructed in restore could also be closed in case of task cancel, for
    // example the data
    // input stream opened for serDe during restore.
    CloseableRegistry cancelStreamRegistryForRestore = new CloseableRegistry();
    backendCloseableRegistry.registerCloseable(cancelStreamRegistryForRestore);
    BackendRestorerProcedure<CheckpointableKeyedStateBackend<K>, KeyedStateHandle> backendRestorer = new BackendRestorerProcedure<>((stateHandles) -> stateBackend.createKeyedStateBackend(environment, environment.getJobID(), operatorIdentifierText, keySerializer, taskInfo.getMaxNumberOfParallelSubtasks(), keyGroupRange, environment.getTaskKvStateRegistry(), ttlTimeProvider, metricGroup, stateHandles, cancelStreamRegistryForRestore, managedMemoryFraction), backendCloseableRegistry, logDescription);
    try {
        return backendRestorer.createAndRestore(prioritizedOperatorSubtaskStates.getPrioritizedManagedKeyedState());
    } finally {
        if (backendCloseableRegistry.unregisterCloseable(cancelStreamRegistryForRestore)) {
            IOUtils.closeQuietly(cancelStreamRegistryForRestore);
        }
    }
}

Example 92 with KeyGroupRange

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

the class StreamOperatorStateHandler method snapshotState.

public OperatorSnapshotFutures snapshotState(CheckpointedStreamOperator streamOperator, Optional<InternalTimeServiceManager<?>> timeServiceManager, String operatorName, long checkpointId, long timestamp, CheckpointOptions checkpointOptions, CheckpointStreamFactory factory, boolean isUsingCustomRawKeyedState) throws CheckpointException {
    KeyGroupRange keyGroupRange = null != keyedStateBackend ? keyedStateBackend.getKeyGroupRange() : KeyGroupRange.EMPTY_KEY_GROUP_RANGE;
    OperatorSnapshotFutures snapshotInProgress = new OperatorSnapshotFutures();
    StateSnapshotContextSynchronousImpl snapshotContext = new StateSnapshotContextSynchronousImpl(checkpointId, timestamp, factory, keyGroupRange, closeableRegistry);
    snapshotState(streamOperator, timeServiceManager, operatorName, checkpointId, timestamp, checkpointOptions, factory, snapshotInProgress, snapshotContext, isUsingCustomRawKeyedState);
    return snapshotInProgress;
}

Example 93 with KeyGroupRange

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

the class HeapStateBackendTestBase method createKeyedBackend.

public <K> HeapKeyedStateBackend<K> createKeyedBackend(TypeSerializer<K> keySerializer, Collection<KeyedStateHandle> stateHandles) throws Exception {
    final KeyGroupRange keyGroupRange = new KeyGroupRange(0, 15);
    final int numKeyGroups = keyGroupRange.getNumberOfKeyGroups();
    ExecutionConfig executionConfig = new ExecutionConfig();
    return new HeapKeyedStateBackendBuilder<>(mock(TaskKvStateRegistry.class), keySerializer, HeapStateBackendTestBase.class.getClassLoader(), numKeyGroups, keyGroupRange, executionConfig, TtlTimeProvider.DEFAULT, LatencyTrackingStateConfig.disabled(), stateHandles, AbstractStateBackend.getCompressionDecorator(executionConfig), TestLocalRecoveryConfig.disabled(), new HeapPriorityQueueSetFactory(keyGroupRange, numKeyGroups, 128), async, new CloseableRegistry()).build();
}

Example 94 with KeyGroupRange

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

the class RocksDBAsyncSnapshotTest method testCleanupOfSnapshotsInFailureCase.

/**
 * Test that the snapshot files are cleaned up in case of a failure during the snapshot
 * procedure.
 */
@Test
public void testCleanupOfSnapshotsInFailureCase() throws Exception {
    long checkpointId = 1L;
    long timestamp = 42L;
    MockEnvironment env = MockEnvironment.builder().build();
    final IOException testException = new IOException("Test exception");
    CheckpointStateOutputStream outputStream = spy(new FailingStream(testException));
    RocksDBStateBackend backend = new RocksDBStateBackend((StateBackend) new MemoryStateBackend());
    backend.setDbStoragePath(temporaryFolder.newFolder().toURI().toString());
    AbstractKeyedStateBackend<Void> keyedStateBackend = backend.createKeyedStateBackend(env, new JobID(), "test operator", VoidSerializer.INSTANCE, 1, new KeyGroupRange(0, 0), null, TtlTimeProvider.DEFAULT, new UnregisteredMetricsGroup(), Collections.emptyList(), new CloseableRegistry());
    try {
        // register a state so that the state backend has to checkpoint something
        keyedStateBackend.getPartitionedState("namespace", StringSerializer.INSTANCE, new ValueStateDescriptor<>("foobar", String.class));
        RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotFuture = keyedStateBackend.snapshot(checkpointId, timestamp, new TestCheckpointStreamFactory(() -> outputStream), CheckpointOptions.forCheckpointWithDefaultLocation());
        try {
            FutureUtils.runIfNotDoneAndGet(snapshotFuture);
            fail("Expected an exception to be thrown here.");
        } catch (ExecutionException e) {
            Assert.assertEquals(testException, e.getCause());
        }
        verify(outputStream).close();
    } finally {
        IOUtils.closeQuietly(keyedStateBackend);
        keyedStateBackend.dispose();
        IOUtils.closeQuietly(env);
    }
}

Example 95 with KeyGroupRange

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

the class RocksIncrementalCheckpointRescalingTest method testScalingDown.

@Test
@SuppressWarnings("unchecked")
public void testScalingDown() throws Exception {
    // -----------------------------------------> test with initial parallelism 3
    // <---------------------------------------
    KeyedOneInputStreamOperatorTestHarness<String, String, Integer>[] harness3 = new KeyedOneInputStreamOperatorTestHarness[3];
    OperatorSubtaskState snapshot3;
    try {
        List<KeyGroupRange> keyGroupPartitions = StateAssignmentOperation.createKeyGroupPartitions(maxParallelism, 3);
        // task's key-group [0, 3], this should trigger the condition to use clip
        KeyGroupRange localKeyGroupRange30 = keyGroupPartitions.get(0);
        Assert.assertEquals(new KeyGroupRange(0, 3), localKeyGroupRange30);
        harness3[0] = getHarnessTest(keySelector, maxParallelism, 3, 0);
        harness3[0].setStateBackend(getStateBackend());
        harness3[0].open();
        // task's key-group [4, 6]
        KeyGroupRange localKeyGroupRange31 = keyGroupPartitions.get(1);
        Assert.assertEquals(new KeyGroupRange(4, 6), localKeyGroupRange31);
        harness3[1] = getHarnessTest(keySelector, maxParallelism, 3, 1);
        harness3[1].setStateBackend(getStateBackend());
        harness3[1].open();
        // task's key-group [7, 9]
        KeyGroupRange localKeyGroupRange32 = keyGroupPartitions.get(2);
        Assert.assertEquals(new KeyGroupRange(7, 9), localKeyGroupRange32);
        harness3[2] = getHarnessTest(keySelector, maxParallelism, 3, 2);
        harness3[2].setStateBackend(getStateBackend());
        harness3[2].open();
        validHarnessResult(harness3[0], 1, records[0], records[1], records[2], records[3]);
        validHarnessResult(harness3[1], 1, records[4], records[5], records[6]);
        validHarnessResult(harness3[2], 1, records[7], records[8], records[9]);
        snapshot3 = AbstractStreamOperatorTestHarness.repackageState(harness3[0].snapshot(0, 0), harness3[1].snapshot(0, 0), harness3[2].snapshot(0, 0));
    } finally {
        closeHarness(harness3);
    }
    // -----------------------------------------> test rescaling from 3 to 2
    // <---------------------------------------
    // init state for new subtask-0
    OperatorSubtaskState initState1 = AbstractStreamOperatorTestHarness.repartitionOperatorState(snapshot3, maxParallelism, 3, 2, 0);
    // init state for new subtask-1
    OperatorSubtaskState initState2 = AbstractStreamOperatorTestHarness.repartitionOperatorState(snapshot3, maxParallelism, 3, 2, 1);
    KeyedOneInputStreamOperatorTestHarness<String, String, Integer>[] harness2 = new KeyedOneInputStreamOperatorTestHarness[3];
    OperatorSubtaskState snapshot2;
    try {
        List<KeyGroupRange> keyGroupPartitions = StateAssignmentOperation.createKeyGroupPartitions(maxParallelism, 2);
        // task's key-group [0, 4]
        // this will choose the state handle generated by harness3[0] to init the target db
        // without any clipping.
        KeyGroupRange localKeyGroupRange20 = keyGroupPartitions.get(0);
        Assert.assertEquals(new KeyGroupRange(0, 4), localKeyGroupRange20);
        harness2[0] = getHarnessTest(keySelector, maxParallelism, 2, 0);
        harness2[0].setStateBackend(getStateBackend());
        harness2[0].setup();
        harness2[0].initializeState(initState1);
        harness2[0].open();
        // task's key-group [5, 9], this will open a empty db, and insert records from two state
        // handles.
        KeyGroupRange localKeyGroupRange21 = keyGroupPartitions.get(1);
        Assert.assertEquals(new KeyGroupRange(5, 9), localKeyGroupRange21);
        harness2[1] = getHarnessTest(keySelector, maxParallelism, 2, 1);
        harness2[1].setStateBackend(getStateBackend());
        harness2[1].setup();
        harness2[1].initializeState(initState2);
        harness2[1].open();
        validHarnessResult(harness2[0], 2, records[0], records[1], records[2], records[3], records[4]);
        validHarnessResult(harness2[1], 2, records[5], records[6], records[7], records[8], records[9]);
        snapshot2 = AbstractStreamOperatorTestHarness.repackageState(harness2[0].snapshot(0, 0), harness2[1].snapshot(0, 0));
    } finally {
        closeHarness(harness2);
    }
    // -----------------------------------------> test rescaling from 2 to 1
    // <---------------------------------------
    // init state for new subtask-0
    initState1 = AbstractStreamOperatorTestHarness.repartitionOperatorState(snapshot2, maxParallelism, 2, 1, 0);
    try (KeyedOneInputStreamOperatorTestHarness<String, String, Integer> harness = getHarnessTest(keySelector, maxParallelism, 1, 0)) {
        // this will choose the state handle generated by harness2[0] to init the target db
        // without any clipping.
        harness.setStateBackend(getStateBackend());
        harness.setup();
        harness.initializeState(initState1);
        harness.open();
        validHarnessResult(harness, 3, records);
    }
}