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

Example 26 with KeyGroupsStateHandle

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

the class CheckpointCoordinatorTest method testRestoreLatestCheckpointFailureWhenParallelismChanges.

/**
	 * Tests that the checkpoint restoration fails if the parallelism of a job vertices with
	 * non-partitioned state has changed.
	 *
	 * @throws Exception
	 */
@Test(expected = IllegalStateException.class)
public void testRestoreLatestCheckpointFailureWhenParallelismChanges() throws Exception {
    final JobID jid = new JobID();
    final long timestamp = System.currentTimeMillis();
    final JobVertexID jobVertexID1 = new JobVertexID();
    final JobVertexID jobVertexID2 = new JobVertexID();
    int parallelism1 = 3;
    int parallelism2 = 2;
    int maxParallelism1 = 42;
    int maxParallelism2 = 13;
    final ExecutionJobVertex jobVertex1 = mockExecutionJobVertex(jobVertexID1, parallelism1, maxParallelism1);
    final ExecutionJobVertex jobVertex2 = mockExecutionJobVertex(jobVertexID2, parallelism2, maxParallelism2);
    List<ExecutionVertex> allExecutionVertices = new ArrayList<>(parallelism1 + parallelism2);
    allExecutionVertices.addAll(Arrays.asList(jobVertex1.getTaskVertices()));
    allExecutionVertices.addAll(Arrays.asList(jobVertex2.getTaskVertices()));
    ExecutionVertex[] arrayExecutionVertices = allExecutionVertices.toArray(new ExecutionVertex[allExecutionVertices.size()]);
    // set up the coordinator and validate the initial state
    CheckpointCoordinator coord = new CheckpointCoordinator(jid, 600000, 600000, 0, Integer.MAX_VALUE, ExternalizedCheckpointSettings.none(), arrayExecutionVertices, arrayExecutionVertices, arrayExecutionVertices, new StandaloneCheckpointIDCounter(), new StandaloneCompletedCheckpointStore(1), null, Executors.directExecutor());
    // trigger the checkpoint
    coord.triggerCheckpoint(timestamp, false);
    assertTrue(coord.getPendingCheckpoints().keySet().size() == 1);
    long checkpointId = Iterables.getOnlyElement(coord.getPendingCheckpoints().keySet());
    CheckpointMetaData checkpointMetaData = new CheckpointMetaData(checkpointId, 0L);
    List<KeyGroupRange> keyGroupPartitions1 = StateAssignmentOperation.createKeyGroupPartitions(maxParallelism1, parallelism1);
    List<KeyGroupRange> keyGroupPartitions2 = StateAssignmentOperation.createKeyGroupPartitions(maxParallelism2, parallelism2);
    for (int index = 0; index < jobVertex1.getParallelism(); index++) {
        ChainedStateHandle<StreamStateHandle> valueSizeTuple = generateStateForVertex(jobVertexID1, index);
        KeyGroupsStateHandle keyGroupState = generateKeyGroupState(jobVertexID1, keyGroupPartitions1.get(index), false);
        SubtaskState checkpointStateHandles = new SubtaskState(valueSizeTuple, null, null, keyGroupState, null);
        AcknowledgeCheckpoint acknowledgeCheckpoint = new AcknowledgeCheckpoint(jid, jobVertex1.getTaskVertices()[index].getCurrentExecutionAttempt().getAttemptId(), checkpointId, new CheckpointMetrics(), checkpointStateHandles);
        coord.receiveAcknowledgeMessage(acknowledgeCheckpoint);
    }
    for (int index = 0; index < jobVertex2.getParallelism(); index++) {
        ChainedStateHandle<StreamStateHandle> state = generateStateForVertex(jobVertexID2, index);
        KeyGroupsStateHandle keyGroupState = generateKeyGroupState(jobVertexID2, keyGroupPartitions2.get(index), false);
        SubtaskState checkpointStateHandles = new SubtaskState(state, null, null, keyGroupState, null);
        AcknowledgeCheckpoint acknowledgeCheckpoint = new AcknowledgeCheckpoint(jid, jobVertex2.getTaskVertices()[index].getCurrentExecutionAttempt().getAttemptId(), checkpointId, new CheckpointMetrics(), checkpointStateHandles);
        coord.receiveAcknowledgeMessage(acknowledgeCheckpoint);
    }
    List<CompletedCheckpoint> completedCheckpoints = coord.getSuccessfulCheckpoints();
    assertEquals(1, completedCheckpoints.size());
    Map<JobVertexID, ExecutionJobVertex> tasks = new HashMap<>();
    int newParallelism1 = 4;
    int newParallelism2 = 3;
    final ExecutionJobVertex newJobVertex1 = mockExecutionJobVertex(jobVertexID1, newParallelism1, maxParallelism1);
    final ExecutionJobVertex newJobVertex2 = mockExecutionJobVertex(jobVertexID2, newParallelism2, maxParallelism2);
    tasks.put(jobVertexID1, newJobVertex1);
    tasks.put(jobVertexID2, newJobVertex2);
    coord.restoreLatestCheckpointedState(tasks, true, false);
    fail("The restoration should have failed because the parallelism of an vertex with " + "non-partitioned state changed.");
}

