How to chain non-blocking action in CompletionStage.exceptionally

Question

I am writing a Play2 application service method in Java that should do the following. Asynchronously call method A, and if that fails, asynchronously call method B.

To illustrate assume this interface for the backend called by the service:

public interface MyBackend {
    CompletionStage<Object> tryWrite(Object foo);
    CompletionStage<Object> tryCleanup(Object foo);
}

So in my service method, I want to return a Future that can complete with these:

• Success of tryWrite completed
• Fail of tryWrite and Success of tryCleanup completed and failing with exception of tryWrite()

(Note: Of course tryWrite() could do any cleanup itself, this is a simplified example to illustrate a problem)

The implementation of a service calling the backend like this seems difficult to me because the CompletionStage.exceptionally() method does not allow Composing.

Version 1:

public class MyServiceImpl {
    public CompletionStage<Object> tryWriteWithCleanup(Object foo) {

        CompletionStage<Object> writeFuture = myBackend.tryWrite(foo)
            .exceptionally((throwable) -> {
                CompletionStage<Object> cleanupFuture = myBackend.tryCleanup(foo);
                throw new RuntimeException(throwable);
        });
        return writeFuture;
    }
}

So version 1 calls tryCleanup(foo) in a non-blocking way, but the CompletionStage returned by tryWriteWithCleanup() will not wait for cleanupFuture to complete. How to change this code to return a future from the service that would also wait for completion of cleanupFuture?

Version 2:

public class MyServiceImpl {
    public CompletionStage<Object> tryWriteWithCleanup(Object foo) {

        final AtomicReference<Throwable> saveException = new AtomicReference<>();
        CompletionStage<Object> writeFuture = myBackend
            .tryWrite(foo)
            .exceptionally(t -> {
                saveException.set(t);
                // continue with cleanup
                return null;
            })
            .thenCompose((nil) -> {
                // if no cleanup necessary, return
                if (saveException.get() == null) {
                    return CompletableFuture.completedFuture(null);
                }
                return CompletionStage<Object> cleanupFuture = myBackend.tryCleanup(foo)
                    .exceptionally(cleanupError -> {
                        // log error
                        return null;
                    })
                    .thenRun(() -> {
                        throw saveException.get();
                    });
        });
        return writeFuture;
    }
}

Version2 uses an external AtomicReference to store the failure, and makes the asynchronous second call in another thenCompose() block, if there was a failure.

All my other attempts to do so ended up so unwieldy that I don't want to paste them here.

Answer 1

You may simply wait for the completion inside the handler:

public CompletionStage<Object> tryWriteWithCleanup(Object foo) {
    return myBackend.tryWrite(foo).exceptionally(throwable -> {
        myBackend.tryCleanup(foo).toCompletableFuture().join();
        throw new CompletionException(throwable);
    });
}

This will defer the completion of the result CompletionStage to the completion of the cleanup stage. Using CompletionException as wrapper will make the wrapping transparent to the caller.

However, it has some drawbacks. While the framework might utilize the thread while waiting or spawn a compensation thread, if it is a worker thread, the blocked thread might be the caller thread if the stage returned by tryWrite happens to be already completed when entering exceptionally. Unfortunately, there is no exceptionallyAsync method. You may use handleAsync instead, but it will complicate the code while still feeling like a kludge.

Further, exceptions thrown by the cleanup may shadow the original failure.

A cleaner solution may be a bit more involved:

public CompletionStage<Object> tryWriteWithCleanup(Object foo) {

    CompletableFuture<Object> writeFuture = new CompletableFuture<>();

    myBackend.tryWrite(foo).whenComplete((obj,throwable) -> {
        if(throwable==null)
            writeFuture.complete(obj);
        else
            myBackend.tryCleanup(foo).whenComplete((x,next) -> {
                try {
                    if(next!=null) throwable.addSuppressed(next);
                }
                finally {
                    writeFuture.completeExceptionally(throwable);
                }
        });
    });
    return writeFuture;
}

This simply creates a CompletableFuture manually, allowing to control its completion, which will happen either directly by the action chained to tryWrite’s stage in the successful case, or by the action chained to the cleanup stage in the exceptional case. Note that the latter takes care about chaining a possible subsequent cleanup exception via addSuppressed.

Answer 2

Unfortunately CompletionStage/CompletableFuture does not provide exception handling API's with composition.

You can work around this though by relying on a handle() with a BiFunction that returns a CompletionStage. This will give you nested stages (CompletionStage<CompletionStage<Object>>) that you can the "unnest" using compose(identity()):

public CompletionStage<Object> tryWriteWithCleanup(Object foo) {
    return myBackend.tryWrite(foo)
            .handle((r, e) -> {
                if (e != null) {
                    return myBackend.tryCleanup(foo)
                            .handle((r2, e2) -> {
                                // Make sure we always return the original exception
                                // but keep track of new exception if any,
                                // as if run in a finally block
                                if (e2 != null) {
                                    e.addSuppressed(e2);
                                }
                                // wrapping in CompletionException  behaves as if
                                // we threw the original exception
                                throw new CompletionException(e);
                            });
                }
                return CompletableFuture.completedFuture(r);
            })
            .thenCompose(Function.identity());
}