optimize request channel to let you peak at the first item in the requestChannel

Author: robertroeser

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

@@ -27,7 +27,10 @@
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
 import reactor.core.Disposable;
-import reactor.core.publisher.*;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.core.publisher.SignalType;
+import reactor.core.publisher.UnicastProcessor;
 
 import java.util.Collections;
 import java.util.Map;
@@ -39,15 +42,17 @@
 import static io.rsocket.frame.FrameHeaderFlyweight.FLAGS_M;
 
 /** Server side RSocket. Receives {@link Frame}s from a {@link RSocketClient} */
-class RSocketServer implements RSocket {
+class RSocketServer implements RequestHandler {
 
   private final DuplexConnection connection;
   private final RSocket requestHandler;
+  private final RequestHandler optimizedRequestHandler;
+  private final boolean hasOptimizedRequestHandler;
   private final Function<Frame, ? extends Payload> frameDecoder;
   private final Consumer<Throwable> errorConsumer;
 
   private final Map<Integer, Subscription> sendingSubscriptions;
-  private final Map<Integer, Processor<Payload,Payload>> channelProcessors;
+  private final Map<Integer, Processor<Payload, Payload>> channelProcessors;
 
   private final UnboundedProcessor<Frame> sendProcessor;
   private KeepAliveHandler keepAliveHandler;
@@ -69,12 +74,22 @@ class RSocketServer implements RSocket {
       Consumer<Throwable> errorConsumer,
       long tickPeriod,
       long ackTimeout) {
+    
+    if (requestHandler instanceof RequestHandler) {
+      this.optimizedRequestHandler = (RequestHandler) requestHandler;
+      this.hasOptimizedRequestHandler = true;
+      this.requestHandler = null;
+    } else {
+      this.hasOptimizedRequestHandler = false;
+      this.requestHandler = requestHandler;
+      this.optimizedRequestHandler = null;
+    }
+
     this.connection = connection;
-    this.requestHandler = requestHandler;
     this.frameDecoder = frameDecoder;
     this.errorConsumer = errorConsumer;
     this.sendingSubscriptions = Collections.synchronizedMap(new IntObjectHashMap<>());
-    this.channelProcessors    = Collections.synchronizedMap(new IntObjectHashMap<>());
+    this.channelProcessors = Collections.synchronizedMap(new IntObjectHashMap<>());
 
     // DO NOT Change the order here. The Send processor must be subscribed to before receiving
     // connections
@@ -116,43 +131,55 @@ class RSocketServer implements RSocket {
   }
 
   private void handleSendProcessorError(Throwable t) {
-    sendingSubscriptions.values().forEach(subscription -> {
-      try {
-        subscription.cancel();
-      } catch (Throwable e) {
-        errorConsumer.accept(e);
-      }
-    });
+    sendingSubscriptions
+        .values()
+        .forEach(
+            subscription -> {
+              try {
+                subscription.cancel();
+              } catch (Throwable e) {
+                errorConsumer.accept(e);
+              }
+            });
 
-    channelProcessors.values().forEach(subscription -> {
-      try {
-        subscription.onError(t);
-      } catch (Throwable e) {
-        errorConsumer.accept(e);
-      }
-    });
+    channelProcessors
+        .values()
+        .forEach(
+            subscription -> {
+              try {
+                subscription.onError(t);
+              } catch (Throwable e) {
+                errorConsumer.accept(e);
+              }
+            });
   }
 
   private void handleSendProcessorCancel(SignalType t) {
     if (SignalType.ON_ERROR == t) {
       return;
     }
 
-    sendingSubscriptions.values().forEach(subscription -> {
-      try {
-        subscription.cancel();
-      } catch (Throwable e) {
-        errorConsumer.accept(e);
-      }
-    });
+    sendingSubscriptions
+        .values()
+        .forEach(
+            subscription -> {
+              try {
+                subscription.cancel();
+              } catch (Throwable e) {
+                errorConsumer.accept(e);
+              }
+            });
 
-    channelProcessors.values().forEach(subscription -> {
-      try {
-        subscription.onComplete();
-      } catch (Throwable e) {
-        errorConsumer.accept(e);
-      }
-    });
+    channelProcessors
+        .values()
+        .forEach(
+            subscription -> {
+              try {
+                subscription.onComplete();
+              } catch (Throwable e) {
+                errorConsumer.accept(e);
+              }
+            });
   }
 
   @Override
@@ -191,6 +218,15 @@ public Flux<Payload> requestChannel(Publisher<Payload> payloads) {
     }
   }
 
+  @Override
+  public Flux<Payload> requestChannel(Payload payload, Publisher<Payload> payloads) {
+    try {
+      return optimizedRequestHandler.requestChannel(payloads);
+    } catch (Throwable t) {
+      return Flux.error(t);
+    }
+  }
+
   @Override
   public Mono<Void> metadataPush(Payload payload) {
     try {
@@ -232,9 +268,7 @@ private synchronized void cleanUpSendingSubscriptions() {
   }
 
   private synchronized void cleanUpChannelProcessors() {
-    channelProcessors
-        .values()
-        .forEach(Processor::onComplete);
+    channelProcessors.values().forEach(Processor::onComplete);
     channelProcessors.clear();
   }
 
@@ -381,7 +415,11 @@ private void handleChannel(int streamId, Payload payload, int initialRequestN) {
     // and any later payload can be processed
     frames.onNext(payload);
 
-    handleStream(streamId, requestChannel(payloads), initialRequestN);
+    if (hasOptimizedRequestHandler) {
+      handleStream(streamId, requestChannel(payload, payloads), initialRequestN);
+    } else {
+      handleStream(streamId, requestChannel(payloads), initialRequestN);
+    }
   }
 
   private void handleKeepAliveFrame(Frame frame) {

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

@@ -0,0 +1,19 @@
+package io.rsocket;
+
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Flux;
+
+/**
+ * Extends the {@link RSocket} that allows an implementor to peek at the first request payload of a channel.
+ */
+public interface RequestHandler extends RSocket {
+  /**
+   * Implement this method to peak at the first payload of the incoming request stream
+   * @param payload First payload in the stream
+   * @param payloads Stream of request payloads.
+   * @return Stream of response payloads.
+   */
+  default Flux<Payload> requestChannel(Payload payload, Publisher<Payload> payloads) {
+    return requestChannel(payloads);
+  }
+}