Example 27 with KeyGroupsStateHandle

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

the class HeapKeyedStateBackendSnapshotMigrationTest method testRestore1_2ToMaster.

/**
	 * [FLINK-5979]
	 *
	 * This test takes a snapshot that was created with Flink 1.2 and tries to restore it in master to check
	 * the backwards compatibility of the serialization format of {@link StateTable}s.
	 */
@Test
public void testRestore1_2ToMaster() throws Exception {
    ClassLoader cl = getClass().getClassLoader();
    URL resource = cl.getResource("heap_keyed_statebackend_1_2.snapshot");
    Preconditions.checkNotNull(resource, "Binary snapshot resource not found!");
    final Integer namespace1 = 1;
    final Integer namespace2 = 2;
    final Integer namespace3 = 3;
    try (final HeapKeyedStateBackend<String> keyedBackend = createKeyedBackend()) {
        final KeyGroupsStateHandle stateHandle;
        try (BufferedInputStream bis = new BufferedInputStream((new FileInputStream(resource.getFile())))) {
            stateHandle = InstantiationUtil.deserializeObject(bis, Thread.currentThread().getContextClassLoader());
        }
        keyedBackend.restore(Collections.singleton(stateHandle));
        final ListStateDescriptor<Long> stateDescr = new ListStateDescriptor<>("my-state", Long.class);
        stateDescr.initializeSerializerUnlessSet(new ExecutionConfig());
        InternalListState<Integer, Long> state = keyedBackend.createListState(IntSerializer.INSTANCE, stateDescr);
        assertEquals(7, keyedBackend.numStateEntries());
        keyedBackend.setCurrentKey("abc");
        state.setCurrentNamespace(namespace1);
        assertEquals(asList(33L, 55L), state.get());
        state.setCurrentNamespace(namespace2);
        assertEquals(asList(22L, 11L), state.get());
        state.setCurrentNamespace(namespace3);
        assertEquals(Collections.singletonList(44L), state.get());
        keyedBackend.setCurrentKey("def");
        state.setCurrentNamespace(namespace1);
        assertEquals(asList(11L, 44L), state.get());
        state.setCurrentNamespace(namespace3);
        assertEquals(asList(22L, 55L, 33L), state.get());
        keyedBackend.setCurrentKey("jkl");
        state.setCurrentNamespace(namespace1);
        assertEquals(asList(11L, 22L, 33L, 44L, 55L), state.get());
        keyedBackend.setCurrentKey("mno");
        state.setCurrentNamespace(namespace3);
        assertEquals(asList(11L, 22L, 33L, 44L, 55L), state.get());
    }
}

Example 28 with KeyGroupsStateHandle

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

the class AbstractStreamOperator method initializeState.

@Override
public final void initializeState(OperatorStateHandles stateHandles) throws Exception {
    Collection<KeyGroupsStateHandle> keyedStateHandlesRaw = null;
    Collection<OperatorStateHandle> operatorStateHandlesRaw = null;
    Collection<OperatorStateHandle> operatorStateHandlesBackend = null;
    boolean restoring = null != stateHandles;
    //TODO we should move the actual initialization of this from StreamTask to this class
    initKeyedState();
    if (getKeyedStateBackend() != null && timeServiceManager == null) {
        timeServiceManager = new InternalTimeServiceManager<>(getKeyedStateBackend().getNumberOfKeyGroups(), getKeyedStateBackend().getKeyGroupRange(), this, getRuntimeContext().getProcessingTimeService());
    }
    if (restoring) {
        //pass directly
        operatorStateHandlesBackend = stateHandles.getManagedOperatorState();
        operatorStateHandlesRaw = stateHandles.getRawOperatorState();
        if (null != getKeyedStateBackend()) {
            //only use the keyed state if it is meant for us (aka head operator)
            keyedStateHandlesRaw = stateHandles.getRawKeyedState();
        }
    }
    checkpointStreamFactory = container.createCheckpointStreamFactory(this);
    initOperatorState(operatorStateHandlesBackend);
    StateInitializationContext initializationContext = new StateInitializationContextImpl(// information whether we restore or start for the first time
    restoring, // access to operator state backend
    operatorStateBackend, // access to keyed state backend
    keyedStateStore, // access to keyed state stream
    keyedStateHandlesRaw, // access to operator state stream
    operatorStateHandlesRaw, // access to register streams for canceling
    getContainingTask().getCancelables());
    initializeState(initializationContext);
    if (restoring) {
        // finally restore the legacy state in case we are
        // migrating from a previous Flink version.
        restoreStreamCheckpointed(stateHandles);
    }
}

