rsocket · mostroverkhov · Jun 17, 2019 · Jun 4, 2019 · Jun 12, 2019 · Jun 13, 2019
diff --git a/gradle.properties b/gradle.properties
@@ -11,4 +11,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-version=1.0.0-RC1-SNAPSHOT
+version=1.0.0-RC1-LEASE-SNAPSHOT
diff --git a/rsocket-core/src/main/java/io/rsocket/RSocketFactory.java b/rsocket-core/src/main/java/io/rsocket/RSocketFactory.java
@@ -31,6 +31,7 @@
 import io.rsocket.internal.ClientSetup;
 import io.rsocket.internal.ServerSetup;
 import io.rsocket.keepalive.KeepAliveHandler;
+import io.rsocket.lease.*;
 import io.rsocket.plugins.DuplexConnectionInterceptor;
 import io.rsocket.plugins.PluginRegistry;
 import io.rsocket.plugins.Plugins;
@@ -40,12 +41,10 @@
 import io.rsocket.transport.ServerTransport;
 import io.rsocket.util.ConnectionUtils;
 import io.rsocket.util.EmptyPayload;
+import io.rsocket.util.MultiSubscriberRSocket;
 import java.time.Duration;
 import java.util.Objects;
-import java.util.function.BiFunction;
-import java.util.function.Consumer;
-import java.util.function.Function;
-import java.util.function.Supplier;
+import java.util.function.*;
 import reactor.core.publisher.Mono;
 
 /** Factory for creating RSocket clients and servers. */
