Propagate traces and log context to non-main threads

jeromevdl commented 3 months ago

Powertools for AWS Lambda (Java) should support multithreading properly, especially but not-only:

with log context and MDC, so that custom keys are not lost between threads
with traces and X-Ray, so that subsegment have the good parent segment when changing thread

Note that these 2 components use ThreadLocal.

What were you trying to accomplish?

Expected Behavior

X-Ray traces should have the good hierarchy: subsegments within threads should be under segment within 'main' thread
MDC keys (for logs) should be kept in all threads.

Current Behavior

Using the new feature (#1620), we can see in X-Ray that thread sub-segments don't have the good parent (handler).
MDC keys are lost between threads

Steps to Reproduce (for bugs)

Use the powertools-examples-batch with SQS
See traces:

Possible Solution

For logs, a temporary solution has been implemented (MultiThreadMDC.

We should think about a more "generic" or "universal" way of managing multithreading in Powertools to not only make logging and tracing work, but also enable users to use multithread more simply.

Initial thoughts:

Having an ThreadAware interface in the common module and a static list of ThreadAware classes.
ThreadAware interface should have at least 2 methods to act before the thread starts and within the thread execution, to be able to copy data from main thread (or perform another task like getting the parent x-ray segment) and reuse it within the thread.
When working with thread, we should iterate over this list and call these methods. Maybe a util class in common could simplify this too.

Environment

Powertools for AWS Lambda (Java) version used: 2.0.0-SNAPSHOT

scottgerring commented 3 months ago

x-ray specific #1670

humanzz commented 3 months ago

@scottgerring on x-ray specifically, and independent of powertools, I do wonder if anything can be done in https://github.com/aws/aws-xray-sdk-java (don't really have thoughts/solutions on this) but I am just saying, that if u guys are thinking abt it, and if u find a nice solution, then maybe we should see if such thing can fit into the xray sdk itself

I know it's not a solution for now, but the moment I came across Scoped Values (https://openjdk.org/jeps/464), I felt like, why are we not in the future, and already logging context/xray is using that 😄

scottgerring commented 3 months ago

@scottgerring on x-ray specifically, and independent of powertools, I do wonder if anything can be done in https://github.com/aws/aws-xray-sdk-java (don't really have thoughts/solutions on this) but I am just saying, that if u guys are thinking abt it, and if u find a nice solution, then maybe we should see if such thing can fit into the xray sdk itself

It looks like X-Ray used to use InheritableThreadLocal (see 1.0.6-beta) to store trace context, which may have just worked. Now it's suggesting providing a SegmentContext impl. Chasing after this I've stumbled across LambdaSegmentContext in the x-ray package itself. Because this is pulling the trace header from the environment, I wonder if providing this, or encouraging our users to use this, will be enough to jam all the subsegments under the current request together?

I know it's not a solution for now, but the moment I came across Scoped Values (https://openjdk.org/jeps/464), I felt like, why are we not in the future, and already logging context/xray is using that 😄

This is super jazzy, but probably not helpful for a few years yet sadly 😄

jeromevdl commented 3 months ago

There's a great comment from @humanzz that outlines some of the issues that happen with thread pooling in general in a Java-Lambda environment.

aws-powertools / powertools-lambda-java