Fix for LoadbalanceTest concurrency issue

Signed-off-by: Rossen Stoyanchev <[email protected]>

Author: rstoyanchev

rsocket-core/src/test/java/io/rsocket/loadbalance/LoadbalanceTest.java

@@ -10,8 +10,11 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.assertj.core.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mockito;
 import org.reactivestreams.Publisher;
@@ -26,10 +29,13 @@
 
 public class LoadbalanceTest {
 
-  @Test
-  public void shouldDeliverAllTheRequestsWithRoundRobinStrategy() {
+  @BeforeEach
+  void setUp() {
     Hooks.onErrorDropped((__) -> {});
+  }
 
+  @Test
+  public void shouldDeliverAllTheRequestsWithRoundRobinStrategy() throws Exception {
     final AtomicInteger counter = new AtomicInteger();
     final ClientTransport mockTransport = Mockito.mock(ClientTransport.class);
     final RSocketConnector rSocketConnectorMock = Mockito.mock(RSocketConnector.class);
@@ -51,12 +57,14 @@ public Mono<Void> fireAndForget(Payload payload) {
       final TestPublisher<List<LoadbalanceTarget>> source = TestPublisher.create();
       final RSocketPool rSocketPool =
           new RSocketPool(rSocketConnectorMock, source, new RoundRobinLoadbalanceStrategy());
+      final CountDownLatch requestLatch = new CountDownLatch(1000);
 
       RaceTestUtils.race(
           () -> {
             for (int j = 0; j < 1000; j++) {
               Mono.defer(() -> rSocketPool.select().fireAndForget(EmptyPayload.INSTANCE))
                   .retry()
+                  .doFinally(s -> requestLatch.countDown())
                   .subscribe();
             }
           },
@@ -75,16 +83,15 @@ public Mono<Void> fireAndForget(Payload payload) {
             }
           });
 
+      requestLatch.await(5, TimeUnit.SECONDS);
       Assertions.assertThat(counter.get()).isEqualTo(1000);
 
       counter.set(0);
     }
   }
 
   @Test
-  public void shouldDeliverAllTheRequestsWithWightedStrategy() {
-    Hooks.onErrorDropped((__) -> {});
-
+  public void shouldDeliverAllTheRequestsWithWeightedStrategy() throws InterruptedException {
     final AtomicInteger counter = new AtomicInteger();
     final ClientTransport mockTransport = Mockito.mock(ClientTransport.class);
     final RSocketConnector rSocketConnectorMock = Mockito.mock(RSocketConnector.class);
@@ -101,13 +108,15 @@ public void shouldDeliverAllTheRequestsWithWightedStrategy() {
               WeightedLoadbalanceStrategy.builder()
                   .weightedStatsResolver(r -> (WeightedStats) r)
                   .build());
+      final CountDownLatch requestLatch = new CountDownLatch(1000);
 
       RaceTestUtils.race(
           () -> {
             for (int j = 0; j < 1000; j++) {
               Mono.defer(() -> rSocketPool.select().fireAndForget(EmptyPayload.INSTANCE))
                   .retry()
-                  .subscribe();
+                  .doFinally(s -> requestLatch.countDown())
+                  .subscribe(aVoid -> {}, Throwable::printStackTrace);
             }
           },
           () -> {
@@ -125,6 +134,7 @@ public void shouldDeliverAllTheRequestsWithWightedStrategy() {
             }
           });
 
+      requestLatch.await(5, TimeUnit.SECONDS);
       Assertions.assertThat(counter.get()).isEqualTo(1000);
 
       counter.set(0);
@@ -133,8 +143,6 @@ public void shouldDeliverAllTheRequestsWithWightedStrategy() {
 
   @Test
   public void ensureRSocketIsCleanedFromThePoolIfSourceRSocketIsDisposed() {
-    Hooks.onErrorDropped((__) -> {});
-
     final AtomicInteger counter = new AtomicInteger();
     final ClientTransport mockTransport = Mockito.mock(ClientTransport.class);
     final RSocketConnector rSocketConnectorMock = Mockito.mock(RSocketConnector.class);
@@ -179,8 +187,6 @@ public Mono<Void> fireAndForget(Payload payload) {
 
   @Test
   public void ensureContextIsPropagatedCorrectlyForRequestChannel() {
-    Hooks.onErrorDropped((__) -> {});
-
     final AtomicInteger counter = new AtomicInteger();
     final ClientTransport mockTransport = Mockito.mock(ClientTransport.class);
     final RSocketConnector rSocketConnectorMock = Mockito.mock(RSocketConnector.class);

rsocket-core/src/test/java/io/rsocket/loadbalance/RoundRobinLoadbalanceStrategyTest.java

@@ -11,6 +11,7 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import org.assertj.core.api.Assertions;
 import org.assertj.core.data.Offset;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mockito;
 import reactor.core.publisher.Hooks;
@@ -19,10 +20,13 @@
 
 public class RoundRobinLoadbalanceStrategyTest {
 
-  @Test
-  public void shouldDeliverValuesProportionally() {
+  @BeforeEach
+  void setUp() {
     Hooks.onErrorDropped((__) -> {});
+  }
 
+  @Test
+  public void shouldDeliverValuesProportionally() {
     final AtomicInteger counter1 = new AtomicInteger();
     final AtomicInteger counter2 = new AtomicInteger();
     final ClientTransport mockTransport = Mockito.mock(ClientTransport.class);
@@ -71,8 +75,6 @@ public Mono<Void> fireAndForget(Payload payload) {
 
   @Test
   public void shouldDeliverValuesToNewlyConnectedSockets() {
-    Hooks.onErrorDropped((__) -> {});
-
     final AtomicInteger counter1 = new AtomicInteger();
     final AtomicInteger counter2 = new AtomicInteger();
     final ClientTransport mockTransport1 = Mockito.mock(ClientTransport.class);
@@ -104,7 +106,7 @@ public Mono<Void> fireAndForget(Payload payload) {
       rSocketPool.select().fireAndForget(EmptyPayload.INSTANCE).subscribe();
     }
 
-    source.next(Arrays.asList(LoadbalanceTarget.from("1", mockTransport1)));
+    source.next(Collections.singletonList(LoadbalanceTarget.from("1", mockTransport1)));
 
     Assertions.assertThat(counter1.get()).isCloseTo(1000, Offset.offset(1));
 
@@ -114,7 +116,7 @@ public Mono<Void> fireAndForget(Payload payload) {
       rSocketPool.select().fireAndForget(EmptyPayload.INSTANCE).subscribe();
     }
 
-    source.next(Arrays.asList(LoadbalanceTarget.from("1", mockTransport1)));
+    source.next(Collections.singletonList(LoadbalanceTarget.from("1", mockTransport1)));
 
     Assertions.assertThat(counter1.get()).isCloseTo(2000, Offset.offset(1));
 
@@ -130,7 +132,7 @@ public Mono<Void> fireAndForget(Payload payload) {
     Assertions.assertThat(counter1.get()).isCloseTo(2500, Offset.offset(1));
     Assertions.assertThat(counter2.get()).isCloseTo(500, Offset.offset(1));
 
-    source.next(Arrays.asList(LoadbalanceTarget.from("2", mockTransport1)));
+    source.next(Collections.singletonList(LoadbalanceTarget.from("2", mockTransport1)));
 
     for (int j = 0; j < 1000; j++) {
       rSocketPool.select().fireAndForget(EmptyPayload.INSTANCE).subscribe();