fixes LoadBalancedRSocketMono to propagate context

Signed-off-by: Oleh Dokuka <[email protected]>

Author: OlegDokuka

rsocket-load-balancer/src/main/java/io/rsocket/client/LoadBalancedRSocketMono.java

@@ -34,14 +34,15 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 import org.reactivestreams.Publisher;
-import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import reactor.core.CoreSubscriber;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.MonoProcessor;
+import reactor.core.publisher.Operators;
+import reactor.util.context.Context;
 import reactor.util.retry.Retry;
 
 /**
@@ -667,26 +668,28 @@ private class WeightedSocket implements LoadBalancerSocketMetrics, RSocket {
     @Override
     public Mono<Payload> requestResponse(Payload payload) {
       return rSocketMono.flatMap(
-          source -> {
-            return Mono.from(
-                subscriber ->
-                    source
-                        .requestResponse(payload)
-                        .subscribe(new LatencySubscriber<>(subscriber, this)));
-          });
+          source ->
+              Mono.from(
+                  subscriber ->
+                      source
+                          .requestResponse(payload)
+                          .subscribe(
+                              new LatencySubscriber<>(
+                                  Operators.toCoreSubscriber(subscriber), this))));
     }
 
     @Override
     public Flux<Payload> requestStream(Payload payload) {
 
       return rSocketMono.flatMapMany(
-          source -> {
-            return Flux.from(
-                subscriber ->
-                    source
-                        .requestStream(payload)
-                        .subscribe(new CountingSubscriber<>(subscriber, this)));
-          });
+          source ->
+              Flux.from(
+                  subscriber ->
+                      source
+                          .requestStream(payload)
+                          .subscribe(
+                              new CountingSubscriber<>(
+                                  Operators.toCoreSubscriber(subscriber), this))));
     }
 
     @Override
@@ -698,7 +701,9 @@ public Mono<Void> fireAndForget(Payload payload) {
                 subscriber ->
                     source
                         .fireAndForget(payload)
-                        .subscribe(new CountingSubscriber<>(subscriber, this)));
+                        .subscribe(
+                            new CountingSubscriber<>(
+                                Operators.toCoreSubscriber(subscriber), this)));
           });
     }
 
@@ -710,21 +715,24 @@ public Mono<Void> metadataPush(Payload payload) {
                 subscriber ->
                     source
                         .metadataPush(payload)
-                        .subscribe(new CountingSubscriber<>(subscriber, this)));
+                        .subscribe(
+                            new CountingSubscriber<>(
+                                Operators.toCoreSubscriber(subscriber), this)));
           });
     }
 
     @Override
     public Flux<Payload> requestChannel(Publisher<Payload> payloads) {
 
       return rSocketMono.flatMapMany(
-          source -> {
-            return Flux.from(
-                subscriber ->
-                    source
-                        .requestChannel(payloads)
-                        .subscribe(new CountingSubscriber<>(subscriber, this)));
-          });
+          source ->
+              Flux.from(
+                  subscriber ->
+                      source
+                          .requestChannel(payloads)
+                          .subscribe(
+                              new CountingSubscriber<>(
+                                  Operators.toCoreSubscriber(subscriber), this))));
     }
 
     synchronized double getPredictedLatency() {
@@ -867,18 +875,23 @@ public long lastTimeUsedMillis() {
      * Subscriber wrapper used for request/response interaction model, measure and collect latency
      * information.
      */
-    private class LatencySubscriber<U> implements Subscriber<U> {
-      private final Subscriber<U> child;
+    private class LatencySubscriber<U> implements CoreSubscriber<U> {
+      private final CoreSubscriber<U> child;
       private final WeightedSocket socket;
       private final AtomicBoolean done;
       private long start;
 
-      LatencySubscriber(Subscriber<U> child, WeightedSocket socket) {
+      LatencySubscriber(CoreSubscriber<U> child, WeightedSocket socket) {
         this.child = child;
         this.socket = socket;
         this.done = new AtomicBoolean(false);
       }
 
+      @Override
+      public Context currentContext() {
+        return child.currentContext();
+      }
+
       @Override
       public void onSubscribe(Subscription s) {
         start = incr();
@@ -931,15 +944,20 @@ public void onComplete() {
      * Subscriber wrapper used for stream like interaction model, it only counts the number of
      * active streams
      */
-    private class CountingSubscriber<U> implements Subscriber<U> {
-      private final Subscriber<U> child;
+    private class CountingSubscriber<U> implements CoreSubscriber<U> {
+      private final CoreSubscriber<U> child;
       private final WeightedSocket socket;
 
-      CountingSubscriber(Subscriber<U> child, WeightedSocket socket) {
+      CountingSubscriber(CoreSubscriber<U> child, WeightedSocket socket) {
         this.child = child;
         this.socket = socket;
       }
 
+      @Override
+      public Context currentContext() {
+        return child.currentContext();
+      }
+
       @Override
       public void onSubscribe(Subscription s) {
         socket.pendingStreams.incrementAndGet();