oh!

Signed-off-by: sezen.leblay <sezen.leblay@datadoghq.com>

Author: sezen-datadog

dd-java-agent/appsec/src/main/java/com/datadog/appsec/ddwaf/WAFModule.java

@@ -404,8 +404,8 @@ public void onDataAvailable(
               // Keep event related span, because it could be ignored in case of
               // reduced datadog sampling rate.
               activeSpan.getLocalRootSpan().setTag(DDTags.MANUAL_KEEP, true);
-            } else if (resultWithData.events) {
-              activeSpan.getLocalRootSpan().setTag(Tags.ASM_KEEP, true);
+            } else if (resultWithData.events) { // sampler still might keep
+              activeSpan.getLocalRootSpan().setTag(DDTags.ASM_SAMPLER_KEEP, true);
             }
             activeSpan
                 .getLocalRootSpan()

dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java

@@ -754,7 +754,7 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) {
           traceSeg.setTagTop(DDTags.MANUAL_KEEP, true);
         } else {
           // Set asm keep in case that root span was not available when events are detected
-          traceSeg.setTagTop(Tags.ASM_KEEP, true);
+          traceSeg.setTagTop(DDTags.ASM_SAMPLER_KEEP, true);
         }
         traceSeg.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM);
         traceSeg.setTagTop("appsec.event", true);

dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/TraceTaggingSmokeTest.groovy

@@ -330,7 +330,8 @@ class TraceTaggingSmokeTest extends AbstractAppSecServerSmokeTest {
     assert appsecJson.triggers != null, "Missing triggers in WAF attack event"
 
     // Should NOT have USER_KEEP sampling priority since keep: false
-    assert rootSpan.metrics.get('_sampling_priority_v1') < 2, "Should not have USER_KEEP sampling priority when keep: false"
+    assert rootSpan.metrics.get('_sampling_priority_v1') < 2,
+    "Should not have USER_KEEP sampling priority when keep: false"
 
   }
 

dd-trace-api/src/main/java/datadog/trace/api/DDTags.java

@@ -59,6 +59,9 @@ public class DDTags {
   /** Manually force tracer to drop the trace */
   public static final String MANUAL_DROP = "manual.drop";
 
+  /** ASM force tracer to keep the trace with sampler priority (SAMPLER_KEEP) */
+  public static final String ASM_SAMPLER_KEEP = "asm.sampler.keep";
+
   public static final String TRACE_START_TIME = "t0";
 
   /* Tags below are for internal use only. */

dd-trace-core/src/main/java/datadog/trace/common/sampling/AsmStandaloneSampler.java

@@ -41,7 +41,7 @@ public <T extends CoreSpan<T>> boolean sample(final T span) {
   public <T extends CoreSpan<T>> void setSamplingPriority(final T span) {
 
     if (shouldSample()) {
-      log.debug("Set SAMPLER_KEEP for span {}", span.getSpanId());
+      log.debug("Set USER_KEEP for span {}", span.getSpanId());
       span.setSamplingPriority(SAMPLER_KEEP, SamplingMechanism.APPSEC);
     } else {
       log.debug("Set SAMPLER_DROP for span {}", span.getSpanId());

dd-trace-core/src/main/java/datadog/trace/core/taginterceptor/TagInterceptor.java

@@ -4,6 +4,7 @@
 import static datadog.trace.api.DDTags.MEASURED;
 import static datadog.trace.api.DDTags.ORIGIN_KEY;
 import static datadog.trace.api.DDTags.SPAN_TYPE;
+import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_KEEP;
 import static datadog.trace.api.sampling.PrioritySampling.USER_DROP;
 import static datadog.trace.api.sampling.PrioritySampling.USER_KEEP;
 import static datadog.trace.bootstrap.instrumentation.api.Tags.HTTP_METHOD;
@@ -108,6 +109,7 @@ public boolean needsIntercept(String tag) {
       case DDTags.MANUAL_KEEP:
       case DDTags.MANUAL_DROP:
       case Tags.ASM_KEEP:
+      case DDTags.ASM_SAMPLER_KEEP:
       case Tags.SAMPLING_PRIORITY:
       case Tags.PROPAGATED_TRACE_SOURCE:
       case Tags.PROPAGATED_DEBUG:
@@ -155,6 +157,12 @@ public boolean interceptTag(DDSpanContext span, String tag, Object value) {
           return true;
         }
         return false;
+      case DDTags.ASM_SAMPLER_KEEP:
+        if (asBoolean(value)) {
+          span.setSamplingPriority(SAMPLER_KEEP, SamplingMechanism.APPSEC);
+          return true;
+        }
+        return false;
       case Tags.SAMPLING_PRIORITY:
         return interceptSamplingPriority(span, value);
       case Tags.PROPAGATED_TRACE_SOURCE: