Add a method to run a block of code with a segment mounted as the cur… (#240)

* Add a method to run a block of code with a segment mounted as the current entity.

* Deprecate setTraceEntity

Author: Anuraag Agrawal

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

@@ -15,6 +15,7 @@
 
 package com.amazonaws.xray;
 
+import com.amazonaws.xray.contexts.SegmentContextExecutors;
 import com.amazonaws.xray.entities.Entity;
 import com.amazonaws.xray.entities.Segment;
 import com.amazonaws.xray.entities.Subsegment;
@@ -186,6 +187,11 @@ public static void clearThreadLocal() {
         globalRecorder.clearThreadLocal();
     }
 
+    /**
+     * @deprecated Use {@link Entity#run(Runnable)} or methods in {@link SegmentContextExecutors} instead of directly setting
+     * the trace entity so it can be restored correctly.
+     */
+    @Deprecated
     public static void setTraceEntity(Entity entity) {
         globalRecorder.setTraceEntity(entity);
     }

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

@@ -17,6 +17,7 @@
 
 import com.amazonaws.xray.contexts.LambdaSegmentContextResolver;
 import com.amazonaws.xray.contexts.SegmentContext;
+import com.amazonaws.xray.contexts.SegmentContextExecutors;
 import com.amazonaws.xray.contexts.SegmentContextResolverChain;
 import com.amazonaws.xray.contexts.ThreadLocalSegmentContextResolver;
 import com.amazonaws.xray.emitters.Emitter;
@@ -753,7 +754,11 @@ private SegmentContext getSegmentContext() {
      *
      * @param entity
      *            the trace entity to set
+     *
+     * @deprecated Use {@link Entity#run(Runnable)} or methods in {@link SegmentContextExecutors} instead of directly setting
+     * the trace entity so it can be restored correctly.
      */
+    @Deprecated
     public void setTraceEntity(@Nullable Entity entity) {
         SegmentContext context = getSegmentContext();
         if (context == null) {

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

@@ -19,7 +19,6 @@
 
 import com.amazonaws.xray.AWSXRay;
 import com.amazonaws.xray.AWSXRayRecorder;
-import com.amazonaws.xray.entities.Entity;
 import com.amazonaws.xray.entities.Segment;
 import java.util.concurrent.Executor;
 import org.checkerframework.checker.nullness.qual.Nullable;
@@ -80,13 +79,7 @@ private SegmentContextExecutor(AWSXRayRecorder recorder, Segment segment) {
 
         @Override
         public void execute(Runnable command) {
-            Entity previous = recorder.getTraceEntity();
-            recorder.setTraceEntity(segment);
-            try {
-                command.run();
-            } finally {
-                recorder.setTraceEntity(previous);
-            }
+            segment.run(command, recorder);
         }
     }
 }

aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/Entity.java

@@ -36,6 +36,26 @@ static String generateId() {
         return AWSXRay.getGlobalRecorder().getIdGenerator().newEntityId();
     }
 
+    /**
+     * Immediately runs the provided {@link Runnable} with this {@link Segment} as the current entity.
+     */
+    default void run(Runnable runnable) {
+        run(runnable, getCreator());
+    }
+
+    /**
+     * Immediately runs the provided {@link Runnable} with this {@link Segment} as the current entity.
+     */
+    default void run(Runnable runnable, AWSXRayRecorder recorder) {
+        Entity previous = recorder.getTraceEntity();
+        recorder.setTraceEntity(this);
+        try {
+            runnable.run();
+        } finally {
+            recorder.setTraceEntity(previous);
+        }
+    }
+
     String getName();
 
     /**

aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/javax/servlet/AWSXRayServletAsyncListener.java

@@ -48,17 +48,13 @@ private AWSXRayRecorder getRecorder() {
 
     private void processEvent(AsyncEvent event) throws IOException {
         AWSXRayRecorder recorder = getRecorder();
-        Entity prior = recorder.getTraceEntity();
-        try {
-            Entity entity = (Entity) event.getSuppliedRequest().getAttribute(ENTITY_ATTRIBUTE_KEY);
-            recorder.setTraceEntity(entity);
+        Entity entity = (Entity) event.getSuppliedRequest().getAttribute(ENTITY_ATTRIBUTE_KEY);
+        entity.run(() -> {
             if (event.getThrowable() != null) {
                 entity.addException(event.getThrowable());
             }
             filter.postFilter(event.getSuppliedRequest(), event.getSuppliedResponse());
-        } finally {
-            recorder.setTraceEntity(prior);
-        }
+        });
     }
 
     @Override

aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/SegmentContextExecutorsTest.java

@@ -39,6 +39,7 @@
 import org.junit.Rule;
 import org.junit.Test;
 import org.mockito.Mock;
+import org.mockito.Spy;
 import org.mockito.junit.MockitoJUnit;
 import org.mockito.junit.MockitoRule;
 
@@ -52,13 +53,13 @@ public class SegmentContextExecutorsTest {
     @Mock
     private volatile AWSXRayRecorder recorder;
 
-    @Mock
+    @Spy
     private volatile Segment current;
 
-    @Mock
+    @Spy
     private volatile Segment manual;
 
-    @Mock
+    @Spy
     private volatile Segment previous;
 
     @BeforeClass

dependencyManagement/build.gradle.kts

@@ -37,7 +37,7 @@ val DEPENDENCY_SETS = listOf(
         ),
         DependencySet(
                 "org.mockito",
-                "2.28.2",
+                "3.6.0",
                 listOf("mockito-all", "mockito-core", "mockito-junit-jupiter")
         )
 )