more fixes

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

Author: OlegDokuka

rsocket-core/src/main/java/io/rsocket/core/RSocketRequester.java

@@ -40,8 +40,6 @@
 import io.rsocket.frame.RequestResponseFrameFlyweight;
 import io.rsocket.frame.RequestStreamFrameFlyweight;
 import io.rsocket.frame.decoder.PayloadDecoder;
-import io.rsocket.internal.FluxSwitchOnFirst;
-import io.rsocket.internal.RateLimitableRequestSubscriber;
 import io.rsocket.internal.SynchronizedIntObjectHashMap;
 import io.rsocket.internal.UnboundedProcessor;
 import io.rsocket.internal.UnicastMonoEmpty;
@@ -63,6 +61,7 @@
 import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
+import reactor.core.publisher.BaseSubscriber;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.SignalType;
@@ -86,7 +85,7 @@ class RSocketRequester implements RSocket {
   private final PayloadDecoder payloadDecoder;
   private final Consumer<Throwable> errorConsumer;
   private final StreamIdSupplier streamIdSupplier;
-  private final IntObjectMap<RateLimitableRequestSubscriber> senders;
+  private final IntObjectMap<Subscription> senders;
   private final IntObjectMap<Processor<Payload, Payload>> receivers;
   private final UnboundedProcessor<ByteBuf> sendProcessor;
   private final RequesterLeaseHandler leaseHandler;
@@ -265,7 +264,6 @@ private Flux<Payload> handleRequestStream(final Payload payload) {
     receivers.put(streamId, receiver);
 
     return receiver
-        .log()
         .doOnRequest(
             new LongConsumer() {
 
@@ -315,29 +313,31 @@ private Flux<Payload> handleChannel(Flux<Payload> request) {
       return Flux.error(err);
     }
 
-    return request.transform(
-        f ->
-            new FluxSwitchOnFirst<>(
-                f,
-                (s, flux) -> {
-                  Payload payload = s.get();
-                  if (payload != null) {
-                    return handleChannel(flux, payload);
-                  } else {
-                    return flux;
-                  }
-                },
-                false));
+    return request.switchOnFirst(
+        (s, flux) -> {
+          Payload payload = s.get();
+          if (payload != null) {
+            return handleChannel(payload, flux.skip(1));
+          } else {
+            return flux;
+          }
+        },
+        false);
   }
 
-  private Flux<? extends Payload> handleChannel(Flux<Payload> inboundFlux, Payload initialPayload) {
+  private Flux<? extends Payload> handleChannel(Payload initialPayload, Flux<Payload> inboundFlux) {
     final UnboundedProcessor<ByteBuf> sendProcessor = this.sendProcessor;
     final AtomicBoolean payloadReleasedFlag = new AtomicBoolean(false);
     final int streamId = streamIdSupplier.nextStreamId(receivers);
 
     final UnicastProcessor<Payload> receiver = UnicastProcessor.create();
-    final RateLimitableRequestSubscriber<Payload> upstreamSubscriber =
-        new RateLimitableRequestSubscriber<Payload>(Queues.SMALL_BUFFER_SIZE) {
+    final BaseSubscriber<Payload> upstreamSubscriber =
+        new BaseSubscriber<Payload>() {
+
+          @Override
+          protected void hookOnSubscribe(Subscription subscription) {
+            // noops
+          }
 
           @Override
           protected void hookOnNext(Payload payload) {
@@ -358,7 +358,7 @@ protected void hookOnComplete() {
           protected void hookOnError(Throwable t) {
             ByteBuf frame = ErrorFrameFlyweight.encode(allocator, streamId, t);
             sendProcessor.onNext(frame);
-            receiver.dispose();
+            receiver.onError(t);
           }
 
           @Override
@@ -380,21 +380,20 @@ public void accept(long n) {
                   senders.put(streamId, upstreamSubscriber);
                   receivers.put(streamId, receiver);
 
-                  inboundFlux.subscribe(upstreamSubscriber);
-
-                  ByteBuf frame =
-                      RequestChannelFrameFlyweight.encode(
-                          allocator,
-                          streamId,
-                          false,
-                          false,
-                          n,
-                          initialPayload.sliceMetadata().retain(),
-                          initialPayload.sliceData().retain());
-
-                  sendProcessor.onNext(frame);
-
+                  inboundFlux.limitRate(Queues.SMALL_BUFFER_SIZE).subscribe(upstreamSubscriber);
                   if (!payloadReleasedFlag.getAndSet(true)) {
+                    ByteBuf frame =
+                        RequestChannelFrameFlyweight.encode(
+                            allocator,
+                            streamId,
+                            false,
+                            false,
+                            n,
+                            initialPayload.sliceMetadata().retain(),
+                            initialPayload.sliceData().retain());
+
+                    sendProcessor.onNext(frame);
+
                     initialPayload.release();
                   }
                 } else {
@@ -408,16 +407,15 @@ public void accept(long n) {
                 upstreamSubscriber.cancel();
               }
             })
+        .doOnComplete(() -> receivers.remove(streamId, receiver))
         .doOnCancel(
             () -> {
               if (!payloadReleasedFlag.getAndSet(true)) {
                 initialPayload.release();
               }
-              if (contains(streamId)) {
+              if (receivers.remove(streamId, receiver)) {
                 sendProcessor.onNext(CancelFrameFlyweight.encode(allocator, streamId));
-                if (receivers.remove(streamId, receiver)) {
-                  upstreamSubscriber.cancel();
-                }
+                upstreamSubscriber.cancel();
               }
             });
   }

rsocket-core/src/main/java/io/rsocket/core/RSocketResponder.java

@@ -27,7 +27,6 @@
 import io.rsocket.exceptions.ApplicationErrorException;
 import io.rsocket.frame.*;
 import io.rsocket.frame.decoder.PayloadDecoder;
-import io.rsocket.internal.RateLimitableRequestSubscriber;
 import io.rsocket.internal.SynchronizedIntObjectHashMap;
 import io.rsocket.internal.UnboundedProcessor;
 import io.rsocket.lease.ResponderLeaseHandler;
@@ -362,16 +361,10 @@ protected void hookFinally(SignalType type) {
   }
 
   private void handleRequestResponse(int streamId, Mono<Payload> response) {
-    response.subscribe(
+    final BaseSubscriber<Payload> subscriber =
         new BaseSubscriber<Payload>() {
           private boolean isEmpty = true;
 
-          @Override
-          protected void hookOnSubscribe(Subscription subscription) {
-            sendingSubscriptions.put(streamId, subscription);
-            subscription.request(Long.MAX_VALUE);
-          }
-
           @Override
           protected void hookOnNext(Payload payload) {
             if (isEmpty) {
@@ -405,18 +398,20 @@ protected void hookOnComplete() {
 
           @Override
           protected void hookFinally(SignalType type) {
-            sendingSubscriptions.remove(streamId);
+            sendingSubscriptions.remove(streamId, this);
           }
-        });
+        };
+
+    sendingSubscriptions.put(streamId, subscriber);
+    response.subscribe(subscriber);
   }
 
   private void handleStream(int streamId, Flux<Payload> response, int initialRequestN) {
-    response.subscribe(
-        new RateLimitableRequestSubscriber<Payload>(Queues.SMALL_BUFFER_SIZE) {
+    final BaseSubscriber<Payload> subscriber =
+        new BaseSubscriber<Payload>() {
 
           @Override
           protected void hookOnSubscribe(Subscription s) {
-            sendingSubscriptions.put(streamId, s);
             s.request(initialRequestN >= Integer.MAX_VALUE ? Long.MAX_VALUE : initialRequestN);
           }
 
@@ -449,7 +444,10 @@ protected void hookOnError(Throwable throwable) {
           protected void hookFinally(SignalType type) {
             sendingSubscriptions.remove(streamId);
           }
-        });
+        };
+
+    sendingSubscriptions.put(streamId, subscriber);
+    response.limitRate(Queues.SMALL_BUFFER_SIZE).subscribe(subscriber);
   }
 
   private void handleChannel(int streamId, Payload payload, int initialRequestN) {
@@ -470,7 +468,7 @@ public void accept(long l) {
                     long n;
                     if (first) {
                       first = false;
-                      n = l - 1;
+                      n = l - 1L;
                     } else {
                       n = l;
                     }

rsocket-core/src/main/java/io/rsocket/frame/FrameUtil.java

@@ -22,6 +22,16 @@ public static String toString(ByteBuf frame) {
         .append(Integer.toBinaryString(FrameHeaderFlyweight.flags(frame)))
         .append(" Length: " + frame.readableBytes());
 
+    if (frameType.hasInitialRequestN()) {
+      payload
+          .append(" InitialRequestN: ")
+          .append(RequestStreamFrameFlyweight.initialRequestN(frame));
+    }
+
+    if (frameType == FrameType.REQUEST_N) {
+      payload.append(" RequestN: ").append(RequestNFrameFlyweight.requestN(frame));
+    }
+
     if (FrameHeaderFlyweight.hasMetadata(frame)) {
       payload.append("\nMetadata:\n");