Example 29 with KeyGroupsStateHandle

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

the class OperatorSnapshotResultTest method testCancelAndCleanup.

/**
	 * Tests that all runnable futures in an OperatorSnapshotResult are properly cancelled and if
	 * the StreamStateHandle result is retrievable that the state handle are discarded.
	 */
@Test
public void testCancelAndCleanup() throws Exception {
    OperatorSnapshotResult operatorSnapshotResult = new OperatorSnapshotResult();
    operatorSnapshotResult.cancel();
    KeyGroupsStateHandle keyedManagedStateHandle = mock(KeyGroupsStateHandle.class);
    RunnableFuture<KeyGroupsStateHandle> keyedStateManagedFuture = mock(RunnableFuture.class);
    when(keyedStateManagedFuture.get()).thenReturn(keyedManagedStateHandle);
    KeyGroupsStateHandle keyedRawStateHandle = mock(KeyGroupsStateHandle.class);
    RunnableFuture<KeyGroupsStateHandle> keyedStateRawFuture = mock(RunnableFuture.class);
    when(keyedStateRawFuture.get()).thenReturn(keyedRawStateHandle);
    OperatorStateHandle operatorManagedStateHandle = mock(OperatorStateHandle.class);
    RunnableFuture<OperatorStateHandle> operatorStateManagedFuture = mock(RunnableFuture.class);
    when(operatorStateManagedFuture.get()).thenReturn(operatorManagedStateHandle);
    OperatorStateHandle operatorRawStateHandle = mock(OperatorStateHandle.class);
    RunnableFuture<OperatorStateHandle> operatorStateRawFuture = mock(RunnableFuture.class);
    when(operatorStateRawFuture.get()).thenReturn(operatorRawStateHandle);
    operatorSnapshotResult = new OperatorSnapshotResult(keyedStateManagedFuture, keyedStateRawFuture, operatorStateManagedFuture, operatorStateRawFuture);
    operatorSnapshotResult.cancel();
    verify(keyedStateManagedFuture).cancel(true);
    verify(keyedStateRawFuture).cancel(true);
    verify(operatorStateManagedFuture).cancel(true);
    verify(operatorStateRawFuture).cancel(true);
    verify(keyedManagedStateHandle).discardState();
    verify(keyedRawStateHandle).discardState();
    verify(operatorManagedStateHandle).discardState();
    verify(operatorRawStateHandle).discardState();
}

Example 30 with KeyGroupsStateHandle

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

the class StreamTaskTest method testAsyncCheckpointingConcurrentCloseAfterAcknowledge.

/**
	 * FLINK-5667
	 *
	 * Tests that a concurrent cancel operation does not discard the state handles of an
	 * acknowledged checkpoint. The situation can only happen if the cancel call is executed
	 * after Environment.acknowledgeCheckpoint() and before the
	 * CloseableRegistry.unregisterClosable() call.
	 */
