Updated transport action sampler

Signed-off-by: Dev Agarwal <[email protected]>
opensearch-project · Feb 7, 2024 · afe700c · afe700c
1 parent 81f4f97
commit afe700c
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/...in/java/org/opensearch/telemetry/tracing/sampler/ProbabilisticTransportActionSampler.java b/...in/java/org/opensearch/telemetry/tracing/sampler/ProbabilisticTransportActionSampler.java
@@ -21,6 +21,7 @@
 import io.opentelemetry.context.Context;
 import io.opentelemetry.sdk.trace.data.LinkData;
 import io.opentelemetry.sdk.trace.samplers.Sampler;
+import io.opentelemetry.sdk.trace.samplers.SamplingDecision;
 import io.opentelemetry.sdk.trace.samplers.SamplingResult;
 
 import static org.opensearch.telemetry.tracing.AttributeNames.TRANSPORT_ACTION;
@@ -71,7 +72,11 @@ public SamplingResult shouldSample(
     ) {
         final String action = attributes.get(AttributeKey.stringKey(TRANSPORT_ACTION));
         if (action != null) {
-            return actionSampler.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);
+            final SamplingResult result = actionSampler.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);
+            if (result.getDecision() != SamplingDecision.DROP && fallbackSampler != null) {
+                return fallbackSampler.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);
+            }
+            return result;
         }
         if (fallbackSampler != null) return fallbackSampler.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);