Support log correlation during entity propagation

aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/SegmentContext.java

@@ -1,8 +1,8 @@
 package com.amazonaws.xray.contexts;
 
+import java.util.Objects;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
-
 import com.amazonaws.xray.AWSXRayRecorder;
 import com.amazonaws.xray.ThreadLocalStorage;
 import com.amazonaws.xray.entities.Entity;
@@ -25,10 +25,20 @@ default Entity getTraceEntity() {
     }
 
     default void setTraceEntity(Entity entity) {
+        if(entity != null && entity.getCreator() != null) {
+            entity.getCreator().getSegmentListeners().stream().filter(Objects::nonNull).forEach(l -> {
+                l.onSetEntity(ThreadLocalStorage.get(), entity);
+            });
+        }
         ThreadLocalStorage.set(entity);
     }
 
     default void clearTraceEntity() {
+        Entity oldEntity = ThreadLocalStorage.get();
+        if(oldEntity != null && oldEntity.getCreator() != null)
+        oldEntity.getCreator().getSegmentListeners().stream().filter(Objects::nonNull).forEach(l -> {
+            l.onClearEntity(oldEntity);
+        });
         ThreadLocalStorage.clear();
     }
 

aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/listeners/SegmentListener.java

@@ -1,5 +1,6 @@
 package com.amazonaws.xray.listeners;
 
+import com.amazonaws.xray.entities.Entity;
 import com.amazonaws.xray.entities.Segment;
 
 /**
@@ -41,4 +42,25 @@ default void beforeEndSegment(Segment segment) {
     default void afterEndSegment(Segment segment) {
 
     }
+
+    /**
+     * onSetEntity is invoked when the SegmentContext is being updated with a new entity.
+     * Both the new entity and the previous entity (or null if unset) are passed.
+     *
+     * @param previousEntity
+     * @param newEntity
+     */
+    default void onSetEntity(Entity previousEntity, Entity newEntity) {
+
+    }
+
+    /**
+     * onClearEntity is invoked just before the SegmentContext is cleared.
+     *
+     * @param previousEntity
+     */
+    default void onClearEntity(Entity previousEntity) {
+
+    }
+
 }

aws-xray-recorder-sdk-log4j/src/main/java/com/amazonaws/xray/log4j/Log4JSegmentListener.java

@@ -1,6 +1,6 @@
 package com.amazonaws.xray.log4j;
 
-import com.amazonaws.xray.entities.Segment;
+import com.amazonaws.xray.entities.Entity;
 import com.amazonaws.xray.listeners.SegmentListener;
 import org.apache.logging.log4j.ThreadContext;
 
@@ -17,24 +17,25 @@ public class Log4JSegmentListener implements SegmentListener {
     /**
      * Maps the AWS-XRAY-TRACE-ID key to the trace ID of the segment that's just been created in the Log4J ThreadContext.
      *
-     * @param segment
+     * @param oldEntity the previous entity or null
+     * @param newEntity the new entity
      * The segment that has just begun
      */
     @Override
-    public void onBeginSegment(Segment segment) {
-        if (segment != null) {
-            ThreadContext.put(TRACE_ID_KEY, TRACE_ID_KEY + ": " + segment.getTraceId().toString());
+    public void onSetEntity(Entity oldEntity, Entity newEntity) {
+        if (newEntity != null && newEntity.getTraceId() != null) {
+            ThreadContext.put(TRACE_ID_KEY, TRACE_ID_KEY + ": " + newEntity.getTraceId().toString());
         }
     }
 
     /**
      * Removes the AWS-XRAY-TRACE-ID key from the ThreadContext upon the completion of each segment.
      *
-     * @param segment
+     * @param entity
      * The segment that has just ended
      */
     @Override
-    public void beforeEndSegment(Segment segment) {
+    public void onClearEntity(Entity entity) {
         ThreadContext.remove(TRACE_ID_KEY);
     }
 }

aws-xray-recorder-sdk-slf4j/src/main/java/com/amazonaws/xray/slf4j/SLF4JSegmentListener.java

@@ -1,6 +1,6 @@
 package com.amazonaws.xray.slf4j;
 
-import com.amazonaws.xray.entities.Segment;
+import com.amazonaws.xray.entities.Entity;
 import com.amazonaws.xray.listeners.SegmentListener;
 import org.slf4j.MDC;
 
@@ -18,26 +18,25 @@ public class SLF4JSegmentListener implements SegmentListener {
     /**
      * Maps the AWS-XRAY-TRACE-ID key to the trace ID of the segment that's just been created in the MDC.
      *
-     * @param segment
+     * @param oldEntity the previous entity or null
+     * @param newEntity the new entity
      * The segment that has just begun
      */
     @Override
-    public void onBeginSegment(Segment segment) {
-
-        // TODO: Check if we are in a Lambda environment (dealing w/ facade segment), and grab the Trace ID from the Header instead if so
-        if (segment != null) {
-            MDC.put(TRACE_ID_KEY, TRACE_ID_KEY + ": " + segment.getTraceId().toString());
+    public void onSetEntity(Entity oldEntity, Entity newEntity) {
+        if (newEntity != null && newEntity.getTraceId() != null) {
+            MDC.put(TRACE_ID_KEY, TRACE_ID_KEY + ": " + newEntity.getTraceId().toString());
         }
     }
 
     /**
      * Removes the AWS-XRAY-TRACE-ID key from the MDC upon the completion of each segment.
      *
-     * @param segment
+     * @param entity
      * The segment that has just ended
      */
     @Override
-    public void beforeEndSegment(Segment segment) {
+    public void onClearEntity(Entity entity) {
         MDC.remove(TRACE_ID_KEY);
     }
 }