From 317eb6c3b9f030d54bdd33d727d63c6d83fc4a5c Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Thu, 27 Mar 2025 10:13:38 +0000 Subject: [PATCH] Simplify the Akka/Pekko checkpointing advice --- .../AkkaActorCellInstrumentation.java | 45 +++++++------------ .../AkkaMailboxInstrumentation.java | 16 +------ .../PekkoActorCellInstrumentation.java | 45 +++++++------------ .../PekkoMailboxInstrumentation.java | 17 +------ 4 files changed, 32 insertions(+), 91 deletions(-) diff --git a/dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java b/dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java index 25edf7bcbc2..7b04b3aef16 100644 --- a/dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java +++ b/dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java @@ -1,10 +1,7 @@ package datadog.trace.instrumentation.akka.concurrent; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint; import static java.util.Collections.singletonMap; import static net.bytebuddy.matcher.ElementMatchers.isMethod; @@ -15,7 +12,6 @@ import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.InstrumentationContext; import datadog.trace.bootstrap.instrumentation.api.AgentScope; -import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils; import datadog.trace.bootstrap.instrumentation.java.concurrent.State; import java.util.Map; @@ -58,40 +54,29 @@ public void methodAdvice(MethodTransformer transformer) { */ public static class InvokeAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static void enter( - @Advice.Argument(value = 0) Envelope envelope, - @Advice.Local(value = "taskScope") AgentScope taskScope) { - checkpointActiveForRollback(); - // note: task scope may be the same as the scope we want to roll back to, - // so we must remember to close it on exit to balance the activation count - taskScope = + public static AgentScope enter(@Advice.Argument(value = 0) Envelope envelope) { + + // do this before checkpointing, as the envelope's task scope may already be active + AgentScope taskScope = AdviceUtils.startTaskScope( InstrumentationContext.get(Envelope.class, State.class), envelope); - // There was a scope created from the envelope, so use that - if (taskScope != null) { - return; - } - AgentSpan activeSpan = activeSpan(); - // If there is no active scope, we can clean all the way to the bottom - if (activeSpan == null) { - return; - } - // If there is a noop span in the active scope, we can clean all the way to this scope - if (activeSpan == noopSpan()) { - return; - } - // Create an active scope with a noop span, and clean all the way to the previous scope - activateSpan(noopSpan()); + + // remember the currently active scope so we can roll back to this point + checkpointActiveForRollback(); + + return taskScope; } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) - public static void exit(@Advice.Local(value = "taskScope") AgentScope taskScope) { + public static void exit(@Advice.Enter AgentScope taskScope) { + + // Clean up any leaking scopes from akka-streams/akka-http etc. + rollbackActiveToCheckpoint(); + + // close envelope's task scope if we previously started it if (taskScope != null) { - // then we have invoked an Envelope and need to mark the work complete taskScope.close(); } - // Clean up any leaking scopes from akka-streams/akka-http etc. - rollbackActiveToCheckpoint(); } } } diff --git a/dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java b/dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java index c20d8d2e021..b9f949ecd0f 100644 --- a/dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java +++ b/dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java @@ -1,10 +1,7 @@ package datadog.trace.instrumentation.akka.concurrent; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint; import static java.util.Collections.singletonList; import static net.bytebuddy.matcher.ElementMatchers.isMethod; @@ -13,7 +10,6 @@ import datadog.trace.agent.tooling.ExcludeFilterProvider; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; -import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter; import java.util.Collection; import java.util.EnumMap; @@ -64,18 +60,8 @@ public void methodAdvice(MethodTransformer transformer) { public static final class SuppressMailboxRunAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void enter() { + // remember the currently active scope so we can roll back to this point checkpointActiveForRollback(); - AgentSpan activeSpan = activeSpan(); - // If there is no active scope, we can clean all the way to the bottom - if (activeSpan == null) { - return; - } - // If there is a noop span in the active scope, we can clean all the way to this scope - if (activeSpan == noopSpan()) { - return; - } - // Create an active scope with a noop span, and clean all the way to the previous scope - activateSpan(noopSpan()); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) diff --git a/dd-java-agent/instrumentation/pekko-concurrent/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java b/dd-java-agent/instrumentation/pekko-concurrent/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java index c5d21d91d9e..9f449a53e63 100644 --- a/dd-java-agent/instrumentation/pekko-concurrent/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java +++ b/dd-java-agent/instrumentation/pekko-concurrent/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java @@ -1,10 +1,7 @@ package datadog.trace.instrumentation.pekko.concurrent; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint; import static java.util.Collections.singletonMap; import static net.bytebuddy.matcher.ElementMatchers.isMethod; @@ -14,7 +11,6 @@ import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.InstrumentationContext; import datadog.trace.bootstrap.instrumentation.api.AgentScope; -import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils; import datadog.trace.bootstrap.instrumentation.java.concurrent.State; import java.util.Map; @@ -58,40 +54,29 @@ public void methodAdvice(MethodTransformer transformer) { */ public static class InvokeAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static void enter( - @Advice.Argument(value = 0) Envelope envelope, - @Advice.Local(value = "taskScope") AgentScope taskScope) { - checkpointActiveForRollback(); - // note: task scope may be the same as the scope we want to roll back to, - // so we must remember to close it on exit to balance the activation count - taskScope = + public static AgentScope enter(@Advice.Argument(value = 0) Envelope envelope) { + + // do this before checkpointing, as the envelope's task scope may already be active + AgentScope taskScope = AdviceUtils.startTaskScope( InstrumentationContext.get(Envelope.class, State.class), envelope); - // There was a scope created from the envelope, so use that - if (taskScope != null) { - return; - } - AgentSpan activeSpan = activeSpan(); - // If there is no active scope, we can clean all the way to the bottom - if (activeSpan == null) { - return; - } - // If there is a noop span in the active scope, we can clean all the way to this scope - if (activeSpan == noopSpan()) { - return; - } - // Create an active scope with a noop span, and clean all the way to the previous scope - activateSpan(noopSpan()); + + // remember the currently active scope so we can roll back to this point + checkpointActiveForRollback(); + + return taskScope; } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) - public static void exit(@Advice.Local(value = "taskScope") AgentScope taskScope) { + public static void exit(@Advice.Enter AgentScope taskScope) { + + // Clean up any leaking scopes from pekko-streams/pekko-http etc. + rollbackActiveToCheckpoint(); + + // close envelope's task scope if we previously started it if (taskScope != null) { - // then we have invoked an Envelope and need to mark the work complete taskScope.close(); } - // Clean up any leaking scopes from pekko-streams/pekko-http etc. - rollbackActiveToCheckpoint(); } } } diff --git a/dd-java-agent/instrumentation/pekko-concurrent/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoMailboxInstrumentation.java b/dd-java-agent/instrumentation/pekko-concurrent/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoMailboxInstrumentation.java index a881c22fb07..ceb432a5fff 100644 --- a/dd-java-agent/instrumentation/pekko-concurrent/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoMailboxInstrumentation.java +++ b/dd-java-agent/instrumentation/pekko-concurrent/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoMailboxInstrumentation.java @@ -1,10 +1,7 @@ package datadog.trace.instrumentation.pekko.concurrent; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint; import static java.util.Collections.singletonList; import static net.bytebuddy.matcher.ElementMatchers.isMethod; @@ -13,7 +10,6 @@ import datadog.trace.agent.tooling.ExcludeFilterProvider; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; -import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter; import java.util.Collection; import java.util.EnumMap; @@ -64,19 +60,8 @@ public void methodAdvice(MethodTransformer transformer) { public static final class SuppressMailboxRunAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void enter() { + // remember the currently active scope so we can roll back to this point checkpointActiveForRollback(); - - AgentSpan activeSpan = activeSpan(); - // If there is no active scope, we can clean all the way to the bottom - if (activeSpan == null) { - return; - } - // If there is a noop span in the active scope, we can clean all the way to this scope - if (activeSpan == noopSpan()) { - return; - } - // Create an active scope with a noop span, and clean all the way to the previous scope - activateSpan(noopSpan()); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)