Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 2 commits into from
Nov 24, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is deprecating this the path forward we're certain we want to take? It is certainly a sharp edge and I definitely agree with recommending using Entity.run() over this. But IMO customers who know their own architecture should still be able to use these low-level APIs if it suits their app and they're comfortable with the risks. Especially if they've implemented a custom SegmentContext interface that overrides the default setTraceEntity() behavior.

All this being said, if the team disagrees and we move forward with this path, should we go ahead and also deprecate clearTraceEntity and getTraceEntity?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

get is required since you definitely need to be able to get the current entity. It's also safe.

clear I debated - but even in OTel it's important in request handlers as a safeguard against leaks. For maximum safety we always clear before creating a root span (segment) since parenting never makes sense. So I think it's ok to keep - it's basically safe (can lead to broken traces potentially/rarely, but definitely not leaks).

set is bad since it requires the restoration maneuvering - it's never better than a Closeable or wrapping a lambda, just more tedious. We haven't added a closeable in this PR yet but I'm still comfortable removing set, it can never provide a good experience I think.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok, thanks for the explainers. Here's to safer usage going forward :)

public static void setTraceEntity(Entity entity) {
globalRecorder.setTraceEntity(entity);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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();

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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
Expand Down
2 changes: 1 addition & 1 deletion dependencyManagement/build.gradle.kts
Original file line number Diff line number Diff line change
Expand Up @@ -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")
)
)
Expand Down