Random failures regarding SegmentNotFoundException in multithreaded application

[NoMemoryError]

I am facing a problem in a multithreaded application where the dynamodb calls randomly fails to create a subsegment resulting in error message "Suppressing AWS X-Ray context missing exception (SegmentNotFoundException): Failed to begin subsegment named 'AmazonDynamoDBv2': segment cannot be found."

I am following the procedure in X-Ray documentation for passing segment context between threads as stated here - https://docs.aws.amazon.com/xray/latest/devguide/xray-sdk-java-multithreading.html and how its implemented in this example - https://docs.aws.amazon.com/xray/latest/devguide/scorekeep-workerthreads.html. In almost 70% cases the instrumentation is successful but the rest fails to find the segment when creating the subsegment.

In terms of implementation I am using custom fork join pool along with Java parallel stream.

forkJoinPool
              .submit(
                  () ->
                      filters.parallelStream()
                          .map(traceableFetcher)
                          .flatMap(Collection::stream)
                          .collect(Collectors.toList()))
              .get();

traceableFetcher gets passed the segment using awsXRayRecorder.getTraceEntity() and the implementation of the traceable class is the following

public class TraceableFunction<T, R> implements Function<T, R> {

  private final Function<T, R> target;

  private final Entity segment;

  @Getter private final MdcSnapshot mdcSnapshot;

  public TraceableFunction(Function<T, R> target, final Entity segment) {
    this.target = target;
    this.segment = segment;
    this.mdcSnapshot = new MdcSnapshot();
  }

  @Override
  public R apply(T t) {
    mdcSnapshot.populateMdc();
    var resultReference = new AtomicReference<R>();

    segment.run(() -> resultReference.set(target.apply(t)));

    mdcSnapshot.cleanMdc();
    return resultReference.get();
  }
}

It seems like in certain cases the ThreadLocalStorage doesn't have the segment or at least it has been cleared. Are you folks aware of any such issues or have any pointers? I am using com.amazonaws:aws-xray-recorder-sdk-core:2.9.1

[wangzlei]

|---parent segment---|
------|--- child subsegment ---|
That solution only woks if child subsegment is generated before parent segment close. If the sub thread starts later than main thread close, will get SegmentNotFoundException issue.

Please try to use SegmentContextExecutors