CompletableFuture与Spring的Sleuth结合工具类

Span span = tracer.currentSpan();
try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
    //执行异步代码
}

import brave.Span;
import brave.Tracer;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.function.BiConsumer;
import java.util.function.BiFunction;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;

public class CompletableFutureWithSpan extends CompletableFuture {
    private final CompletableFuture completableFuture;
    private final Tracer tracer;

    CompletableFutureWithSpan(CompletableFuture completableFuture, Tracer tracer) {
        this.completableFuture = completableFuture;
        this.tracer = tracer;
    }

    private static  CompletableFutureWithSpan from(CompletableFuture completableFuture, Tracer tracer) {
        return new CompletableFutureWithSpan(completableFuture, tracer);
    }

    public static  CompletableFutureWithSpan supplyAsync(Supplier supplier, Tracer tracer) {
        Span span = tracer.currentSpan();
        return from(CompletableFuture.supplyAsync(() -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                return supplier.get();
            }
        }), tracer);
    }

    public static  CompletableFutureWithSpan supplyAsync(Supplier supplier, Tracer tracer, Executor executor) {
        Span span = tracer.currentSpan();
        return from(CompletableFuture.supplyAsync(() -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                return supplier.get();
            }
        }, executor), tracer);
    }

    public static CompletableFuture runAsync(Runnable runnable, Tracer tracer) {
        Span span = tracer.currentSpan();
        return from(CompletableFuture.runAsync(() -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                runnable.run();
            }
        }), tracer);
    }

    public static CompletableFuture runAsync(Runnable runnable, Tracer tracer, Executor executor) {
        Span span = tracer.currentSpan();
        return from(CompletableFuture.runAsync(() -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                runnable.run();
            }
        }, executor), tracer);
    }


    @Override
    public  CompletableFutureWithSpan thenApplyAsync(Function fn) {
        Span span = tracer.currentSpan();
        return from(completableFuture.thenApplyAsync(t -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                return fn.apply(t);
            }
        }), this.tracer);
    }

    @Override
    public  CompletableFutureWithSpan thenApplyAsync(Function fn, Executor executor) {
        Span span = tracer.currentSpan();
        return from(completableFuture.thenApplyAsync(t -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                return fn.apply(t);
            }
        }, executor), this.tracer);
    }

    @Override
    public CompletableFutureWithSpan thenAcceptAsync(Consumer action) {
        Span span = tracer.currentSpan();
        return from(completableFuture.thenAcceptAsync(t -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                action.accept(t);
            }
        }), this.tracer);
    }

    @Override
    public CompletableFutureWithSpan thenAcceptAsync(Consumer action, Executor executor) {
        Span span = tracer.currentSpan();
        return from(completableFuture.thenAcceptAsync(t -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                action.accept(t);
            }
        }, executor), this.tracer);
    }

    @Override
    public CompletableFutureWithSpan thenRunAsync(Runnable action) {
        Span span = tracer.currentSpan();
        return from(completableFuture.thenRunAsync(() -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                action.run();
            }
        }), this.tracer);
    }

    @Override
    public CompletableFutureWithSpan thenRunAsync(Runnable action, Executor executor) {
        Span span = tracer.currentSpan();
        return from(completableFuture.thenRunAsync(() -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                action.run();
            }
        }, executor), this.tracer);
    }

    @Override
    public  CompletableFutureWithSpan thenCombineAsync(CompletionStage other, BiFunction fn) {
        Span span = tracer.currentSpan();
        return from(completableFuture.thenCombineAsync(other, (t, u) -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                return fn.apply(t, u);
            }
        }), this.tracer);
    }

    @Override
    public  CompletableFutureWithSpan thenCombineAsync(CompletionStage other, BiFunction fn, Executor executor) {
        Span span = tracer.currentSpan();
        return from(completableFuture.thenCombineAsync(other, (t, u) -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                return fn.apply(t, u);
            }
        }, executor), this.tracer);
    }

    @Override
    public  CompletableFutureWithSpan thenAcceptBothAsync(CompletionStage other, BiConsumer action) {
        Span span = tracer.currentSpan();
        return from(completableFuture.thenAcceptBothAsync(other, (t, u) -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                action.accept(t, u);
            }
        }), this.tracer);
    }

    @Override
    public  CompletableFutureWithSpan thenAcceptBothAsync(CompletionStage other, BiConsumer action, Executor executor) {
        Span span = tracer.currentSpan();
        return from(completableFuture.thenAcceptBothAsync(other, (t, u) -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                action.accept(t, u);
            }
        }, executor), this.tracer);
    }

    @Override
    public CompletableFutureWithSpan runAfterBothAsync(CompletionStage other, Runnable action) {
        Span span = tracer.currentSpan();
        return from(completableFuture.runAfterBothAsync(other, () -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                action.run();
            }
        }), this.tracer);
    }

    @Override
    public CompletableFutureWithSpan runAfterBothAsync(CompletionStage other, Runnable action, Executor executor) {
        Span span = tracer.currentSpan();
        return from(completableFuture.runAfterBothAsync(other, () -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                action.run();
            }
        }, executor), this.tracer);
    }

    @Override
    public  CompletableFutureWithSpan applyToEitherAsync(CompletionStage other, Function fn) {
        Span span = tracer.currentSpan();
        return from(completableFuture.applyToEitherAsync(other, t -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                return fn.apply(t);
            }
        }), this.tracer);
    }

    @Override
    public  CompletableFutureWithSpan applyToEitherAsync(CompletionStage other, Function fn, Executor executor) {
        Span span = tracer.currentSpan();
        return from(completableFuture.applyToEitherAsync(other, t -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                return fn.apply(t);
            }
        }, executor), this.tracer);
    }

    @Override
    public CompletableFutureWithSpan acceptEitherAsync(CompletionStage other, Consumer action) {
        Span span = tracer.currentSpan();
        return from(completableFuture.acceptEitherAsync(other, t -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                action.accept(t);
            }
        }), this.tracer);
    }

    @Override
    public CompletableFutureWithSpan acceptEitherAsync(CompletionStage other, Consumer action, Executor executor) {
        Span span = tracer.currentSpan();
        return from(completableFuture.acceptEitherAsync(other, t -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                action.accept(t);
            }
        }, executor), this.tracer);
    }

    @Override
    public CompletableFutureWithSpan runAfterEitherAsync(CompletionStage other, Runnable action) {
        Span span = tracer.currentSpan();
        return from(completableFuture.runAfterEitherAsync(other, () -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                action.run();
            }
        }), this.tracer);
    }

    @Override
    public CompletableFutureWithSpan runAfterEitherAsync(CompletionStage other, Runnable action, Executor executor) {
        Span span = tracer.currentSpan();
        return from(completableFuture.runAfterEitherAsync(other, () -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                action.run();
            }
        }, executor), this.tracer);
    }

    @Override
    public  CompletableFutureWithSpan thenComposeAsync(Function> fn) {
        Span span = tracer.currentSpan();
        return from(completableFuture.thenComposeAsync(t -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                return fn.apply(t);
            }
        }), this.tracer);
    }

    @Override
    public  CompletableFutureWithSpan thenComposeAsync(Function> fn, Executor executor) {
        Span span = tracer.currentSpan();
        return from(completableFuture.thenComposeAsync(t -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                return fn.apply(t);
            }
        }, executor), this.tracer);
    }

    @Override
    public CompletableFutureWithSpan whenCompleteAsync(BiConsumer action) {
        Span span = tracer.currentSpan();
        return from(completableFuture.whenCompleteAsync((t, throwable) -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                action.accept(t, throwable);
            }
        }), this.tracer);
    }

    @Override
    public CompletableFutureWithSpan whenCompleteAsync(BiConsumer action, Executor executor) {
        Span span = tracer.currentSpan();
        return from(completableFuture.whenCompleteAsync((t, throwable) -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                action.accept(t, throwable);
            }
        }, executor), this.tracer);
    }

    @Override
    public  CompletableFutureWithSpan handleAsync(BiFunction fn) {
        Span span = tracer.currentSpan();
        return from(completableFuture.handleAsync((t, throwable) -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                return fn.apply(t, throwable);
            }
        }), this.tracer);
    }

    @Override
    public  CompletableFutureWithSpan handleAsync(BiFunction fn, Executor executor) {
        Span span = tracer.currentSpan();
        return from(completableFuture.handleAsync((t, throwable) -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                return fn.apply(t, throwable);
            }
        }, executor), this.tracer);
    }

    @Override
    public CompletableFuture completeAsync(Supplier supplier, Executor executor) {
        Span span = tracer.currentSpan();
        return completableFuture.completeAsync(() -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                return supplier.get();
            }
        }, executor);
    }

    @Override
    public CompletableFuture completeAsync(Supplier supplier) {
        Span span = tracer.currentSpan();
        return completableFuture.completeAsync(() -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                return supplier.get();
            }
        });
    }

    @Override
    public boolean isDone() {
        return completableFuture.isDone();
    }

    @Override
    public T get() throws InterruptedException, ExecutionException {
        return completableFuture.get();
    }

    @Override
    public T get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
        return completableFuture.get(timeout, unit);
    }

    @Override
    public T join() {
        return completableFuture.join();
    }

    @Override
    public T getNow(T valueIfAbsent) {
        return completableFuture.getNow(valueIfAbsent);
    }

    @Override
    public boolean complete(T value) {
        return completableFuture.complete(value);
    }

    @Override
    public boolean completeExceptionally(Throwable ex) {
        return completableFuture.completeExceptionally(ex);
    }

    @Override
    public  CompletableFuture thenApply(Function fn) {
        return completableFuture.thenApply(fn);
    }

    @Override
    public CompletableFuture thenAccept(Consumer action) {
        return completableFuture.thenAccept(action);
    }

    @Override
    public CompletableFuture thenRun(Runnable action) {
        return completableFuture.thenRun(action);
    }

    @Override
    public  CompletableFuture thenCombine(CompletionStage other, BiFunction fn) {
        return completableFuture.thenCombine(other, fn);
    }

    @Override
    public  CompletableFuture thenAcceptBoth(CompletionStage other, BiConsumer action) {
        return completableFuture.thenAcceptBoth(other, action);
    }

    @Override
    public CompletableFuture runAfterBoth(CompletionStage other, Runnable action) {
        return completableFuture.runAfterBoth(other, action);
    }

    @Override
    public  CompletableFuture applyToEither(CompletionStage other, Function fn) {
        return completableFuture.applyToEither(other, fn);
    }

    @Override
    public CompletableFuture acceptEither(CompletionStage other, Consumer action) {
        return completableFuture.acceptEither(other, action);
    }

    @Override
    public CompletableFuture runAfterEither(CompletionStage other, Runnable action) {
        return completableFuture.runAfterEither(other, action);
    }

    @Override
    public  CompletableFuture thenCompose(Function> fn) {
        return completableFuture.thenCompose(fn);
    }

    @Override
    public CompletableFuture whenComplete(BiConsumer action) {
        return completableFuture.whenComplete(action);
    }

    @Override
    public  CompletableFuture handle(BiFunction fn) {
        return completableFuture.handle(fn);
    }

    @Override
    public CompletableFuture toCompletableFuture() {
        return completableFuture.toCompletableFuture();
    }

    @Override
    public CompletableFuture exceptionally(Function fn) {
        return completableFuture.exceptionally(fn);
    }

    @Override
    public boolean cancel(boolean mayInterruptIfRunning) {
        return completableFuture.cancel(mayInterruptIfRunning);
    }

    @Override
    public boolean isCancelled() {
        return completableFuture.isCancelled();
    }

    @Override
    public boolean isCompletedExceptionally() {
        return completableFuture.isCompletedExceptionally();
    }

    @Override
    public void obtrudeValue(T value) {
        completableFuture.obtrudeValue(value);
    }

    @Override
    public void obtrudeException(Throwable ex) {
        completableFuture.obtrudeException(ex);
    }

    @Override
    public int getNumberOfDependents() {
        return completableFuture.getNumberOfDependents();
    }

    @Override
    public String toString() {
        return completableFuture.toString();
    }

    @Override
    public  CompletableFuture newIncompleteFuture() {
        return completableFuture.newIncompleteFuture();
    }

    @Override
    public Executor defaultExecutor() {
        return completableFuture.defaultExecutor();
    }

    @Override
    public CompletableFuture copy() {
        return completableFuture.copy();
    }

    @Override
    public CompletionStage minimalCompletionStage() {
        return completableFuture.minimalCompletionStage();
    }

    @Override
    public CompletableFuture orTimeout(long timeout, TimeUnit unit) {
        return completableFuture.orTimeout(timeout, unit);
    }

    @Override
    public CompletableFuture completeOnTimeout(T value, long timeout, TimeUnit unit) {
        return completableFuture.completeOnTimeout(value, timeout, unit);
    }
}

    @Override
    public CompletableFuture exceptionallyCompose(Function> fn) {
        return completableFuture.exceptionallyCompose(fn);
    }
     @Override
    public CompletableFutureWithSpan exceptionallyAsync(Function fn) {
        Span span = tracer.currentSpan();
        return from(completableFuture.exceptionallyAsync(throwable -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                return fn.apply(throwable);
            }
        }), this.tracer);
    }

    @Override
    public CompletableFutureWithSpan exceptionallyAsync(Function fn, Executor executor) {
        Span span = tracer.currentSpan();
        return from(completableFuture.exceptionallyAsync(throwable -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                return fn.apply(throwable);
            }
        }, executor), this.tracer);
    }

    @Override
    public CompletableFutureWithSpan exceptionallyComposeAsync(Function> fn) {
        Span span = tracer.currentSpan();
        return from(completableFuture.exceptionallyComposeAsync(throwable -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                return fn.apply(throwable);
            }
        }), this.tracer);
    }

    @Override
    public CompletableFutureWithSpan exceptionallyComposeAsync(Function> fn, Executor executor) {
        Span span = tracer.currentSpan();
        return from(completableFuture.exceptionallyComposeAsync(throwable -> {
            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
                return fn.apply(throwable);
            }
        }, executor), this.tracer);
    }

Mono.fromFuture(CompletableFutureWithSpan)