Unregister sessions on server close

4 files changed, 24 insertions, 17 deletions

diff --git a/container-messagebus/src/main/java/com/yahoo/messagebus/shared/ServerSession.java b/container-messagebus/src/main/java/com/yahoo/messagebus/shared/ServerSession.java
index 56713815c7a..4cd4a776292 100644
--- a/container-messagebus/src/main/java/com/yahoo/messagebus/shared/ServerSession.java
+++ b/container-messagebus/src/main/java/com/yahoo/messagebus/shared/ServerSession.java
@@ -10,13 +10,16 @@ import com.yahoo.messagebus.Reply;
  */
 public interface ServerSession extends SharedResource {
 
-    public MessageHandler getMessageHandler();
+    MessageHandler getMessageHandler();
 
-    public void setMessageHandler(MessageHandler msgHandler);
+    void setMessageHandler(MessageHandler msgHandler);
 
-    public void sendReply(Reply reply);
+    void sendReply(Reply reply);
 
-    public String connectionSpec();
+    String connectionSpec();
+
+    String name();
+
+    void close();
 
-    public String name();
 }
diff --git a/container-messagebus/src/main/java/com/yahoo/messagebus/shared/SharedDestinationSession.java b/container-messagebus/src/main/java/com/yahoo/messagebus/shared/SharedDestinationSession.java
index 7da164757cd..5a9cd39c5b4 100644
--- a/container-messagebus/src/main/java/com/yahoo/messagebus/shared/SharedDestinationSession.java
+++ b/container-messagebus/src/main/java/com/yahoo/messagebus/shared/SharedDestinationSession.java
@@ -77,6 +77,11 @@ public class SharedDestinationSession extends AbstractResource implements Messag
     }
 
     @Override
+    public void close() {
+        session.destroy();
+    }
+
+    @Override
     protected void destroy() {
         log.log(Level.FINE, "Destroying shared destination session.");
         session.destroy();
diff --git a/container-messagebus/src/main/java/com/yahoo/messagebus/shared/SharedIntermediateSession.java b/container-messagebus/src/main/java/com/yahoo/messagebus/shared/SharedIntermediateSession.java
index 5c9fab46e34..64cc1aaf510 100644
--- a/container-messagebus/src/main/java/com/yahoo/messagebus/shared/SharedIntermediateSession.java
+++ b/container-messagebus/src/main/java/com/yahoo/messagebus/shared/SharedIntermediateSession.java
@@ -96,6 +96,11 @@ public class SharedIntermediateSession extends AbstractResource
     }
 
     @Override
+    public void close() {
+        session.destroy();
+    }
+
+    @Override
     protected void destroy() {
         log.log(Level.FINE, "Destroying shared intermediate session.");
         session.destroy();
diff --git a/container-messagebus/src/test/java/com/yahoo/messagebus/jdisc/MbusServerTestCase.java b/container-messagebus/src/test/java/com/yahoo/messagebus/jdisc/MbusServerTestCase.java
index 9d45d2e7abf..6ebb41c4ab7 100644
--- a/container-messagebus/src/test/java/com/yahoo/messagebus/jdisc/MbusServerTestCase.java
+++ b/container-messagebus/src/test/java/com/yahoo/messagebus/jdisc/MbusServerTestCase.java
@@ -331,9 +331,7 @@ public class MbusServerTestCase {
         int refCount = 1;
 
         @Override
-        public void sendReply(Reply reply) {
-
-        }
+        public void sendReply(Reply reply) { }
 
         @Override
         public MessageHandler getMessageHandler() {
@@ -341,9 +339,7 @@ public class MbusServerTestCase {
         }
 
         @Override
-        public void setMessageHandler(MessageHandler msgHandler) {
-
-        }
+        public void setMessageHandler(MessageHandler msgHandler) { }
 
         @Override
         public String connectionSpec() {
@@ -356,14 +352,12 @@ public class MbusServerTestCase {
         }
 
         @Override
+        public void close() { }
+
+        @Override
         public ResourceReference refer() {
             ++refCount;
-            return new ResourceReference() {
-                @Override
-                public void close() {
-                    --refCount;
-                }
-            };
+            return () -> --refCount;
         }
 
         @Override