Cleanup operation samplers if absent from strategy response

Signed-off-by: jung <jung@uber.com>
jaegertracing · Oct 2, 2019 · 8a5720e · 8a5720e
1 parent 21ea7af
commit 8a5720e
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 9 deletions.
diff --git a/jaeger-core/src/main/java/io/jaegertracing/internal/samplers/PerOperationSampler.java b/jaeger-core/src/main/java/io/jaegertracing/internal/samplers/PerOperationSampler.java
@@ -18,13 +18,15 @@
 import io.jaegertracing.internal.samplers.http.PerOperationSamplingParameters;
 import io.jaegertracing.spi.Sampler;
 import java.util.HashMap;
+import java.util.Map;
 import lombok.AccessLevel;
 import lombok.AllArgsConstructor;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
 import lombok.ToString;
 import lombok.extern.slf4j.Slf4j;
 
+
 /**
  * Computes {@link #sample(String, long)} using the name of the operation, and maintains a specific
  * {@link GuaranteedThroughputSampler} instance for each operation.
@@ -36,7 +38,7 @@
 @Getter(AccessLevel.PACKAGE) //Visible for testing
 public class PerOperationSampler implements Sampler {
   private final int maxOperations;
-  private final HashMap<String, GuaranteedThroughputSampler> operationNameToSampler;
+  private Map<String, GuaranteedThroughputSampler> operationNameToSampler;
   private ProbabilisticSampler defaultSampler;
   private double lowerBound;
 
@@ -51,9 +53,9 @@ public PerOperationSampler(int maxOperations, OperationSamplingParameters strate
   /**
    * Updates the GuaranteedThroughputSampler for each operation
    * @param strategies The parameters for operation sampling
-   * @return true iff any samplers were updated
+   * @return true if any samplers were updated
    */
-  public synchronized boolean update(OperationSamplingParameters strategies) {
+  public synchronized boolean update(final OperationSamplingParameters strategies) {
     boolean isUpdated = false;
 
     lowerBound = strategies.getDefaultLowerBoundTracesPerSecond();
@@ -64,23 +66,28 @@ public synchronized boolean update(OperationSamplingParameters strategies) {
       isUpdated = true;
     }
 
+    Map<String, GuaranteedThroughputSampler> newOpsSamplers = new HashMap<String, GuaranteedThroughputSampler>();
+    //add or update operation samples using given strategies
     for (PerOperationSamplingParameters strategy : strategies.getPerOperationStrategies()) {
       String operation = strategy.getOperation();
       double samplingRate = strategy.getProbabilisticSampling().getSamplingRate();
       GuaranteedThroughputSampler sampler = operationNameToSampler.get(operation);
       if (sampler != null) {
         isUpdated = sampler.update(samplingRate, lowerBound) || isUpdated;
+        newOpsSamplers.put(operation, sampler);
       } else {
-        if (operationNameToSampler.size() < maxOperations) {
+        if (newOpsSamplers.size() < maxOperations) {
           sampler = new GuaranteedThroughputSampler(samplingRate, lowerBound);
-          operationNameToSampler.put(operation, sampler);
+          newOpsSamplers.put(operation, sampler);
           isUpdated = true;
         } else {
           log.info("Exceeded the maximum number of operations({}) for per operations sampling",
               maxOperations);
         }
       }
     }
+
+    operationNameToSampler = newOpsSamplers;
     return isUpdated;
   }
 

diff --git a/jaeger-core/src/test/java/io/jaegertracing/internal/samplers/PerOperationSamplerTest.java b/jaeger-core/src/test/java/io/jaegertracing/internal/samplers/PerOperationSamplerTest.java
@@ -154,8 +154,8 @@ public void testUpdateIgnoreGreaterThanMax() {
     undertest.update(new OperationSamplingParameters(DEFAULT_SAMPLING_PROBABILITY,
                                                      DEFAULT_LOWER_BOUND_TRACES_PER_SECOND, parametersList));
 
-    assertEquals(1, operationToSamplers.size());
-    assertNotNull(operationToSamplers.get(OPERATION));
+    assertEquals(1, undertest.getOperationNameToSampler().size());
+    assertNotNull(undertest.getOperationNameToSampler().get(OPERATION));
   }
 
   @Test
@@ -169,10 +169,29 @@ public void testUpdateAddOperation() {
                                                      DEFAULT_LOWER_BOUND_TRACES_PER_SECOND,
                                                      parametersList));
 
-    assertEquals(1, operationToSamplers.size());
+    assertEquals(1, undertest.getOperationNameToSampler().size());
     assertEquals(new GuaranteedThroughputSampler(SAMPLING_RATE,
                                                  DEFAULT_LOWER_BOUND_TRACES_PER_SECOND),
-                 operationToSamplers.get(OPERATION));
+            undertest.getOperationNameToSampler().get(OPERATION));
   }
 
+  @Test
+  public void testAbsentOperationIsRemoved() {
+    String absentOp = "ShouldBeRemoved";
+    operationToSamplers.put(absentOp, mock(GuaranteedThroughputSampler.class));
+    PerOperationSamplingParameters perOperationSamplingParameters1 =
+            new PerOperationSamplingParameters(OPERATION, new ProbabilisticSamplingStrategy(SAMPLING_RATE));
+    List<PerOperationSamplingParameters> parametersList = new ArrayList<>();
+    parametersList.add(perOperationSamplingParameters1);
+
+    undertest.update(new OperationSamplingParameters(DEFAULT_SAMPLING_PROBABILITY,
+            DEFAULT_LOWER_BOUND_TRACES_PER_SECOND,
+            parametersList));
+
+    assertEquals(1, undertest.getOperationNameToSampler().size());
+    assertEquals(new GuaranteedThroughputSampler(SAMPLING_RATE,
+                    DEFAULT_LOWER_BOUND_TRACES_PER_SECOND),
+            undertest.getOperationNameToSampler().get(OPERATION));
+    assertFalse(undertest.getOperationNameToSampler().containsKey(absentOp));
+  }
 }