Minor refactor based on commits

devagarwal1803 · Sep 21, 2023 · ae6783c · ae6783c
1 parent 3fba81b
commit ae6783c
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 17 deletions.
diff --git a/...try-otel/src/main/java/org/opensearch/telemetry/tracing/sampler/ProbabilisticSampler.java b/...try-otel/src/main/java/org/opensearch/telemetry/tracing/sampler/ProbabilisticSampler.java
@@ -17,6 +17,8 @@
 
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.sdk.trace.data.LinkData;
@@ -63,12 +65,12 @@ Sampler getSampler(Attributes attributes) {
 
         final String action = attributes.get(AttributeKey.stringKey(TRANSPORT_ACTION));
         if (action != null && telemetrySettings.isActionSamplingOverrideSet(action)) {
-            if (actionSamplerNeedsUpdate(action)) {
+            if (isActionSamplerUpdateRequired(action)) {
                 synchronized (this) {
                     updateActionSampler(action);
                 }
-                return actionBasedSampler.get(action);
             }
+            return actionBasedSampler.get(action);
         }
 
         return defaultSampler;
@@ -78,7 +80,7 @@ private boolean isSamplingRatioChanged(double newSamplingRatio) {
         return Double.compare(this.samplingRatio, newSamplingRatio) != 0;
     }
 
-    private boolean actionSamplerNeedsUpdate(String action) {
+    private boolean isActionSamplerUpdateRequired(String action) {
         return (!actionBasedSampler.containsKey(action)
             || actionBasedSamplingProbability.get(action) != telemetrySettings.getActionSamplingProbability(action));
     }
@@ -109,7 +111,12 @@ public SamplingResult shouldSample(
         Attributes attributes,
         List<LinkData> parentLinks
     ) {
-        return getSampler(attributes).shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);
+        // Use action sampler only if it is root span & override is present for the action
+        SpanContext parentSpanContext = Span.fromContext(parentContext).getSpanContext();
+        if (!parentSpanContext.isValid()) {
+            return getSampler(attributes).shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);
+        }
+        return defaultSampler.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);
     }
 
     @Override

diff --git a/...tel/src/test/java/org/opensearch/telemetry/tracing/sampler/ProbabilisticSamplerTests.java b/...tel/src/test/java/org/opensearch/telemetry/tracing/sampler/ProbabilisticSamplerTests.java
@@ -82,27 +82,30 @@ public void testGetSamplerWithCustomActionSamplingRatio() {
 
         Map<String, Double> filters = new HashMap<>();
         // Setting 100% sampling for indexing request
-        filters.put("indexing", 1.00);
+        filters.put("internal:cluster/coordination/join", 1.00);
         // Setting 50% sampling rate for query request
-        filters.put("query", 0.5);
+        filters.put("cluster:monitor/nodes/stats[n]", 0.5);
         telemetrySettings.setActionSamplingProbability(filters);
 
-        attributes.put(AttributeKey.stringKey("action"), "indexing");
+        attributes.put(AttributeKey.stringKey("action"), "internal:cluster/coordination/join");
         probabilisticSampler.getSampler(attributes);
 
-        // Validates sampling probability for general request.
+        // Validates sampling probability for dummy request.
         assertEquals(0.02, probabilisticSampler.getActionSamplingRatio("dummy"), 0.0d);
-        // Validates sampling probability for indexing request as override is present for it.
-        assertEquals(1.00, probabilisticSampler.getActionSamplingRatio("indexing"), 0.0d);
-        // Validates sampling probability for query request as override is present for it.
-        assertEquals(0.02, probabilisticSampler.getActionSamplingRatio("query"), 0.0d);
-
-        filters.put("indexing", 0.30);
+        // Validates sampling probability for cluster coordination request as override is present for it.
+        assertEquals(1.00, probabilisticSampler.getActionSamplingRatio("internal:cluster/coordination/join"), 0.0d);
+        // Validates sampling probability for nodes stats request as override is present for it.
+        assertEquals(0.02, probabilisticSampler.getActionSamplingRatio("cluster:monitor/nodes/stats[n]"), 0.0d);
+        // Validates sampling probability for nodes stats second request.
+        assertEquals(0.30, probabilisticSampler.getActionSamplingRatio("cluster:monitor/nodes/stats[n]"), 0.0d);
+
+        filters.put("internal:cluster/coordination/join", 0.30);
         telemetrySettings.setActionSamplingProbability(filters);
         // Need to call getSampler() to update the value of tracerHeadSamplerSamplingRatio
         probabilisticSampler.getSampler(attributes);
-        // Validates sampling probability for indexing request as override is present for it.
-        assertEquals(0.30, probabilisticSampler.getActionSamplingRatio("indexing"), 0.0d);
+
+        // Validates sampling probability for cluster coordination as override is present for it.
+        assertEquals(0.30, probabilisticSampler.getActionSamplingRatio("internal:cluster/coordination/join"), 0.0d);
     }
 
 }
diff --git a/server/src/main/java/org/opensearch/telemetry/TelemetrySettings.java b/server/src/main/java/org/opensearch/telemetry/TelemetrySettings.java
@@ -46,7 +46,7 @@ public class TelemetrySettings {
      * Probability of action based sampler
      */
     public static final Setting.AffixSetting<Double> TRACER_SAMPLER_ACTION_PROBABILITY = Setting.affixKeySetting(
-        "telemetry.tracer.sampler.",
+        "telemetry.tracer.action.sampler.",
         "probability",
         (ns, key) -> Setting.doubleSetting(key, 0.00d, 0.00d, 1.00d, Setting.Property.Dynamic, Setting.Property.NodeScope)
     );