Keep the current sampler when the Jaeger sampling strategy is unchanged

sdk-extensions/jaeger-remote-sampler/src/main/java/io/opentelemetry/sdk/extension/trace/jaeger/sampler/JaegerRemoteSampler.java

@@ -6,6 +6,7 @@
 package io.opentelemetry.sdk.extension.trace.jaeger.sampler;
 
 import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.internal.GuardedBy;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.sdk.common.CompletableResultCode;
@@ -48,6 +49,11 @@ public final class JaegerRemoteSampler implements Sampler {
   private final GrpcSender grpcSender;
   private final URI endpoint;
   private final int pollingIntervalMs;
+  private final Object lock = new Object();
+
+  @GuardedBy("lock")
+  @Nullable
+  private SamplingStrategyResponse lastStrategyResponse;
 
   JaegerRemoteSampler(
       GrpcSender grpcSender,
@@ -98,7 +104,15 @@ private void onResponse(GrpcResponse grpcResponse) {
       try {
         SamplingStrategyResponse strategyResponse =
             SamplingStrategyResponseUnMarshaler.read(grpcResponse.getResponseMessage());
-        sampler = updateSampler(strategyResponse);
+        synchronized (lock) {
+          // Keep the current sampler when the strategy hasn't changed, so stateful samplers such
+          // as RateLimitingSampler don't lose their state on every poll.
+          if (strategyResponse.equals(lastStrategyResponse)) {
+            return;
+          }
+          sampler = updateSampler(strategyResponse);
+          lastStrategyResponse = strategyResponse;
+        }
       } catch (IOException e) {
         logger.log(Level.WARNING, "Failed to unmarshal strategy response", e);
       }

sdk-extensions/jaeger-remote-sampler/src/main/java/io/opentelemetry/sdk/extension/trace/jaeger/sampler/SamplingStrategyResponse.java

@@ -8,22 +8,37 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import javax.annotation.Nullable;
 
-class SamplingStrategyResponse {
+final class SamplingStrategyResponse {
 
   enum SamplingStrategyType {
     PROBABILISTIC,
     RATE_LIMITING,
     UNRECOGNIZED,
   }
 
-  static class RateLimitingSamplingStrategy {
+  static final class RateLimitingSamplingStrategy {
     final int maxTracesPerSecond;
 
     private RateLimitingSamplingStrategy(Builder builder) {
       this.maxTracesPerSecond = builder.maxTracesPerSecond;
     }
 
+    @Override
+    public boolean equals(@Nullable Object o) {
+      if (!(o instanceof RateLimitingSamplingStrategy)) {
+        return false;
+      }
+      RateLimitingSamplingStrategy that = (RateLimitingSamplingStrategy) o;
+      return maxTracesPerSecond == that.maxTracesPerSecond;
+    }
+
+    @Override
+    public int hashCode() {
+      return Integer.hashCode(maxTracesPerSecond);
+    }
+
     static class Builder {
       private int maxTracesPerSecond;
 
@@ -38,13 +53,27 @@ RateLimitingSamplingStrategy build() {
     }
   }
 
-  static class ProbabilisticSamplingStrategy {
+  static final class ProbabilisticSamplingStrategy {
     final double samplingRate;
 
     ProbabilisticSamplingStrategy(Builder builder) {
       this.samplingRate = builder.samplingRate;
     }
 
+    @Override
+    public boolean equals(@Nullable Object o) {
+      if (!(o instanceof ProbabilisticSamplingStrategy)) {
+        return false;
+      }
+      ProbabilisticSamplingStrategy that = (ProbabilisticSamplingStrategy) o;
+      return Double.doubleToLongBits(samplingRate) == Double.doubleToLongBits(that.samplingRate);
+    }
+
+    @Override
+    public int hashCode() {
+      return Double.hashCode(samplingRate);
+    }
+
     static class Builder {
       private double samplingRate;
 
@@ -59,7 +88,7 @@ ProbabilisticSamplingStrategy build() {
     }
   }
 
-  static class PerOperationSamplingStrategies {
+  static final class PerOperationSamplingStrategies {
     final double defaultSamplingProbability;
     final double defaultLowerBoundTracesPerSecond;
     final double defaultUpperBoundTracesPerSecond;
@@ -72,6 +101,30 @@ private PerOperationSamplingStrategies(Builder builder) {
       this.strategies = Collections.unmodifiableList(builder.strategies);
     }
 
+    @Override
+    public boolean equals(@Nullable Object o) {
+      if (!(o instanceof PerOperationSamplingStrategies)) {
+        return false;
+      }
+      PerOperationSamplingStrategies that = (PerOperationSamplingStrategies) o;
+      return Double.doubleToLongBits(defaultSamplingProbability)
+              == Double.doubleToLongBits(that.defaultSamplingProbability)
+          && Double.doubleToLongBits(defaultLowerBoundTracesPerSecond)
+              == Double.doubleToLongBits(that.defaultLowerBoundTracesPerSecond)
+          && Double.doubleToLongBits(defaultUpperBoundTracesPerSecond)
+              == Double.doubleToLongBits(that.defaultUpperBoundTracesPerSecond)
+          && strategies.equals(that.strategies);
+    }
+
+    @Override
+    public int hashCode() {
+      int result = Double.hashCode(defaultSamplingProbability);
+      result = 31 * result + Double.hashCode(defaultLowerBoundTracesPerSecond);
+      result = 31 * result + Double.hashCode(defaultUpperBoundTracesPerSecond);
+      result = 31 * result + strategies.hashCode();
+      return result;
+    }
+
     static class Builder {
       private double defaultSamplingProbability;
       private double defaultLowerBoundTracesPerSecond;
@@ -104,7 +157,7 @@ PerOperationSamplingStrategies build() {
     }
   }
 
-  static class OperationSamplingStrategy {
+  static final class OperationSamplingStrategy {
     final String operation;
     final ProbabilisticSamplingStrategy probabilisticSamplingStrategy;
 
@@ -113,6 +166,21 @@ private OperationSamplingStrategy(Builder builder) {
       this.probabilisticSamplingStrategy = builder.probabilisticSamplingStrategy;
     }
 
+    @Override
+    public boolean equals(@Nullable Object o) {
+      if (!(o instanceof OperationSamplingStrategy)) {
+        return false;
+      }
+      OperationSamplingStrategy that = (OperationSamplingStrategy) o;
+      return operation.equals(that.operation)
+          && probabilisticSamplingStrategy.equals(that.probabilisticSamplingStrategy);
+    }
+
+    @Override
+    public int hashCode() {
+      return 31 * operation.hashCode() + probabilisticSamplingStrategy.hashCode();
+    }
+
     static class Builder {
       private String operation = "";
       private ProbabilisticSamplingStrategy probabilisticSamplingStrategy =
@@ -147,6 +215,27 @@ private SamplingStrategyResponse(Builder builder) {
     this.perOperationSamplingStrategies = builder.perOperationSamplingStrategies;
   }
 
+  @Override
+  public boolean equals(@Nullable Object o) {
+    if (!(o instanceof SamplingStrategyResponse)) {
+      return false;
+    }
+    SamplingStrategyResponse that = (SamplingStrategyResponse) o;
+    return strategyType == that.strategyType
+        && rateLimitingSamplingStrategy.equals(that.rateLimitingSamplingStrategy)
+        && probabilisticSamplingStrategy.equals(that.probabilisticSamplingStrategy)
+        && perOperationSamplingStrategies.equals(that.perOperationSamplingStrategies);
+  }
+
+  @Override
+  public int hashCode() {
+    int result = strategyType.hashCode();
+    result = 31 * result + rateLimitingSamplingStrategy.hashCode();
+    result = 31 * result + probabilisticSamplingStrategy.hashCode();
+    result = 31 * result + perOperationSamplingStrategies.hashCode();
+    return result;
+  }
+
   static class Builder {
     private SamplingStrategyType samplingStrategyType = SamplingStrategyType.UNRECOGNIZED;
     private RateLimitingSamplingStrategy rateLimitingSamplingStrategy =

sdk-extensions/jaeger-remote-sampler/src/test/java/io/opentelemetry/sdk/extension/trace/jaeger/sampler/JaegerRemoteSamplerTest.java

@@ -38,6 +38,7 @@
 import java.util.Collections;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Stream;
 import javax.annotation.Nullable;
 import javax.net.ssl.KeyManager;
@@ -71,6 +72,8 @@ class JaegerRemoteSamplerTest {
   private static final ConcurrentLinkedQueue<Sampling.SamplingStrategyResponse> responses =
       new ConcurrentLinkedQueue<>();
 
+  private static final AtomicInteger requestCount = new AtomicInteger();
+
   private static void addGrpcError(GrpcStatusCode code, @Nullable String message) {
     grpcErrors.add(new ArmeriaStatusException(code.getValue(), message));
   }
@@ -104,6 +107,7 @@ public void getSamplingStrategy(
                                     io.opentelemetry.sdk.extension.trace.jaeger.proto.api_v2
                                         .Sampling.SamplingStrategyResponse>
                                 responseObserver) {
+                          requestCount.incrementAndGet();
                           ArmeriaStatusException grpcError = grpcErrors.poll();
                           if (grpcError != null) {
                             responseObserver.onError(grpcError);
@@ -142,6 +146,7 @@ public void getSamplingStrategy(
   public void before() {
     grpcErrors.clear();
     responses.clear();
+    requestCount.set(0);
   }
 
   @Test
@@ -449,6 +454,33 @@ void perOperationSampling_zeroRateLimiting() {
     }
   }
 
+  @Test
+  void unchangedStrategy_keepsSampler() {
+    try (JaegerRemoteSampler sampler =
+        JaegerRemoteSampler.builder()
+            .setEndpoint(server.httpUri().toString())
+            .setServiceName(SERVICE_NAME)
+            .setPollingInterval(50, TimeUnit.MILLISECONDS)
+            .build()) {
+      await().untilAsserted(samplerIsType(sampler, RateLimitingSampler.class));
+      Sampler initial = sampler.getSampler();
+
+      // Further polls return the same strategy, so the stateful sampler must be retained.
+      int polls = requestCount.get();
+      await().untilAsserted(() -> assertThat(requestCount.get()).isGreaterThan(polls + 2));
+      assertThat(sampler.getSampler()).isSameAs(initial);
+
+      // A changed strategy still replaces the sampler.
+      responses.add(
+          Sampling.SamplingStrategyResponse.newBuilder()
+              .setStrategyType(SamplingStrategyType.RATE_LIMITING)
+              .setRateLimitingSampling(
+                  RateLimitingSamplingStrategy.newBuilder().setMaxTracesPerSecond(5).build())
+              .build());
+      await().untilAsserted(() -> assertThat(sampler.getSampler()).isNotSameAs(initial));
+    }
+  }
+
   @Test
   void internal_error_server_response() {
     addGrpcError(GrpcStatusCode.INTERNAL, "internal error");

sdk-extensions/jaeger-remote-sampler/src/test/java/io/opentelemetry/sdk/extension/trace/jaeger/sampler/SamplingStrategyResponseTest.java

@@ -0,0 +1,29 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.sdk.extension.trace.jaeger.sampler;
+
+import io.opentelemetry.sdk.extension.trace.jaeger.sampler.SamplingStrategyResponse.OperationSamplingStrategy;
+import io.opentelemetry.sdk.extension.trace.jaeger.sampler.SamplingStrategyResponse.PerOperationSamplingStrategies;
+import io.opentelemetry.sdk.extension.trace.jaeger.sampler.SamplingStrategyResponse.ProbabilisticSamplingStrategy;
+import io.opentelemetry.sdk.extension.trace.jaeger.sampler.SamplingStrategyResponse.RateLimitingSamplingStrategy;
+import nl.jqno.equalsverifier.EqualsVerifier;
+import nl.jqno.equalsverifier.Warning;
+import org.junit.jupiter.api.Test;
+
+class SamplingStrategyResponseTest {
+
+  @Test
+  void equalsVerifier() {
+    // The builders always supply a value, so the fields are never null in practice.
+    EqualsVerifier.forClass(RateLimitingSamplingStrategy.class).verify();
+    EqualsVerifier.forClass(ProbabilisticSamplingStrategy.class).verify();
+    EqualsVerifier.forClass(OperationSamplingStrategy.class).suppress(Warning.NULL_FIELDS).verify();
+    EqualsVerifier.forClass(PerOperationSamplingStrategies.class)
+        .suppress(Warning.NULL_FIELDS)
+        .verify();
+    EqualsVerifier.forClass(SamplingStrategyResponse.class).suppress(Warning.NULL_FIELDS).verify();
+  }
+}