Skip to content

Commit 830f9db

Browse files
mccullsmcculls
committed
WIP
1 parent 372ceb0 commit 830f9db

File tree

9 files changed

+64
-134
lines changed

9 files changed

+64
-134
lines changed

‎dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java‎

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,17 @@
11
package datadog.trace.instrumentation.akka.concurrent;
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint;
64
import static java.util.Collections.singletonMap;
75
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
86

97
import akka.dispatch.Envelope;
108
import com.google.auto.service.AutoService;
9+
import datadog.context.Context;
1110
import datadog.trace.agent.tooling.Instrumenter;
1211
import datadog.trace.agent.tooling.InstrumenterModule;
1312
import datadog.trace.bootstrap.InstrumentationContext;
1413
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
14+
import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge;
1515
import datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils;
1616
import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
1717
import java.util.Map;
@@ -54,24 +54,25 @@ public void methodAdvice(MethodTransformer transformer) {
5454
*/
5555
public static class InvokeAdvice {
5656
@Advice.OnMethodEnter(suppress = Throwable.class)
57-
public static AgentScope enter(@Advice.Argument(value = 0) Envelope envelope) {
57+
public static Context enter(
58+
@Advice.Argument(value = 0) Envelope envelope,
59+
@Advice.Local("taskScope") AgentScope taskScope) {
5860

5961
// do this before checkpointing, as the envelope's task scope may already be active
60-
AgentScope taskScope =
62+
taskScope =
6163
AdviceUtils.startTaskScope(
6264
InstrumentationContext.get(Envelope.class, State.class), envelope);
6365

64-
// remember the currently active scope so we can roll back to this point
65-
checkpointActiveForRollback();
66-
67-
return taskScope;
66+
// Use swap to checkpoint the active context so we can roll back to this point
67+
return Java8BytecodeBridge.getCurrentContext().swap();
6868
}
6969

7070
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
71-
public static void exit(@Advice.Enter AgentScope taskScope) {
71+
public static void exit(
72+
@Advice.Enter Context checkpoint, @Advice.Local("taskScope") AgentScope taskScope) {
7273

7374
// Clean up any leaking scopes from akka-streams/akka-http etc.
74-
rollbackActiveToCheckpoint();
75+
checkpoint.swap();
7576

7677
// close envelope's task scope if we previously started it
7778
if (taskScope != null) {

‎dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,15 @@
11
package datadog.trace.instrumentation.akka.concurrent;
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint;
64
import static java.util.Collections.singletonList;
75
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
86

97
import com.google.auto.service.AutoService;
8+
import datadog.context.Context;
109
import datadog.trace.agent.tooling.ExcludeFilterProvider;
1110
import datadog.trace.agent.tooling.Instrumenter;
1211
import datadog.trace.agent.tooling.InstrumenterModule;
12+
import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge;
1313
import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter;
1414
import java.util.Collection;
1515
import java.util.EnumMap;
@@ -59,15 +59,15 @@ public void methodAdvice(MethodTransformer transformer) {
5959
*/
6060
public static final class SuppressMailboxRunAdvice {
6161
@Advice.OnMethodEnter(suppress = Throwable.class)
62-
public static void enter() {
63-
// remember the currently active scope so we can roll back to this point
64-
checkpointActiveForRollback();
62+
public static Context enter() {
63+
// Use swap to checkpoint the active context so we can roll back to this point
64+
return Java8BytecodeBridge.getCurrentContext().swap();
6565
}
6666

6767
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
68-
public static void exit() {
68+
public static void exit(@Advice.Enter Context checkpoint) {
6969
// Clean up any leaking scopes from akka-streams/akka-http etc.
70-
rollbackActiveToCheckpoint();
70+
checkpoint.swap();
7171
}
7272
}
7373
}

‎dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,11 @@
11
package datadog.trace.instrumentation.pekko.concurrent;
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint;
64
import static java.util.Collections.singletonMap;
75
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
86

97
import com.google.auto.service.AutoService;
8+
import datadog.context.Context;
109
import datadog.trace.agent.tooling.Instrumenter;
1110
import datadog.trace.agent.tooling.InstrumenterModule;
1211
import datadog.trace.bootstrap.InstrumentationContext;
@@ -54,24 +53,25 @@ public void methodAdvice(MethodTransformer transformer) {
5453
*/
5554
public static class InvokeAdvice {
5655
@Advice.OnMethodEnter(suppress = Throwable.class)
57-
public static AgentScope enter(@Advice.Argument(value = 0) Envelope envelope) {
56+
public static Context enter(
57+
@Advice.Argument(value = 0) Envelope envelope,
58+
@Advice.Local("taskScope") AgentScope taskScope) {
5859

5960
// do this before checkpointing, as the envelope's task scope may already be active
60-
AgentScope taskScope =
61+
taskScope =
6162
AdviceUtils.startTaskScope(
6263
InstrumentationContext.get(Envelope.class, State.class), envelope);
6364

64-
// remember the currently active scope so we can roll back to this point
65-
checkpointActiveForRollback();
66-
67-
return taskScope;
65+
// Use swap to checkpoint the active context so we can roll back to this point
66+
return Context.current().swap();
6867
}
6968

7069
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
71-
public static void exit(@Advice.Enter AgentScope taskScope) {
70+
public static void exit(
71+
@Advice.Enter Context checkpoint, @Advice.Local("taskScope") AgentScope taskScope) {
7272

7373
// Clean up any leaking scopes from pekko-streams/pekko-http etc.
74-
rollbackActiveToCheckpoint();
74+
checkpoint.swap();
7575

7676
// close envelope's task scope if we previously started it
7777
if (taskScope != null) {

‎dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoMailboxInstrumentation.java‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,11 @@
11
package datadog.trace.instrumentation.pekko.concurrent;
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint;
64
import static java.util.Collections.singletonList;
75
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
86

97
import com.google.auto.service.AutoService;
8+
import datadog.context.Context;
109
import datadog.trace.agent.tooling.ExcludeFilterProvider;
1110
import datadog.trace.agent.tooling.Instrumenter;
1211
import datadog.trace.agent.tooling.InstrumenterModule;
@@ -59,15 +58,15 @@ public void methodAdvice(MethodTransformer transformer) {
5958
*/
6059
public static final class SuppressMailboxRunAdvice {
6160
@Advice.OnMethodEnter(suppress = Throwable.class)
62-
public static void enter() {
63-
// remember the currently active scope so we can roll back to this point
64-
checkpointActiveForRollback();
61+
public static Context enter() {
62+
// Use swap to checkpoint the active context so we can roll back to this point
63+
return Context.current().swap();
6564
}
6665

6766
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
68-
public static void exit() {
67+
public static void exit(@Advice.Enter Context checkpoint) {
6968
// Clean up any leaking scopes from pekko-streams/pekko-http etc.
70-
rollbackActiveToCheckpoint();
69+
checkpoint.swap();
7170
}
7271
}
7372
}

‎dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java‎

Lines changed: 0 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1183,16 +1183,6 @@ public AgentSpan activeSpan() {
11831183
return scopeManager.activeSpan();
11841184
}
11851185

1186-
@Override
1187-
public void checkpointActiveForRollback() {
1188-
this.scopeManager.checkpointActiveForRollback();
1189-
}
1190-
1191-
@Override
1192-
public void rollbackActiveToCheckpoint() {
1193-
this.scopeManager.rollbackActiveToCheckpoint();
1194-
}
1195-
11961186
@Override
11971187
public void closeActive() {
11981188
AgentScope activeScope = this.scopeManager.active();

‎dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScope.java‎

Lines changed: 0 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,6 @@ class ContinuableScope implements AgentScope {
2121

2222
private boolean asyncPropagating;
2323

24-
private short checkpointCount = 0;
25-
2624
private final byte source;
2725

2826
private short referenceCount = 1;
@@ -140,19 +138,6 @@ public final String toString() {
140138
return super.toString() + "->" + context;
141139
}
142140

143-
public void checkpoint() {
144-
checkpointCount++;
145-
}
146-
147-
public boolean rollback() {
148-
if (checkpointCount > 0) {
149-
checkpointCount--;
150-
return false; // stop rollback at checkpoint
151-
} else {
152-
return true;
153-
}
154-
}
155-
156141
public final void beforeActivated() {
157142
AgentSpan span = span();
158143
if (span == null) {

‎dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScopeManager.java‎

Lines changed: 0 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -284,24 +284,6 @@ public AgentScope active() {
284284
return scopeStack().active();
285285
}
286286

287-
public void checkpointActiveForRollback() {
288-
ContinuableScope active = scopeStack().active();
289-
if (active != null) {
290-
active.checkpoint();
291-
}
292-
}
293-
294-
public void rollbackActiveToCheckpoint() {
295-
ContinuableScope active;
296-
while ((active = scopeStack().active()) != null) {
297-
if (active.rollback()) {
298-
active.close();
299-
} else {
300-
break; // stop at the most recent checkpointed scope
301-
}
302-
}
303-
}
304-
305287
public AgentSpan activeSpan() {
306288
final ContinuableScope active = scopeStack().active();
307289
return active == null ? null : active.span();

‎dd-trace-core/src/test/groovy/datadog/trace/core/scopemanager/ScopeManagerTest.groovy‎

Lines changed: 30 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,37 +1,37 @@
11
package datadog.trace.core.scopemanager
22

3+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopContinuation
4+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan
5+
import static datadog.trace.core.scopemanager.EVENT.ACTIVATE
6+
import static datadog.trace.core.scopemanager.EVENT.CLOSE
7+
import static datadog.trace.test.util.GCUtils.awaitGC
8+
39
import datadog.context.Context
410
import datadog.context.ContextKey
5-
import datadog.trace.test.util.ThreadUtils
11+
import datadog.context.ContextScope
612
import datadog.trace.api.DDTraceId
713
import datadog.trace.api.Stateful
814
import datadog.trace.api.interceptor.MutableSpan
915
import datadog.trace.api.interceptor.TraceInterceptor
1016
import datadog.trace.api.scopemanager.ExtendedScopeListener
17+
import datadog.trace.api.scopemanager.ScopeListener
1118
import datadog.trace.bootstrap.instrumentation.api.AgentScope
1219
import datadog.trace.bootstrap.instrumentation.api.AgentSpan
1320
import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration
1421
import datadog.trace.common.writer.ListWriter
15-
import datadog.trace.api.scopemanager.ScopeListener
1622
import datadog.trace.context.TraceScope
1723
import datadog.trace.core.CoreTracer
1824
import datadog.trace.core.DDSpan
1925
import datadog.trace.core.test.DDCoreSpecification
20-
import spock.lang.Shared
21-
26+
import datadog.trace.test.util.ThreadUtils
2227
import java.lang.ref.WeakReference
2328
import java.util.concurrent.ExecutorService
2429
import java.util.concurrent.Executors
2530
import java.util.concurrent.Future
2631
import java.util.concurrent.TimeUnit
2732
import java.util.concurrent.atomic.AtomicInteger
2833
import java.util.concurrent.atomic.AtomicReference
29-
30-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopContinuation
31-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan
32-
import static datadog.trace.core.scopemanager.EVENT.ACTIVATE
33-
import static datadog.trace.core.scopemanager.EVENT.CLOSE
34-
import static datadog.trace.test.util.GCUtils.awaitGC
34+
import spock.lang.Shared
3535

3636
enum EVENT {
3737
ACTIVATE, CLOSE
@@ -1072,59 +1072,66 @@ class ScopeManagerTest extends DDCoreSpecification {
10721072
scopeManager.activeSpan() == null
10731073
}
10741074

1075-
def "rollback stops at most recent checkpoint"() {
1075+
def "swap can be used to checkpoint and rollback"() {
10761076
when:
1077+
def span0 = tracer.buildSpan("test0", "test0").start()
10771078
def span1 = tracer.buildSpan("test1", "test1").start()
10781079
def span2 = tracer.buildSpan("test2", "test2").start()
10791080
def span3 = tracer.buildSpan("test3", "test3").start()
10801081
then:
10811082
scopeManager.activeSpan() == null
10821083

10831084
when:
1084-
tracer.checkpointActiveForRollback()
1085+
ContextScope scope0 = tracer.activateSpan(span0)
1086+
Context c0 = Context.current().swap()
10851087
tracer.activateSpan(span1)
1086-
tracer.checkpointActiveForRollback()
1088+
Context c1 = Context.current().swap()
10871089
tracer.activateSpan(span2)
1088-
tracer.checkpointActiveForRollback()
1090+
Context c2 = Context.current().swap()
10891091
tracer.activateSpan(span1)
1090-
tracer.checkpointActiveForRollback()
1092+
Context c3 = Context.current().swap()
10911093
tracer.activateSpan(span2)
1092-
tracer.checkpointActiveForRollback()
1094+
Context c4 = Context.current().swap()
10931095
tracer.activateSpan(span2)
1094-
tracer.checkpointActiveForRollback()
1096+
Context c5 = Context.current().swap()
10951097
tracer.activateSpan(span1)
10961098
tracer.activateSpan(span2)
10971099
tracer.activateSpan(span3)
10981100
then:
10991101
scopeManager.activeSpan() == span3
11001102

11011103
when:
1102-
tracer.rollbackActiveToCheckpoint()
1104+
c5.swap()
11031105
then:
11041106
scopeManager.activeSpan() == span2
11051107

11061108
when:
1107-
tracer.rollbackActiveToCheckpoint()
1109+
c4.swap()
11081110
then:
11091111
scopeManager.activeSpan() == span2
11101112

11111113
when:
1112-
tracer.rollbackActiveToCheckpoint()
1114+
c3.swap()
11131115
then:
11141116
scopeManager.activeSpan() == span1
11151117

11161118
when:
1117-
tracer.rollbackActiveToCheckpoint()
1119+
c2.swap()
11181120
then:
11191121
scopeManager.activeSpan() == span2
11201122

11211123
when:
1122-
tracer.rollbackActiveToCheckpoint()
1124+
c1.swap()
11231125
then:
11241126
scopeManager.activeSpan() == span1
11251127

11261128
when:
1127-
tracer.rollbackActiveToCheckpoint()
1129+
c0.swap()
1130+
then:
1131+
scopeManager.activeSpan() == span0
1132+
1133+
when:
1134+
scope0.close()
11281135
then:
11291136
scopeManager.activeSpan() == null
11301137
}

0 commit comments

Comments
 (0)