Overlapping leases and redistribution on new connects (#188)

#### Problem

This PR addresses two problems:

- Typically users of `FairLeaseDistributor` would supply the same value for lease TTL and the period in which leases are distributed. In such cases, there would be a time period when new lease has not arrived and old lease is expired. This isn't a good reflection of server's intent as the server can handle the new requests but the lease has not yet reached the client.
- New clients have to wait for the next distribution tick to get leases.

#### Modified

- Modified `FairLeaseDistributor` to increase user supplied lease TTL by 10% so that leases overlap. A server can still reject requests if it is overloaded.
- Redistribute leases whenever a new client connects.

#### Result

Better leasing behavior as seen by the clients.

Author: NiteshKant

reactivesocket-core/src/main/java/io/reactivesocket/lease/DefaultLeaseEnforcingSocket.java

@@ -18,6 +18,8 @@
 
 import io.reactivesocket.ReactiveSocket;
 import io.reactivesocket.reactivestreams.extensions.Px;
+import io.reactivesocket.exceptions.RejectedException;
+import io.reactivesocket.reactivestreams.extensions.Px;
 import io.reactivesocket.reactivestreams.extensions.internal.Cancellable;
 import io.reactivesocket.reactivestreams.extensions.internal.subscribers.Subscribers;
 import org.reactivestreams.Publisher;
@@ -30,11 +32,23 @@ public class DefaultLeaseEnforcingSocket extends DefaultLeaseHonoringSocket impl
     private final LeaseDistributor leaseDistributor;
     private volatile Consumer<Lease> leaseSender;
     private Cancellable distributorCancellation;
+    @SuppressWarnings("rawtypes")
+    private final Px rejectError;
 
     public DefaultLeaseEnforcingSocket(ReactiveSocket delegate, LeaseDistributor leaseDistributor,
-                                       LongSupplier currentTimeSupplier) {
+                                       LongSupplier currentTimeSupplier, boolean clientHonorsLeases) {
         super(delegate, currentTimeSupplier);
         this.leaseDistributor = leaseDistributor;
+        if (!clientHonorsLeases) {
+            rejectError = Px.error(new RejectedException("Server overloaded."));
+        } else {
+            rejectError = null;
+        }
+    }
+
+    public DefaultLeaseEnforcingSocket(ReactiveSocket delegate, LeaseDistributor leaseDistributor,
+                                       LongSupplier currentTimeSupplier) {
+        this(delegate, leaseDistributor, currentTimeSupplier, true);
     }
 
     public DefaultLeaseEnforcingSocket(ReactiveSocket delegate, LeaseDistributor leaseDistributor) {
@@ -66,6 +80,11 @@ public Publisher<Void> close() {
                  });
     }
 
+    @Override
+    protected <T> Publisher<T> rejectError() {
+        return null == rejectError ? super.rejectError() : rejectError;
+    }
+
     /**
      * A distributor of leases for an instance of {@link LeaseEnforcingSocket}.
      */

reactivesocket-core/src/main/java/io/reactivesocket/lease/DefaultLeaseHonoringSocket.java

