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

Example 96 with KeyGroupRange

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

the class RocksIncrementalCheckpointRescalingTest method testScalingUp.

@Test
@SuppressWarnings("unchecked")
public void testScalingUp() throws Exception {
    // -----------------------------------------> test with initial parallelism 1
    // <---------------------------------------
    OperatorSubtaskState snapshot;
    try (KeyedOneInputStreamOperatorTestHarness<String, String, Integer> harness = getHarnessTest(keySelector, maxParallelism, 1, 0)) {
        harness.setStateBackend(getStateBackend());
        harness.open();
        validHarnessResult(harness, 1, records);
        snapshot = harness.snapshot(0, 0);
    }
    // -----------------------------------------> test rescaling from 1 to 2
    // <---------------------------------------
    // init state for new subtask-0
    OperatorSubtaskState initState1 = AbstractStreamOperatorTestHarness.repartitionOperatorState(snapshot, maxParallelism, 1, 2, 0);
    // init state for new subtask-1
    OperatorSubtaskState initState2 = AbstractStreamOperatorTestHarness.repartitionOperatorState(snapshot, maxParallelism, 1, 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]
        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]
        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 3
    // <---------------------------------------
    // init state for new subtask-0
    initState1 = AbstractStreamOperatorTestHarness.repartitionOperatorState(snapshot2, maxParallelism, 2, 3, 0);
    // init state for new subtask-1
    initState2 = AbstractStreamOperatorTestHarness.repartitionOperatorState(snapshot2, maxParallelism, 2, 3, 1);
    // init state for new subtask-2
    OperatorSubtaskState initState3 = AbstractStreamOperatorTestHarness.repartitionOperatorState(snapshot2, maxParallelism, 2, 3, 2);
    KeyedOneInputStreamOperatorTestHarness<String, String, Integer>[] harness3 = new KeyedOneInputStreamOperatorTestHarness[3];
    try {
        List<KeyGroupRange> keyGroupPartitions = StateAssignmentOperation.createKeyGroupPartitions(maxParallelism, 3);
        // task's key-group [0, 3]
        // this will choose the state handle to harness2[0] to init the target db with clipping.
        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].setup();
        harness3[0].initializeState(initState1);
        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].setup();
        harness3[1].initializeState(initState2);
        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].setup();
        harness3[2].initializeState(initState3);
        harness3[2].open();
        validHarnessResult(harness3[0], 3, records[0], records[1], records[2], records[3]);
        validHarnessResult(harness3[1], 3, records[4], records[5], records[6]);
        validHarnessResult(harness3[2], 3, records[7], records[8], records[9]);
    } finally {
        closeHarness(harness3);
    }
}

Example 97 with KeyGroupRange

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

the class RocksDBIncrementalCheckpointUtilsTest method testChooseTheBestStateHandleForInitial.

@Test
public void testChooseTheBestStateHandleForInitial() {
    List<KeyedStateHandle> keyedStateHandles = new ArrayList<>(3);
    KeyedStateHandle keyedStateHandle1 = mock(KeyedStateHandle.class);
    when(keyedStateHandle1.getKeyGroupRange()).thenReturn(new KeyGroupRange(0, 3));
    keyedStateHandles.add(keyedStateHandle1);
    KeyedStateHandle keyedStateHandle2 = mock(KeyedStateHandle.class);
    when(keyedStateHandle2.getKeyGroupRange()).thenReturn(new KeyGroupRange(4, 7));
    keyedStateHandles.add(keyedStateHandle2);
    KeyedStateHandle keyedStateHandle3 = mock(KeyedStateHandle.class);
    when(keyedStateHandle3.getKeyGroupRange()).thenReturn(new KeyGroupRange(8, 12));
    keyedStateHandles.add(keyedStateHandle3);
    // this should choose no one handle.
    Assert.assertNull(RocksDBIncrementalCheckpointUtils.chooseTheBestStateHandleForInitial(keyedStateHandles, new KeyGroupRange(3, 5)));
    // this should choose keyedStateHandle2, because keyedStateHandle2's key-group range
    // satisfies the overlap fraction demand.
    Assert.assertEquals(keyedStateHandle2, RocksDBIncrementalCheckpointUtils.chooseTheBestStateHandleForInitial(keyedStateHandles, new KeyGroupRange(3, 6)));
    // both keyedStateHandle2 & keyedStateHandle3's key-group range satisfies the overlap
    // fraction, but keyedStateHandle3's key group range is better.
    Assert.assertEquals(keyedStateHandle3, RocksDBIncrementalCheckpointUtils.chooseTheBestStateHandleForInitial(keyedStateHandles, new KeyGroupRange(5, 12)));
    // The intersect key group number of keyedStateHandle2 & keyedStateHandle3's with [4, 11]
    // are 4. But the over fraction of keyedStateHandle2 is better.
    Assert.assertEquals(keyedStateHandle2, RocksDBIncrementalCheckpointUtils.chooseTheBestStateHandleForInitial(keyedStateHandles, new KeyGroupRange(4, 11)));
    // both keyedStateHandle2 & keyedStateHandle3's key-group range are covered by [3, 12],
    // but this should choose the keyedStateHandle3, because keyedStateHandle3's key-group is
    // bigger than keyedStateHandle2.
    Assert.assertEquals(keyedStateHandle3, RocksDBIncrementalCheckpointUtils.chooseTheBestStateHandleForInitial(keyedStateHandles, new KeyGroupRange(3, 12)));
}

