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.

Sign up for GitHub

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

Simplify the Akka/Pekko checkpointing advice #8625

Merged
merged 1 commit into from
Mar 27, 2025
Merged

Simplify the Akka/Pekko checkpointing advice #8625

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
45 changes: 15 additions & 30 deletions ...main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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();
}
}
}
16 changes: 1 addition & 15 deletions ...c/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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)
Expand Down
45 changes: 15 additions & 30 deletions ...in/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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();
}
}
}
17 changes: 1 addition & 16 deletions ...main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoMailboxInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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)
Expand Down