Enhance log probes to honor debug session tags

dd-java-agent/agent-debugger/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/DebuggerContext.java

@@ -21,8 +21,26 @@ public class DebuggerContext {
   private static final ThreadLocal<Boolean> IN_PROBE = ThreadLocal.withInitial(() -> Boolean.FALSE);
 
   public enum SkipCause {
-    RATE,
-    CONDITION
+    RATE {
+      @Override
+      public String tag() {
+        return "cause:rate";
+      }
+    },
+    CONDITION {
+      @Override
+      public String tag() {
+        return "cause:condition";
+      }
+    },
+    DEBUG_SESSION_DISABLED {
+      @Override
+      public String tag() {
+        return "cause:debug session disabled";
+      }
+    };
+
+    public abstract String tag();
   }
 
   public interface ProbeResolver {

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/DebugSessionStatus.java

@@ -0,0 +1,11 @@
+package com.datadog.debugger.probe;
+
+public enum DebugSessionStatus {
+  NONE,
+  ACTIVE,
+  DISABLED;
+
+  public boolean isDisabled() {
+    return this == DISABLED;
+  }
+}

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/LogProbe.java

@@ -32,14 +32,21 @@
 import datadog.trace.bootstrap.debugger.ProbeImplementation;
 import datadog.trace.bootstrap.debugger.ProbeRateLimiter;
 import datadog.trace.bootstrap.debugger.util.TimeoutChecker;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer.TracerAPI;
+import datadog.trace.core.DDSpan;
+import datadog.trace.core.DDSpanContext;
 import java.io.IOException;
 import java.lang.reflect.ParameterizedType;
 import java.time.Duration;
 import java.time.temporal.ChronoUnit;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.function.Consumer;
 import org.slf4j.Logger;
@@ -51,6 +58,30 @@ public class LogProbe extends ProbeDefinition implements Sampled {
   private static final Limits LIMITS = new Limits(1, 3, 8192, 5);
   private static final int LOG_MSG_LIMIT = 8192;
 
+  private static Map<String, String> getDebugSessions() {
+    HashMap<String, String> sessions = new HashMap<>();
+    TracerAPI tracer = AgentTracer.get();
+    if (tracer != null) {
+      AgentSpan span = tracer.activeSpan();
+      if (span instanceof DDSpan) {
+        DDSpanContext context = (DDSpanContext) span.context();
+        String debug = context.getPropagationTags().getDebugPropagation();
+        if (debug != null) {
+          String[] entries = debug.split(",");
+          for (String entry : entries) {
+            if (!entry.contains(":")) {
+              sessions.put("*", entry);
+            } else {
+              String[] values = entry.split(":");
+              sessions.put(values[0], values[1]);
+            }
+          }
+        }
+      }
+    }
+    return sessions;
+  }
+
   /** Stores part of a templated message either a str or an expression */
   public static class Segment {
     private final String str;
@@ -503,10 +534,16 @@ private void sample(LogStatus logStatus, MethodLocation methodLocation) {
     if (!MethodLocation.isSame(methodLocation, evaluateAt)) {
       return;
     }
-    boolean sampled = ProbeRateLimiter.tryProbe(id);
+    boolean sampled =
+        !logStatus.getDebugSessionStatus().isDisabled() && ProbeRateLimiter.tryProbe(id);
     logStatus.setSampled(sampled);
     if (!sampled) {
-      DebuggerAgent.getSink().skipSnapshot(id, DebuggerContext.SkipCause.RATE);
+      DebuggerAgent.getSink()
+          .skipSnapshot(
+              id,
+              logStatus.getDebugSessionStatus().isDisabled()
+                  ? DebuggerContext.SkipCause.DEBUG_SESSION_DISABLED
+                  : DebuggerContext.SkipCause.RATE);
     }
   }
 
@@ -673,6 +710,10 @@ public boolean hasCondition() {
     return probeCondition != null;
   }
 
+  protected String getDebugSessionId() {
+    return getTagMap().get("session_id");
+  }
+
   @Override
   public CapturedContext.Status createStatus() {
     return new LogStatus(this);
@@ -685,6 +726,7 @@ public static class LogStatus extends CapturedContext.Status {
         new LogStatus(ProbeImplementation.UNKNOWN, true);
 
     private boolean condition = true;
+    private final DebugSessionStatus debugSessionStatus;
     private boolean hasLogTemplateErrors;
     private boolean hasConditionErrors;
     private boolean sampled = true;
@@ -693,10 +735,11 @@ public static class LogStatus extends CapturedContext.Status {
 
     public LogStatus(ProbeImplementation probeImplementation) {
       super(probeImplementation);
+      debugSessionStatus = debugSessionStatus();
     }
 
     private LogStatus(ProbeImplementation probeImplementation, boolean condition) {
-      super(probeImplementation);
+      this(probeImplementation);
       this.condition = condition;
     }
 
@@ -711,7 +754,11 @@ public boolean isCapturing() {
     }
 
     public boolean shouldSend() {
-      return sampled && condition && !hasConditionErrors;
+      DebugSessionStatus status = getDebugSessionStatus();
+      // an ACTIVE status overrides the sampling as the sampling decision was made by the trigger
+      // probe
+      return status == DebugSessionStatus.ACTIVE
+          || !status.isDisabled() && sampled && condition && !hasConditionErrors;
     }
 
     public boolean shouldReportError() {
@@ -722,6 +769,26 @@ public boolean getCondition() {
       return condition;
     }
 
+    public DebugSessionStatus getDebugSessionStatus() {
+      return debugSessionStatus;
+    }
+
+    private DebugSessionStatus debugSessionStatus() {
+      if (probeImplementation instanceof LogProbe) {
+        LogProbe definition = (LogProbe) probeImplementation;
+        Map<String, String> sessions = getDebugSessions();
+        String sessionId = definition.getDebugSessionId();
+        if (sessionId == null) {
+          return DebugSessionStatus.NONE;
+        }
+        return "1".equals(sessions.get(sessionId)) || "1".equals(sessions.get("*"))
+            ? DebugSessionStatus.ACTIVE
+            : DebugSessionStatus.DISABLED;
+      }
+
+      return DebugSessionStatus.NONE;
+    }
+
     public void setCondition(boolean value) {
       this.condition = value;
     }
@@ -765,6 +832,29 @@ public boolean isForceSampling() {
     public void setForceSampling(boolean forceSampling) {
       this.forceSampling = forceSampling;
     }
+
+    @Override
+    public String toString() {
+      return "LogStatus{"
+          + ", probeId="
+          + probeImplementation.getId()
+          + ", condition="
+          + condition
+          + ", debugSessionStatus="
+          + debugSessionStatus
+          + ", forceSampling="
+          + forceSampling
+          + ", hasConditionErrors="
+          + hasConditionErrors
+          + ", hasLogTemplateErrors="
+          + hasLogTemplateErrors
+          + ", message='"
+          + message
+          + '\''
+          + ", sampled="
+          + sampled
+          + '}';
+    }
   }
 
   @Generated

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/ProbeDefinition.java

@@ -93,6 +93,9 @@ public String concatTags() {
   }
 
   public Map<String, String> getTagMap() {
+    if (tagMap.isEmpty() && tags != null) {
+      initTagMap(tagMap, tags);
+    }
     return Collections.unmodifiableMap(tagMap);
   }
 

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/TriggerProbe.java

@@ -59,6 +59,15 @@ public InstrumentationResult.Status instrument(
         .instrument();
   }
 
+  public String getSessionId() {
+    return sessionId;
+  }
+
+  public TriggerProbe setSessionId(String sessionId) {
+    this.sessionId = sessionId;
+    return this;
+  }
+
   public Sampling getSampling() {
     return sampling;
   }
@@ -110,7 +119,7 @@ private void decorateTags() {
     AgentTracer.TracerAPI tracerAPI = AgentTracer.get();
 
     AgentSpan agentSpan = tracerAPI.activeSpan().getLocalRootSpan();
-    agentSpan.setTag(Tags.PROPAGATED_DEBUG, "1");
+    agentSpan.setTag(Tags.PROPAGATED_DEBUG, sessionId + ":1");
     agentSpan.setTag(format("_dd.ld.probe_id.%s", probeId.getId()), true);
   }
 
@@ -148,8 +157,20 @@ public int hashCode() {
 
   @Override
   public String toString() {
-    return format(
-        "TriggerProbe{id='%s', where=%s, sampling=%s, probeCondition=%s}",
-        id, where, sampling, probeCondition);
+    return String.format(
+        "TriggerProbe{id='%s', sessionId='%s', evaluateAt=%s, language='%s', location=%s, probeCondition=%s, probeId=%s,"
+            + " sampling=%s, tagMap=%s, tags=%s, version=%d, where=%s}",
+        id,
+        sessionId,
+        evaluateAt,
+        language,
+        location,
+        probeCondition,
+        probeId,
+        sampling,
+        tagMap,
+        Arrays.toString(tags),
+        version,
+        where);
   }
 }

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/sink/DebuggerSink.java

@@ -4,7 +4,7 @@
 import com.datadog.debugger.uploader.BatchUploader;
 import com.datadog.debugger.util.DebuggerMetrics;
 import datadog.trace.api.Config;
-import datadog.trace.bootstrap.debugger.DebuggerContext;
+import datadog.trace.bootstrap.debugger.DebuggerContext.SkipCause;
 import datadog.trace.bootstrap.debugger.ProbeId;
 import datadog.trace.util.AgentTaskScheduler;
 import java.util.List;
@@ -223,20 +223,8 @@ private void reportError(ProbeId probeId, DiagnosticMessage msg) {
   }
 
   /** Notifies the snapshot was skipped for one of the SkipCause reason */
-  public void skipSnapshot(String probeId, DebuggerContext.SkipCause cause) {
-    String causeTag;
-    switch (cause) {
-      case RATE:
-        causeTag = "cause:rate";
-        break;
-      case CONDITION:
-        causeTag = "cause:condition";
-        break;
-      default:
-        throw new IllegalArgumentException("Unknown cause: " + cause);
-    }
-    String probeIdTag = "probe_id:" + probeId;
-    debuggerMetrics.incrementCounter(PREFIX + "skip", causeTag, probeIdTag);
+  public void skipSnapshot(String probeId, SkipCause cause) {
+    debuggerMetrics.incrementCounter(PREFIX + "skip", cause.tag(), "probe_id:" + probeId);
   }
 
   long getCurrentLowRateFlushInterval() {

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturingTestBase.java

@@ -352,14 +352,7 @@ public static LogProbe.Builder createProbeBuilder(
   protected TestSnapshotListener installProbes(
       Configuration configuration, ProbeDefinition... probes) {
 
-    config = mock(Config.class);
-    when(config.isDebuggerEnabled()).thenReturn(true);
-    when(config.isDebuggerClassFileDumpEnabled()).thenReturn(true);
-    when(config.isDebuggerVerifyByteCode()).thenReturn(false);
-    when(config.getFinalDebuggerSnapshotUrl())
-        .thenReturn("http://localhost:8126/debugger/v1/input");
-    when(config.getFinalDebuggerSymDBUrl()).thenReturn("http://localhost:8126/symdb/v1/input");
-    when(config.getDebuggerCodeOriginMaxUserFrames()).thenReturn(20);
+    config = mockConfig();
     instrumentationListener = new MockInstrumentationListener();
     probeStatusSink = mock(ProbeStatusSink.class);
 
@@ -404,6 +397,19 @@ protected TestSnapshotListener installProbes(
     return listener;
   }
 
+  public static Config mockConfig() {
+    Config config = mock(Config.class);
+    when(config.isDebuggerEnabled()).thenReturn(true);
+    when(config.isDebuggerClassFileDumpEnabled()).thenReturn(true);
+    when(config.isDebuggerVerifyByteCode()).thenReturn(false);
+    when(config.getFinalDebuggerSnapshotUrl())
+        .thenReturn("http://localhost:8126/debugger/v1/input");
+    when(config.getFinalDebuggerSymDBUrl()).thenReturn("http://localhost:8126/symdb/v1/input");
+    when(config.getDebuggerCodeOriginMaxUserFrames()).thenReturn(20);
+
+    return config;
+  }
+
   public static ProbeImplementation resolver(
       String encodedId,
       Collection<LogProbe> logProbes,