CompletableFuture与Spring的Sleuth结合工具类

系列目录：

1. Spring WebFlux运用中的思考与对比
2. CompletableFuture与Spring的Sleuth结合工具类
3. CommpetableFuture使用anyOf过程中的一些优化思考
4. 结合CompletableFuture与Spring的Sleuth结合工具类与allOf以及anyOf

本文基于JDK 11 and JDK 12

按照上一篇内容的分析，我们想在异步代码保留原有的spanId和traceId需要在异步调用前，使用：

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

每次使用CompletableFuture都要这么写的话，太麻烦了，所以，我们继承，使用代理的设计模式，将所有的Async方法都覆盖：

对于JDK11：

1import brave.Span;
2import brave.Tracer;
3
4import java.util.concurrent.CompletableFuture;
5import java.util.concurrent.CompletionStage;
6import java.util.concurrent.ExecutionException;
7import java.util.concurrent.Executor;
8import java.util.concurrent.TimeUnit;
9import java.util.concurrent.TimeoutException;
10import java.util.function.BiConsumer;
11import java.util.function.BiFunction;
12import java.util.function.Consumer;
13import java.util.function.Function;
14import java.util.function.Supplier;
15
16public class CompletableFutureWithSpan<T> extends CompletableFuture<T> {
17    private final CompletableFuture<T> completableFuture;
18    private final Tracer tracer;
19
20    CompletableFutureWithSpan(CompletableFuture<T> completableFuture, Tracer tracer) {
21        this.completableFuture = completableFuture;
22        this.tracer = tracer;
23    }
24
25    private static <T> CompletableFutureWithSpan<T> from(CompletableFuture<T> completableFuture, Tracer tracer) {
26        return new CompletableFutureWithSpan(completableFuture, tracer);
27    }
28
29    public static <U> CompletableFutureWithSpan<U> supplyAsync(Supplier<U> supplier, Tracer tracer) {
30        Span span = tracer.currentSpan();
31        return from(CompletableFuture.supplyAsync(() -> {
32            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
33                return supplier.get();
34            }
35        }), tracer);
36    }
37
38    public static <U> CompletableFutureWithSpan<U> supplyAsync(Supplier<U> supplier, Tracer tracer, Executor executor) {
39        Span span = tracer.currentSpan();
40        return from(CompletableFuture.supplyAsync(() -> {
41            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
42                return supplier.get();
43            }
44        }, executor), tracer);
45    }
46
47    public static CompletableFuture<Void> runAsync(Runnable runnable, Tracer tracer) {
48        Span span = tracer.currentSpan();
49        return from(CompletableFuture.runAsync(() -> {
50            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
51                runnable.run();
52            }
53        }), tracer);
54    }
55
56    public static CompletableFuture<Void> runAsync(Runnable runnable, Tracer tracer, Executor executor) {
57        Span span = tracer.currentSpan();
58        return from(CompletableFuture.runAsync(() -> {
59            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
60                runnable.run();
61            }
62        }, executor), tracer);
63    }
64
65
66    @Override
67    public <U> CompletableFutureWithSpan<U> thenApplyAsync(Function<? super T, ? extends U> fn) {
68        Span span = tracer.currentSpan();
69        return from(completableFuture.thenApplyAsync(t -> {
70            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
71                return fn.apply(t);
72            }
73        }), this.tracer);
74    }
75
76    @Override
77    public <U> CompletableFutureWithSpan<U> thenApplyAsync(Function<? super T, ? extends U> fn, Executor executor) {
78        Span span = tracer.currentSpan();
79        return from(completableFuture.thenApplyAsync(t -> {
80            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
81                return fn.apply(t);
82            }
83        }, executor), this.tracer);
84    }
85
86    @Override
87    public CompletableFutureWithSpan<Void> thenAcceptAsync(Consumer<? super T> action) {
88        Span span = tracer.currentSpan();
89        return from(completableFuture.thenAcceptAsync(t -> {
90            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
91                action.accept(t);
92            }
93        }), this.tracer);
94    }
95
96    @Override
97    public CompletableFutureWithSpan<Void> thenAcceptAsync(Consumer<? super T> action, Executor executor) {
98        Span span = tracer.currentSpan();
99        return from(completableFuture.thenAcceptAsync(t -> {
100            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
101                action.accept(t);
102            }
103        }, executor), this.tracer);
104    }
105
106    @Override
107    public CompletableFutureWithSpan<Void> thenRunAsync(Runnable action) {
108        Span span = tracer.currentSpan();
109        return from(completableFuture.thenRunAsync(() -> {
110            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
111                action.run();
112            }
113        }), this.tracer);
114    }
115
116    @Override
117    public CompletableFutureWithSpan<Void> thenRunAsync(Runnable action, Executor executor) {
118        Span span = tracer.currentSpan();
119        return from(completableFuture.thenRunAsync(() -> {
120            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
121                action.run();
122            }
123        }, executor), this.tracer);
124    }
125
126    @Override
127    public <U, V> CompletableFutureWithSpan<V> thenCombineAsync(CompletionStage<? extends U> other, BiFunction<? super T, ? super U, ? extends V> fn) {
128        Span span = tracer.currentSpan();
129        return from(completableFuture.thenCombineAsync(other, (t, u) -> {
130            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
131                return fn.apply(t, u);
132            }
133        }), this.tracer);
134    }
135
136    @Override
137    public <U, V> CompletableFutureWithSpan<V> thenCombineAsync(CompletionStage<? extends U> other, BiFunction<? super T, ? super U, ? extends V> fn, Executor executor) {
138        Span span = tracer.currentSpan();
139        return from(completableFuture.thenCombineAsync(other, (t, u) -> {
140            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
141                return fn.apply(t, u);
142            }
143        }, executor), this.tracer);
144    }
145
146    @Override
147    public <U> CompletableFutureWithSpan<Void> thenAcceptBothAsync(CompletionStage<? extends U> other, BiConsumer<? super T, ? super U> action) {
148        Span span = tracer.currentSpan();
149        return from(completableFuture.thenAcceptBothAsync(other, (t, u) -> {
150            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
151                action.accept(t, u);
152            }
153        }), this.tracer);
154    }
155
156    @Override
157    public <U> CompletableFutureWithSpan<Void> thenAcceptBothAsync(CompletionStage<? extends U> other, BiConsumer<? super T, ? super U> action, Executor executor) {
158        Span span = tracer.currentSpan();
159        return from(completableFuture.thenAcceptBothAsync(other, (t, u) -> {
160            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
161                action.accept(t, u);
162            }
163        }, executor), this.tracer);
164    }
165
166    @Override
167    public CompletableFutureWithSpan<Void> runAfterBothAsync(CompletionStage<?> other, Runnable action) {
168        Span span = tracer.currentSpan();
169        return from(completableFuture.runAfterBothAsync(other, () -> {
170            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
171                action.run();
172            }
173        }), this.tracer);
174    }
175
176    @Override
177    public CompletableFutureWithSpan<Void> runAfterBothAsync(CompletionStage<?> other, Runnable action, Executor executor) {
178        Span span = tracer.currentSpan();
179        return from(completableFuture.runAfterBothAsync(other, () -> {
180            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
181                action.run();
182            }
183        }, executor), this.tracer);
184    }
185
186    @Override
187    public <U> CompletableFutureWithSpan<U> applyToEitherAsync(CompletionStage<? extends T> other, Function<? super T, U> fn) {
188        Span span = tracer.currentSpan();
189        return from(completableFuture.applyToEitherAsync(other, t -> {
190            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
191                return fn.apply(t);
192            }
193        }), this.tracer);
194    }
195
196    @Override
197    public <U> CompletableFutureWithSpan<U> applyToEitherAsync(CompletionStage<? extends T> other, Function<? super T, U> fn, Executor executor) {
198        Span span = tracer.currentSpan();
199        return from(completableFuture.applyToEitherAsync(other, t -> {
200            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
201                return fn.apply(t);
202            }
203        }, executor), this.tracer);
204    }
205
206    @Override
207    public CompletableFutureWithSpan<Void> acceptEitherAsync(CompletionStage<? extends T> other, Consumer<? super T> action) {
208        Span span = tracer.currentSpan();
209        return from(completableFuture.acceptEitherAsync(other, t -> {
210            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
211                action.accept(t);
212            }
213        }), this.tracer);
214    }
215
216    @Override
217    public CompletableFutureWithSpan<Void> acceptEitherAsync(CompletionStage<? extends T> other, Consumer<? super T> action, Executor executor) {
218        Span span = tracer.currentSpan();
219        return from(completableFuture.acceptEitherAsync(other, t -> {
220            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
221                action.accept(t);
222            }
223        }, executor), this.tracer);
224    }
225
226    @Override
227    public CompletableFutureWithSpan<Void> runAfterEitherAsync(CompletionStage<?> other, Runnable action) {
228        Span span = tracer.currentSpan();
229        return from(completableFuture.runAfterEitherAsync(other, () -> {
230            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
231                action.run();
232            }
233        }), this.tracer);
234    }
235
236    @Override
237    public CompletableFutureWithSpan<Void> runAfterEitherAsync(CompletionStage<?> other, Runnable action, Executor executor) {
238        Span span = tracer.currentSpan();
239        return from(completableFuture.runAfterEitherAsync(other, () -> {
240            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
241                action.run();
242            }
243        }, executor), this.tracer);
244    }
245
246    @Override
247    public <U> CompletableFutureWithSpan<U> thenComposeAsync(Function<? super T, ? extends CompletionStage<U>> fn) {
248        Span span = tracer.currentSpan();
249        return from(completableFuture.thenComposeAsync(t -> {
250            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
251                return fn.apply(t);
252            }
253        }), this.tracer);
254    }
255
256    @Override
257    public <U> CompletableFutureWithSpan<U> thenComposeAsync(Function<? super T, ? extends CompletionStage<U>> fn, Executor executor) {
258        Span span = tracer.currentSpan();
259        return from(completableFuture.thenComposeAsync(t -> {
260            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
261                return fn.apply(t);
262            }
263        }, executor), this.tracer);
264    }
265
266    @Override
267    public CompletableFutureWithSpan<T> whenCompleteAsync(BiConsumer<? super T, ? super Throwable> action) {
268        Span span = tracer.currentSpan();
269        return from(completableFuture.whenCompleteAsync((t, throwable) -> {
270            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
271                action.accept(t, throwable);
272            }
273        }), this.tracer);
274    }
275
276    @Override
277    public CompletableFutureWithSpan<T> whenCompleteAsync(BiConsumer<? super T, ? super Throwable> action, Executor executor) {
278        Span span = tracer.currentSpan();
279        return from(completableFuture.whenCompleteAsync((t, throwable) -> {
280            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
281                action.accept(t, throwable);
282            }
283        }, executor), this.tracer);
284    }
285
286    @Override
287    public <U> CompletableFutureWithSpan<U> handleAsync(BiFunction<? super T, Throwable, ? extends U> fn) {
288        Span span = tracer.currentSpan();
289        return from(completableFuture.handleAsync((t, throwable) -> {
290            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
291                return fn.apply(t, throwable);
292            }
293        }), this.tracer);
294    }
295
296    @Override
297    public <U> CompletableFutureWithSpan<U> handleAsync(BiFunction<? super T, Throwable, ? extends U> fn, Executor executor) {
298        Span span = tracer.currentSpan();
299        return from(completableFuture.handleAsync((t, throwable) -> {
300            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
301                return fn.apply(t, throwable);
302            }
303        }, executor), this.tracer);
304    }
305
306    @Override
307    public CompletableFuture<T> completeAsync(Supplier<? extends T> supplier, Executor executor) {
308        Span span = tracer.currentSpan();
309        return completableFuture.completeAsync(() -> {
310            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
311                return supplier.get();
312            }
313        }, executor);
314    }
315
316    @Override
317    public CompletableFuture<T> completeAsync(Supplier<? extends T> supplier) {
318        Span span = tracer.currentSpan();
319        return completableFuture.completeAsync(() -> {
320            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
321                return supplier.get();
322            }
323        });
324    }
325
326    @Override
327    public boolean isDone() {
328        return completableFuture.isDone();
329    }
330
331    @Override
332    public T get() throws InterruptedException, ExecutionException {
333        return completableFuture.get();
334    }
335
336    @Override
337    public T get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
338        return completableFuture.get(timeout, unit);
339    }
340
341    @Override
342    public T join() {
343        return completableFuture.join();
344    }
345
346    @Override
347    public T getNow(T valueIfAbsent) {
348        return completableFuture.getNow(valueIfAbsent);
349    }
350
351    @Override
352    public boolean complete(T value) {
353        return completableFuture.complete(value);
354    }
355
356    @Override
357    public boolean completeExceptionally(Throwable ex) {
358        return completableFuture.completeExceptionally(ex);
359    }
360
361    @Override
362    public <U> CompletableFuture<U> thenApply(Function<? super T, ? extends U> fn) {
363        return completableFuture.thenApply(fn);
364    }
365
366    @Override
367    public CompletableFuture<Void> thenAccept(Consumer<? super T> action) {
368        return completableFuture.thenAccept(action);
369    }
370
371    @Override
372    public CompletableFuture<Void> thenRun(Runnable action) {
373        return completableFuture.thenRun(action);
374    }
375
376    @Override
377    public <U, V> CompletableFuture<V> thenCombine(CompletionStage<? extends U> other, BiFunction<? super T, ? super U, ? extends V> fn) {
378        return completableFuture.thenCombine(other, fn);
379    }
380
381    @Override
382    public <U> CompletableFuture<Void> thenAcceptBoth(CompletionStage<? extends U> other, BiConsumer<? super T, ? super U> action) {
383        return completableFuture.thenAcceptBoth(other, action);
384    }
385
386    @Override
387    public CompletableFuture<Void> runAfterBoth(CompletionStage<?> other, Runnable action) {
388        return completableFuture.runAfterBoth(other, action);
389    }
390
391    @Override
392    public <U> CompletableFuture<U> applyToEither(CompletionStage<? extends T> other, Function<? super T, U> fn) {
393        return completableFuture.applyToEither(other, fn);
394    }
395
396    @Override
397    public CompletableFuture<Void> acceptEither(CompletionStage<? extends T> other, Consumer<? super T> action) {
398        return completableFuture.acceptEither(other, action);
399    }
400
401    @Override
402    public CompletableFuture<Void> runAfterEither(CompletionStage<?> other, Runnable action) {
403        return completableFuture.runAfterEither(other, action);
404    }
405
406    @Override
407    public <U> CompletableFuture<U> thenCompose(Function<? super T, ? extends CompletionStage<U>> fn) {
408        return completableFuture.thenCompose(fn);
409    }
410
411    @Override
412    public CompletableFuture<T> whenComplete(BiConsumer<? super T, ? super Throwable> action) {
413        return completableFuture.whenComplete(action);
414    }
415
416    @Override
417    public <U> CompletableFuture<U> handle(BiFunction<? super T, Throwable, ? extends U> fn) {
418        return completableFuture.handle(fn);
419    }
420
421    @Override
422    public CompletableFuture<T> toCompletableFuture() {
423        return completableFuture.toCompletableFuture();
424    }
425
426    @Override
427    public CompletableFuture<T> exceptionally(Function<Throwable, ? extends T> fn) {
428        return completableFuture.exceptionally(fn);
429    }
430
431    @Override
432    public boolean cancel(boolean mayInterruptIfRunning) {
433        return completableFuture.cancel(mayInterruptIfRunning);
434    }
435
436    @Override
437    public boolean isCancelled() {
438        return completableFuture.isCancelled();
439    }
440
441    @Override
442    public boolean isCompletedExceptionally() {
443        return completableFuture.isCompletedExceptionally();
444    }
445
446    @Override
447    public void obtrudeValue(T value) {
448        completableFuture.obtrudeValue(value);
449    }
450
451    @Override
452    public void obtrudeException(Throwable ex) {
453        completableFuture.obtrudeException(ex);
454    }
455
456    @Override
457    public int getNumberOfDependents() {
458        return completableFuture.getNumberOfDependents();
459    }
460
461    @Override
462    public String toString() {
463        return completableFuture.toString();
464    }
465
466    @Override
467    public <U> CompletableFuture<U> newIncompleteFuture() {
468        return completableFuture.newIncompleteFuture();
469    }
470
471    @Override
472    public Executor defaultExecutor() {
473        return completableFuture.defaultExecutor();
474    }
475
476    @Override
477    public CompletableFuture<T> copy() {
478        return completableFuture.copy();
479    }
480
481    @Override
482    public CompletionStage<T> minimalCompletionStage() {
483        return completableFuture.minimalCompletionStage();
484    }
485
486    @Override
487    public CompletableFuture<T> orTimeout(long timeout, TimeUnit unit) {
488        return completableFuture.orTimeout(timeout, unit);
489    }
490
491    @Override
492    public CompletableFuture<T> completeOnTimeout(T value, long timeout, TimeUnit unit) {
493        return completableFuture.completeOnTimeout(value, timeout, unit);
494    }
495}
496

