From ce542024548d9ff732bcb962844f23a57c4d0d48 Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Fri, 20 Nov 2020 12:39:10 +0900 Subject: [PATCH 1/2] Add a method to run a block of code with a segment mounted as the current entity. --- .../contexts/SegmentContextExecutors.java | 9 +-------- .../com/amazonaws/xray/entities/Segment.java | 20 +++++++++++++++++++ .../contexts/SegmentContextExecutorsTest.java | 7 ++++--- dependencyManagement/build.gradle.kts | 2 +- 4 files changed, 26 insertions(+), 12 deletions(-) diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/SegmentContextExecutors.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/SegmentContextExecutors.java index d27d2663..64de0634 100644 --- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/SegmentContextExecutors.java +++ b/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); } } } diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/Segment.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/Segment.java index b85d0667..f62a188e 100644 --- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/Segment.java +++ b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/Segment.java @@ -26,6 +26,26 @@ static Segment noOp(TraceID traceId, AWSXRayRecorder recorder) { return new NoOpSegment(traceId, recorder); } + /** + * 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); + } + } + /** * Ends the segment. Sets the end time to the current time. Sets inProgress to false. * diff --git a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/SegmentContextExecutorsTest.java b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/SegmentContextExecutorsTest.java index 8c17494c..b82d8b24 100644 --- a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/SegmentContextExecutorsTest.java +++ b/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 diff --git a/dependencyManagement/build.gradle.kts b/dependencyManagement/build.gradle.kts index 8b457023..8b00e46f 100644 --- a/dependencyManagement/build.gradle.kts +++ b/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") ) ) From 90944f74fafb0755b424f4778a77cfb1ac97a517 Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Fri, 20 Nov 2020 13:01:18 +0900 Subject: [PATCH 2/2] Deprecate setTraceEntity --- .../main/java/com/amazonaws/xray/AWSXRay.java | 6 ++++++ .../com/amazonaws/xray/AWSXRayRecorder.java | 5 +++++ .../com/amazonaws/xray/entities/Entity.java | 20 +++++++++++++++++++ .../com/amazonaws/xray/entities/Segment.java | 20 ------------------- .../servlet/AWSXRayServletAsyncListener.java | 10 +++------- 5 files changed, 34 insertions(+), 27 deletions(-) diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/AWSXRay.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/AWSXRay.java index 10da6ef7..d2f1f45f 100644 --- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/AWSXRay.java +++ b/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); } diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/AWSXRayRecorder.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/AWSXRayRecorder.java index 11fc9dc6..ca9be2da 100644 --- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/AWSXRayRecorder.java +++ b/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) { diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/Entity.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/Entity.java index 02c720c6..3cbc81ed 100644 --- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/Entity.java +++ b/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(); /** diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/Segment.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/Segment.java index f62a188e..b85d0667 100644 --- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/Segment.java +++ b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/Segment.java @@ -26,26 +26,6 @@ static Segment noOp(TraceID traceId, AWSXRayRecorder recorder) { return new NoOpSegment(traceId, recorder); } - /** - * 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); - } - } - /** * Ends the segment. Sets the end time to the current time. Sets inProgress to false. * diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/javax/servlet/AWSXRayServletAsyncListener.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/javax/servlet/AWSXRayServletAsyncListener.java index 3335846b..36c39925 100644 --- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/javax/servlet/AWSXRayServletAsyncListener.java +++ b/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