ensures local server awaits all connections are close before termination notification

Signed-off-by: Oleh Dokuka <[email protected]>
Signed-off-by: Oleh Dokuka <[email protected]>
Signed-off-by: OlegDokuka <[email protected]>

Author: OlegDokuka

rsocket-transport-local/src/main/java/io/rsocket/transport/local/LocalServerTransport.java

@@ -17,12 +17,15 @@
 package io.rsocket.transport.local;
 
 import io.rsocket.Closeable;
+import io.rsocket.DuplexConnection;
 import io.rsocket.transport.ClientTransport;
 import io.rsocket.transport.ServerTransport;
 import java.util.Objects;
+import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.stream.Collectors;
 import reactor.core.Scannable;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.Sinks;
@@ -34,7 +37,7 @@
  */
 public final class LocalServerTransport implements ServerTransport<Closeable> {
 
-  private static final ConcurrentMap<String, ConnectionAcceptor> registry =
+  private static final ConcurrentMap<String, ServerCloseableAcceptor> registry =
       new ConcurrentHashMap<>();
 
   private final String name;
@@ -72,7 +75,10 @@ public static LocalServerTransport createEphemeral() {
    */
   public static void dispose(String name) {
     Objects.requireNonNull(name, "name must not be null");
-    registry.remove(name);
+    ServerCloseableAcceptor sca = registry.remove(name);
+    if (sca != null) {
+      sca.dispose();
+    }
   }
 
   /**
@@ -107,34 +113,55 @@ public Mono<Closeable> start(ConnectionAcceptor acceptor) {
     Objects.requireNonNull(acceptor, "acceptor must not be null");
     return Mono.create(
         sink -> {
-          ServerCloseable closeable = new ServerCloseable(name, acceptor);
-          if (registry.putIfAbsent(name, acceptor) != null) {
-            throw new IllegalStateException("name already registered: " + name);
+          ServerCloseableAcceptor closeable = new ServerCloseableAcceptor(name, acceptor);
+          if (registry.putIfAbsent(name, closeable) != null) {
+            sink.error(new IllegalStateException("name already registered: " + name));
           }
           sink.success(closeable);
         });
   }
 
-  static class ServerCloseable implements Closeable {
+  @SuppressWarnings({"ReactorTransformationOnMonoVoid", "CallingSubscribeInNonBlockingScope"})
+  static class ServerCloseableAcceptor implements ConnectionAcceptor, Closeable {
 
     private final LocalSocketAddress address;
 
     private final ConnectionAcceptor acceptor;
 
-    private final Sinks.Empty<Void> onClose = Sinks.empty();
+    private final Set<DuplexConnection> activeConnections = ConcurrentHashMap.newKeySet();
+
+    private final Sinks.Empty<Void> onClose = Sinks.unsafe().empty();
 
-    ServerCloseable(String name, ConnectionAcceptor acceptor) {
+    ServerCloseableAcceptor(String name, ConnectionAcceptor acceptor) {
       Objects.requireNonNull(name, "name must not be null");
       this.address = new LocalSocketAddress(name);
       this.acceptor = acceptor;
     }
 
+    @Override
+    public Mono<Void> apply(DuplexConnection duplexConnection) {
+      activeConnections.add(duplexConnection);
+      duplexConnection
+          .onClose()
+          .doFinally(__ -> activeConnections.remove(duplexConnection))
+          .subscribe();
+      return acceptor.apply(duplexConnection);
+    }
+
     @Override
     public void dispose() {
-      if (!registry.remove(address.getName(), acceptor)) {
-        throw new AssertionError();
+      if (!registry.remove(address.getName(), this)) {
+        // already disposed
+        return;
       }
-      onClose.tryEmitEmpty();
+
+      Mono.whenDelayError(
+              activeConnections
+                  .stream()
+                  .peek(DuplexConnection::dispose)
+                  .map(DuplexConnection::onClose)
+                  .collect(Collectors.toList()))
+          .subscribe(null, onClose::tryEmitError, onClose::tryEmitEmpty);
     }
 
     @Override

rsocket-transport-local/src/test/java/io/rsocket/transport/local/LocalClientTransportTest.java

@@ -19,9 +19,10 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatNullPointerException;
 
+import io.rsocket.Closeable;
+import java.time.Duration;
 import org.junit.jupiter.api.DisplayName;
 import org.junit.jupiter.api.Test;
-import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
 final class LocalClientTransportTest {
@@ -31,12 +32,20 @@ final class LocalClientTransportTest {
   void connect() {
     LocalServerTransport serverTransport = LocalServerTransport.createEphemeral();
 
-    serverTransport
-        .start(duplexConnection -> Mono.empty())
-        .flatMap(closeable -> LocalClientTransport.create(serverTransport.getName()).connect())
-        .as(StepVerifier::create)
-        .expectNextCount(1)
-        .verifyComplete();
+    Closeable closeable =
+        serverTransport.start(duplexConnection -> duplexConnection.receive().then()).block();
+
+    try {
+      LocalClientTransport.create(serverTransport.getName())
+          .connect()
+          .doOnNext(d -> d.receive().subscribe())
+          .as(StepVerifier::create)
+          .expectNextCount(1)
+          .verifyComplete();
+    } finally {
+      closeable.dispose();
+      closeable.onClose().block(Duration.ofSeconds(5));
+    }
   }
 
   @DisplayName("generates error if server not started")

rsocket-transport-local/src/test/java/io/rsocket/transport/local/LocalServerTransportTest.java

@@ -96,11 +96,16 @@ void named() {
   @DisplayName("starts local server transport")
   @Test
   void start() {
-    LocalServerTransport.createEphemeral()
-        .start(duplexConnection -> Mono.empty())
-        .as(StepVerifier::create)
-        .expectNextCount(1)
-        .verifyComplete();
+    LocalServerTransport ephemeral = LocalServerTransport.createEphemeral();
+    try {
+      ephemeral
+          .start(duplexConnection -> Mono.empty())
+          .as(StepVerifier::create)
+          .expectNextCount(1)
+          .verifyComplete();
+    } finally {
+      LocalServerTransport.dispose(ephemeral.getName());
+    }
   }
 
   @DisplayName("start throws NullPointerException with null acceptor")