Ensure atomicity when removing last registered session owner

2 files changed, 8 insertions, 12 deletions

diff --git a/messagebus/src/main/java/com/yahoo/messagebus/network/NetworkMultiplexer.java b/messagebus/src/main/java/com/yahoo/messagebus/network/NetworkMultiplexer.java
index d36579d5be1..959fff488b1 100644
--- a/messagebus/src/main/java/com/yahoo/messagebus/network/NetworkMultiplexer.java
+++ b/messagebus/src/main/java/com/yahoo/messagebus/network/NetworkMultiplexer.java
@@ -59,7 +59,7 @@ public class NetworkMultiplexer implements NetworkOwner {
                     net.registerSession(session);
             }
             else if (owners.contains(owner))
-                throw new IllegalArgumentException("Session '" + session + "' with owner '" + owner + "' already registered with this");
+                throw new IllegalArgumentException("Session '" + session + "' with owner '" + owner + "' already registered with " +  this);
 
             owners.push(owner);
             return owners;
@@ -68,12 +68,12 @@ public class NetworkMultiplexer implements NetworkOwner {
 
     public void unregisterSession(String session, NetworkOwner owner, boolean broadcast) {
         sessions.computeIfPresent(session, (name, owners) -> {
-            owners.remove(owner);
-            if (owners.isEmpty()) {
+            if (owners.size() == 1 && owners.contains(owner)) {
                 if (broadcast)
                     net.unregisterSession(session);
                 return null;
             }
+            owners.remove(owner);
             return owners;
         });
     }
@@ -103,7 +103,7 @@ public class NetworkMultiplexer implements NetworkOwner {
     /** Attach the network owner to this, allowing this to forward messages to it. */
     public void attach(NetworkOwner owner) {
         if (owners.contains(owner))
-            throw new IllegalArgumentException(owner + " is already attached to this");
+            throw new IllegalArgumentException(owner + " is already attached to " + this);
 
         owners.add(owner);
     }
@@ -111,7 +111,7 @@ public class NetworkMultiplexer implements NetworkOwner {
     /** Detach the network owner from this, no longer allowing messages to it, and shutting down this is ownerless. */
     public void detach(NetworkOwner owner) {
         if ( ! owners.remove(owner))
-            throw new IllegalArgumentException(owner + " not attached to this");
+            throw new IllegalArgumentException(owner + " not attached to " + this);
 
         destroyIfOwnerless();
     }
@@ -137,12 +137,7 @@ public class NetworkMultiplexer implements NetworkOwner {
 
     @Override
     public String toString() {
-        return "NetworkMultiplexer{" +
-               "net=" + net +
-               ", owners=" + owners +
-               ", sessions=" + sessions +
-               ", destructible=" + disowned +
-               '}';
+        return "network multiplexer with owners: " + owners + ", sessions: " + sessions + " and destructible: " + disowned.get();
     }
 
 }
\ No newline at end of file
diff --git a/messagebus/src/test/java/com/yahoo/messagebus/network/NetworkMultiplexerTest.java b/messagebus/src/test/java/com/yahoo/messagebus/network/NetworkMultiplexerTest.java
index 60da1b49d90..57ded9f750d 100644
--- a/messagebus/src/test/java/com/yahoo/messagebus/network/NetworkMultiplexerTest.java
+++ b/messagebus/src/test/java/com/yahoo/messagebus/network/NetworkMultiplexerTest.java
@@ -43,7 +43,8 @@ public class NetworkMultiplexerTest {
             fail("Illegal to register same session multiple times with the same owner");
         }
         catch (IllegalArgumentException expected) {
-            assertEquals("Session 's1' with owner 'mock owner' already registered with this", expected.getMessage());
+            assertEquals("Session 's1' with owner 'mock owner' already registered with network multiplexer with owners: [mock owner], sessions: {s1=[mock owner]} and destructible: false",
+                         expected.getMessage());
         }
         assertEquals(Set.of("s1"), net.registered);