Fix sampling rate recorded for dependencies (#1582)

agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/instrumentation/sdk/BytecodeUtilImpl.java

@@ -27,6 +27,7 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import com.google.common.base.Strings;
+import com.microsoft.applicationinsights.agent.Exporter;
 import com.microsoft.applicationinsights.agent.bootstrap.BytecodeUtil.BytecodeUtilDelegate;
 import com.microsoft.applicationinsights.agent.internal.Global;
 import com.microsoft.applicationinsights.agent.internal.sampling.SamplingScoreGeneratorV2;
@@ -220,20 +221,31 @@ public void logErrorOnce(Throwable t) {
 
     private static void track(Telemetry telemetry) {
         SpanContext context = Span.current().getSpanContext();
+        double samplingPercentage;
         if (context.isValid()) {
-            String traceId = context.getTraceId();
-            String spanId = context.getSpanId();
-            telemetry.getContext().getOperation().setId(traceId);
-            telemetry.getContext().getOperation().setParentId(spanId);
-        }
-        double samplingPercentage = Global.getSamplingPercentage();
-        if (sample(telemetry, samplingPercentage)) {
-            if (telemetry instanceof SupportSampling && samplingPercentage != 100) {
-                ((SupportSampling) telemetry).setSamplingPercentage(samplingPercentage);
+            if (!context.isSampled()) {
+                // sampled out
+                return;
             }
-            // this is not null because sdk instrumentation is not added until Global.setTelemetryClient() is called
-            checkNotNull(Global.getTelemetryClient()).track(telemetry);
+            telemetry.getContext().getOperation().setId(context.getTraceId());
+            telemetry.getContext().getOperation().setParentId(context.getSpanId());
+            samplingPercentage =
+                    Exporter.getSamplingPercentage(context.getTraceState(), Global.getSamplingPercentage(), false);
+        } else {
+            // sampling is done using the configured sampling percentage
+            samplingPercentage = Global.getSamplingPercentage();
+            if (!sample(telemetry, samplingPercentage)) {
+                // sampled out
+                return;
+            }
+        }
+        // sampled in
+
+        if (telemetry instanceof SupportSampling && samplingPercentage != 100) {
+            ((SupportSampling) telemetry).setSamplingPercentage(samplingPercentage);
         }
+        // this is not null because sdk instrumentation is not added until Global.setTelemetryClient() is called
+        checkNotNull(Global.getTelemetryClient()).track(telemetry);
     }
 
     private static boolean sample(Telemetry telemetry, double samplingPercentage) {

agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/propagator/DelegatingPropagator.java

@@ -3,6 +3,11 @@
 import java.util.Collection;
 import javax.annotation.Nullable;
 
+import com.microsoft.applicationinsights.agent.Exporter;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.api.trace.TraceFlags;
+import io.opentelemetry.api.trace.TraceState;
 import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.propagation.TextMapGetter;
@@ -21,8 +26,13 @@ public static DelegatingPropagator getInstance() {
     }
 
     public void setUpStandardDelegate() {
+
+        // TODO when should we enable baggage propagation?
+        // currently using modified W3CTraceContextPropagator because "ai-internal-sp" trace state
+        // shouldn't be sent over the wire (at least not yet, and not with that name)
+
         // important that W3CTraceContextPropagator is last, so it will take precedence if both sets of headers are present
-        delegate = TextMapPropagator.composite(AiLegacyPropagator.getInstance(), W3CTraceContextPropagator.getInstance());
+        delegate = TextMapPropagator.composite(AiLegacyPropagator.getInstance(), new ModifiedW3CTraceContextPropagator());
     }
 
     @Override
@@ -39,4 +49,74 @@ public <C> void inject(Context context, @Nullable C carrier, TextMapSetter<C> se
     public <C> Context extract(Context context, @Nullable C carrier, TextMapGetter<C> getter) {
         return delegate.extract(context, carrier, getter);
     }
+
+    private static class ModifiedW3CTraceContextPropagator implements TextMapPropagator {
+
+        private final TextMapPropagator delegate = W3CTraceContextPropagator.getInstance();
+
+        @Override
+        public Collection<String> fields() {
+            return delegate.fields();
+        }
+
+        @Override
+        public <C> void inject(Context context, @Nullable C carrier, TextMapSetter<C> setter) {
+            // do not propagate sampling percentage downstream YET
+            SpanContext spanContext = Span.fromContext(context).getSpanContext();
+            // sampling percentage should always be present, so no need to optimize with checking if present
+            TraceState traceState = spanContext.getTraceState();
+            TraceState updatedTraceState;
+            if (traceState.size() == 1 && traceState.get(Exporter.SAMPLING_PERCENTAGE_TRACE_STATE) != null) {
+                // this is a common case, worth optimizing
+                updatedTraceState = TraceState.getDefault();
+            } else {
+                updatedTraceState = traceState.toBuilder()
+                        .remove(Exporter.SAMPLING_PERCENTAGE_TRACE_STATE)
+                        .build();
+            }
+            SpanContext updatedSpanContext = new ModifiedSpanContext(spanContext, updatedTraceState);
+            delegate.inject(Context.root().with(Span.wrap(updatedSpanContext)), carrier, setter);
+        }
+
+        @Override
+        public <C> Context extract(Context context, @Nullable C carrier, TextMapGetter<C> getter) {
+            return delegate.extract(context, carrier, getter);
+        }
+    }
+
+    private static class ModifiedSpanContext implements SpanContext {
+
+        private final SpanContext delegate;
+        private final TraceState traceState;
+
+        private ModifiedSpanContext(SpanContext delegate, TraceState traceState) {
+            this.delegate = delegate;
+            this.traceState = traceState;
+        }
+
+        @Override
+        public String getTraceId() {
+            return delegate.getTraceId();
+        }
+
+        @Override
+        public String getSpanId() {
+            return delegate.getSpanId();
+        }
+
+        @Override
+        public TraceFlags getTraceFlags() {
+            return delegate.getTraceFlags();
+        }
+
+        @Override
+        public TraceState getTraceState() {
+            return traceState;
+        }
+
+        @Override
+        public boolean isRemote() {
+            return delegate.isRemote();
+        }
+    }
 }

agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/sampling/AiSampler.java

@@ -1,7 +1,6 @@
 package com.microsoft.applicationinsights.agent.internal.sampling;
 
 import java.util.List;
-import javax.annotation.Nullable;
 
 import com.microsoft.applicationinsights.agent.internal.sampling.SamplingOverrides.MatcherGroup;
 import io.opentelemetry.api.common.Attributes;
@@ -27,7 +26,7 @@ class AiSampler implements Sampler {
     //
     // failure to follow this pattern can result in unexpected / incorrect computation of values in the portal
     private final double defaultSamplingPercentage;
-    private final SamplingResult defaultRecordAndSampleResult;
+    private final SamplingResult recordAndSampleAndAddTraceStateIfMissing;
 
     private final SamplingOverrides samplingOverrides;
 
@@ -38,12 +37,13 @@ class AiSampler implements Sampler {
     // samplingPercentage is still used in BehaviorIfNoMatchingOverrides.RECORD_AND_SAMPLE
     // to set an approximate value for the span attribute "applicationinsights.internal.sampling_percentage"
     //
-    // in the future the sampling percentage (or its invert "count") will be
+    // in the future the sampling percentage (or its inverse "count") will be
     // carried down by trace state to set the accurate value
     AiSampler(double samplingPercentage, SamplingOverrides samplingOverrides,
               BehaviorIfNoMatchingOverrides behaviorIfNoMatchingOverrides) {
         this.defaultSamplingPercentage = samplingPercentage;
-        defaultRecordAndSampleResult = SamplingOverrides.getRecordAndSampleResult(defaultSamplingPercentage);
+        recordAndSampleAndAddTraceStateIfMissing =
+                SamplingOverrides.getRecordAndSampleAndAddTraceStateIfMissing(samplingPercentage);
 
         this.samplingOverrides = samplingOverrides;
 
@@ -53,7 +53,7 @@ class AiSampler implements Sampler {
     }
 
     @Override
-    public SamplingResult shouldSample(@Nullable Context parentContext,
+    public SamplingResult shouldSample(Context parentContext,
                                        String traceId,
                                        String name,
                                        SpanKind spanKind,
@@ -63,23 +63,27 @@ public SamplingResult shouldSample(@Nullable Context parentContext,
         MatcherGroup override = samplingOverrides.getOverride(attributes);
 
         if (override != null) {
-            return getSamplingResult(override.getPercentage(), override.getRecordAndSampleResult(), traceId, name);
+            return getSamplingResult(override.getPercentage(), override.getRecordAndSampleAndOverwriteTraceState(), traceId, name);
         }
 
         switch (behaviorIfNoMatchingOverrides) {
             case RECORD_AND_SAMPLE:
-                return defaultRecordAndSampleResult;
+                // this is used for localParentSampled and remoteParentSampled
+                // (note: currently sampling percentage portion of trace state is not propagated,
+                //        so it will always be missing in the remoteParentSampled case)
+                return recordAndSampleAndAddTraceStateIfMissing;
             case USE_DEFAULT_SAMPLING_PERCENTAGE:
-                return getSamplingResult(defaultSamplingPercentage, defaultRecordAndSampleResult, traceId, name);
+                // this is used for root sampler
+                return getSamplingResult(defaultSamplingPercentage, recordAndSampleAndAddTraceStateIfMissing, traceId, name);
             default:
                 throw new IllegalStateException("Unexpected BehaviorIfNoMatchingOverrides: " + behaviorIfNoMatchingOverrides);
         }
     }
 
-    private SamplingResult getSamplingResult(double percentage, SamplingResult recordAndSampleResult, String traceId, String name) {
+    private SamplingResult getSamplingResult(double percentage, SamplingResult sampledSamplingResult, String traceId, String name) {
         if (percentage == 100) {
             // optimization, no need to calculate score in this case
-            return recordAndSampleResult;
+            return sampledSamplingResult;
         }
         if (percentage == 0) {
             // optimization, no need to calculate score in this case
@@ -89,7 +93,7 @@ private SamplingResult getSamplingResult(double percentage, SamplingResult recor
             logger.debug("Item {} sampled out", name);
             return dropDecision;
         }
-        return recordAndSampleResult;
+        return sampledSamplingResult;
     }
 
     @Override

agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/sampling/Samplers.java

@@ -13,6 +13,8 @@ public static Sampler getSampler(double samplingPercentage, Configuration config
                 AiSampler.BehaviorIfNoMatchingOverrides.RECORD_AND_SAMPLE);
         // ignoreRemoteParentNotSampled is currently needed
         // because .NET SDK always propagates trace flags "00" (not sampled)
+        // NOTE: once we start propagating sampling percentage over the wire, we can use that to know that we can
+        // respect upstream decision for remoteParentNotSampled
         Sampler remoteParentNotSampled = config.preview.ignoreRemoteParentNotSampled ? rootSampler : Sampler.alwaysOff();
         return Sampler.parentBasedBuilder(rootSampler)
                 .setRemoteParentNotSampled(remoteParentNotSampled)

agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/sampling/SamplingOverrides.java

@@ -1,5 +1,7 @@
 package com.microsoft.applicationinsights.agent.internal.sampling;
 
+import java.math.BigDecimal;
+import java.math.MathContext;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.function.Predicate;
@@ -11,6 +13,7 @@
 import com.microsoft.applicationinsights.agent.internal.wasbootstrap.configuration.Configuration.SamplingOverrideAttribute;
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.trace.TraceState;
 import io.opentelemetry.sdk.trace.samplers.SamplingDecision;
 import io.opentelemetry.sdk.trace.samplers.SamplingResult;
 
@@ -35,37 +38,100 @@ MatcherGroup getOverride(Attributes attributes) {
         return null;
     }
 
-    static SamplingResult getRecordAndSampleResult(double percentage) {
-        Attributes alwaysOnAttributes;
-        if (percentage != 100) {
-            alwaysOnAttributes = Attributes.of(Exporter.AI_SAMPLING_PERCENTAGE_KEY, percentage);
-        } else {
-            // the exporter assumes 100 when the AI_SAMPLING_PERCENTAGE_KEY attribute is not present
-            alwaysOnAttributes = Attributes.empty();
+    static SamplingResult getRecordAndSampleAndOverwriteTraceState(double samplingPercentage) {
+        return new TraceStateUpdatingSamplingResult(SamplingDecision.RECORD_AND_SAMPLE, toRoundedString(samplingPercentage), true);
+    }
+
+    static SamplingResult getRecordAndSampleAndAddTraceStateIfMissing(double samplingPercentage) {
+        return new TraceStateUpdatingSamplingResult(SamplingDecision.RECORD_AND_SAMPLE, toRoundedString(samplingPercentage), false);
+    }
+
+    // TODO write test for
+    //  * 33.33333333333
+    //  * 66.66666666666
+    //  * 1.123456
+    //  * 50.0
+    //  * 1.0
+    //  * 0
+    //  * 0.001
+    //  * 0.000001
+    // 5 digit of precision, and remove any trailing zeros beyond the decimal point
+    private static String toRoundedString(double percentage) {
+        BigDecimal bigDecimal = new BigDecimal(percentage);
+        bigDecimal = bigDecimal.round(new MathContext(5));
+        String formatted = bigDecimal.toString();
+        double dv = bigDecimal.doubleValue();
+        if (dv > 0 && dv < 1) {
+            while (formatted.endsWith("0")) {
+                formatted = formatted.substring(0, formatted.length() - 1);
+            }
+        }
+        return formatted;
+    }
+
+    private static final class TraceStateUpdatingSamplingResult implements SamplingResult {
+
+        private final SamplingDecision decision;
+        private final String samplingPercentage;
+        private final TraceState traceState;
+        private final boolean overwriteExisting;
+
+        private TraceStateUpdatingSamplingResult(SamplingDecision decision, String samplingPercentage,
+                                                 boolean overwriteExisting) {
+            this.decision = decision;
+            this.samplingPercentage = samplingPercentage;
+            this.overwriteExisting = overwriteExisting;
+            traceState = TraceState.builder().put(Exporter.SAMPLING_PERCENTAGE_TRACE_STATE, samplingPercentage).build();
+        }
+
+        @Override
+        public SamplingDecision getDecision() {
+            return decision;
+        }
+
+        @Override
+        public Attributes getAttributes() {
+            return Attributes.empty();
+        }
+
+        @Override
+        public TraceState getUpdatedTraceState(TraceState parentTraceState) {
+            if (parentTraceState.isEmpty()) {
+                return traceState;
+            }
+            String existingSamplingPercentage = parentTraceState.get(Exporter.SAMPLING_PERCENTAGE_TRACE_STATE);
+            if (samplingPercentage.equals(existingSamplingPercentage)) {
+                return parentTraceState;
+            }
+            if (existingSamplingPercentage != null && !overwriteExisting) {
+                return parentTraceState;
+            }
+            return parentTraceState.toBuilder()
+                    .put(Exporter.SAMPLING_PERCENTAGE_TRACE_STATE, samplingPercentage)
+                    .build();
         }
-        return SamplingResult.create(SamplingDecision.RECORD_AND_SAMPLE, alwaysOnAttributes);
     }
 
     static class MatcherGroup {
         private final List<Predicate<Attributes>> predicates;
         private final double percentage;
-        private final SamplingResult recordAndSampleResult;
+        private final SamplingResult recordAndSampleAndOverwriteTraceState;
 
         private MatcherGroup(SamplingOverride override) {
             predicates = new ArrayList<>();
             for (SamplingOverrideAttribute attribute : override.attributes) {
                 predicates.add(toPredicate(attribute));
             }
             percentage = override.percentage;
-            recordAndSampleResult = SamplingOverrides.getRecordAndSampleResult(percentage);
+            recordAndSampleAndOverwriteTraceState = SamplingOverrides.getRecordAndSampleAndOverwriteTraceState(percentage);
         }
 
         double getPercentage() {
             return percentage;
         }
 
-        SamplingResult getRecordAndSampleResult() {
-            return recordAndSampleResult;
+        SamplingResult getRecordAndSampleAndOverwriteTraceState() {
+            return recordAndSampleAndOverwriteTraceState;
         }
 
         private boolean matches(Attributes attributes) {