class SchedulerImpl implements Scheduler {

    @Override
    public <T> ScheduledCompletableFuture<T> before(CompletableFuture<T> promise, Duration timeout) {
        final AtomicBoolean done = new AtomicBoolean();
        final Consumer<Runnable> runOnce = runnable -> {
            if (!done.getAndSet(true)) {
                runnable.run();
            }
        };
        final ScheduledCompletableFutureOnce<T> wrappedPromise = new ScheduledCompletableFutureOnce<>();
        Callable<T> callable = () -> {
            wrappedPromise.completeExceptionally(new TimeoutException());
            return null;
        };
        final ScheduledCompletableFutureOnce<T> afterPromise = afterInternal(wrappedPromise, callable, timeout);
        wrappedPromise.exceptionally(e -> {
            if (e instanceof CancellationException) {
                // Also cancel the scheduled timer if returned completable future is cancelled.
                afterPromise.cancel(true);
            }
            return null;
        });

        promise.thenAccept(p -> runOnce.accept(() -> wrappedPromise.complete(p))) //
                .exceptionally(ex -> {
                    runOnce.accept(() -> wrappedPromise.completeExceptionally(ex));
                    return null;
                });
        return wrappedPromise;
    }

}