Migrate simple try-with-resources use of AgentScope to ContextScope

dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java

@@ -11,6 +11,7 @@
 import com.squareup.moshi.Moshi;
 import com.squareup.moshi.Types;
 import datadog.communication.http.OkHttpUtils;
+import datadog.context.ContextScope;
 import datadog.trace.api.Config;
 import datadog.trace.api.aiguard.AIGuard;
 import datadog.trace.api.aiguard.AIGuard.AIGuardAbortError;
@@ -26,7 +27,6 @@
 import datadog.trace.api.aiguard.noop.NoOpEvaluator;
 import datadog.trace.api.gateway.RequestContext;
 import datadog.trace.api.telemetry.WafMetricCollector;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.ClientIpAddressData;
@@ -281,7 +281,7 @@ public Evaluation evaluate(final List<Message> messages, final Options options)
       // sure client IP tags were populated.
       copyAnomalyDetectionTags(span, localRootSpan);
     }
-    try (final AgentScope scope = tracer.activateSpan(span)) {
+    try (final ContextScope scope = tracer.activateSpan(span)) {
       final Message last = messages.get(messages.size() - 1);
       if (isToolCall(last)) {
         span.setTag(TARGET_TAG, "tool");

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/httpurlconnection/HttpUrlState.java

@@ -5,8 +5,8 @@
 import static datadog.trace.bootstrap.instrumentation.httpurlconnection.HttpUrlConnectionDecorator.DECORATE;
 import static datadog.trace.bootstrap.instrumentation.httpurlconnection.HttpUrlConnectionDecorator.HTTP_URL_CONNECTION;
 
+import datadog.context.ContextScope;
 import datadog.trace.bootstrap.ContextStore;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import java.net.HttpURLConnection;
 
@@ -18,7 +18,7 @@ public class HttpUrlState {
 
   public AgentSpan start(final HttpURLConnection connection) {
     span = startSpan(HTTP_URL_CONNECTION.toString(), DECORATE.operationName());
-    try (final AgentScope scope = activateSpan(span)) {
+    try (final ContextScope scope = activateSpan(span)) {
       DECORATE.afterStart(span);
       DECORATE.onRequest(span, connection);
       return span;
@@ -39,7 +39,7 @@ public void finish() {
 
   public void finishSpan(
       final HttpURLConnection connection, final int responseCode, final Throwable throwable) {
-    try (final AgentScope scope = activateSpan(span)) {
+    try (final ContextScope scope = activateSpan(span)) {
       if (responseCode > 0) {
         // safe to access response data as 'responseCode' is set
         DECORATE.onResponse(span, connection);
@@ -62,7 +62,7 @@ public void finishSpan(final HttpURLConnection connection, final int responseCod
      * (e.g. breaks getOutputStream).
      */
     if (responseCode > 0) {
-      try (final AgentScope scope = activateSpan(span)) {
+      try (final ContextScope scope = activateSpan(span)) {
         // safe to access response data as 'responseCode' is set
         DECORATE.onResponse(span, connection);
         DECORATE.beforeFinish(span);

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/Wrapper.java

@@ -5,6 +5,7 @@
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.exclude;
 
+import datadog.context.ContextScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import java.util.concurrent.RunnableFuture;
 
@@ -43,7 +44,7 @@ public Wrapper(T delegate, AgentScope.Continuation continuation) {
 
   @Override
   public void run() {
-    try (AgentScope scope = activate()) {
+    try (ContextScope scope = activate()) {
       delegate.run();
     }
   }

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/probe/LogProbeTest.java

@@ -16,14 +16,14 @@
 import com.datadog.debugger.sink.DebuggerSink;
 import com.datadog.debugger.sink.ProbeStatusSink;
 import com.datadog.debugger.sink.Snapshot;
+import datadog.context.ContextScope;
 import datadog.trace.api.Config;
 import datadog.trace.api.IdGenerationStrategy;
 import datadog.trace.bootstrap.debugger.CapturedContext;
 import datadog.trace.bootstrap.debugger.EvaluationError;
 import datadog.trace.bootstrap.debugger.MethodLocation;
 import datadog.trace.bootstrap.debugger.ProbeId;
 import datadog.trace.bootstrap.debugger.ProbeRateLimiter;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer.TracerAPI;
@@ -136,7 +136,7 @@ private int runTrace(TracerAPI tracer, boolean captureSnapshot, Integer line, St
     if (sessionId != null) {
       span.setTag(Tags.PROPAGATED_DEBUG, sessionId + ":1");
     }
-    try (AgentScope scope = tracer.activateManualSpan(span)) {
+    try (ContextScope scope = tracer.activateManualSpan(span)) {
       Builder builder =
           createLog("Budget testing").probeId(ProbeId.newId()).captureSnapshot(captureSnapshot);
       if (sessionId != null) {
@@ -177,7 +177,7 @@ private boolean fillSnapshot(DebugSessionStatus status) {
         CoreTracer.builder().idGenerationStrategy(IdGenerationStrategy.fromName("random")).build();
     AgentTracer.registerIfAbsent(tracer);
     AgentSpan span = tracer.startSpan("log probe debug session testing", "test span");
-    try (AgentScope scope = tracer.activateManualSpan(span)) {
+    try (ContextScope scope = tracer.activateManualSpan(span)) {
       if (status == DebugSessionStatus.ACTIVE) {
         span.setTag(Tags.PROPAGATED_DEBUG, DEBUG_SESSION_ID + ":1");
       } else if (status == DebugSessionStatus.DISABLED) {

dd-java-agent/agent-debugger/src/test/resources/com/datadog/debugger/CapturedSnapshot20.java

@@ -1,7 +1,7 @@
 package com.datadog.debugger;
 
+import datadog.context.ContextScope;
 import datadog.trace.agent.tooling.TracerInstaller;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.core.CoreTracer;
@@ -25,7 +25,7 @@ public class CapturedSnapshot20 {
   public static int main(String arg) {
     AgentTracer.TracerAPI tracerAPI = AgentTracer.get();
     AgentSpan span = tracerAPI.buildSpan("dynamic-instrumentation", "process").start();
-    try (AgentScope scope = tracerAPI.activateManualSpan(span)) {
+    try (ContextScope scope = tracerAPI.activateManualSpan(span)) {
       if (arg.equals("exception") || arg.equals("illegal")) {
         return new CapturedSnapshot20().processWithException(arg);
       }

dd-java-agent/agent-debugger/src/test/resources/com/datadog/debugger/CapturedSnapshot21.java

@@ -1,7 +1,7 @@
 package com.datadog.debugger;
 
+import datadog.context.ContextScope;
 import datadog.trace.agent.tooling.TracerInstaller;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.core.CoreTracer;
@@ -25,7 +25,7 @@ public class CapturedSnapshot21 {
   public static int main(String arg) {
     AgentTracer.TracerAPI tracerAPI = AgentTracer.get();
     AgentSpan span = tracerAPI.buildSpan("dynamic-instrumentation", "rootProcess").start();
-    try (AgentScope scope = tracerAPI.activateManualSpan(span)) {
+    try (ContextScope scope = tracerAPI.activateManualSpan(span)) {
       return new CapturedSnapshot21().rootProcess(arg);
     } finally {
       span.finish();
@@ -35,7 +35,7 @@ public static int main(String arg) {
   private int rootProcess(String arg) {
     AgentTracer.TracerAPI tracerAPI = AgentTracer.get();
     AgentSpan span = tracerAPI.buildSpan("dynamic-instrumentation", "process1").start();
-    try (AgentScope scope = tracerAPI.activateManualSpan(span)) {
+    try (ContextScope scope = tracerAPI.activateManualSpan(span)) {
       return process1(arg) + 1;
     } finally {
       span.finish();
@@ -45,7 +45,7 @@ private int rootProcess(String arg) {
   private int process1(String arg) {
     AgentTracer.TracerAPI tracerAPI = AgentTracer.get();
     AgentSpan span = tracerAPI.buildSpan("dynamic-instrumentation", "process2").start();
-    try (AgentScope scope = tracerAPI.activateManualSpan(span)) {
+    try (ContextScope scope = tracerAPI.activateManualSpan(span)) {
       return process2(arg) + 1;
     } finally {
       span.finish();
@@ -55,7 +55,7 @@ private int process1(String arg) {
   private int process2(String arg) {
     AgentTracer.TracerAPI tracerAPI = AgentTracer.get();
     AgentSpan span = tracerAPI.buildSpan("dynamic-instrumentation", "process3").start();
-    try (AgentScope scope = tracerAPI.activateManualSpan(span)) {
+    try (ContextScope scope = tracerAPI.activateManualSpan(span)) {
       return process3(arg) + 1;
     } finally {
       span.finish();

dd-java-agent/agent-debugger/src/test/resources/com/datadog/debugger/CapturedSnapshot28.java

@@ -1,7 +1,7 @@
 package com.datadog.debugger;
 
+import datadog.context.ContextScope;
 import datadog.trace.agent.tooling.TracerInstaller;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.core.CoreTracer;
@@ -26,7 +26,7 @@ public class CapturedSnapshot28 {
   public static int main(String arg) {
     AgentTracer.TracerAPI tracerAPI = AgentTracer.get();
     AgentSpan span = tracerAPI.buildSpan("dynamic-instrumentation", "process").start();
-    try (AgentScope scope = tracerAPI.activateManualSpan(span)) {
+    try (ContextScope scope = tracerAPI.activateManualSpan(span)) {
       return new CapturedSnapshot28().process(arg);
     } finally {
       span.finish();

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/Reporter.java

@@ -7,6 +7,7 @@
 import com.datadog.iast.model.Vulnerability;
 import com.datadog.iast.model.VulnerabilityBatch;
 import com.datadog.iast.taint.TaintedObjects;
+import datadog.context.ContextScope;
 import datadog.trace.api.Config;
 import datadog.trace.api.ProductTraceSource;
 import datadog.trace.api.gateway.RequestContext;
@@ -61,7 +62,7 @@ public void report(@Nullable final AgentSpan span, @Nonnull final Vulnerability
     }
     if (span == null) {
       final AgentSpan newSpan = startNewSpan();
-      try (final AgentScope autoClosed = tracer().activateManualSpan(newSpan)) {
+      try (final ContextScope autoClosed = tracer().activateManualSpan(newSpan)) {
         vulnerability.updateSpan(newSpan);
         reportVulnerability(newSpan, vulnerability);
       } finally {

dd-java-agent/instrumentation/aerospike-4.0/src/main/java/datadog/trace/instrumentation/aerospike4/TracingListener.java

@@ -17,7 +17,7 @@
 import com.aerospike.client.listener.RecordListener;
 import com.aerospike.client.listener.RecordSequenceListener;
 import com.aerospike.client.listener.WriteListener;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.context.ContextScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope.Continuation;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import java.util.List;
@@ -73,7 +73,7 @@ public void onSuccess(final Key key, final boolean exists) {
     clientSpan.finish();
 
     if (listener != null) {
-      try (final AgentScope scope = continuation.activate()) {
+      try (final ContextScope scope = continuation.activate()) {
         if (listener instanceof ExistsListener) {
           ((ExistsListener) listener).onSuccess(key, exists);
         } else if (listener instanceof DeleteListener) {
@@ -91,7 +91,7 @@ public void onSuccess(final Key[] keys, final boolean[] exists) {
     clientSpan.finish();
 
     if (listener != null) {
-      try (final AgentScope scope = continuation.activate()) {
+      try (final ContextScope scope = continuation.activate()) {
         ((ExistsArrayListener) listener).onSuccess(keys, exists);
       }
     } else {
@@ -105,7 +105,7 @@ public void onSuccess(final Key key, final Record record) {
     clientSpan.finish();
 
     if (listener != null) {
-      try (final AgentScope scope = continuation.activate()) {
+      try (final ContextScope scope = continuation.activate()) {
         ((RecordListener) listener).onSuccess(key, record);
       }
     } else {
@@ -119,7 +119,7 @@ public void onSuccess(final Key[] keys, final Record[] records) {
     clientSpan.finish();
 
     if (listener != null) {
-      try (final AgentScope scope = continuation.activate()) {
+      try (final ContextScope scope = continuation.activate()) {
         ((RecordArrayListener) listener).onSuccess(keys, records);
       }
     } else {
@@ -133,7 +133,7 @@ public void onSuccess(final List<BatchRead> records) {
     clientSpan.finish();
 
     if (listener != null) {
-      try (final AgentScope scope = continuation.activate()) {
+      try (final ContextScope scope = continuation.activate()) {
         ((BatchListListener) listener).onSuccess(records);
       }
     } else {
@@ -147,7 +147,7 @@ public void onSuccess(final Key key) {
     clientSpan.finish();
 
     if (listener != null) {
-      try (final AgentScope scope = continuation.activate()) {
+      try (final ContextScope scope = continuation.activate()) {
         ((WriteListener) listener).onSuccess(key);
       }
     } else {
@@ -161,7 +161,7 @@ public void onSuccess(final Key key, final Object obj) {
     clientSpan.finish();
 
     if (listener != null) {
-      try (final AgentScope scope = continuation.activate()) {
+      try (final ContextScope scope = continuation.activate()) {
         ((ExecuteListener) listener).onSuccess(key, obj);
       }
     } else {
@@ -175,7 +175,7 @@ public void onSuccess() {
     clientSpan.finish();
 
     if (listener != null) {
-      try (final AgentScope scope = continuation.activate()) {
+      try (final ContextScope scope = continuation.activate()) {
         if (listener instanceof ExistsSequenceListener) {
           ((ExistsSequenceListener) listener).onSuccess();
         } else if (listener instanceof RecordSequenceListener) {
@@ -196,7 +196,7 @@ public void onFailure(final AerospikeException error) {
     clientSpan.finish();
 
     if (listener != null) {
-      try (final AgentScope scope = continuation.activate()) {
+      try (final ContextScope scope = continuation.activate()) {
         if (listener instanceof ExistsListener) {
           ((ExistsListener) listener).onFailure(error);
         } else if (listener instanceof ExistsSequenceListener) {

dd-java-agent/instrumentation/akka/akka-actor-2.5/src/test/java/LinearTask.java

@@ -2,7 +2,7 @@
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
 
 import akka.dispatch.forkjoin.RecursiveTask;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.context.ContextScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 
 public class LinearTask extends RecursiveTask<Integer> {
@@ -32,7 +32,7 @@ protected Integer compute() {
     } else {
       int next = parent + 1;
       AgentSpan span = startSpan("test", Integer.toString(next));
-      try (AgentScope scope = activateSpan(span)) {
+      try (ContextScope scope = activateSpan(span)) {
         LinearTask child = new LinearTask(next, depth);
         return child.fork().join();
       } finally {

dd-java-agent/instrumentation/apache-httpclient/apache-httpasyncclient-4.0/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/TraceContinuedFutureCallback.java

@@ -3,6 +3,7 @@
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopContinuation;
 import static datadog.trace.instrumentation.apachehttpasyncclient.ApacheHttpAsyncClientDecorator.DECORATE;
 
+import datadog.context.ContextScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import org.apache.http.concurrent.FutureCallback;
@@ -35,7 +36,7 @@ public void completed(final T result) {
     if (parentContinuation == noopContinuation()) {
       completeDelegate(result);
     } else {
-      try (final AgentScope scope = parentContinuation.activate()) {
+      try (final ContextScope scope = parentContinuation.activate()) {
         completeDelegate(result);
       }
     }
@@ -51,7 +52,7 @@ public void failed(final Exception ex) {
     if (parentContinuation == noopContinuation()) {
       failDelegate(ex);
     } else {
-      try (final AgentScope scope = parentContinuation.activate()) {
+      try (final ContextScope scope = parentContinuation.activate()) {
         failDelegate(ex);
       }
     }
@@ -66,7 +67,7 @@ public void cancelled() {
     if (parentContinuation == noopContinuation()) {
       cancelDelegate();
     } else {
-      try (final AgentScope scope = parentContinuation.activate()) {
+      try (final ContextScope scope = parentContinuation.activate()) {
         cancelDelegate();
       }
     }

dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/TraceContinuedFutureCallback.java

@@ -3,6 +3,7 @@
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopContinuation;
 import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.DECORATE;
 
+import datadog.context.ContextScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import javax.annotation.Nullable;
@@ -38,7 +39,7 @@ public void completed(final T result) {
     if (parentContinuation == noopContinuation()) {
       completeDelegate(result);
     } else {
-      try (final AgentScope scope = parentContinuation.activate()) {
+      try (final ContextScope scope = parentContinuation.activate()) {
         completeDelegate(result);
       }
     }
@@ -54,7 +55,7 @@ public void failed(final Exception ex) {
     if (parentContinuation == noopContinuation()) {
       failDelegate(ex);
     } else {
-      try (final AgentScope scope = parentContinuation.activate()) {
+      try (final ContextScope scope = parentContinuation.activate()) {
         failDelegate(ex);
       }
     }
@@ -69,7 +70,7 @@ public void cancelled() {
     if (parentContinuation == noopContinuation()) {
       cancelDelegate();
     } else {
-      try (final AgentScope scope = parentContinuation.activate()) {
+      try (final ContextScope scope = parentContinuation.activate()) {
         cancelDelegate();
       }
     }