@@ -21,19 +21,25 @@
 import io.reactivesocket.exceptions.RejectedException;
 import io.reactivesocket.reactivestreams.extensions.Px;
 import org.reactivestreams.Publisher;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.LongSupplier;
 
 public class DefaultLeaseHonoringSocket implements LeaseHonoringSocket {
 
+    private static final Logger logger = LoggerFactory.getLogger(DefaultLeaseHonoringSocket.class);
+
     private volatile Lease currentLease;
     private final ReactiveSocket delegate;
     private final LongSupplier currentTimeSupplier;
     private final AtomicInteger remainingQuota;
 
     @SuppressWarnings("ThrowableInstanceNeverThrown")
     private static final RejectedException rejectedException = new RejectedException("Lease exhausted.");
+    @SuppressWarnings("rawtypes")
+    private static final Px rejectedPx = Px.error(rejectedException);
 
     public DefaultLeaseHonoringSocket(ReactiveSocket delegate, LongSupplier currentTimeSupplier) {
         this.delegate = delegate;
@@ -55,7 +61,7 @@ public void accept(Lease lease) {
     public Publisher<Void> fireAndForget(Payload payload) {
         return Px.defer(() -> {
             if (!checkLease()) {
-                return Px.error(rejectedException);
+                return rejectError();
             }
             return delegate.fireAndForget(payload);
         });
@@ -65,7 +71,7 @@ public Publisher<Void> fireAndForget(Payload payload) {
     public Publisher<Payload> requestResponse(Payload payload) {
         return Px.defer(() -> {
             if (!checkLease()) {
-                return Px.error(rejectedException);
+                return rejectError();
             }
             return delegate.requestResponse(payload);
         });
@@ -75,7 +81,7 @@ public Publisher<Payload> requestResponse(Payload payload) {
     public Publisher<Payload> requestStream(Payload payload) {
         return Px.defer(() -> {
             if (!checkLease()) {
-                return Px.error(rejectedException);
+                return rejectError();
             }
             return delegate.requestStream(payload);
         });
@@ -85,7 +91,7 @@ public Publisher<Payload> requestStream(Payload payload) {
     public Publisher<Payload> requestSubscription(Payload payload) {
         return Px.defer(() -> {
             if (!checkLease()) {
-                return Px.error(rejectedException);
+                return rejectError();
             }
             return delegate.requestSubscription(payload);
         });
@@ -95,7 +101,7 @@ public Publisher<Payload> requestSubscription(Payload payload) {
     public Publisher<Payload> requestChannel(Publisher<Payload> payloads) {
         return Px.defer(() -> {
             if (!checkLease()) {
-                return Px.error(rejectedException);
+                return rejectError();
             }
             return delegate.requestChannel(payloads);
         });
@@ -105,7 +111,7 @@ public Publisher<Payload> requestChannel(Publisher<Payload> payloads) {
     public Publisher<Void> metadataPush(Payload payload) {
         return Px.defer(() -> {
             if (!checkLease()) {
-                return Px.error(rejectedException);
+                return rejectError();
             }
             return delegate.metadataPush(payload);
         });
@@ -126,7 +132,20 @@ public Publisher<Void> onClose() {
         return delegate.onClose();
     }
 
+    @SuppressWarnings("unchecked")
+    protected  <T> Publisher<T> rejectError() {
+        return rejectedPx;
+    }
+
     private boolean checkLease() {
-        return remainingQuota.getAndDecrement() > 0 && !currentLease.isExpired(currentTimeSupplier.getAsLong());
+        boolean allow = remainingQuota.getAndDecrement() > 0 && !currentLease.isExpired(currentTimeSupplier.getAsLong());
+        if (!allow) {
+            if (logger.isDebugEnabled()) {
+                logger.debug("Lease expired. Lease: " + currentLease + ", remaining quota: "
+                             + Math.max(0, remainingQuota.get()) + ", current time (ms) "
+                             + currentTimeSupplier.getAsLong());
+            }
+        }
+        return allow;
     }
 }

reactivesocket-core/src/main/java/io/reactivesocket/lease/FairLeaseDistributor.java

@@ -38,17 +38,31 @@ public final class FairLeaseDistributor implements DefaultLeaseEnforcingSocket.L
     private Subscription ticksSubscription;
     private volatile boolean startTicks;
     private final IntSupplier capacitySupplier;
-    private final int leaseTTL;
+    private final int leaseTTLMillis;
     private final Publisher<Long> leaseDistributionTicks;
-    private volatile int remainingPermits;
+    private final boolean redistributeOnConnect;
 
-    public FairLeaseDistributor(IntSupplier capacitySupplier, int leaseTTL, Publisher<Long> leaseDistributionTicks) {
+    public FairLeaseDistributor(IntSupplier capacitySupplier, int leaseTTLMillis,
+                                Publisher<Long> leaseDistributionTicks, boolean redistributeOnConnect) {
         this.capacitySupplier = capacitySupplier;
-        this.leaseTTL = leaseTTL;
+        /*
+         * If lease TTL is exactly the same as the period of replenishment, then there would be a time period when new
+         * lease has not arrived and old lease is expired. This isn't a good reflection of server's intent as the server
+         * can handle the new requests but the lease has not yet reached the client. So, having TTL slightly more
+         * than distribution period (accomodating for network lag) is more representative of server's intent. OTOH, if
+         * server isn't ready, it can always reject a request.
+         */
+        this.leaseTTLMillis = (int) (leaseTTLMillis * 1.1);
         this.leaseDistributionTicks = leaseDistributionTicks;
+        this.redistributeOnConnect = redistributeOnConnect;
         activeRecipients = new LinkedBlockingQueue<>();
     }
 
+    public FairLeaseDistributor(IntSupplier capacitySupplier, int leaseTTLMillis,
+                                Publisher<Long> leaseDistributionTicks) {
+        this(capacitySupplier, leaseTTLMillis, leaseDistributionTicks, true);
+    }
+
     /**
      * Shutdown this distributor. No more leases will be provided to the registered sockets.
      */
@@ -68,12 +82,23 @@ public void shutdown() {
     @Override
     public Cancellable registerSocket(Consumer<Lease> leaseConsumer) {
         activeRecipients.add(leaseConsumer);
+        boolean _started;
         synchronized (this) {
+            _started = startTicks;
             if (!startTicks) {
                 startTicks();
                 startTicks = true;
             }
         }
+
+        if (_started && redistributeOnConnect) {
+            /*
+             * This is a way to make sure that the clients that arrive in the middle of a distribution period, do not
+             * have to wait for the next tick to arrive.
+             */
+            distribute(capacitySupplier.getAsInt());
+        }
+
         return new CancellableImpl() {
             @Override
             protected void onCancel() {
@@ -86,20 +111,19 @@ private void distribute(int permits) {
         if (activeRecipients.isEmpty()) {
             return;
         }
-        remainingPermits -= permits;
         int recipients = activeRecipients.size();
         int budget = permits / recipients;
 
         // it would be more fair to randomized the distribution of extra
         int extra = permits - budget * recipients;
-        Lease budgetLease = new LeaseImpl(budget, leaseTTL, Frame.NULL_BYTEBUFFER);
+        Lease budgetLease = new LeaseImpl(budget, leaseTTLMillis, Frame.NULL_BYTEBUFFER);
         for (Consumer<Lease> recipient: activeRecipients) {
             Lease leaseToSend = budgetLease;
             int n = budget;
             if (extra > 0) {
                 n += 1;
                 extra -= 1;
-                leaseToSend = new LeaseImpl(n, leaseTTL, Frame.NULL_BYTEBUFFER);
+                leaseToSend = new LeaseImpl(n, leaseTTLMillis, Frame.NULL_BYTEBUFFER);
             }
             recipient.accept(leaseToSend);
         }
@@ -109,8 +133,7 @@ private void startTicks() {
         Px.from(leaseDistributionTicks)
           .doOnSubscribe(subscription -> ticksSubscription = subscription)
           .doOnNext(aLong -> {
-              remainingPermits = capacitySupplier.getAsInt();
-              distribute(remainingPermits);
+              distribute(capacitySupplier.getAsInt());
           })
           .ignore()
           .subscribe();

reactivesocket-core/src/test/java/io/reactivesocket/lease/FairLeaseDistributorTest.java

@@ -42,7 +42,7 @@ public void testRegisterCancel() throws Exception {
         assertThat("Unexpected leases received.", rule.leases, hasSize(1));
         Lease lease = rule.leases.remove(0);
         assertThat("Unexpected permits", lease.getAllowedRequests(), is(rule.permits));
-        assertThat("Unexpected ttl", lease.getTtl(), is(rule.ttl));
+        rule.assertTTL(lease);
         cancel.cancel();
         rule.ticks.onNext(1L);
         assertThat("Unexpected leases received post cancellation.", rule.leases, is(empty()));
@@ -73,8 +73,25 @@ public void testTwoSocketsAndCancel() throws Exception {
         assertThat("Unexpected leases received.", rule.leases, hasSize(1));
     }
 
+    @Test(timeout = 10000)
+    public void testRedistribute() throws Exception {
+        rule.permits = 2;
+        rule.redistributeLeasesOnConnect();
+
+        Cancellable cancel1 = rule.distributor.registerSocket(rule);
+        rule.distributor.registerSocket(rule);
+
+        assertThat("Unexpected leases received.", rule.leases, hasSize(2));
+        rule.assertLease(rule.permits/2);
+        rule.assertLease(rule.permits/2);
+        cancel1.cancel();
+        rule.ticks.onNext(1L);
+        assertThat("Unexpected leases received.", rule.leases, hasSize(1));
+    }
+
     public static class DistributorRule extends ExternalResource implements Consumer<Lease> {
 
+        private boolean redistributeOnConnect;
         private FairLeaseDistributor distributor;
         private int permits;
         private int ttl;
@@ -86,20 +103,29 @@ public Statement apply(final Statement base, Description description) {
             return new Statement() {
                 @Override
                 public void evaluate() throws Throwable {
-                    ticks = PublishProcessor.create();
-                    if (0 == permits) {
-                        permits = 1;
-                    }
-                    if (0 == ttl) {
-                        ttl = 10;
-                    }
-                    distributor = new FairLeaseDistributor(() -> permits, ttl, ticks);
-                    leases = new CopyOnWriteArrayList<>();
+                    init();
                     base.evaluate();
                 }
             };
         }
 
+        protected void init() {
+            ticks = PublishProcessor.create();
+            if (0 == permits) {
+                permits = 1;
+            }
+            if (0 == ttl) {
+                ttl = 10;
+            }
+            distributor = new FairLeaseDistributor(() -> permits, ttl, ticks, redistributeOnConnect);
+            leases = new CopyOnWriteArrayList<>();
+        }
+
+        public void redistributeLeasesOnConnect() {
+            redistributeOnConnect = true;
+            init();
+        }
+
         @Override
         public void accept(Lease lease) {
             leases.add(lease);
@@ -108,7 +134,11 @@ public void accept(Lease lease) {
         public void assertLease(int expectedPermits) {
             Lease lease = leases.remove(0);
             assertThat("Unexpected permits", lease.getAllowedRequests(), is(expectedPermits));
-            assertThat("Unexpected ttl", lease.getTtl(), is(ttl));
+            assertTTL(lease);
+        }
+
+        protected void assertTTL(Lease lease) {
+            assertThat("Unexpected ttl", lease.getTtl(), is((int)(ttl * 1.1)));
         }
     }
 }