feat: client code adjustments and sessionAffinity hints

klopfdreh · klopfdreh · commit ac68dd7e5703 · 2025-03-05T14:30:34.000+01:00
Signed-off-by: Tobias Soloschenko &lt;tsoloschenko@apache.org&gt;
diff --git a/README.md b/README.md
@@ -117,6 +117,33 @@ This installation includes Prometheus and Grafana as well.
 1. `kubectl apply -f scripts/kubernetes/prometheus/`
 3. `kubectl apply -f scripts/kubernetes/grafana`
 
+Note: Within container orchestration platforms you have to ensure that all interactions of the client and server are bound to the same Pod the initial request was sent to. This can be achieved by using a `sessionAffinity` which is set to ClientIP
+
+Example service definition:
+```yaml
+kind: Service
+apiVersion: v1
+metadata:
+ labels:
+   app: prometheus-rsocket-proxy-server
+ # ...
+spec:
+  ports:
+    - name: https
+      protocol: TCP
+      port: 8443
+      targetPort: https
+    - name: rsocket
+      protocol: TCP
+      port: 7001
+      targetPort: rsocket
+  selector:
+    app: prometheus-rsocket-proxy-server
+  type: ClusterIP
+  sessionAffinity: ClientIP
+```
+
+
 ## Expected performance
 
 A 3-pod deployment easily handles 1,000 connected application instances each serving 1,000 distinct time series with <1vCPU and <3Gi RAM total on GKE.
