CompletableFutureTest.java

@Test
  void sharedCompletableFuture() throws Exception {

    Supplier<String> longTokenFetchProcess = () ->
        testing.runWithSpan("tokenFetch", () -> {
          try {
            Thread.sleep(5000L);
          } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RuntimeException(e);
          }
          return "token";
        });

    // Originally this CompletableFuture would be a result of java.net.http.HttpClient::sendAsync
    // A supplier to ensure the call will be executed only if needed
    Supplier<CompletableFuture<String>> singleTokenFetch =  () -> CompletableFuture.supplyAsync(
        longTokenFetchProcess
    );

    ConcurrentHashMap<String, CompletableFuture<String>> fetchStorage = new ConcurrentHashMap<>();
    CountDownLatch sync = new CountDownLatch(2);

    Runnable c1 = makeController("controller1", sync, fetchStorage, singleTokenFetch);
    Runnable c2 = makeController("controller2", sync, fetchStorage, singleTokenFetch);

    ExecutorService controllerExecutors = Executors.newFixedThreadPool(2);
    controllerExecutors.submit(c1);
    controllerExecutors.submit(c2);
    controllerExecutors.shutdown();
    controllerExecutors.awaitTermination(10, TimeUnit.SECONDS);

    List<List<SpanData>> traces = testing.waitForTraces(2);

    Set<String> c1Spans = traces.get(0)
        .stream()
        .map(SpanData::getName)
        .collect(Collectors.toSet());
    Set<String> c2Spans = traces.get(1)
        .stream()
        .map(SpanData::getName)
        .collect(Collectors.toSet());

    assertThat("tokenFetch").satisfiesAnyOf(
        span -> assertThat(c1Spans).contains(span),
        span -> assertThat(c2Spans).contains(span)
    );
    assertThat(c2Spans).contains("controller2", "controller2-end");
    assertThat(c1Spans).contains("controller1", "controller1-end");
  }

  private static Runnable makeController(String name, CountDownLatch sync,
      ConcurrentHashMap<String, CompletableFuture<String>> fetchStorage,
      Supplier<CompletableFuture<String>> singleTokenFetch) {
    return () -> {
      Tracer tracer = GlobalOpenTelemetry.get().getTracer("test");
      Span span = tracer.spanBuilder(name).startSpan();
      try(Scope scope = span.makeCurrent()){ // root of the trace
        sync.countDown();
        try {
          sync.await();
        } catch (InterruptedException e) {
          throw new RuntimeException(e);
        }
        // either another span in the trace or not, depending if fetch is needed
        fetchStorage.computeIfAbsent("COMMON_TOKEN", _ignored -> singleTokenFetch.get()).join();
        // ending span
        Span span2 = tracer.spanBuilder(name + "-end").startSpan();
        try(Scope scope2 = span.makeCurrent()){
          //nop
        }
        span2.end();
      }
      span.end();
    };
  }