@@ -92,6 +91,8 @@ default <T extends Closeable> Start<T> transport(ServerTransport<T> transport) {
   }
 
   public static class ClientRSocketFactory implements ClientTransportAcceptor {
+    private static final String CLIENT_TAG = "client";
+
     private Supplier<Function<RSocket, RSocket>> acceptor =
         () -> rSocket -> new AbstractRSocket() {};
 
@@ -115,13 +116,17 @@ public static class ClientRSocketFactory implements ClientTransportAcceptor {
     private boolean resumeCleanupStoreOnKeepAlive;
     private Supplier<ByteBuf> resumeTokenSupplier = ResumeFrameFlyweight::generateResumeToken;
     private Function<? super ByteBuf, ? extends ResumableFramesStore> resumeStoreFactory =
-        token -> new InMemoryResumableFramesStore("client", 100_000);
+        token -> new InMemoryResumableFramesStore(CLIENT_TAG, 100_000);
     private Duration resumeSessionDuration = Duration.ofMinutes(2);
     private Duration resumeStreamTimeout = Duration.ofSeconds(10);
     private Supplier<ResumeStrategy> resumeStrategySupplier =
         () ->
             new ExponentialBackoffResumeStrategy(Duration.ofSeconds(1), Duration.ofSeconds(16), 2);
 
+    private boolean multiSubscriberRequester = true;
+    private boolean leaseEnabled;
+    private Supplier<Leases<?>> leasesSupplier = Leases::new;
+
     private ByteBufAllocator allocator = ByteBufAllocator.DEFAULT;
 
     public ClientRSocketFactory byteBufAllocator(ByteBufAllocator allocator) {
@@ -205,6 +210,22 @@ public ClientRSocketFactory metadataMimeType(String metadataMimeType) {
       return this;
     }
 
+    public ClientRSocketFactory lease(Supplier<Leases<? extends LeaseStats>> leasesSupplier) {
+      this.leaseEnabled = true;
+      this.leasesSupplier = Objects.requireNonNull(leasesSupplier);
+      return this;
+    }
+
+    public ClientRSocketFactory lease() {
+      this.leaseEnabled = true;
+      return this;
+    }
+
+    public ClientRSocketFactory singleSubscriberRequester() {
+      this.multiSubscriberRequester = false;
+      return this;
+    }
+
     public ClientRSocketFactory resume() {
       this.resumeEnabled = true;
       return this;
@@ -302,7 +323,14 @@ public Mono<RSocket> start() {
                   ClientServerInputMultiplexer multiplexer =
                       new ClientServerInputMultiplexer(wrappedConnection, plugins, true);
 
-                  RSocketRequester rSocketRequester =
+                  boolean isLeaseEnabled = leaseEnabled;
+                  Leases<?> leases = leasesSupplier.get();
+                  RequesterLeaseHandler requesterLeaseHandler =
+                      isLeaseEnabled
+                          ? new RequesterLeaseHandler.Impl(CLIENT_TAG, leases.receiver())
+                          : RequesterLeaseHandler.None;
+
+                  RSocket rSocketRequester =
                       new RSocketRequester(
                           allocator,
                           multiplexer.asClientConnection(),
@@ -311,12 +339,17 @@ public Mono<RSocket> start() {
                           StreamIdSupplier.clientSupplier(),
                           keepAliveTickPeriod(),
                           keepAliveTimeout(),
-                          keepAliveHandler);
+                          keepAliveHandler,
+                          requesterLeaseHandler);
+
+                  if (multiSubscriberRequester) {
+                    rSocketRequester = new MultiSubscriberRSocket(rSocketRequester);
+                  }
 
                   ByteBuf setupFrame =
                       SetupFrameFlyweight.encode(
                           allocator,
-                          false,
+                          isLeaseEnabled,
                           keepAliveTickPeriod(),
                           keepAliveTimeout(),
                           resumeToken,
@@ -337,13 +370,20 @@ public Mono<RSocket> start() {
 
                   RSocket wrappedRSocketHandler = plugins.applyResponder(rSocketHandler);
 
-                  RSocketResponder rSocketResponder =
+                  ResponderLeaseHandler responderLeaseHandler =
+                      isLeaseEnabled
+                          ? new ResponderLeaseHandler.Impl<>(
+                              CLIENT_TAG, allocator, leases.sender(), errorConsumer, leases.stats())
+                          : ResponderLeaseHandler.None;
+
+                  RSocket rSocketResponder =
                       new RSocketResponder(
                           allocator,
                           multiplexer.asServerConnection(),
                           wrappedRSocketHandler,
                           payloadDecoder,
-                          errorConsumer);
+                          errorConsumer,
+                          responderLeaseHandler);
 
                   return wrappedConnection.sendOne(setupFrame).thenReturn(wrappedRSocketRequester);
                 });
@@ -382,16 +422,23 @@ private Mono<DuplexConnection> newConnection() {
   }
 
   public static class ServerRSocketFactory {
+    private static final String SERVER_TAG = "server";
+
     private SocketAcceptor acceptor;
     private PayloadDecoder payloadDecoder = PayloadDecoder.DEFAULT;
     private Consumer<Throwable> errorConsumer = Throwable::printStackTrace;
     private int mtu = 0;
     private PluginRegistry plugins = new PluginRegistry(Plugins.defaultPlugins());
+
     private boolean resumeSupported;
     private Duration resumeSessionDuration = Duration.ofSeconds(120);
     private Duration resumeStreamTimeout = Duration.ofSeconds(10);
     private Function<? super ByteBuf, ? extends ResumableFramesStore> resumeStoreFactory =
-        token -> new InMemoryResumableFramesStore("server", 100_000);
+        token -> new InMemoryResumableFramesStore(SERVER_TAG, 100_000);
+
+    private boolean multiSubscriberRequester = true;
+    private boolean leaseEnabled;
+    private Supplier<Leases<?>> leasesSupplier = Leases::new;
 
     private ByteBufAllocator allocator = ByteBufAllocator.DEFAULT;
     private boolean resumeCleanupStoreOnKeepAlive;
@@ -450,6 +497,22 @@ public ServerRSocketFactory errorConsumer(Consumer<Throwable> errorConsumer) {
       return this;
     }
 
+    public ServerRSocketFactory lease(Supplier<Leases<?>> leasesSupplier) {
+      this.leaseEnabled = true;
+      this.leasesSupplier = Objects.requireNonNull(leasesSupplier);
+      return this;
+    }
+
+    public ServerRSocketFactory lease() {
+      this.leaseEnabled = true;
+      return this;
+    }
+
+    public ServerRSocketFactory singleSubscriberRequester() {
+      this.multiSubscriberRequester = false;
+      return this;
+    }
+
     public ServerRSocketFactory resume() {
       this.resumeSupported = true;
       return this;
@@ -541,13 +604,31 @@ private Mono<Void> acceptSetup(
                     multiplexer.dispose();
                   });
         }
+
+        boolean isLeaseEnabled = leaseEnabled;
+
+        if (SetupFrameFlyweight.honorLease(setupFrame) && !isLeaseEnabled) {
+          return sendError(multiplexer, new InvalidSetupException("lease is not supported"))
+              .doFinally(
+                  signalType -> {
+                    setupFrame.release();
+                    multiplexer.dispose();
+                  });
+        }
+
         return serverSetup.acceptRSocketSetup(
             setupFrame,
             multiplexer,
             (keepAliveHandler, wrappedMultiplexer) -> {
               ConnectionSetupPayload setupPayload = ConnectionSetupPayload.create(setupFrame);
 
-              RSocketRequester rSocketRequester =
+              Leases<?> leases = leasesSupplier.get();
+              RequesterLeaseHandler requesterLeaseHandler =
+                  isLeaseEnabled
+                      ? new RequesterLeaseHandler.Impl(SERVER_TAG, leases.receiver())
+                      : RequesterLeaseHandler.None;
+
+              RSocket rSocketRequester =
                   new RSocketRequester(
                       allocator,
                       wrappedMultiplexer.asServerConnection(),
@@ -556,8 +637,12 @@ private Mono<Void> acceptSetup(
                       StreamIdSupplier.serverSupplier(),
                       setupPayload.keepAliveInterval(),
                       setupPayload.keepAliveMaxLifetime(),
-                      keepAliveHandler);
+                      keepAliveHandler,
+                      requesterLeaseHandler);
 
+              if (multiSubscriberRequester) {
+                rSocketRequester = new MultiSubscriberRSocket(rSocketRequester);
+              }
               RSocket wrappedRSocketRequester = plugins.applyRequester(rSocketRequester);
 
               return acceptor
@@ -568,13 +653,24 @@ private Mono<Void> acceptSetup(
                       rSocketHandler -> {
                         RSocket wrappedRSocketHandler = plugins.applyResponder(rSocketHandler);
 
-                        RSocketResponder rSocketResponder =
+                        ResponderLeaseHandler responderLeaseHandler =
+                            isLeaseEnabled
+                                ? new ResponderLeaseHandler.Impl<>(
+                                    SERVER_TAG,
+                                    allocator,
+                                    leases.sender(),
+                                    errorConsumer,
+                                    leases.stats())
+                                : ResponderLeaseHandler.None;
+
+                        RSocket rSocketResponder =
                             new RSocketResponder(
                                 allocator,
                                 wrappedMultiplexer.asClientConnection(),
                                 wrappedRSocketHandler,
                                 payloadDecoder,
-                                errorConsumer);
+                                errorConsumer,
+                                responderLeaseHandler);
                       })
                   .doFinally(signalType -> setupPayload.release())
                   .then();