@Test
public void testAsyncCheckpointingConcurrentCloseAfterAcknowledge() throws Exception {
    final long checkpointId = 42L;
    final long timestamp = 1L;
    final OneShotLatch acknowledgeCheckpointLatch = new OneShotLatch();
    final OneShotLatch completeAcknowledge = new OneShotLatch();
    TaskInfo mockTaskInfo = mock(TaskInfo.class);
    when(mockTaskInfo.getTaskNameWithSubtasks()).thenReturn("foobar");
    when(mockTaskInfo.getIndexOfThisSubtask()).thenReturn(0);
    Environment mockEnvironment = mock(Environment.class);
    when(mockEnvironment.getTaskInfo()).thenReturn(mockTaskInfo);
    doAnswer(new Answer() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            acknowledgeCheckpointLatch.trigger();
            // block here so that we can issue the concurrent cancel call
            completeAcknowledge.await();
            return null;
        }
    }).when(mockEnvironment).acknowledgeCheckpoint(anyLong(), any(CheckpointMetrics.class), any(SubtaskState.class));
    StreamTask<?, AbstractStreamOperator<?>> streamTask = mock(StreamTask.class, Mockito.CALLS_REAL_METHODS);
    CheckpointMetaData checkpointMetaData = new CheckpointMetaData(checkpointId, timestamp);
    streamTask.setEnvironment(mockEnvironment);
    StreamOperator<?> streamOperator = mock(StreamOperator.class, withSettings().extraInterfaces(StreamCheckpointedOperator.class));
    KeyGroupsStateHandle managedKeyedStateHandle = mock(KeyGroupsStateHandle.class);
    KeyGroupsStateHandle rawKeyedStateHandle = mock(KeyGroupsStateHandle.class);
    OperatorStateHandle managedOperatorStateHandle = mock(OperatorStateHandle.class);
    OperatorStateHandle rawOperatorStateHandle = mock(OperatorStateHandle.class);
    OperatorSnapshotResult operatorSnapshotResult = new OperatorSnapshotResult(new DoneFuture<>(managedKeyedStateHandle), new DoneFuture<>(rawKeyedStateHandle), new DoneFuture<>(managedOperatorStateHandle), new DoneFuture<>(rawOperatorStateHandle));
    when(streamOperator.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class))).thenReturn(operatorSnapshotResult);
    StreamOperator<?>[] streamOperators = { streamOperator };
    OperatorChain<Void, AbstractStreamOperator<Void>> operatorChain = mock(OperatorChain.class);
    when(operatorChain.getAllOperators()).thenReturn(streamOperators);
    StreamStateHandle streamStateHandle = mock(StreamStateHandle.class);
    CheckpointStreamFactory.CheckpointStateOutputStream outStream = mock(CheckpointStreamFactory.CheckpointStateOutputStream.class);
    when(outStream.closeAndGetHandle()).thenReturn(streamStateHandle);
    CheckpointStreamFactory mockStreamFactory = mock(CheckpointStreamFactory.class);
    when(mockStreamFactory.createCheckpointStateOutputStream(anyLong(), anyLong())).thenReturn(outStream);
    AbstractStateBackend mockStateBackend = mock(AbstractStateBackend.class);
    when(mockStateBackend.createStreamFactory(any(JobID.class), anyString())).thenReturn(mockStreamFactory);
    Whitebox.setInternalState(streamTask, "isRunning", true);
    Whitebox.setInternalState(streamTask, "lock", new Object());
    Whitebox.setInternalState(streamTask, "operatorChain", operatorChain);
    Whitebox.setInternalState(streamTask, "cancelables", new CloseableRegistry());
    Whitebox.setInternalState(streamTask, "asyncOperationsThreadPool", Executors.newFixedThreadPool(1));
    Whitebox.setInternalState(streamTask, "configuration", new StreamConfig(new Configuration()));
    Whitebox.setInternalState(streamTask, "stateBackend", mockStateBackend);
    streamTask.triggerCheckpoint(checkpointMetaData, CheckpointOptions.forFullCheckpoint());
    acknowledgeCheckpointLatch.await();
    ArgumentCaptor<SubtaskState> subtaskStateCaptor = ArgumentCaptor.forClass(SubtaskState.class);
    // check that the checkpoint has been completed
    verify(mockEnvironment).acknowledgeCheckpoint(eq(checkpointId), any(CheckpointMetrics.class), subtaskStateCaptor.capture());
    SubtaskState subtaskState = subtaskStateCaptor.getValue();
    // check that the subtask state contains the expected state handles
    assertEquals(managedKeyedStateHandle, subtaskState.getManagedKeyedState());
    assertEquals(rawKeyedStateHandle, subtaskState.getRawKeyedState());
    assertEquals(new ChainedStateHandle<>(Collections.singletonList(managedOperatorStateHandle)), subtaskState.getManagedOperatorState());
    assertEquals(new ChainedStateHandle<>(Collections.singletonList(rawOperatorStateHandle)), subtaskState.getRawOperatorState());
    // check that the state handles have not been discarded
    verify(managedKeyedStateHandle, never()).discardState();
    verify(rawKeyedStateHandle, never()).discardState();
    verify(managedOperatorStateHandle, never()).discardState();
    verify(rawOperatorStateHandle, never()).discardState();
    streamTask.cancel();
    completeAcknowledge.trigger();
    // canceling the stream task after it has acknowledged the checkpoint should not discard
    // the state handles
    verify(managedKeyedStateHandle, never()).discardState();
    verify(rawKeyedStateHandle, never()).discardState();
    verify(managedOperatorStateHandle, never()).discardState();
    verify(rawOperatorStateHandle, never()).discardState();
}