Example 98 with KeyGroupRange

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

the class RocksDBIncrementalCheckpointUtilsTest method testClipDBWithKeyGroupRange.

@Test
public void testClipDBWithKeyGroupRange() throws Exception {
    testClipDBWithKeyGroupRangeHelper(new KeyGroupRange(0, 1), new KeyGroupRange(0, 2), 1);
    testClipDBWithKeyGroupRangeHelper(new KeyGroupRange(0, 1), new KeyGroupRange(0, 1), 1);
    testClipDBWithKeyGroupRangeHelper(new KeyGroupRange(0, 1), new KeyGroupRange(1, 2), 1);
    testClipDBWithKeyGroupRangeHelper(new KeyGroupRange(0, 1), new KeyGroupRange(2, 4), 1);
    testClipDBWithKeyGroupRangeHelper(new KeyGroupRange(Byte.MAX_VALUE - 15, Byte.MAX_VALUE), new KeyGroupRange(Byte.MAX_VALUE - 10, Byte.MAX_VALUE), 1);
    testClipDBWithKeyGroupRangeHelper(new KeyGroupRange(Short.MAX_VALUE - 15, Short.MAX_VALUE), new KeyGroupRange(Short.MAX_VALUE - 10, Short.MAX_VALUE), 2);
    testClipDBWithKeyGroupRangeHelper(new KeyGroupRange(Byte.MAX_VALUE - 15, Byte.MAX_VALUE - 1), new KeyGroupRange(Byte.MAX_VALUE - 10, Byte.MAX_VALUE), 1);
    testClipDBWithKeyGroupRangeHelper(new KeyGroupRange(Short.MAX_VALUE - 15, Short.MAX_VALUE - 1), new KeyGroupRange(Short.MAX_VALUE - 10, Short.MAX_VALUE), 2);
}

Example 99 with KeyGroupRange

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

the class RocksDBStateBackendConfigTest method testUseTempDirectories.

// ------------------------------------------------------------------------
// RocksDB local file automatic from temp directories
// ------------------------------------------------------------------------
/**
 * This tests whether the RocksDB backends uses the temp directories that are provided from the
 * {@link Environment} when no db storage path is set.
 */
@Test
public void testUseTempDirectories() throws Exception {
    String checkpointPath = tempFolder.newFolder().toURI().toString();
    RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);
    File dir1 = tempFolder.newFolder();
    assertNull(rocksDbBackend.getDbStoragePaths());
    final MockEnvironment env = getMockEnvironment(dir1);
    RocksDBKeyedStateBackend<Integer> keyedBackend = (RocksDBKeyedStateBackend<Integer>) rocksDbBackend.createKeyedStateBackend(env, env.getJobID(), "test_op", IntSerializer.INSTANCE, 1, new KeyGroupRange(0, 0), env.getTaskKvStateRegistry(), TtlTimeProvider.DEFAULT, new UnregisteredMetricsGroup(), Collections.emptyList(), new CloseableRegistry());
    try {
        File instanceBasePath = keyedBackend.getInstanceBasePath();
        assertThat(instanceBasePath.getAbsolutePath(), startsWith(dir1.getAbsolutePath()));
    } finally {
        IOUtils.closeQuietly(keyedBackend);
        keyedBackend.dispose();
        env.close();
    }
}

Example 100 with KeyGroupRange

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

the class CheckpointCoordinatorTest method testSharedStateRegistrationOnRestore.

