Refresh sockets on select in LoadBalancedRSocketMono before checking for active sockets (#623)

jareddellitt · robertroeser · commit 1b87e1f0121a · 2019-04-18T13:02:38.000-07:00
Signed-off-by: Jared Dellitt &lt;jared@dwolla.com&gt;
diff --git a/rsocket-load-balancer/src/main/java/io/rsocket/client/LoadBalancedRSocketMono.java b/rsocket-load-balancer/src/main/java/io/rsocket/client/LoadBalancedRSocketMono.java
@@ -374,10 +374,11 @@ public synchronized double availability() {
   }
 
   private synchronized RSocket select() {
+    refreshSockets();
+
     if (activeSockets.isEmpty()) {
       return FAILING_REACTIVE_SOCKET;
     }
-    refreshSockets();
 
     int size = activeSockets.size();
     if (size == 1) {
diff --git a/rsocket-load-balancer/src/test/java/io/rsocket/client/LoadBalancedRSocketMonoTest.java b/rsocket-load-balancer/src/test/java/io/rsocket/client/LoadBalancedRSocketMonoTest.java
@@ -19,31 +19,24 @@
 import io.rsocket.Payload;
 import io.rsocket.RSocket;
 import io.rsocket.client.filter.RSocketSupplier;
-import io.rsocket.util.EmptyPayload;
-import java.net.InetSocketAddress;
-import java.net.SocketAddress;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
-import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.CompletableFuture;
 import java.util.function.Function;
 import org.junit.Assert;
 import org.junit.Test;
 import org.mockito.Mockito;
 import org.reactivestreams.Publisher;
-import org.reactivestreams.Subscriber;
-import org.reactivestreams.Subscription;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 public class LoadBalancedRSocketMonoTest {
 
   @Test(timeout = 10_000L)
   public void testNeverSelectFailingFactories() throws InterruptedException {
-    InetSocketAddress local0 = InetSocketAddress.createUnresolved("localhost", 7000);
-    InetSocketAddress local1 = InetSocketAddress.createUnresolved("localhost", 7001);
-
     TestingRSocket socket = new TestingRSocket(Function.identity());
-    RSocketSupplier failing = failingClient(local0);
+    RSocketSupplier failing = failingClient();
     RSocketSupplier succeeding = succeedingFactory(socket);
     List<RSocketSupplier> factories = Arrays.asList(failing, succeeding);
 
@@ -52,9 +45,6 @@ public void testNeverSelectFailingFactories() throws InterruptedException {
 
   @Test(timeout = 10_000L)
   public void testNeverSelectFailingSocket() throws InterruptedException {
-    InetSocketAddress local0 = InetSocketAddress.createUnresolved("localhost", 7000);
-    InetSocketAddress local1 = InetSocketAddress.createUnresolved("localhost", 7001);
-
     TestingRSocket socket = new TestingRSocket(Function.identity());
     TestingRSocket failingSocket =
         new TestingRSocket(Function.identity()) {
@@ -76,6 +66,33 @@ public double availability() {
     testBalancer(clients);
   }
 
+  @Test(timeout = 10_000L)
+  public void testRefreshesSocketsOnSelectBeforeReturningFailedAfterNewFactoriesDelivered() {
+    TestingRSocket socket = new TestingRSocket(Function.identity());
+
+    CompletableFuture<RSocketSupplier> laterSupplier = new CompletableFuture<>();
+    Flux<List<RSocketSupplier>> factories =
+        Flux.create(
+            s -> {
+              s.next(Collections.emptyList());
+
+              laterSupplier.handle(
+                  (RSocketSupplier result, Throwable t) -> {
+                    s.next(Collections.singletonList(result));
+                    return null;
+                  });
+            });
+
+    LoadBalancedRSocketMono balancer = LoadBalancedRSocketMono.create(factories);
+
+    Assert.assertEquals(0.0, balancer.availability(), 0);
+
+    laterSupplier.complete(succeedingFactory(socket));
+    balancer.rSocketMono.block();
+
+    Assert.assertEquals(1.0, balancer.availability(), 0);
+  }
+
   private void testBalancer(List<RSocketSupplier> factories) throws InterruptedException {
     Publisher<List<RSocketSupplier>> src =
         s -> {
@@ -92,39 +109,6 @@ private void testBalancer(List<RSocketSupplier> factories) throws InterruptedExc
     Flux.range(0, 100).flatMap(i -> balancer).blockLast();
   }
 
-  private void makeAcall(RSocket balancer) throws InterruptedException {
-    CountDownLatch latch = new CountDownLatch(1);
-
-    balancer
-        .requestResponse(EmptyPayload.INSTANCE)
-        .subscribe(
-            new Subscriber<Payload>() {
-              @Override
-              public void onSubscribe(Subscription s) {
-                s.request(1L);
-              }
-
-              @Override
-              public void onNext(Payload payload) {
-                System.out.println("Successfully receiving a response");
-              }
-
-              @Override
-              public void onError(Throwable t) {
-                t.printStackTrace();
-                Assert.assertTrue(false);
-                latch.countDown();
-              }
-
-              @Override
-              public void onComplete() {
-                latch.countDown();
-              }
-            });
-
-    latch.await();
-  }
-
   private static RSocketSupplier succeedingFactory(RSocket socket) {
     RSocketSupplier mock = Mockito.mock(RSocketSupplier.class);
 
@@ -135,7 +119,7 @@ private static RSocketSupplier succeedingFactory(RSocket socket) {
     return mock;
   }
 
-  private static RSocketSupplier failingClient(SocketAddress sa) {
+  private static RSocketSupplier failingClient() {
     RSocketSupplier mock = Mockito.mock(RSocketSupplier.class);
 
     Mockito.when(mock.availability()).thenReturn(0.0);

Original file line number	Diff line number	Diff line change
`@@ -374,10 +374,11 @@ public synchronized double availability() {`
`374`	`374`	`}`
`375`	`375`
`376`	`376`	`private synchronized RSocket select() {`
	`377`	`+ refreshSockets();`
	`378`	`+`
`377`	`379`	`if (activeSockets.isEmpty()) {`
`378`	`380`	`return FAILING_REACTIVE_SOCKET;`
`379`	`381`	`}`
`380`		`- refreshSockets();`
`381`	`382`
`382`	`383`	`int size = activeSockets.size();`
`383`	`384`	`if (size == 1) {`