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

Example 41 with KeyGroupRange

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

the class InternalTimerServiceImplTest method testSnapshotAndRebalancingRestore.

private void testSnapshotAndRebalancingRestore(int snapshotVersion) throws Exception {
    @SuppressWarnings("unchecked") Triggerable<Integer, String> mockTriggerable = mock(Triggerable.class);
    TestKeyContext keyContext = new TestKeyContext();
    TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
    final PriorityQueueSetFactory queueFactory = createQueueFactory();
    InternalTimerServiceImpl<Integer, String> timerService = createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, queueFactory);
    int midpoint = testKeyGroupRange.getStartKeyGroup() + (testKeyGroupRange.getEndKeyGroup() - testKeyGroupRange.getStartKeyGroup()) / 2;
    // get two sub key-ranges so that we can restore two ranges separately
    KeyGroupRange subKeyGroupRange1 = new KeyGroupRange(testKeyGroupRange.getStartKeyGroup(), midpoint);
    KeyGroupRange subKeyGroupRange2 = new KeyGroupRange(midpoint + 1, testKeyGroupRange.getEndKeyGroup());
    // get two different keys, one per sub range
    int key1 = getKeyInKeyGroupRange(subKeyGroupRange1, maxParallelism);
    int key2 = getKeyInKeyGroupRange(subKeyGroupRange2, maxParallelism);
    keyContext.setCurrentKey(key1);
    timerService.registerProcessingTimeTimer("ciao", 10);
    timerService.registerEventTimeTimer("hello", 10);
    keyContext.setCurrentKey(key2);
    timerService.registerEventTimeTimer("ciao", 10);
    timerService.registerProcessingTimeTimer("hello", 10);
    assertEquals(2, timerService.numProcessingTimeTimers());
    assertEquals(1, timerService.numProcessingTimeTimers("hello"));
    assertEquals(1, timerService.numProcessingTimeTimers("ciao"));
    assertEquals(2, timerService.numEventTimeTimers());
    assertEquals(1, timerService.numEventTimeTimers("hello"));
    assertEquals(1, timerService.numEventTimeTimers("ciao"));
    // one map per sub key-group range
    Map<Integer, byte[]> snapshot1 = new HashMap<>();
    Map<Integer, byte[]> snapshot2 = new HashMap<>();
    for (Integer keyGroupIndex : testKeyGroupRange) {
        try (ByteArrayOutputStream outStream = new ByteArrayOutputStream()) {
            InternalTimersSnapshot<Integer, String> timersSnapshot = timerService.snapshotTimersForKeyGroup(keyGroupIndex);
            InternalTimersSnapshotReaderWriters.getWriterForVersion(snapshotVersion, timersSnapshot, timerService.getKeySerializer(), timerService.getNamespaceSerializer()).writeTimersSnapshot(new DataOutputViewStreamWrapper(outStream));
            if (subKeyGroupRange1.contains(keyGroupIndex)) {
                snapshot1.put(keyGroupIndex, outStream.toByteArray());
            } else if (subKeyGroupRange2.contains(keyGroupIndex)) {
                snapshot2.put(keyGroupIndex, outStream.toByteArray());
            } else {
                throw new IllegalStateException("Key-Group index doesn't belong to any sub range.");
            }
        }
    }
    // from now on we need everything twice. once per sub key-group range
    @SuppressWarnings("unchecked") Triggerable<Integer, String> mockTriggerable1 = mock(Triggerable.class);
    @SuppressWarnings("unchecked") Triggerable<Integer, String> mockTriggerable2 = mock(Triggerable.class);
    TestKeyContext keyContext1 = new TestKeyContext();
    TestKeyContext keyContext2 = new TestKeyContext();
    TestProcessingTimeService processingTimeService1 = new TestProcessingTimeService();
    TestProcessingTimeService processingTimeService2 = new TestProcessingTimeService();
    InternalTimerServiceImpl<Integer, String> timerService1 = restoreTimerService(snapshot1, snapshotVersion, mockTriggerable1, keyContext1, processingTimeService1, subKeyGroupRange1, queueFactory);
    InternalTimerServiceImpl<Integer, String> timerService2 = restoreTimerService(snapshot2, snapshotVersion, mockTriggerable2, keyContext2, processingTimeService2, subKeyGroupRange2, queueFactory);
    processingTimeService1.setCurrentTime(10);
    timerService1.advanceWatermark(10);
    verify(mockTriggerable1, times(1)).onProcessingTime(anyInternalTimer());
    verify(mockTriggerable1, times(1)).onProcessingTime(eq(new TimerHeapInternalTimer<>(10, key1, "ciao")));
    verify(mockTriggerable1, never()).onProcessingTime(eq(new TimerHeapInternalTimer<>(10, key2, "hello")));
    verify(mockTriggerable1, times(1)).onEventTime(anyInternalTimer());
    verify(mockTriggerable1, times(1)).onEventTime(eq(new TimerHeapInternalTimer<>(10, key1, "hello")));
    verify(mockTriggerable1, never()).onEventTime(eq(new TimerHeapInternalTimer<>(10, key2, "ciao")));
    assertEquals(0, timerService1.numEventTimeTimers());
    processingTimeService2.setCurrentTime(10);
    timerService2.advanceWatermark(10);
    verify(mockTriggerable2, times(1)).onProcessingTime(anyInternalTimer());
    verify(mockTriggerable2, never()).onProcessingTime(eq(new TimerHeapInternalTimer<>(10, key1, "ciao")));
    verify(mockTriggerable2, times(1)).onProcessingTime(eq(new TimerHeapInternalTimer<>(10, key2, "hello")));
    verify(mockTriggerable2, times(1)).onEventTime(anyInternalTimer());
    verify(mockTriggerable2, never()).onEventTime(eq(new TimerHeapInternalTimer<>(10, key1, "hello")));
    verify(mockTriggerable2, times(1)).onEventTime(eq(new TimerHeapInternalTimer<>(10, key2, "ciao")));
    assertEquals(0, timerService2.numEventTimeTimers());
}

