Java Code Examples for org.apache.flink.runtime.executiongraph.utils.SimpleAckingTaskManagerGateway#CheckpointConsumer

@Test
public void testTriggerCheckpointBeforePreviousOneCompleted() throws Exception {
	// create some mock Execution vertices that receive the checkpoint trigger messages
	final ExecutionAttemptID attemptID = new ExecutionAttemptID();
	final AtomicInteger taskManagerCheckpointTriggeredTimes = new AtomicInteger(0);
	final SimpleAckingTaskManagerGateway.CheckpointConsumer checkpointConsumer =
		(executionAttemptID,
		jobId, checkpointId,
		timestamp,
		checkpointOptions,
		advanceToEndOfEventTime) -> taskManagerCheckpointTriggeredTimes.incrementAndGet();
	ExecutionVertex vertex = mockExecutionVertex(attemptID, checkpointConsumer);

	// set up the coordinator and validate the initial state
	CheckpointCoordinator checkpointCoordinator = createCheckpointCoordinator(vertex);

	checkpointCoordinator.startCheckpointScheduler();
	// start a periodic checkpoint first
	final CompletableFuture<CompletedCheckpoint> onCompletionPromise1 =
		triggerPeriodicCheckpoint(checkpointCoordinator);
	assertTrue(checkpointCoordinator.isTriggering());
	assertEquals(0, checkpointCoordinator.getTriggerRequestQueue().size());
	final CompletableFuture<CompletedCheckpoint> onCompletionPromise2 =
		triggerPeriodicCheckpoint(checkpointCoordinator);
	// another trigger before the prior one finished

	assertTrue(checkpointCoordinator.isTriggering());
	assertEquals(1, checkpointCoordinator.getTriggerRequestQueue().size());

	manuallyTriggeredScheduledExecutor.triggerAll();
	assertFalse(onCompletionPromise1.isCompletedExceptionally());
	assertFalse(onCompletionPromise2.isCompletedExceptionally());
	assertFalse(checkpointCoordinator.isTriggering());
	assertEquals(0, checkpointCoordinator.getTriggerRequestQueue().size());
	assertEquals(2, taskManagerCheckpointTriggeredTimes.get());
}
 

@Test
public void testTriggerCheckpointRequestQueuedWithFailure() throws Exception {
	// create some mock Execution vertices that receive the checkpoint trigger messages
	final ExecutionAttemptID attemptID = new ExecutionAttemptID();
	final AtomicInteger taskManagerCheckpointTriggeredTimes = new AtomicInteger(0);
	final SimpleAckingTaskManagerGateway.CheckpointConsumer checkpointConsumer =
		(executionAttemptID,
		jobId, checkpointId,
		timestamp,
		checkpointOptions,
		advanceToEndOfEventTime) -> taskManagerCheckpointTriggeredTimes.incrementAndGet();
	ExecutionVertex vertex = mockExecutionVertex(attemptID, checkpointConsumer);

	// set up the coordinator and validate the initial state
	CheckpointCoordinator checkpointCoordinator =
		new CheckpointCoordinatorBuilder()
			.setTasks(new ExecutionVertex[] { vertex })
			.setCheckpointIDCounter(new UnstableCheckpointIDCounter(id -> id == 0))
			.setTimer(manuallyTriggeredScheduledExecutor)
			.build();

	checkpointCoordinator.startCheckpointScheduler();
	// start a periodic checkpoint first
	final CompletableFuture<CompletedCheckpoint> onCompletionPromise1 =
		triggerNonPeriodicCheckpoint(checkpointCoordinator);
	assertTrue(checkpointCoordinator.isTriggering());
	assertEquals(0, checkpointCoordinator.getTriggerRequestQueue().size());

	// another trigger before the prior one finished
	final CompletableFuture<CompletedCheckpoint> onCompletionPromise2 =
		triggerNonPeriodicCheckpoint(checkpointCoordinator);

	// another trigger before the first one finished
	final CompletableFuture<CompletedCheckpoint> onCompletionPromise3 =
		triggerNonPeriodicCheckpoint(checkpointCoordinator);
	assertTrue(checkpointCoordinator.isTriggering());
	assertEquals(2, checkpointCoordinator.getTriggerRequestQueue().size());

	manuallyTriggeredScheduledExecutor.triggerAll();
	// the first triggered checkpoint fails by design through UnstableCheckpointIDCounter
	assertTrue(onCompletionPromise1.isCompletedExceptionally());
	assertFalse(onCompletionPromise2.isCompletedExceptionally());
	assertFalse(onCompletionPromise3.isCompletedExceptionally());
	assertFalse(checkpointCoordinator.isTriggering());
	assertEquals(0, checkpointCoordinator.getTriggerRequestQueue().size());
	assertEquals(2, taskManagerCheckpointTriggeredTimes.get());
}
 

