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;
@@ -19,6 +20,7 @@
 import io.opentelemetry.sdk.trace.samplers.SamplingResult;
 import java.io.IOException;
 import java.net.URI;
+import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
@@ -48,6 +50,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 byte[] lastResponseMessage;
 
   JaegerRemoteSampler(
       GrpcSender grpcSender,
@@ -95,10 +102,19 @@ private void onResponse(GrpcResponse grpcResponse) {
       if (isShutdown.get()) {
         return;
       }
+      byte[] responseMessage = grpcResponse.getResponseMessage();
       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 (Arrays.equals(responseMessage, lastResponseMessage)) {
+            return;
+          }
+          SamplingStrategyResponse strategyResponse =
+              SamplingStrategyResponseUnMarshaler.read(responseMessage);
+          sampler = updateSampler(strategyResponse);
+          lastResponseMessage = responseMessage;
+        }
       } catch (IOException e) {
         logger.log(Level.WARNING, "Failed to unmarshal strategy response", e);
       }

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");