suggestions

Author: amarziali

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

@@ -57,10 +57,9 @@ public void methodAdvice(MethodTransformer transformer) {
    */
   public static class InvokeAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void enter(
+    public static Context enter(
         @Advice.Argument(value = 0) Envelope envelope,
-        @Advice.Local("taskScope") AgentScope taskScope,
-        @Advice.Local("checkpointContext") Context checkpointContext) {
+        @Advice.Local("taskScope") AgentScope taskScope) {
 
       // do this before checkpointing, as the envelope's task scope may already be active
       taskScope =
@@ -70,15 +69,15 @@ public static void enter(
       if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
         // remember the currently active scope so we can roll back to this point
         checkpointActiveForRollback();
+        return null;
       } else {
-        checkpointContext = getCurrentContext().swap();
+        return getCurrentContext().swap();
       }
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void exit(
-        @Advice.Local("taskScope") AgentScope taskScope,
-        @Advice.Local("checkpointContext") Context checkpointContext) {
+        @Advice.Local("taskScope") AgentScope taskScope, @Advice.Enter Context checkpointContext) {
 
       if (checkpointContext == null) {
         // Clean up any leaking scopes from akka-streams/akka-http etc.

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

@@ -57,10 +57,9 @@ public void methodAdvice(MethodTransformer transformer) {
    */
   public static class InvokeAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void enter(
+    public static Context enter(
         @Advice.Argument(value = 0) Envelope envelope,
-        @Advice.Local("taskScope") AgentScope taskScope,
-        @Advice.Local("checkpointContext") Context checkpointContext) {
+        @Advice.Local("taskScope") AgentScope taskScope) {
 
       // do this before checkpointing, as the envelope's task scope may already be active
       taskScope =
@@ -70,15 +69,15 @@ public static void enter(
       if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
         // remember the currently active scope so we can roll back to this point
         checkpointActiveForRollback();
+        return null;
       } else {
-        checkpointContext = getCurrentContext().swap();
+        return getCurrentContext().swap();
       }
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void exit(
-        @Advice.Local("taskScope") AgentScope taskScope,
-        @Advice.Local("checkpointContext") Context checkpointContext) {
+        @Advice.Local("taskScope") AgentScope taskScope, @Advice.Enter Context checkpointContext) {
 
       if (checkpointContext == null) {
         // Clean up any leaking scopes from pekko-streams/pekko-http etc.