X-Ray and async multithreading: segment cannot be found

Question

I need to run multiple async methods from the main thread, which should be able to finish even after the main thread has returned a response.

I've added a method which does the async call and handles the X-Ray based things, which works if I call it only once. However, if I call it again from the same thread, then X-Ray starts throwing .a.x.e.SegmentNotFoundException: Failed to begin subsegment named '30baae07-3a63-4036-807d-dcb0fdbac4eb': segment cannot be found.

The async method:

public static void runAsync(Runnable runnable) {
    CompletableFuture.runAsync(() -> {
        try {
            Subsegment subsegment = AWSXRay.beginSubsegment(UUID.randomUUID().toString());

            runnable.run();

            GLOBAL_RECORDER.endSubsegment(subsegment);
        } catch (Exception e) {
            log.error("Failed to run async method", e);
        }
    });
}

And its usage:

    AsyncUtil.runAsync(() -> alerterService.refreshDataA(...));

    // Here is a call of another unrelated method

    AsyncUtil.runAsync(() -> qualifierService.refreshDataB(...));

The source of this code is taken from https://github.com/aws/aws-xray-sdk-java/issues/227. I've looked through multiple sources about how to make this work and this seemed to be the correct way. Yet it does not work in my case. If I remove one of the async calls, it works as intended, so the issue must be that those 2 async methods are run in parallel alongside the main thread.

Do you know how to make this work without running the two async methods in a series using just 1 thread?

Please note that I'm very new to X-Ray and this is the first time I'm working with it, even though it's been in our project for a while.

EDIT: After changing it to work the way as @lei-wang responded, the issue still persisted.

public static void runAsync(Runnable runnable) {
    Entity traceEntity = Objects.requireNonNull(AWSXRay.getTraceEntity());

    new Thread(() -> {
        try {
            AWSXRay.beginSegment("Async segment", traceEntity.getTraceId(), traceEntity.getId());

            CompletableFuture
                    .runAsync(runnable)
                    .join();

            AWSXRay.endSegment();
        } catch (Exception e) {
            log.error("Failed to run async method", e);
        }
    }).start();
}

It's probably good to note that in this case it can't find the segment named as the service which is used inside the runnable, eg. "Failed to begin subsegment named 'alerter-service': segment cannot be found."

Answer 1

XRay Subsegment has to be under a living Segment, the issue because 2 runAsyn get started after the main thread finished, so Subsegments cannot find parent Segment. The workaround is to create a new thread to manage these 2 async methods, as the parent Segment of 2 Subsegments:

public static CompletableFuture runAsync(Runnable runnable) {
 return CompletableFuture.runAsync(().....)
}

// from segment in main thread.
traceid = AWSXRay.currentTraceId();
parentid = AWSXRay.currentEntityId();

new Thread() {
  void run() {
    AWSXRay.beginSegment("Async management", traceid, parentid);
    CompletableFuture future = AsyncUtil.runAsync();
    future.join(); // wait for asyn complete
    AWSXRay.endSegment();
  }
}

Or you can beginSegment instead of beginSubsegment in your async thread directly.

Answer 2

My test code

private static void testAsync() throws InterruptedException {
    AWSXRay.beginSegment("testAsync");
    runAsyncUtil(() -> System.out.println("1"));
    runAsyncUtil(() -> System.out.println("2"));
    AWSXRay.endSegment();
  }

  private static void runAsyncUtil(Runnable runnable) {
    CompletableFuture.runAsync(() -> {
      AWSXRay.beginSubsegment("sub");
      runnable.run();
      AWSXRay.endSubsegment();
    }, SegmentContextExecutors.newSegmentContextExecutor());
  }