Java Code Examples for org.apache.flink.runtime.concurrent.ManuallyTriggeredScheduledExecutor#triggerAll()

/**
 * This test makes sure that the checkpoint is already registered by the time.
 * that the hooks are called
 */
@Test
public void ensureRegisteredAtHookTime() throws Exception {
	final String id = "id";

	// create the checkpoint coordinator
	final JobID jid = new JobID();
	final ExecutionAttemptID execId = new ExecutionAttemptID();
	final ExecutionVertex ackVertex = mockExecutionVertex(execId);
	final ManuallyTriggeredScheduledExecutor manuallyTriggeredScheduledExecutor =
		new ManuallyTriggeredScheduledExecutor();
	final CheckpointCoordinator cc = instantiateCheckpointCoordinator(
		jid, manuallyTriggeredScheduledExecutor, ackVertex);

	final MasterTriggerRestoreHook<Void> hook = mockGeneric(MasterTriggerRestoreHook.class);
	when(hook.getIdentifier()).thenReturn(id);
	when(hook.triggerCheckpoint(anyLong(), anyLong(), any(Executor.class))).thenAnswer(
			new Answer<CompletableFuture<Void>>() {

				@Override
				public CompletableFuture<Void> answer(InvocationOnMock invocation) throws Throwable {
					assertEquals(1, cc.getNumberOfPendingCheckpoints());

					long checkpointId = (Long) invocation.getArguments()[0];
					assertNotNull(cc.getPendingCheckpoints().get(checkpointId));
					return null;
				}
			}
	);

	cc.addMasterHook(hook);

	// trigger a checkpoint
	final CompletableFuture<CompletedCheckpoint> checkpointFuture = cc.triggerCheckpoint(false);
	manuallyTriggeredScheduledExecutor.triggerAll();
	assertFalse(checkpointFuture.isCompletedExceptionally());
}
 

/**
 * Tests that idle but not releasable task managers will not be released even if timed out before it can be.
 */
@Test
public void testTaskManagerNotReleasedBeforeItCanBe() throws Exception {
	final CompletableFuture<InstanceID> releaseFuture = new CompletableFuture<>();
	final ResourceActions resourceManagerActions = new TestingResourceActionsBuilder()
		.setReleaseResourceConsumer((instanceID, e) -> releaseFuture.complete(instanceID))
		.build();
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceID resourceID = ResourceID.generate();

	final AtomicBoolean canBeReleased = new AtomicBoolean(false);
	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
		.setCanBeReleasedSupplier(canBeReleased::get)
		.createTestingTaskExecutorGateway();
	final TaskExecutorConnection taskManagerConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);

	final SlotID slotId = new SlotID(resourceID, 0);
	final ResourceProfile resourceProfile = new ResourceProfile(1.0, 1);
	final SlotStatus slotStatus = new SlotStatus(slotId, resourceProfile);
	final SlotReport slotReport = new SlotReport(slotStatus);

	final ManuallyTriggeredScheduledExecutor mainThreadExecutor = new ManuallyTriggeredScheduledExecutor();

	try (SlotManager slotManager = SlotManagerBuilder.newBuilder()
		.setScheduledExecutor(mainThreadExecutor)
		.setTaskManagerTimeout(Time.milliseconds(0L))
		.build()) {

		slotManager.start(resourceManagerId, mainThreadExecutor, resourceManagerActions);

		mainThreadExecutor.execute(() -> slotManager.registerTaskManager(taskManagerConnection, slotReport));

		// now it can not be released yet
		canBeReleased.set(false);
		mainThreadExecutor.execute(slotManager::checkTaskManagerTimeouts);
		mainThreadExecutor.triggerAll();
		assertFalse(releaseFuture.isDone());

		// now it can and should be released
		canBeReleased.set(true);
		mainThreadExecutor.execute(slotManager::checkTaskManagerTimeouts);
		mainThreadExecutor.triggerAll();
		assertThat(releaseFuture.get(), is(equalTo(taskManagerConnection.getInstanceID())));
	}
}
 