对于JDK12： 额外覆盖如下几个方法：

1    @Override
2    public CompletableFuture<T> exceptionallyCompose(Function<Throwable, ? extends CompletionStage<T>> fn) {
3        return completableFuture.exceptionallyCompose(fn);
4    }
5     @Override
6    public CompletableFutureWithSpan<T> exceptionallyAsync(Function<Throwable, ? extends T> fn) {
7        Span span = tracer.currentSpan();
8        return from(completableFuture.exceptionallyAsync(throwable -> {
9            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
10                return fn.apply(throwable);
11            }
12        }), this.tracer);
13    }
14
15    @Override
16    public CompletableFutureWithSpan<T> exceptionallyAsync(Function<Throwable, ? extends T> fn, Executor executor) {
17        Span span = tracer.currentSpan();
18        return from(completableFuture.exceptionallyAsync(throwable -> {
19            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
20                return fn.apply(throwable);
21            }
22        }, executor), this.tracer);
23    }
24
25    @Override
26    public CompletableFutureWithSpan<T> exceptionallyComposeAsync(Function<Throwable, ? extends CompletionStage<T>> fn) {
27        Span span = tracer.currentSpan();
28        return from(completableFuture.exceptionallyComposeAsync(throwable -> {
29            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
30                return fn.apply(throwable);
31            }
32        }), this.tracer);
33    }
34
35    @Override
36    public CompletableFutureWithSpan<T> exceptionallyComposeAsync(Function<Throwable, ? extends CompletionStage<T>> fn, Executor executor) {
37        Span span = tracer.currentSpan();
38        return from(completableFuture.exceptionallyComposeAsync(throwable -> {
39            try (Tracer.SpanInScope cleared = tracer.withSpanInScope(span)) {
40                return fn.apply(throwable);
41            }
42        }, executor), this.tracer);
43    }

这样使用和CompletableFuture完全一样，并且：

Mono.fromFuture(CompletableFutureWithSpan)

也是没没有问题的