RSocketServer simplifications

Author: rdegnan

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

@@ -228,14 +228,14 @@ private void handleFrame(Frame frame) {
           handleStream(streamId, requestStream(frameDecoder.apply(frame)), initialRequestN(frame));
           break;
         case REQUEST_CHANNEL:
-          handleChannel(streamId, frame);
-          break;
-        case PAYLOAD:
-          // TODO: Hook in receiving socket.
+          handleChannel(streamId, frameDecoder.apply(frame), initialRequestN(frame));
           break;
         case METADATA_PUSH:
           metadataPush(frameDecoder.apply(frame));
           break;
+        case PAYLOAD:
+          // TODO: Hook in receiving socket.
+          break;
         case LEASE:
           // Lease must not be received here as this is the server end of the socket which sends
           // leases.
@@ -282,18 +282,17 @@ private void handleFrame(Frame frame) {
 
   private void handleFireAndForget(int streamId, Mono<Void> result) {
     result
+        .doOnSubscribe(subscription -> sendingSubscriptions.put(streamId, subscription))
         .doFinally(signalType -> sendingSubscriptions.remove(streamId))
-        .subscribe(
-            null,
-            errorConsumer,
-            null,
-            subscription -> sendingSubscriptions.put(streamId, subscription));
+        .subscribe(null, errorConsumer);
   }
 
   private void handleRequestResponse(int streamId, Mono<Payload> response) {
     response
+        .switchIfEmpty(Mono.fromCallable(() -> Frame.PayloadFrame.from(streamId, FrameType.COMPLETE)))
         .doOnSubscribe(subscription -> sendingSubscriptions.put(streamId, subscription))
-        .map(
+        .doFinally(signalType -> sendingSubscriptions.remove(streamId))
+        .subscribe(
             payload -> {
               int flags = FLAGS_C;
               if (payload.hasMetadata()) {
@@ -302,34 +301,35 @@ private void handleRequestResponse(int streamId, Mono<Payload> response) {
               final Frame frame =
                   Frame.PayloadFrame.from(streamId, FrameType.NEXT_COMPLETE, payload, flags);
               payload.release();
-              return frame;
-            })
-        .switchIfEmpty(Mono.fromCallable(() -> Frame.PayloadFrame.from(streamId, FrameType.COMPLETE)))
-        .doFinally(signalType -> sendingSubscriptions.remove(streamId))
-        .subscribe(sendProcessor::onNext, t -> handleError(streamId, t));
+              sendProcessor.onNext(frame);
+            },
+            t -> handleError(streamId, t));
   }
 
   private void handleStream(int streamId, Flux<Payload> response, int initialRequestN) {
     response
-        .map(
-            payload -> {
-              final Frame frame = Frame.PayloadFrame.from(streamId, FrameType.NEXT, payload);
-              payload.release();
-              return frame;
-            })
-        .concatWith(Mono.fromCallable(() -> Frame.PayloadFrame.from(streamId, FrameType.COMPLETE)))
         .transform(
             frameFlux -> {
-              LimitableRequestPublisher<Frame> frames = LimitableRequestPublisher.wrap(frameFlux);
-              sendingSubscriptions.put(streamId, frames);
-              frames.increaseRequestLimit(initialRequestN);
-              return frames;
+              LimitableRequestPublisher<Payload> payloads = LimitableRequestPublisher.wrap(frameFlux);
+              sendingSubscriptions.put(streamId, payloads);
+              payloads.increaseRequestLimit(initialRequestN);
+              return payloads;
             })
         .doFinally(signalType -> sendingSubscriptions.remove(streamId))
-        .subscribe(sendProcessor::onNext, t -> handleError(streamId, t));
+        .subscribe(
+            payload -> {
+              final Frame frame = Frame.PayloadFrame.from(streamId, FrameType.NEXT, payload);
+              payload.release();
+              sendProcessor.onNext(frame);
+            },
+            t -> handleError(streamId, t),
+            () -> {
+              final Frame frame = Frame.PayloadFrame.from(streamId, FrameType.COMPLETE);
+              sendProcessor.onNext(frame);
+            });
   }
 
-  private void handleChannel(int streamId, Frame firstFrame) {
+  private void handleChannel(int streamId, Payload payload, int initialRequestN) {
     UnicastProcessor<Payload> frames = UnicastProcessor.create();
     channelProcessors.put(streamId, frames);
 
@@ -343,9 +343,9 @@ private void handleChannel(int streamId, Frame firstFrame) {
     // not chained, as the payload should be enqueued in the Unicast processor before this method
     // returns
     // and any later payload can be processed
-    frames.onNext(frameDecoder.apply(firstFrame));
+    frames.onNext(payload);
 
-    handleStream(streamId, requestChannel(payloads), initialRequestN(firstFrame));
+    handleStream(streamId, requestChannel(payloads), initialRequestN);
   }
 
   private void handleKeepAliveFrame(Frame frame) {