/**
 * Tests that a failure while storing a completed checkpoint in the completed checkpoint store
 * will properly fail the originating pending checkpoint and clean upt the completed checkpoint.
 */
@Test
public void testFailingCompletedCheckpointStoreAdd() throws Exception {
	JobID jid = new JobID();

	final ManuallyTriggeredScheduledExecutor manuallyTriggeredScheduledExecutor =
		new ManuallyTriggeredScheduledExecutor();

	final ExecutionAttemptID executionAttemptId = new ExecutionAttemptID();
	final ExecutionVertex vertex = CheckpointCoordinatorTestingUtils.mockExecutionVertex(executionAttemptId);

	// set up the coordinator and validate the initial state
	CheckpointCoordinator coord =
		new CheckpointCoordinatorBuilder()
			.setJobId(jid)
			.setTasks(new ExecutionVertex[] { vertex })
			.setCompletedCheckpointStore(new FailingCompletedCheckpointStore())
			.setTimer(manuallyTriggeredScheduledExecutor)
			.build();

	coord.triggerCheckpoint(false);

	manuallyTriggeredScheduledExecutor.triggerAll();

	assertEquals(1, coord.getNumberOfPendingCheckpoints());

	PendingCheckpoint pendingCheckpoint = coord.getPendingCheckpoints().values().iterator().next();

	assertFalse(pendingCheckpoint.isDiscarded());

	final long checkpointId = coord.getPendingCheckpoints().keySet().iterator().next();

	KeyedStateHandle managedKeyedHandle = mock(KeyedStateHandle.class);
	KeyedStateHandle rawKeyedHandle = mock(KeyedStateHandle.class);
	OperatorStateHandle managedOpHandle = mock(OperatorStreamStateHandle.class);
	OperatorStateHandle rawOpHandle = mock(OperatorStreamStateHandle.class);
	InputChannelStateHandle inputChannelStateHandle = new InputChannelStateHandle(new InputChannelInfo(0, 1), mock(StreamStateHandle.class), Collections.singletonList(1L));
	ResultSubpartitionStateHandle resultSubpartitionStateHandle = new ResultSubpartitionStateHandle(new ResultSubpartitionInfo(0, 1), mock(StreamStateHandle.class), Collections.singletonList(1L));

	final OperatorSubtaskState operatorSubtaskState = spy(new OperatorSubtaskState(
		managedOpHandle,
		rawOpHandle,
		managedKeyedHandle,
		rawKeyedHandle,
		StateObjectCollection.singleton(inputChannelStateHandle),
		StateObjectCollection.singleton(resultSubpartitionStateHandle)));

	TaskStateSnapshot subtaskState = spy(new TaskStateSnapshot());
	subtaskState.putSubtaskStateByOperatorID(new OperatorID(), operatorSubtaskState);

	when(subtaskState.getSubtaskStateByOperatorID(OperatorID.fromJobVertexID(vertex.getJobvertexId()))).thenReturn(operatorSubtaskState);

	AcknowledgeCheckpoint acknowledgeMessage = new AcknowledgeCheckpoint(jid, executionAttemptId, checkpointId, new CheckpointMetrics(), subtaskState);

	try {
		coord.receiveAcknowledgeMessage(acknowledgeMessage, "Unknown location");
		fail("Expected a checkpoint exception because the completed checkpoint store could not " +
			"store the completed checkpoint.");
	} catch (CheckpointException e) {
		// ignore because we expected this exception
	}

	// make sure that the pending checkpoint has been discarded after we could not complete it
	assertTrue(pendingCheckpoint.isDiscarded());

	// make sure that the subtask state has been discarded after we could not complete it.
	verify(operatorSubtaskState).discardState();
	verify(operatorSubtaskState.getManagedOperatorState().iterator().next()).discardState();
	verify(operatorSubtaskState.getRawOperatorState().iterator().next()).discardState();
	verify(operatorSubtaskState.getManagedKeyedState().iterator().next()).discardState();
	verify(operatorSubtaskState.getRawKeyedState().iterator().next()).discardState();
	verify(operatorSubtaskState.getInputChannelState().iterator().next().getDelegate()).discardState();
	verify(operatorSubtaskState.getResultSubpartitionState().iterator().next().getDelegate()).discardState();
}
 