diff --git a/client/src/main/java/io/micrometer/prometheus/rsocket/PrometheusRSocketClient.java b/client/src/main/java/io/micrometer/prometheus/rsocket/PrometheusRSocketClient.java
@@ -72,14 +72,15 @@ public class PrometheusRSocketClient {
    *
    * @param registryAndScrape the registry and scrape meter
    * @param transport         the client transport
-   * @param retry             the retry configuration
+   * @param reconnectRetry    the reconnectRetry configuration
    * @param onKeyReceived     the callback if a key has been received
    */
   private PrometheusRSocketClient(MeterRegistryAndScrape<?> registryAndScrape,
                                   ClientTransport transport,
-                                  Retry retry,
+                                  Retry reconnectRetry,
                                   Runnable onKeyReceived) {
-    this(registryAndScrape, transport, retry, Duration.ofSeconds(5), onKeyReceived);
+    this(registryAndScrape, transport, reconnectRetry, Retry.backoff(6, Duration.ofMillis(100))
+        .maxBackoff(Duration.ofSeconds(5)), Duration.ofSeconds(5), onKeyReceived);
   }
 
   /**
@@ -93,6 +94,7 @@ private PrometheusRSocketClient(MeterRegistryAndScrape<?> registryAndScrape,
    */
   private PrometheusRSocketClient(MeterRegistryAndScrape<?> registryAndScrape,
                                   ClientTransport transport,
+                                  Retry reconnectRetry,
                                   Retry retry,
                                   Duration timeout,
                                   Runnable onKeyReceived) {
@@ -103,7 +105,7 @@ private PrometheusRSocketClient(MeterRegistryAndScrape<?> registryAndScrape,
         .reconnect(new Retry() {
           @Override
           public Publisher<?> generateCompanion(Flux<RetrySignal> retrySignals) {
-            return retry.generateCompanion(retrySignals
+            return reconnectRetry.generateCompanion(retrySignals
                 .doOnNext(retrySignal -> {
                       Throwable failure = retrySignal.failure();
                       DistributionSummary.builder("prometheus.connection.retry")
@@ -118,35 +120,57 @@ public Publisher<?> generateCompanion(Flux<RetrySignal> retrySignals) {
           }
         })
         .acceptor((payload, r) -> {
+          LOGGER.trace("Acceptor for requestResponse and fireAndForget has been set up.");
           this.sendingSocket = r;
           return Mono.just(new RSocket() {
             @Override
             public Mono<Payload> requestResponse(Payload payload) {
-              PublicKey key = decodePublicKey(payload.getData());
-              latestKey.set(key);
-              onKeyReceived.run();
-              return Mono.fromCallable(() -> scrapePayload(key));
+              try {
+                LOGGER.trace("Received public key from Prometheus proxy in requestResponse.");
+                PublicKey key = decodePublicKey(payload.getData());
+                latestKey.set(key);
+                onKeyReceived.run();
+                return Mono.fromCallable(() -> scrapePayload(key));
+              } finally {
+                payload.release();
+              }
             }
 
             @Override
             public Mono<Void> fireAndForget(Payload payload) {
-              latestKey.set(decodePublicKey(payload.getData()));
-              onKeyReceived.run();
-              return Mono.empty();
+              try {
+                LOGGER.trace("Received public key from Prometheus proxy in fireAndForget.");
+                latestKey.set(decodePublicKey(payload.getData()));
+                onKeyReceived.run();
+                return Mono.empty();
+              } finally {
+                payload.release();
+              }
             }
           });
         })
         .connect(transport)
-        .doOnError(t -> Counter.builder("prometheus.connection.error")
-            .baseUnit("errors")
-            .tag("exception", t.getClass().getSimpleName() == null ? t.getClass().getName() : t.getClass().getSimpleName())
-            .register(registryAndScrape.registry)
-            .increment())
+        .doOnError(t -> {
+          LOGGER.trace("Failed to connect to Prometheus proxy.", t);
+          Counter.builder("prometheus.connection.error")
+              .baseUnit("errors")
+              .tag("exception", t.getClass().getSimpleName() == null ? t.getClass().getName() : t.getClass().getSimpleName())
+              .register(registryAndScrape.registry)
+              .increment();
+        })
+        .doOnSuccess(connection -> {
+          LOGGER.trace("Successfully established connection to Prometheus proxy.");
+          Counter.builder("prometheus.connection.success")
+              .baseUnit("connections")
+              .register(registryAndScrape.registry)
+              .increment();
+        })
         .doOnNext(connection -> this.connection = connection)
         .flatMap(socket -> socket.onClose()
-            .map(v -> 1) // https://github.com/rsocket/rsocket-java/issues/819
+            .then(Mono.fromCallable(() -> 1)) // https://github.com/rsocket/rsocket-java/issues/819
             .onErrorReturn(1))
         .repeat(() -> !requestedDisconnect)
+        .retryWhen(retry)
         .subscribe();
   }
 
@@ -299,9 +323,12 @@ public static class Builder {
     private MeterRegistryAndScrape<?> registryAndScrape;
     private final ClientTransport clientTransport;
 
-    private Retry retry = Retry.backoff(Long.MAX_VALUE, Duration.ofSeconds(10))
+    private Retry reconnectRetry = Retry.backoff(Long.MAX_VALUE, Duration.ofSeconds(10))
         .maxBackoff(Duration.ofMinutes(10));
 
+    private Retry retry = Retry.backoff(6, Duration.ofMillis(100))
+        .maxBackoff(Duration.ofSeconds(5));
+
     private Duration timeout = Duration.ofSeconds(5);
 
     private Runnable onKeyReceived = () -> {
@@ -312,6 +339,17 @@ <M extends MeterRegistry> Builder(M registry, Supplier<String> scrape, ClientTra
       this.clientTransport = clientTransport;
     }
 
+    /**
+     * Configures the reconnectRetry for {@link PrometheusRSocketClient}.
+     *
+     * @param reconnectRetry the reconnectRetry configuration
+     * @return the {@link Builder}
+     */
+    public Builder reconnectRetry(Retry reconnectRetry) {
+      this.reconnectRetry = reconnectRetry;
+      return this;
+    }
+
     /**
      * Configures the retry for {@link PrometheusRSocketClient}.
      *
@@ -366,6 +404,7 @@ public PrometheusRSocketClient connect(Duration timeout) {
       return new PrometheusRSocketClient(
           registryAndScrape,
           clientTransport,
+          reconnectRetry,
           retry,
           timeout,
           () -> {
@@ -395,6 +434,7 @@ public PrometheusRSocketClient connectBlockingly(Duration timeout) {
       PrometheusRSocketClient client = new PrometheusRSocketClient(
           registryAndScrape,
           clientTransport,
+          reconnectRetry,
           retry,
           timeout,
           () -> {
diff --git a/starter-spring/src/main/java/io/micrometer/prometheus/rsocket/autoconfigure/PrometheusRSocketClientAutoConfiguration.java b/starter-spring/src/main/java/io/micrometer/prometheus/rsocket/autoconfigure/PrometheusRSocketClientAutoConfiguration.java
@@ -28,6 +28,8 @@
 import org.springframework.context.annotation.Bean;
 import reactor.util.retry.Retry;
 
+import java.time.Duration;
+
 @AutoConfiguration
 @AutoConfigureAfter(PrometheusMetricsExportAutoConfiguration.class)
 @ConditionalOnBean(PrometheusMeterRegistry.class)
@@ -39,8 +41,10 @@ public class PrometheusRSocketClientAutoConfiguration {
   @Bean(destroyMethod = "pushAndCloseBlockingly")
   PrometheusRSocketClient prometheusRSocketClient(PrometheusMeterRegistry meterRegistry, PrometheusRSocketClientProperties properties) {
     return PrometheusRSocketClient.build(meterRegistry, properties.createClientTransport())
-        .retry(Retry.backoff(properties.getMaxRetries(), properties.getFirstBackoff())
-            .maxBackoff(properties.getMaxBackoff()))
+        .reconnectRetry(Retry.backoff(properties.getReconnectRetry().getMaxRetries(), properties.getReconnectRetry().getFirstBackoff())
+            .maxBackoff(properties.getReconnectRetry().getMaxBackoff()))
+        .retry(Retry.backoff(properties.getRetry().getMaxRetries(), properties.getRetry().getFirstBackoff())
+            .maxBackoff(properties.getRetry().getMaxBackoff()))
         .timeout(properties.getTimeout())
         .connectBlockingly();
   }
diff --git a/starter-spring/src/main/java/io/micrometer/prometheus/rsocket/autoconfigure/PrometheusRSocketClientProperties.java b/starter-spring/src/main/java/io/micrometer/prometheus/rsocket/autoconfigure/PrometheusRSocketClientProperties.java
@@ -20,9 +20,11 @@
 import io.rsocket.transport.netty.client.TcpClientTransport;
 import io.rsocket.transport.netty.client.WebsocketClientTransport;
 import org.springframework.boot.context.properties.ConfigurationProperties;
+import org.springframework.boot.context.properties.NestedConfigurationProperty;
 import org.springframework.boot.convert.DurationUnit;
 import reactor.netty.http.client.HttpClient;
 import reactor.netty.tcp.TcpClient;
+import reactor.util.retry.Retry;
 
 import java.time.Duration;
 import java.time.temporal.ChronoUnit;
@@ -41,19 +43,98 @@ public class PrometheusRSocketClientProperties {
   private int port = 7001;
 
   /**
-   * The maximum number of connection attempts to make.
+   * Retry configuration for the reconnect.
    */
-  private long maxRetries = Long.MAX_VALUE;
+  @NestedConfigurationProperty
+  private ReconnectRetry reconnectRetry = new ReconnectRetry();
 
   /**
-   * The first connection attempt backoff delay to apply, then grow exponentially.
+   * Retry configuration for the subscription.
    */
-  private Duration firstBackoff = Duration.ofSeconds(10);
+  @NestedConfigurationProperty
+  private Retry retry = new Retry();
 
-  /**
-   * The maximum connection attempt delay to apply despite exponential growth.
-   */
-  private Duration maxBackoff = Duration.ofMinutes(10);
+  public static class ReconnectRetry {
+    /**
+     * The maximum number of connection attempts to make.
+     */
+    private long maxRetries = Long.MAX_VALUE;
+
+    /**
+     * The first connection attempt backoff delay to apply, then grow exponentially.
+     */
+    private Duration firstBackoff = Duration.ofSeconds(10);
+
+    /**
+     * The maximum connection attempt delay to apply despite exponential growth.
+     */
+    private Duration maxBackoff = Duration.ofMinutes(10);
+
+    public long getMaxRetries() {
+      return maxRetries;
+    }
+
+    public void setMaxRetries(long maxRetries) {
+      this.maxRetries = maxRetries;
+    }
+
+    public Duration getFirstBackoff() {
+      return firstBackoff;
+    }
+
+    public void setFirstBackoff(Duration firstBackoff) {
+      this.firstBackoff = firstBackoff;
+    }
+
+    public Duration getMaxBackoff() {
+      return maxBackoff;
+    }
+
+    public void setMaxBackoff(Duration maxBackoff) {
+      this.maxBackoff = maxBackoff;
+    }
+  }
+
+  public static class Retry {
+    /**
+     * The maximum number of connection attempts to make.
+     */
+    private long maxRetries = 6;
+
+    /**
+     * The first connection attempt backoff delay to apply, then grow exponentially.
+     */
+    private Duration firstBackoff = Duration.ofMillis(100);
+
+    /**
+     * The maximum connection attempt delay to apply despite exponential growth.
+     */
+    private Duration maxBackoff = Duration.ofSeconds(5);
+
+    public long getMaxRetries() {
+      return maxRetries;
+    }
+
+    public void setMaxRetries(long maxRetries) {
+      this.maxRetries = maxRetries;
+    }
+
+    public Duration getFirstBackoff() {
+      return firstBackoff;
+    }
+
+    public void setFirstBackoff(Duration firstBackoff) {
+      this.firstBackoff = firstBackoff;
+    }
+
+    public Duration getMaxBackoff() {
+      return maxBackoff;
+    }
+
+    public void setMaxBackoff(Duration maxBackoff) {
+      this.maxBackoff = maxBackoff;
+    }
+  }
 
   /**
    * RSocket transport protocol.
@@ -71,30 +152,6 @@ public class PrometheusRSocketClientProperties {
   @DurationUnit(ChronoUnit.SECONDS)
   private Duration timeout = Duration.ofSeconds(5);
 
-  public long getMaxRetries() {
-    return maxRetries;
-  }
-
-  public void setMaxRetries(long maxRetries) {
-    this.maxRetries = maxRetries;
-  }
-
-  public Duration getFirstBackoff() {
-    return firstBackoff;
-  }
-
-  public void setFirstBackoff(Duration firstBackoff) {
-    this.firstBackoff = firstBackoff;
-  }
-
-  public Duration getMaxBackoff() {
-    return maxBackoff;
-  }
-
-  public void setMaxBackoff(Duration maxBackoff) {
-    this.maxBackoff = maxBackoff;
-  }
-
   public String getHost() {
     return host;
   }
@@ -135,6 +192,22 @@ public void setTimeout(Duration timeout) {
     this.timeout = timeout;
   }
 
+  public ReconnectRetry getReconnectRetry() {
+    return reconnectRetry;
+  }
+
+  public void setReconnectRetry(ReconnectRetry reconnectRetry) {
+    this.reconnectRetry = reconnectRetry;
+  }
+
+  public Retry getRetry() {
+    return retry;
+  }
+
+  public void setRetry(Retry retry) {
+    this.retry = retry;
+  }
+
   ClientTransport createClientTransport() {
     final TcpClient tcpClient = TcpClient.create().host(this.host).port(this.port);
     return this.transport.create(this.secure ? tcpClient.secure() : tcpClient);
