Use refer(context) for mbus too.

6 files changed, 6 insertions, 6 deletions

diff --git a/container-messagebus/src/main/java/com/yahoo/container/jdisc/messagebus/SessionCache.java b/container-messagebus/src/main/java/com/yahoo/container/jdisc/messagebus/SessionCache.java
index 354e5fbfe12..1db288d59dd 100644
--- a/container-messagebus/src/main/java/com/yahoo/container/jdisc/messagebus/SessionCache.java
+++ b/container-messagebus/src/main/java/com/yahoo/container/jdisc/messagebus/SessionCache.java
@@ -177,7 +177,7 @@ public final class SessionCache extends AbstractComponent {
                     sessionReference = References.fromResource(session);
                 } else {
                     try {
-                        sessionReference = session.refer();
+                        sessionReference = session.refer(this);
                         logReuse(session);
                     } catch (final IllegalStateException e) {
                         session = createAndStore(registry, p, key);
diff --git a/container-messagebus/src/main/java/com/yahoo/messagebus/jdisc/MbusClient.java b/container-messagebus/src/main/java/com/yahoo/messagebus/jdisc/MbusClient.java
index 85eb703d4cd..7fa294422d6 100644
--- a/container-messagebus/src/main/java/com/yahoo/messagebus/jdisc/MbusClient.java
+++ b/container-messagebus/src/main/java/com/yahoo/messagebus/jdisc/MbusClient.java
@@ -41,7 +41,7 @@ public final class MbusClient extends AbstractResource implements ClientProvider
     @Inject
     public MbusClient(ClientSession session) {
         this.session = session;
-        this.sessionReference = session.refer();
+        this.sessionReference = session.refer(this);
         thread = new Thread(new SenderTask(), "mbus-client-" + threadId.getAndIncrement());
     }
 
diff --git a/container-messagebus/src/main/java/com/yahoo/messagebus/jdisc/MbusServer.java b/container-messagebus/src/main/java/com/yahoo/messagebus/jdisc/MbusServer.java
index 5f1963cffe6..e7b1fc3e71d 100644
--- a/container-messagebus/src/main/java/com/yahoo/messagebus/jdisc/MbusServer.java
+++ b/container-messagebus/src/main/java/com/yahoo/messagebus/jdisc/MbusServer.java
@@ -44,7 +44,7 @@ public final class MbusServer extends AbstractResource implements ServerProvider
         this.session = session;
         uri = URI.create("mbus://localhost/" + session.name());
         session.setMessageHandler(this);
-        sessionReference = session.refer();
+        sessionReference = session.refer(this);
     }
 
     @Override
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 c4abe963d89..e0ab9dbff4f 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
@@ -29,7 +29,7 @@ public class SharedDestinationSession extends AbstractResource implements Messag
     SharedDestinationSession(SharedMessageBus mbus, DestinationSessionParams params) {
         this.msgHandler.set(params.getMessageHandler());
         this.session = mbus.messageBus().createDetachedDestinationSession(params.setMessageHandler(this));
-        this.mbusReference = mbus.refer();
+        this.mbusReference = mbus.refer(this);
     }
 
     public DestinationSession session() {
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 7565aa6e625..06c6192cb7a 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
@@ -37,7 +37,7 @@ public class SharedIntermediateSession extends AbstractResource
         this.msgHandler.set(params.getMessageHandler());
         this.session = mbus.messageBus().createDetachedIntermediateSession(params.setReplyHandler(this)
                                                                                  .setMessageHandler(this));
-        this.mbusReference = mbus.refer();
+        this.mbusReference = mbus.refer(this);
     }
 
     public IntermediateSession session() {
diff --git a/container-messagebus/src/main/java/com/yahoo/messagebus/shared/SharedSourceSession.java b/container-messagebus/src/main/java/com/yahoo/messagebus/shared/SharedSourceSession.java
index a185922545e..016fb8975a7 100644
--- a/container-messagebus/src/main/java/com/yahoo/messagebus/shared/SharedSourceSession.java
+++ b/container-messagebus/src/main/java/com/yahoo/messagebus/shared/SharedSourceSession.java
@@ -27,7 +27,7 @@ public class SharedSourceSession extends AbstractResource implements ClientSessi
             throw new IllegalArgumentException("Reply handler must be null.");
         }
         this.session = mbus.messageBus().createSourceSession(params.setReplyHandler(this));
-        this.mbusReference = mbus.refer();
+        this.mbusReference = mbus.refer(this);
     }
 
     public SourceSession session() {