SavepointCoordinator extends CheckpointCoordinator { public Future<String> triggerSavepoint(long timestamp) throws Exception { final Promise<String> promise = new scala.concurrent.impl.Promise.DefaultPromise<>(); try { final long checkpointId = getAndIncrementCheckpointId(); if (checkpointId == -1) { throw new IllegalStateException("Failed to get checkpoint Id"); } if (savepointPromises.put(checkpointId, promise) == null) { boolean success = false; try { success = triggerCheckpoint(timestamp, checkpointId); } finally { if (!success) { savepointPromises.remove(checkpointId); promise.failure(new Exception("Failed to trigger savepoint")); } } } else { throw new IllegalStateException("Duplicate checkpoint ID"); } } catch (Throwable t) { promise.failure(new Exception("Failed to trigger savepoint", t)); } return promise.future(); } SavepointCoordinator(
			JobID jobId,
			long baseInterval,
			long checkpointTimeout,
			ExecutionVertex[] tasksToTrigger,
			ExecutionVertex[] tasksToWaitFor,
			ExecutionVertex[] tasksToCommitTo,
			ClassLoader userClassLoader,
			CheckpointIDCounter checkpointIDCounter,
			StateStore<CompletedCheckpoint> savepointStore,
			CheckpointStatsTracker statsTracker); String getSavepointRestorePath(); Future<String> triggerSavepoint(long timestamp); void restoreSavepoint(
			Map<JobVertexID, ExecutionJobVertex> tasks,
			String savepointPath); @Override ActorGateway createActivatorDeactivator(
			ActorSystem actorSystem,
			UUID leaderSessionID);  }
@Test public void testSimpleTriggerSavepoint() throws Exception { JobID jobId = new JobID(); long checkpointTimeout = 60 * 1000; long timestamp = 1272635; ExecutionVertex[] vertices = new ExecutionVertex[] { mockExecutionVertex(jobId), mockExecutionVertex(jobId) }; MockCheckpointIdCounter checkpointIdCounter = new MockCheckpointIdCounter(); HeapStateStore<CompletedCheckpoint> savepointStore = new HeapStateStore<>(); SavepointCoordinator coordinator = createSavepointCoordinator( jobId, checkpointTimeout, vertices, vertices, vertices, checkpointIdCounter, savepointStore); Future<String> savepointPathFuture = coordinator.triggerSavepoint(timestamp); assertFalse(savepointPathFuture.isCompleted()); long checkpointId = checkpointIdCounter.getLastReturnedCount(); assertEquals(0, checkpointId); for (ExecutionVertex vertex : vertices) { verifyTriggerCheckpoint(vertex, checkpointId, timestamp); } PendingCheckpoint pendingCheckpoint = coordinator.getPendingCheckpoints() .get(checkpointId); verifyPendingCheckpoint(pendingCheckpoint, jobId, checkpointId, timestamp, 0, 2, 0, false, false); for (ExecutionVertex vertex : vertices) { coordinator.receiveAcknowledgeMessage(new AcknowledgeCheckpoint( jobId, vertex.getCurrentExecutionAttempt().getAttemptId(), checkpointId, createSerializedStateHandle(vertex), 0)); } assertTrue(pendingCheckpoint.isDiscarded()); assertEquals(0, coordinator.getSuccessfulCheckpoints().size()); for (ExecutionVertex vertex : vertices) { verifyNotifyCheckpointComplete(vertex, checkpointId, timestamp); } assertTrue(savepointPathFuture.isCompleted()); String savepointPath = Await.result(savepointPathFuture, FiniteDuration.Zero()); CompletedCheckpoint savepoint = savepointStore.getState(savepointPath); verifySavepoint(savepoint, jobId, checkpointId, timestamp, vertices); assertEquals(0, getSavepointPromises(coordinator).size()); coordinator.shutdown(); }
@Test public void testAbortSavepointIfTriggerTasksNotExecuted() throws Exception { JobID jobId = new JobID(); ExecutionVertex[] triggerVertices = new ExecutionVertex[] { mock(ExecutionVertex.class), mock(ExecutionVertex.class) }; ExecutionVertex[] ackVertices = new ExecutionVertex[] { mockExecutionVertex(jobId), mockExecutionVertex(jobId) }; SavepointCoordinator coordinator = createSavepointCoordinator( jobId, 60 * 1000, triggerVertices, ackVertices, new ExecutionVertex[] {}, new MockCheckpointIdCounter(), new HeapStateStore<CompletedCheckpoint>()); Future<String> savepointPathFuture = coordinator.triggerSavepoint(1238123); assertTrue(savepointPathFuture.isCompleted()); try { Await.result(savepointPathFuture, FiniteDuration.Zero()); fail("Did not throw expected Exception after shutdown"); } catch (Exception ignored) { } assertEquals(0, getSavepointPromises(coordinator).size()); coordinator.shutdown(); }
@Test public void testAbortSavepointIfTriggerTasksAreFinished() throws Exception { JobID jobId = new JobID(); ExecutionVertex[] triggerVertices = new ExecutionVertex[] { mockExecutionVertex(jobId), mockExecutionVertex(jobId, ExecutionState.FINISHED) }; ExecutionVertex[] ackVertices = new ExecutionVertex[] { mockExecutionVertex(jobId), mockExecutionVertex(jobId) }; SavepointCoordinator coordinator = createSavepointCoordinator( jobId, 60 * 1000, triggerVertices, ackVertices, new ExecutionVertex[] {}, new MockCheckpointIdCounter(), new HeapStateStore<CompletedCheckpoint>()); Future<String> savepointPathFuture = coordinator.triggerSavepoint(1238123); assertTrue(savepointPathFuture.isCompleted()); try { Await.result(savepointPathFuture, FiniteDuration.Zero()); fail("Did not throw expected Exception after shutdown"); } catch (Exception ignored) { } assertEquals(0, getSavepointPromises(coordinator).size()); coordinator.shutdown(); }
@Test public void testAbortSavepointIfAckTasksAreNotExecuted() throws Exception { JobID jobId = new JobID(); ExecutionVertex[] triggerVertices = new ExecutionVertex[] { mockExecutionVertex(jobId), mockExecutionVertex(jobId) }; ExecutionVertex[] ackVertices = new ExecutionVertex[] { mock(ExecutionVertex.class), mock(ExecutionVertex.class) }; SavepointCoordinator coordinator = createSavepointCoordinator( jobId, 60 * 1000, triggerVertices, ackVertices, new ExecutionVertex[] {}, new MockCheckpointIdCounter(), new HeapStateStore<CompletedCheckpoint>()); Future<String> savepointPathFuture = coordinator.triggerSavepoint(1238123); assertTrue(savepointPathFuture.isCompleted()); try { Await.result(savepointPathFuture, FiniteDuration.Zero()); fail("Did not throw expected Exception after shutdown"); } catch (Exception ignored) { } assertEquals(0, getSavepointPromises(coordinator).size()); coordinator.shutdown(); }
@Test public void testAbortOnCheckpointTimeout() throws Exception { JobID jobId = new JobID(); ExecutionVertex[] vertices = new ExecutionVertex[] { mockExecutionVertex(jobId), mockExecutionVertex(jobId) }; ExecutionVertex commitVertex = mockExecutionVertex(jobId); MockCheckpointIdCounter checkpointIdCounter = new MockCheckpointIdCounter(); SavepointCoordinator coordinator = createSavepointCoordinator( jobId, 20, vertices, vertices, new ExecutionVertex[] { commitVertex }, checkpointIdCounter, new HeapStateStore<CompletedCheckpoint>()); Future<String> savepointPathFuture = coordinator.triggerSavepoint(12731273); assertFalse(savepointPathFuture.isCompleted()); long checkpointId = checkpointIdCounter.getLastReturnedCount(); coordinator.receiveAcknowledgeMessage(new AcknowledgeCheckpoint( jobId, vertices[0].getCurrentExecutionAttempt().getAttemptId(), checkpointId, createSerializedStateHandle(vertices[0]), 0)); PendingCheckpoint pendingCheckpoint = coordinator.getPendingCheckpoints() .get(checkpointId); assertFalse(pendingCheckpoint.isDiscarded()); Deadline deadline = FiniteDuration.apply(5, "s").fromNow(); while (deadline.hasTimeLeft() && !pendingCheckpoint.isDiscarded() && coordinator.getNumberOfPendingCheckpoints() > 0) { Thread.sleep(250); } assertTrue(pendingCheckpoint.isDiscarded()); assertEquals(0, coordinator.getNumberOfPendingCheckpoints()); assertEquals(0, coordinator.getNumberOfRetainedSuccessfulCheckpoints()); verify(commitVertex, times(0)).sendMessageToCurrentExecution( any(NotifyCheckpointComplete.class), any(ExecutionAttemptID.class)); assertTrue(savepointPathFuture.isCompleted()); try { Await.result(savepointPathFuture, FiniteDuration.Zero()); fail("Did not throw expected Exception after timeout"); } catch (Exception ignored) { } assertEquals(0, getSavepointPromises(coordinator).size()); coordinator.shutdown(); }
@Test public void testAbortSavepointsOnShutdown() throws Exception { JobID jobId = new JobID(); ExecutionVertex[] vertices = new ExecutionVertex[] { mockExecutionVertex(jobId), mockExecutionVertex(jobId) }; SavepointCoordinator coordinator = createSavepointCoordinator( jobId, 60 * 1000, vertices, vertices, vertices, new MockCheckpointIdCounter(), new HeapStateStore<CompletedCheckpoint>()); List<Future<String>> savepointPathFutures = new ArrayList<>(); savepointPathFutures.add(coordinator.triggerSavepoint(12731273)); savepointPathFutures.add(coordinator.triggerSavepoint(12731273 + 123)); for (Future<String> future : savepointPathFutures) { assertFalse(future.isCompleted()); } coordinator.shutdown(); for (Future<String> future : savepointPathFutures) { assertTrue(future.isCompleted()); try { Await.result(future, FiniteDuration.Zero()); fail("Did not throw expected Exception after shutdown"); } catch (Exception ignored) { } } assertEquals(0, getSavepointPromises(coordinator).size()); }
@Test public void testAbortSavepointOnStateStoreFailure() throws Exception { JobID jobId = new JobID(); ExecutionJobVertex jobVertex = mockExecutionJobVertex(jobId, new JobVertexID(), 4); HeapStateStore<CompletedCheckpoint> savepointStore = spy( new HeapStateStore<CompletedCheckpoint>()); SavepointCoordinator coordinator = createSavepointCoordinator( jobId, 60 * 1000, jobVertex.getTaskVertices(), jobVertex.getTaskVertices(), new ExecutionVertex[] {}, new MockCheckpointIdCounter(), savepointStore); doThrow(new Exception("TestException")) .when(savepointStore).putState(any(CompletedCheckpoint.class)); Future<String> savepointPathFuture = coordinator.triggerSavepoint(1231273123); for (ExecutionVertex vertex : jobVertex.getTaskVertices()) { ExecutionAttemptID attemptId = vertex.getCurrentExecutionAttempt().getAttemptId(); coordinator.receiveAcknowledgeMessage(new AcknowledgeCheckpoint( jobId, attemptId, 0, createSerializedStateHandle(vertex), 0)); } try { Await.result(savepointPathFuture, FiniteDuration.Zero()); fail("Did not throw expected Exception after rollback with savepoint store failure."); } catch (Exception ignored) { } assertEquals(0, getSavepointPromises(coordinator).size()); coordinator.shutdown(); }
@Test public void testAbortSavepointIfSubsumed() throws Exception { JobID jobId = new JobID(); long checkpointTimeout = 60 * 1000; long[] timestamps = new long[] { 1272635, 1272635 + 10 }; long[] checkpointIds = new long[2]; ExecutionVertex[] vertices = new ExecutionVertex[] { mockExecutionVertex(jobId), mockExecutionVertex(jobId) }; MockCheckpointIdCounter checkpointIdCounter = new MockCheckpointIdCounter(); HeapStateStore<CompletedCheckpoint> savepointStore = new HeapStateStore<>(); SavepointCoordinator coordinator = createSavepointCoordinator( jobId, checkpointTimeout, vertices, vertices, vertices, checkpointIdCounter, savepointStore); List<Future<String>> savepointPathFutures = new ArrayList<>(); savepointPathFutures.add(coordinator.triggerSavepoint(timestamps[0])); checkpointIds[0] = checkpointIdCounter.getLastReturnedCount(); savepointPathFutures.add(coordinator.triggerSavepoint(timestamps[1])); checkpointIds[1] = checkpointIdCounter.getLastReturnedCount(); for (Future<String> future : savepointPathFutures) { assertFalse(future.isCompleted()); } for (ExecutionVertex vertex : vertices) { verifyTriggerCheckpoint(vertex, checkpointIds[0], timestamps[0]); verifyTriggerCheckpoint(vertex, checkpointIds[1], timestamps[1]); } PendingCheckpoint[] pendingCheckpoints = new PendingCheckpoint[] { coordinator.getPendingCheckpoints().get(checkpointIds[0]), coordinator.getPendingCheckpoints().get(checkpointIds[1]) }; verifyPendingCheckpoint(pendingCheckpoints[0], jobId, checkpointIds[0], timestamps[0], 0, 2, 0, false, false); verifyPendingCheckpoint(pendingCheckpoints[1], jobId, checkpointIds[1], timestamps[1], 0, 2, 0, false, false); for (ExecutionVertex vertex : vertices) { coordinator.receiveAcknowledgeMessage(new AcknowledgeCheckpoint( jobId, vertex.getCurrentExecutionAttempt().getAttemptId(), checkpointIds[1], createSerializedStateHandle(vertex), 0)); } coordinator.receiveAcknowledgeMessage(new AcknowledgeCheckpoint( jobId, vertices[0].getCurrentExecutionAttempt().getAttemptId(), checkpointIds[0], createSerializedStateHandle(vertices[0]), 0)); assertTrue(pendingCheckpoints[0].isDiscarded()); assertTrue(pendingCheckpoints[1].isDiscarded()); assertEquals(0, coordinator.getSuccessfulCheckpoints().size()); for (ExecutionVertex vertex : vertices) { verifyNotifyCheckpointComplete(vertex, checkpointIds[1], timestamps[1]); } CompletedCheckpoint[] savepoints = new CompletedCheckpoint[2]; String[] savepointPaths = new String[2]; assertTrue(savepointPathFutures.get(0).isCompleted()); try { savepointPaths[0] = Await.result(savepointPathFutures.get(0), FiniteDuration.Zero()); fail("Did not throw expected exception"); } catch (Exception ignored) { } assertTrue(savepointPathFutures.get(1).isCompleted()); savepointPaths[1] = Await.result(savepointPathFutures.get(1), FiniteDuration.Zero()); savepoints[1] = savepointStore.getState(savepointPaths[1]); verifySavepoint(savepoints[1], jobId, checkpointIds[1], timestamps[1], vertices); assertEquals(0, getSavepointPromises(coordinator).size()); coordinator.shutdown(); }
@Test public void testShutdownDoesNotCleanUpCompletedCheckpointsWithFileSystemStore() throws Exception { JobID jobId = new JobID(); long checkpointTimeout = 60 * 1000; long timestamp = 1272635; ExecutionVertex[] vertices = new ExecutionVertex[] { mockExecutionVertex(jobId), mockExecutionVertex(jobId) }; MockCheckpointIdCounter checkpointIdCounter = new MockCheckpointIdCounter(); final File tmpDir = CommonTestUtils.createTempDirectory(); try { FileSystemStateStore<CompletedCheckpoint> savepointStore = new FileSystemStateStore<>( tmpDir.toURI().toString(), "sp-"); SavepointCoordinator coordinator = createSavepointCoordinator( jobId, checkpointTimeout, vertices, vertices, vertices, checkpointIdCounter, savepointStore); Future<String> savepointPathFuture = coordinator.triggerSavepoint(timestamp); assertFalse(savepointPathFuture.isCompleted()); long checkpointId = checkpointIdCounter.getLastReturnedCount(); assertEquals(0, checkpointId); for (ExecutionVertex vertex : vertices) { verifyTriggerCheckpoint(vertex, checkpointId, timestamp); } PendingCheckpoint pendingCheckpoint = coordinator.getPendingCheckpoints() .get(checkpointId); verifyPendingCheckpoint(pendingCheckpoint, jobId, checkpointId, timestamp, 0, 2, 0, false, false); for (ExecutionVertex vertex : vertices) { coordinator.receiveAcknowledgeMessage(new AcknowledgeCheckpoint( jobId, vertex.getCurrentExecutionAttempt().getAttemptId(), checkpointId, createSerializedStateHandle(vertex), 0)); } assertTrue(pendingCheckpoint.isDiscarded()); assertEquals(0, coordinator.getSuccessfulCheckpoints().size()); for (ExecutionVertex vertex : vertices) { verifyNotifyCheckpointComplete(vertex, checkpointId, timestamp); } assertTrue(savepointPathFuture.isCompleted()); String savepointPath = Await.result(savepointPathFuture, FiniteDuration.Zero()); assertEquals(0, getSavepointPromises(coordinator).size()); coordinator.shutdown(); CompletedCheckpoint savepoint = savepointStore.getState(savepointPath); verifySavepoint(savepoint, jobId, checkpointId, timestamp, vertices); } finally { FileUtils.deleteDirectory(tmpDir); } }