Example 42 with KeyGroupRange

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

the class InternalTimerServiceImplTest method testKeyGroupStartIndexSetting.

@Test
public void testKeyGroupStartIndexSetting() {
    int startKeyGroupIdx = 7;
    int endKeyGroupIdx = 21;
    KeyGroupRange testKeyGroupList = new KeyGroupRange(startKeyGroupIdx, endKeyGroupIdx);
    TestKeyContext keyContext = new TestKeyContext();
    TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
    InternalTimerServiceImpl<Integer, String> service = createInternalTimerService(testKeyGroupList, keyContext, processingTimeService, IntSerializer.INSTANCE, StringSerializer.INSTANCE, createQueueFactory());
    Assert.assertEquals(startKeyGroupIdx, service.getLocalKeyGroupRangeStartIdx());
}

Example 43 with KeyGroupRange

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

the class InternalTimeServiceManagerImpl method create.

/**
 * A factory method for creating the {@link InternalTimeServiceManagerImpl}.
 *
 * <p><b>IMPORTANT:</b> Keep in sync with {@link InternalTimeServiceManager.Provider}.
 */
public static <K> InternalTimeServiceManagerImpl<K> create(CheckpointableKeyedStateBackend<K> keyedStateBackend, ClassLoader userClassloader, KeyContext keyContext, ProcessingTimeService processingTimeService, Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates) throws Exception {
    final KeyGroupRange keyGroupRange = keyedStateBackend.getKeyGroupRange();
    final InternalTimeServiceManagerImpl<K> timeServiceManager = new InternalTimeServiceManagerImpl<>(keyGroupRange, keyContext, keyedStateBackend, processingTimeService);
    // and then initialize the timer services
    for (KeyGroupStatePartitionStreamProvider streamProvider : rawKeyedStates) {
        int keyGroupIdx = streamProvider.getKeyGroupId();
        Preconditions.checkArgument(keyGroupRange.contains(keyGroupIdx), "Key Group " + keyGroupIdx + " does not belong to the local range.");
        timeServiceManager.restoreStateForKeyGroup(streamProvider.getStream(), keyGroupIdx, userClassloader);
    }
    return timeServiceManager;
}

Example 44 with KeyGroupRange

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

the class StateChangelogStorageTest method testWriteAndRead.

@Test
public void testWriteAndRead() throws Exception {
    KeyGroupRange kgRange = KeyGroupRange.of(0, 5);
    Map<Integer, List<byte[]>> appendsByKeyGroup = generateAppends(kgRange, 10, 20);
    try (StateChangelogStorage<T> client = getFactory();
        StateChangelogWriter<T> writer = client.createWriter(new OperatorID().toString(), kgRange)) {
        SequenceNumber prev = writer.initialSequenceNumber();
        for (Map.Entry<Integer, List<byte[]>> entry : appendsByKeyGroup.entrySet()) {
            Integer group = entry.getKey();
            List<byte[]> appends = entry.getValue();
            for (byte[] bytes : appends) {
                writer.append(group, bytes);
            }
        }
        T handle = writer.persist(prev).get();
        StateChangelogHandleReader<T> reader = client.createReader();
        assertByteMapsEqual(appendsByKeyGroup, extract(handle, reader));
    }
}

Example 45 with KeyGroupRange

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

the class InterruptSensitiveRestoreTest method createTask.

