TTL does not replay on CompletableFuture sync path
Problem description
CompletableFuture.Completion extends ForkJoinTask, so we expect TTL to replay the captured context when executing completion stages. The problem is that after completion, CompletableFuture may greedily execute its dependent stages via the tryFire method, bypassing the ForkJoinTask.doExec interface. Because TTL only instruments doExec, these dependent stages execute without the replay/restore logic and therefore read the wrong context. This issue arises whenever we use non-async methods, e.g. CompletableFuture.thenCompose, CompletableFuture.thenApply, and CompletableFuture.handle.
Here is a unit test showing the problem
@Test
void test() {
var context = new TransmittableThreadLocal<String>();
var pool = Executors.newSingleThreadExecutor();
context.set("C1");
var base = new CompletableFuture<String>();
var chained = base.thenApply(s -> context.get());
pool.execute(() -> { context.set("C2"); base.complete(""); });
String res = chained.join();
assertEquals("C1", res); // expected: <C1> but was: <C2>
}Potential solution
Treat CompletableFuture.Completion subclasses as a special case of ForkJoinTask. We should not rely on doExec instrumentation since its workload may be called via tryFire. We need to instrument the constructor and modify the stored closure:
From:
class UniCompose<T,V> {
Function<? super T, ? extends CompletionStage<V>> fn;
UniCompose(Executor executor, CompletableFuture<V> dep,
CompletableFuture<T> src,
Function<? super T, ? extends CompletionStage<V>> fn) {
super(executor, dep, src);
this.fn = fn;
}
}To:
class UniCompose<T,V> {
Function<? super T, ? extends CompletionStage<V>> fn;
UniCompose(Executor executor, CompletableFuture<V> dep,
CompletableFuture<T> src,
Function<? super T, ? extends CompletionStage<V>> fn) {
super(executor, dep, src);
this.fn = arg -> {
if (isReplayRedundant()) {
return fn.apply(arg); // Fast path
}
var backup = replay(captured);
try {
return fn.apply(arg);
} finally {
restore(backup);
}
};
}
}Source: alibaba/transmittable-thread-local