@Test
public void testHooksAreCalledOnTrigger() throws Exception {
	final String id1 = "id1";
	final String id2 = "id2";

	final String state1 = "the-test-string-state";
	final byte[] state1serialized = new StringSerializer().serialize(state1);

	final long state2 = 987654321L;
	final byte[] state2serialized = new LongSerializer().serialize(state2);

	final MasterTriggerRestoreHook<String> statefulHook1 = mockGeneric(MasterTriggerRestoreHook.class);
	when(statefulHook1.getIdentifier()).thenReturn(id1);
	when(statefulHook1.createCheckpointDataSerializer()).thenReturn(new StringSerializer());
	when(statefulHook1.triggerCheckpoint(anyLong(), anyLong(), any(Executor.class)))
			.thenReturn(CompletableFuture.completedFuture(state1));

	final MasterTriggerRestoreHook<Long> statefulHook2 = mockGeneric(MasterTriggerRestoreHook.class);
	when(statefulHook2.getIdentifier()).thenReturn(id2);
	when(statefulHook2.createCheckpointDataSerializer()).thenReturn(new LongSerializer());
	when(statefulHook2.triggerCheckpoint(anyLong(), anyLong(), any(Executor.class)))
			.thenReturn(CompletableFuture.completedFuture(state2));

	final MasterTriggerRestoreHook<Void> statelessHook = mockGeneric(MasterTriggerRestoreHook.class);
	when(statelessHook.getIdentifier()).thenReturn("some-id");

	// create the checkpoint coordinator
	final JobID jid = new JobID();
	final ExecutionAttemptID execId = new ExecutionAttemptID();
	final ExecutionVertex ackVertex = mockExecutionVertex(execId);
	final ManuallyTriggeredScheduledExecutor manuallyTriggeredScheduledExecutor =
		new ManuallyTriggeredScheduledExecutor();
	final CheckpointCoordinator cc = instantiateCheckpointCoordinator(
		jid, manuallyTriggeredScheduledExecutor, ackVertex);

	cc.addMasterHook(statefulHook1);
	cc.addMasterHook(statelessHook);
	cc.addMasterHook(statefulHook2);

	// trigger a checkpoint
	final CompletableFuture<CompletedCheckpoint> checkpointFuture = cc.triggerCheckpoint(false);
	manuallyTriggeredScheduledExecutor.triggerAll();
	assertFalse(checkpointFuture.isCompletedExceptionally());
	assertEquals(1, cc.getNumberOfPendingCheckpoints());

	verify(statefulHook1, times(1)).triggerCheckpoint(anyLong(), anyLong(), any(Executor.class));
	verify(statefulHook2, times(1)).triggerCheckpoint(anyLong(), anyLong(), any(Executor.class));
	verify(statelessHook, times(1)).triggerCheckpoint(anyLong(), anyLong(), any(Executor.class));

	final long checkpointId = cc.getPendingCheckpoints().values().iterator().next().getCheckpointId();
	cc.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(jid, execId, checkpointId), "Unknown location");
	assertEquals(0, cc.getNumberOfPendingCheckpoints());

	assertEquals(1, cc.getNumberOfRetainedSuccessfulCheckpoints());
	final CompletedCheckpoint chk = cc.getCheckpointStore().getLatestCheckpoint(false);

	final Collection<MasterState> masterStates = chk.getMasterHookStates();
	assertEquals(2, masterStates.size());

	for (MasterState ms : masterStates) {
		if (ms.name().equals(id1)) {
			assertArrayEquals(state1serialized, ms.bytes());
			assertEquals(StringSerializer.VERSION, ms.version());
		}
		else if (ms.name().equals(id2)) {
			assertArrayEquals(state2serialized, ms.bytes());
			assertEquals(LongSerializer.VERSION, ms.version());
		}
		else {
			fail("unrecognized state name: " + ms.name());
		}
	}
}