// ------------------------------------------------------------------------
// Utilities
// ------------------------------------------------------------------------
private static Task createTask(StreamConfig streamConfig, Configuration taskConfig, StreamStateHandle state, int mode) throws IOException {
    ShuffleEnvironment<?, ?> shuffleEnvironment = new NettyShuffleEnvironmentBuilder().build();
    Collection<KeyedStateHandle> keyedStateFromBackend = Collections.emptyList();
    Collection<KeyedStateHandle> keyedStateFromStream = Collections.emptyList();
    Collection<OperatorStateHandle> operatorStateBackend = Collections.emptyList();
    Collection<OperatorStateHandle> operatorStateStream = Collections.emptyList();
    Map<String, OperatorStateHandle.StateMetaInfo> operatorStateMetadata = new HashMap<>(1);
    OperatorStateHandle.StateMetaInfo metaInfo = new OperatorStateHandle.StateMetaInfo(new long[] { 0 }, OperatorStateHandle.Mode.SPLIT_DISTRIBUTE);
    operatorStateMetadata.put(DefaultOperatorStateBackend.DEFAULT_OPERATOR_STATE_NAME, metaInfo);
    KeyGroupRangeOffsets keyGroupRangeOffsets = new KeyGroupRangeOffsets(new KeyGroupRange(0, 0));
    Collection<OperatorStateHandle> operatorStateHandles = Collections.singletonList(new OperatorStreamStateHandle(operatorStateMetadata, state));
    List<KeyedStateHandle> keyedStateHandles = Collections.singletonList(new KeyGroupsStateHandle(keyGroupRangeOffsets, state));
    switch(mode) {
        case OPERATOR_MANAGED:
            operatorStateBackend = operatorStateHandles;
            break;
        case OPERATOR_RAW:
            operatorStateStream = operatorStateHandles;
            break;
        case KEYED_MANAGED:
            keyedStateFromBackend = keyedStateHandles;
            break;
        case KEYED_RAW:
            keyedStateFromStream = keyedStateHandles;
            break;
        default:
            throw new IllegalArgumentException();
    }
    OperatorSubtaskState operatorSubtaskState = OperatorSubtaskState.builder().setManagedOperatorState(new StateObjectCollection<>(operatorStateBackend)).setRawOperatorState(new StateObjectCollection<>(operatorStateStream)).setManagedKeyedState(new StateObjectCollection<>(keyedStateFromBackend)).setRawKeyedState(new StateObjectCollection<>(keyedStateFromStream)).build();
    JobVertexID jobVertexID = new JobVertexID();
    OperatorID operatorID = OperatorID.fromJobVertexID(jobVertexID);
    streamConfig.setOperatorID(operatorID);
    TaskStateSnapshot stateSnapshot = new TaskStateSnapshot();
    stateSnapshot.putSubtaskStateByOperatorID(operatorID, operatorSubtaskState);
    JobManagerTaskRestore taskRestore = new JobManagerTaskRestore(1L, stateSnapshot);
    JobInformation jobInformation = new JobInformation(new JobID(), "test job name", new SerializedValue<>(new ExecutionConfig()), new Configuration(), Collections.emptyList(), Collections.emptyList());
    TaskInformation taskInformation = new TaskInformation(jobVertexID, "test task name", 1, 1, SourceStreamTask.class.getName(), taskConfig);
    TestTaskStateManager taskStateManager = TestTaskStateManager.builder().setReportedCheckpointId(taskRestore.getRestoreCheckpointId()).setJobManagerTaskStateSnapshotsByCheckpointId(Collections.singletonMap(taskRestore.getRestoreCheckpointId(), taskRestore.getTaskStateSnapshot())).build();
    return new Task(jobInformation, taskInformation, new ExecutionAttemptID(), new AllocationID(), 0, 0, Collections.<ResultPartitionDeploymentDescriptor>emptyList(), Collections.<InputGateDeploymentDescriptor>emptyList(), mock(MemoryManager.class), mock(IOManager.class), shuffleEnvironment, new KvStateService(new KvStateRegistry(), null, null), mock(BroadcastVariableManager.class), new TaskEventDispatcher(), ExternalResourceInfoProvider.NO_EXTERNAL_RESOURCES, taskStateManager, mock(TaskManagerActions.class), mock(InputSplitProvider.class), mock(CheckpointResponder.class), new NoOpTaskOperatorEventGateway(), new TestGlobalAggregateManager(), TestingClassLoaderLease.newBuilder().build(), new FileCache(new String[] { EnvironmentInformation.getTemporaryFileDirectory() }, VoidPermanentBlobService.INSTANCE), new TestingTaskManagerRuntimeInfo(), UnregisteredMetricGroups.createUnregisteredTaskMetricGroup(), new NoOpResultPartitionConsumableNotifier(), mock(PartitionProducerStateChecker.class), mock(Executor.class));
}