Minor refactor based on comments

Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com>
devagarwal1803 · Sep 25, 2023 · b0aa9f9 · b0aa9f9
1 parent ffc297c
commit b0aa9f9
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 19 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
@@ -37,6 +37,8 @@ public class ProbabilisticSampler implements Sampler {
     private final Map<String, Sampler> actionBasedSampler;
     private final Map<String, Double> actionBasedSamplingProbability;
 
+    // private final Map<String, Pair<String, Sampler>> myMap = new HashMap<>();
+
     /**
      * Constructor
      *
@@ -55,14 +57,7 @@ public ProbabilisticSampler(TelemetrySettings telemetrySettings) {
      * @param attributes Telemetry attributes
      */
     Sampler getSampler(Attributes attributes) {
-        double newSamplingRatio = telemetrySettings.getSamplingProbability();
-        if (isSamplingRatioChanged(newSamplingRatio)) {
-            synchronized (this) {
-                this.samplingRatio = newSamplingRatio;
-                defaultSampler = Sampler.traceIdRatioBased(samplingRatio);
-            }
-        }
-
+        // Evaluate if sampling is overridden at action level
         final String action = attributes.get(AttributeKey.stringKey(TRANSPORT_ACTION));
         if (action != null && telemetrySettings.isActionSamplingOverrideSet(action)) {
             if (isActionSamplerUpdateRequired(action)) {
@@ -71,9 +66,17 @@ Sampler getSampler(Attributes attributes) {
                 }
             }
             return actionBasedSampler.get(action);
+        } else {
+            // Update default sampling if no override is present for action
+            double newSamplingRatio = telemetrySettings.getSamplingProbability();
+            if (isSamplingRatioChanged(newSamplingRatio)) {
+                synchronized (this) {
+                    this.samplingRatio = newSamplingRatio;
+                    defaultSampler = Sampler.traceIdRatioBased(samplingRatio);
+                }
+            }
+            return defaultSampler;
         }
-
-        return defaultSampler;
     }
 
     private boolean isSamplingRatioChanged(double newSamplingRatio) {
@@ -82,7 +85,7 @@ private boolean isSamplingRatioChanged(double newSamplingRatio) {
 
     private boolean isActionSamplerUpdateRequired(String action) {
         return (!actionBasedSampler.containsKey(action)
-            || actionBasedSamplingProbability.get(action) != telemetrySettings.getActionSamplingProbability(action));
+            || (actionBasedSamplingProbability.get(action) != telemetrySettings.getActionSamplingProbability(action)));
     }
 
     private void updateActionSampler(String action) {
@@ -96,10 +99,7 @@ private void updateActionSampler(String action) {
     }
 
     double getActionSamplingRatio(String action) {
-        if (actionBasedSamplingProbability.containsKey(action)) {
-            return actionBasedSamplingProbability.get(action);
-        }
-        return samplingRatio;
+        return actionBasedSamplingProbability.get(action);
     }
 
     @Override

diff --git a/server/src/main/java/org/opensearch/telemetry/TelemetrySettings.java b/server/src/main/java/org/opensearch/telemetry/TelemetrySettings.java
@@ -116,9 +116,6 @@ public boolean isActionSamplingOverrideSet(String action) {
      * @return double value of sampling probability for that action
      */
     public double getActionSamplingProbability(String action) {
-        if (affixSamplingProbability.containsKey(action)) {
-            return this.affixSamplingProbability.get(action);
-        }
-        return samplingProbability;
+        return this.affixSamplingProbability.get(action);
     }
 }