Update sampleRate in DSC (#4158)

* wip

* wip2

* wip3

* format + api

* cleanup

* revert change to demo

* make baggage final again

* remove ObjectToString suppressing

* remove outdated comment

* remove getPropagationContext from Scopes again

* remove noop baggage and propagation context again

* fix outbox sender; test; cleanup api file

* Update sampleRate in DSC

* wip

* wip2

* wip3

* format + api

* cleanup

* revert change to demo

* make baggage final again

* remove ObjectToString suppressing

* remove outdated comment

* remove getPropagationContext from Scopes again

* remove noop baggage and propagation context again

* fix outbox sender; test; cleanup api file

* review changes

* Format code

* changelog

* fix build

---------

Co-authored-by: Sentry Github Bot <bot+github-bot@sentry.io>

Author: adinauer

CHANGELOG.md

@@ -15,6 +15,8 @@
 - (Internal) Add API to filter native debug images based on stacktrace addresses ([#4089](https://github.com/getsentry/sentry-java/pull/4089))
 - Propagate sampling random value ([#4153](https://github.com/getsentry/sentry-java/pull/4153))
   - The random value used for sampling traces is now sent to Sentry and attached to the `baggage` header on outgoing requests
+- Update `sampleRate` that is sent to Sentry and attached to the `baggage` header on outgoing requests ([#4158](https://github.com/getsentry/sentry-java/pull/4158))
+  - If the SDK uses its `sampleRate` or `tracesSampler` callback, it now updates the `sampleRate` in Dynamic Sampling Context.
 
 ### Fixes
 

sentry-opentelemetry/sentry-opentelemetry-core/src/main/java/io/sentry/opentelemetry/OtelSentrySpanProcessor.java

@@ -83,6 +83,7 @@ public void onStart(final @NotNull Context parentContext, final @NotNull ReadWri
               new SentryId(traceId), sentrySpanId, sentryParentSpanId, baggage, sampled);
 
       baggage = propagationContext.getBaggage();
+      baggage.setValuesFromSamplingDecision(samplingDecision);
 
       updatePropagationContext(scopes, propagationContext);
     }

sentry/api/sentry.api

@@ -33,6 +33,7 @@ public final class io/sentry/Baggage {
 	public fun <init> (Lio/sentry/Baggage;)V
 	public fun <init> (Lio/sentry/ILogger;)V
 	public fun <init> (Ljava/util/Map;Ljava/lang/String;ZZLio/sentry/ILogger;)V
+	public fun forceSetSampleRate (Ljava/lang/String;)V
 	public fun freeze ()V
 	public static fun fromEvent (Lio/sentry/SentryEvent;Lio/sentry/SentryOptions;)Lio/sentry/Baggage;
 	public static fun fromHeader (Ljava/lang/String;)Lio/sentry/Baggage;
@@ -70,6 +71,7 @@ public final class io/sentry/Baggage {
 	public fun setTraceId (Ljava/lang/String;)V
 	public fun setTransaction (Ljava/lang/String;)V
 	public fun setUserId (Ljava/lang/String;)V
+	public fun setValuesFromSamplingDecision (Lio/sentry/TracesSamplingDecision;)V
 	public fun setValuesFromScope (Lio/sentry/IScope;Lio/sentry/SentryOptions;)V
 	public fun setValuesFromTransaction (Lio/sentry/protocol/SentryId;Lio/sentry/protocol/SentryId;Lio/sentry/SentryOptions;Lio/sentry/TracesSamplingDecision;Ljava/lang/String;Lio/sentry/protocol/TransactionNameSource;)V
 	public fun toHeaderString (Ljava/lang/String;)Ljava/lang/String;

sentry/src/main/java/io/sentry/Baggage.java

@@ -367,6 +367,11 @@ public void setSampleRate(final @Nullable String sampleRate) {
     set(DSCKeys.SAMPLE_RATE, sampleRate);
   }
 
+  @ApiStatus.Internal
+  public void forceSetSampleRate(final @Nullable String sampleRate) {
+    set(DSCKeys.SAMPLE_RATE, sampleRate, true);
+  }
+
   @ApiStatus.Internal
   public @Nullable String getSampleRand() {
     return get(DSCKeys.SAMPLE_RAND);
@@ -399,7 +404,18 @@ public void setReplayId(final @Nullable String replayId) {
 
   @ApiStatus.Internal
   public void set(final @NotNull String key, final @Nullable String value) {
-    if (mutable) {
+    set(key, value, false);
+  }
+
+  /**
+   * Sets / updates a value
+   *
+   * @param key key
+   * @param value value to set
+   * @param force ignores mutability of this baggage and sets the value anyways
+   */
+  private void set(final @NotNull String key, final @Nullable String value, final boolean force) {
+    if (mutable || force) {
       this.keyValues.put(key, value);
     }
   }
@@ -439,6 +455,25 @@ public void setValuesFromTransaction(
     }
     setSampleRate(sampleRateToString(sampleRate(samplingDecision)));
     setSampled(StringUtils.toString(sampled(samplingDecision)));
+    setSampleRand(sampleRateToString(sampleRand(samplingDecision))); // TODO check
+  }
+
+  @ApiStatus.Internal
+  public void setValuesFromSamplingDecision(
+      final @Nullable TracesSamplingDecision samplingDecision) {
+    if (samplingDecision == null) {
+      return;
+    }
+
+    setSampled(StringUtils.toString(sampled(samplingDecision)));
+
+    if (samplingDecision.getSampleRand() != null) {
+      setSampleRand(sampleRateToString(sampleRand(samplingDecision)));
+    }
+
+    if (samplingDecision.getSampleRate() != null) {
+      forceSetSampleRate(sampleRateToString(sampleRate(samplingDecision)));
+    }
   }
 
   @ApiStatus.Internal
@@ -466,6 +501,14 @@ public void setValuesFromScope(
     return samplingDecision.getSampleRate();
   }
 
+  private static @Nullable Double sampleRand(@Nullable TracesSamplingDecision samplingDecision) {
+    if (samplingDecision == null) {
+      return null;
+    }
+
+    return samplingDecision.getSampleRand();
+  }
+
   private static @Nullable String sampleRateToString(@Nullable Double sampleRateAsDouble) {
     if (!SampleRateUtils.isValidTracesSampleRate(sampleRateAsDouble, false)) {
       return null;

sentry/src/main/java/io/sentry/SpanContext.java

@@ -91,10 +91,10 @@ public SpanContext(
     this.spanId = Objects.requireNonNull(spanId, "spanId is required");
     this.op = Objects.requireNonNull(operation, "operation is required");
     this.parentSpanId = parentSpanId;
-    this.samplingDecision = samplingDecision;
     this.description = description;
     this.status = status;
     this.origin = origin;
+    setSamplingDecision(samplingDecision);
   }
 
   /**
@@ -106,7 +106,7 @@ public SpanContext(final @NotNull SpanContext spanContext) {
     this.traceId = spanContext.traceId;
     this.spanId = spanContext.spanId;
     this.parentSpanId = spanContext.parentSpanId;
-    this.samplingDecision = spanContext.samplingDecision;
+    setSamplingDecision(spanContext.samplingDecision);
     this.op = spanContext.op;
     this.description = spanContext.description;
     this.status = spanContext.status;
@@ -209,6 +209,9 @@ public void setSampled(final @Nullable Boolean sampled, final @Nullable Boolean
   @ApiStatus.Internal
   public void setSamplingDecision(final @Nullable TracesSamplingDecision samplingDecision) {
     this.samplingDecision = samplingDecision;
+    if (this.baggage != null) {
+      this.baggage.setValuesFromSamplingDecision(this.samplingDecision);
+    }
   }
 
   public @Nullable String getOrigin() {

sentry/src/test/java/io/sentry/BaggageTest.kt

@@ -355,6 +355,18 @@ class BaggageTest {
         assertEquals("sentry-trace_id=a,sentry-transaction=sentryTransaction", baggage.toHeaderString(null))
     }
 
+    @Test
+    fun `if header contains sentry values baggage is marked as shouldFreeze`() {
+        val baggage = Baggage.fromHeader("sentry-trace_id=a,sentry-transaction=sentryTransaction", logger)
+        assertTrue(baggage.isShouldFreeze)
+    }
+
+    @Test
+    fun `if header does not contain sentry values baggage is not marked as shouldFreeze`() {
+        val baggage = Baggage.fromHeader("a=b", logger)
+        assertFalse(baggage.isShouldFreeze)
+    }
+
     @Test
     fun `value may contain = sign`() {
         val baggage = Baggage(logger)
@@ -560,6 +572,44 @@ class BaggageTest {
     }
 
     @Test
+    fun `sets values from traces sampling decision`() {
+        val baggage = Baggage.fromHeader("a=b,c=d")
+        baggage.setValuesFromSamplingDecision(TracesSamplingDecision(true, 0.021, 0.025))
+
+        assertEquals("true", baggage.sampled)
+        assertEquals("0.021", baggage.sampleRate)
+        assertEquals("0.025", baggage.sampleRand)
+    }
+
+    @Test
+    fun `handles null traces sampling decision`() {
+        val baggage = Baggage.fromHeader("a=b,c=d")
+        baggage.setValuesFromSamplingDecision(null)
+    }
+
+    @Test
+    fun `sets values from traces sampling decision only if non null`() {
+        val baggage = Baggage.fromHeader("a=b,c=d")
+        baggage.setValuesFromSamplingDecision(TracesSamplingDecision(true, 0.021, 0.025))
+        baggage.setValuesFromSamplingDecision(TracesSamplingDecision(false, null, null))
+
+        assertEquals("false", baggage.sampled)
+        assertEquals("0.021", baggage.sampleRate)
+        assertEquals("0.025", baggage.sampleRand)
+    }
+
+    @Test
+    fun `replaces only sample rate if already frozen`() {
+        val baggage = Baggage.fromHeader("a=b,c=d")
+        baggage.setValuesFromSamplingDecision(TracesSamplingDecision(true, 0.021, 0.025))
+        baggage.freeze()
+        baggage.setValuesFromSamplingDecision(TracesSamplingDecision(false, 0.121, 0.125))
+
+        assertEquals("true", baggage.sampled)
+        assertEquals("0.121", baggage.sampleRate)
+        assertEquals("0.025", baggage.sampleRand)
+    }
+
     fun `sample rate can be retrieved as double`() {
         val baggage = Baggage.fromHeader("a=b,c=d")
         baggage.sampleRate = "0.1"

sentry/src/test/java/io/sentry/SpanContextTest.kt

@@ -19,4 +19,15 @@ class SpanContextTest {
         trace.setTag("tagName", "tagValue")
         assertEquals("tagValue", trace.tags["tagName"])
     }
+
+    @Test
+    fun `updates sampling decision on baggage`() {
+        val trace = SpanContext("op")
+        trace.baggage = Baggage.fromHeader("a=b")
+        trace.samplingDecision = TracesSamplingDecision(true, 0.1, 0.2)
+
+        assertEquals("true", trace.baggage?.sampled)
+        assertEquals("0.1", trace.baggage?.sampleRate)
+        assertEquals("0.2", trace.baggage?.sampleRand)
+    }
 }