@Override
    void onFailure(Throwable cause, CompletableFuture<Map<String, String>> failureLabels) {
        if (hasPendingStateTransition) {
            // the error handling remains the same independent of how many tasks have failed
            // we don't want to initiate the same state transition multiple times, so we exit early
            // this could also be achieved via Context#runIfState, but that'd spam the logs
            return;
        }
        hasPendingStateTransition = true;

        FutureUtils.assertNoException(
                internalSavepointFuture.handle(
                        (savepoint, savepointError) -> {
                            // if savepointError is null then the savepoint has been created
                            // successfully, but the job failed while committing side effects,
                            // so we enrich the exception for the user
                            final Throwable ex =
                                    savepointError != null
                                            ? cause
                                            : new StopWithSavepointStoppingException(
                                                    savepoint, getJobId(), cause);
                            operationFailureCause = ex;
                            FailureResultUtil.restartOrFail(
                                    context.howToHandleFailure(ex, failureLabels), context, this);
                            return null;
                        }));
    }
@Test
    void testOnFailureWaitsForSavepointCompletion() throws Exception {
        try (MockStopWithSavepointContext ctx = new MockStopWithSavepointContext()) {
            CheckpointScheduling mockStopWithSavepointOperations = new MockCheckpointScheduling();
            CompletableFuture<String> savepointFuture = new CompletableFuture<>();
            StateTrackingMockExecutionGraph executionGraph = new StateTrackingMockExecutionGraph();
            StopWithSavepoint sws =
                    createStopWithSavepoint(
                            ctx, mockStopWithSavepointOperations, executionGraph, savepointFuture);
            ctx.setStopWithSavepoint(sws);

            ctx.setHowToHandleFailure(failure -> FailureResult.canRestart(failure, Duration.ZERO));

            sws.onFailure(
                    new Exception("task failure"),
                    CompletableFuture.completedFuture(Collections.emptyMap()));
            // this is a sanity check that we haven't scheduled a state transition
            ctx.triggerExecutors();

            ctx.setExpectRestarting(assertNonNull());
            savepointFuture.complete(SAVEPOINT_PATH);
            ctx.triggerExecutors();
        }
    }