@Test
public void testTriggerCheckpointRequestCancelled() throws Exception {
	// create some mock Execution vertices that receive the checkpoint trigger messages
	final ExecutionAttemptID attemptID = new ExecutionAttemptID();
	final AtomicInteger taskManagerCheckpointTriggeredTimes = new AtomicInteger(0);
	final SimpleAckingTaskManagerGateway.CheckpointConsumer checkpointConsumer =
		(executionAttemptID,
		jobId, checkpointId,
		timestamp,
		checkpointOptions,
		advanceToEndOfEventTime) -> taskManagerCheckpointTriggeredTimes.incrementAndGet();
	ExecutionVertex vertex = mockExecutionVertex(attemptID, checkpointConsumer);

	// set up the coordinator and validate the initial state
	CheckpointCoordinator checkpointCoordinator = createCheckpointCoordinator(vertex);

	final CompletableFuture<String> masterHookCheckpointFuture = new CompletableFuture<>();
	checkpointCoordinator.addMasterHook(new TestingMasterHook(masterHookCheckpointFuture));
	checkpointCoordinator.startCheckpointScheduler();
	final CompletableFuture<CompletedCheckpoint> onCompletionPromise =
		triggerPeriodicCheckpoint(checkpointCoordinator);

	// checkpoint trigger will not finish since master hook checkpoint is not finished yet
	manuallyTriggeredScheduledExecutor.triggerAll();
	assertTrue(checkpointCoordinator.isTriggering());

	// trigger cancellation
	manuallyTriggeredScheduledExecutor.triggerNonPeriodicScheduledTasks();
	assertTrue(checkpointCoordinator.isTriggering());

	try {
		onCompletionPromise.get();
		fail("Should not reach here");
	} catch (ExecutionException e) {
		final Optional<CheckpointException> checkpointExceptionOptional =
			ExceptionUtils.findThrowable(e, CheckpointException.class);
		assertTrue(checkpointExceptionOptional.isPresent());
		assertEquals(CheckpointFailureReason.CHECKPOINT_EXPIRED,
			checkpointExceptionOptional.get().getCheckpointFailureReason());
	}

	// continue triggering
	masterHookCheckpointFuture.complete("finish master hook");

	manuallyTriggeredScheduledExecutor.triggerAll();
	assertFalse(checkpointCoordinator.isTriggering());
	// it doesn't really trigger task manager to do checkpoint
	assertEquals(0, taskManagerCheckpointTriggeredTimes.get());
	assertEquals(0, checkpointCoordinator.getTriggerRequestQueue().size());
}
 

@Test
public void testTriggerCheckpointSnapshotMasterHookFailed() throws Exception {
	// create some mock Execution vertices that receive the checkpoint trigger messages
	final ExecutionAttemptID attemptID = new ExecutionAttemptID();
	final AtomicInteger taskManagerCheckpointTriggeredTimes = new AtomicInteger(0);
	final SimpleAckingTaskManagerGateway.CheckpointConsumer checkpointConsumer =
		(executionAttemptID,
		jobId, checkpointId,
		timestamp,
		checkpointOptions,
		advanceToEndOfEventTime) -> taskManagerCheckpointTriggeredTimes.incrementAndGet();
	ExecutionVertex vertex = mockExecutionVertex(attemptID, checkpointConsumer);

	// set up the coordinator and validate the initial state
	CheckpointCoordinator checkpointCoordinator = createCheckpointCoordinator(vertex);

	final CompletableFuture<String> masterHookCheckpointFuture = new CompletableFuture<>();
	checkpointCoordinator.addMasterHook(new TestingMasterHook(masterHookCheckpointFuture));
	checkpointCoordinator.startCheckpointScheduler();
	final CompletableFuture<CompletedCheckpoint> onCompletionPromise =
		triggerPeriodicCheckpoint(checkpointCoordinator);

	// checkpoint trigger will not finish since master hook checkpoint is not finished yet
	manuallyTriggeredScheduledExecutor.triggerAll();
	assertTrue(checkpointCoordinator.isTriggering());

	// continue triggering
	masterHookCheckpointFuture.completeExceptionally(new Exception("by design"));

	manuallyTriggeredScheduledExecutor.triggerAll();
	assertFalse(checkpointCoordinator.isTriggering());

	try {
		onCompletionPromise.get();
		fail("Should not reach here");
	} catch (ExecutionException e) {
		final Optional<CheckpointException> checkpointExceptionOptional =
			ExceptionUtils.findThrowable(e, CheckpointException.class);
		assertTrue(checkpointExceptionOptional.isPresent());
		assertEquals(CheckpointFailureReason.TRIGGER_CHECKPOINT_FAILURE,
			checkpointExceptionOptional.get().getCheckpointFailureReason());
	}
	// it doesn't really trigger task manager to do checkpoint
	assertEquals(0, taskManagerCheckpointTriggeredTimes.get());
	assertEquals(0, checkpointCoordinator.getTriggerRequestQueue().size());
}