@Test
public void testSharedStateRegistrationOnRestore() throws Exception {
    JobVertexID jobVertexID1 = new JobVertexID();
    int parallelism1 = 2;
    int maxParallelism1 = 4;
    ExecutionGraph graph = new CheckpointCoordinatorTestingUtils.CheckpointExecutionGraphBuilder().addJobVertex(jobVertexID1, parallelism1, maxParallelism1).build();
    ExecutionJobVertex jobVertex1 = graph.getJobVertex(jobVertexID1);
    List<CompletedCheckpoint> checkpoints = Collections.emptyList();
    SharedStateRegistry firstInstance = SharedStateRegistry.DEFAULT_FACTORY.create(org.apache.flink.util.concurrent.Executors.directExecutor(), checkpoints);
    final EmbeddedCompletedCheckpointStore store = new EmbeddedCompletedCheckpointStore(10, checkpoints, firstInstance);
    // set up the coordinator and validate the initial state
    final CheckpointCoordinatorBuilder coordinatorBuilder = new CheckpointCoordinatorBuilder().setExecutionGraph(graph).setTimer(manuallyTriggeredScheduledExecutor);
    final CheckpointCoordinator coordinator = coordinatorBuilder.setCompletedCheckpointStore(store).build();
    final int numCheckpoints = 3;
    List<KeyGroupRange> keyGroupPartitions1 = StateAssignmentOperation.createKeyGroupPartitions(maxParallelism1, parallelism1);
    for (int i = 0; i < numCheckpoints; ++i) {
        performIncrementalCheckpoint(graph.getJobID(), coordinator, jobVertex1, keyGroupPartitions1, i);
    }
    List<CompletedCheckpoint> completedCheckpoints = coordinator.getSuccessfulCheckpoints();
    assertEquals(numCheckpoints, completedCheckpoints.size());
    int sharedHandleCount = 0;
    List<Map<StateHandleID, StreamStateHandle>> sharedHandlesByCheckpoint = new ArrayList<>(numCheckpoints);
    for (int i = 0; i < numCheckpoints; ++i) {
        sharedHandlesByCheckpoint.add(new HashMap<>(2));
    }
    int cp = 0;
    for (CompletedCheckpoint completedCheckpoint : completedCheckpoints) {
        for (OperatorState taskState : completedCheckpoint.getOperatorStates().values()) {
            for (OperatorSubtaskState subtaskState : taskState.getStates()) {
                for (KeyedStateHandle keyedStateHandle : subtaskState.getManagedKeyedState()) {
                    // test we are once registered with the current registry
                    verify(keyedStateHandle, times(1)).registerSharedStates(firstInstance, completedCheckpoint.getCheckpointID());
                    IncrementalRemoteKeyedStateHandle incrementalKeyedStateHandle = (IncrementalRemoteKeyedStateHandle) keyedStateHandle;
                    sharedHandlesByCheckpoint.get(cp).putAll(incrementalKeyedStateHandle.getSharedState());
                    for (StreamStateHandle streamStateHandle : incrementalKeyedStateHandle.getSharedState().values()) {
                        assertTrue(!(streamStateHandle instanceof PlaceholderStreamStateHandle));
                        verify(streamStateHandle, never()).discardState();
                        ++sharedHandleCount;
                    }
                    for (StreamStateHandle streamStateHandle : incrementalKeyedStateHandle.getPrivateState().values()) {
                        verify(streamStateHandle, never()).discardState();
                    }
                    verify(incrementalKeyedStateHandle.getMetaStateHandle(), never()).discardState();
                }
                verify(subtaskState, never()).discardState();
            }
        }
        ++cp;
    }
    // 2 (parallelism) x (1 (CP0) + 2 (CP1) + 2 (CP2)) = 10
    assertEquals(10, sharedHandleCount);
    // discard CP0
    store.removeOldestCheckpoint();
    // CP1
    for (Map<StateHandleID, StreamStateHandle> cpList : sharedHandlesByCheckpoint) {
        for (StreamStateHandle streamStateHandle : cpList.values()) {
            verify(streamStateHandle, never()).discardState();
        }
    }
    // shutdown the store
    store.shutdown(JobStatus.SUSPENDED, new CheckpointsCleaner());
    // restore the store
    Set<ExecutionJobVertex> tasks = new HashSet<>();
    tasks.add(jobVertex1);
    assertEquals(JobStatus.SUSPENDED, store.getShutdownStatus().orElse(null));
    SharedStateRegistry secondInstance = SharedStateRegistry.DEFAULT_FACTORY.create(org.apache.flink.util.concurrent.Executors.directExecutor(), store.getAllCheckpoints());
    final EmbeddedCompletedCheckpointStore secondStore = new EmbeddedCompletedCheckpointStore(10, store.getAllCheckpoints(), secondInstance);
    final CheckpointCoordinator secondCoordinator = coordinatorBuilder.setCompletedCheckpointStore(secondStore).build();
    assertTrue(secondCoordinator.restoreLatestCheckpointedStateToAll(tasks, false));
    // validate that all shared states are registered again after the recovery.
    cp = 0;
    for (CompletedCheckpoint completedCheckpoint : completedCheckpoints) {
        for (OperatorState taskState : completedCheckpoint.getOperatorStates().values()) {
            for (OperatorSubtaskState subtaskState : taskState.getStates()) {
                for (KeyedStateHandle keyedStateHandle : subtaskState.getManagedKeyedState()) {
                    VerificationMode verificationMode;
                    // test we are once registered with the new registry
                    if (cp > 0) {
                        verificationMode = times(1);
                    } else {
                        verificationMode = never();
                    }
                    // check that all are registered with the new registry
                    verify(keyedStateHandle, verificationMode).registerSharedStates(secondInstance, completedCheckpoint.getCheckpointID());
                }
            }
        }
        ++cp;
    }
    // discard CP1
    secondStore.removeOldestCheckpoint();
    // we expect that all shared state from CP0 is no longer referenced and discarded. CP2 is
    // still live and also
    // references the state from CP1, so we expect they are not discarded.
    verifyDiscard(sharedHandlesByCheckpoint, cpId -> cpId == 0 ? times(1) : never());
    // discard CP2
    secondStore.removeOldestCheckpoint();
    // still expect shared state not to be discarded because it may be used in later checkpoints
    verifyDiscard(sharedHandlesByCheckpoint, cpId -> cpId == 1 ? never() : atLeast(0));
}