From 86c6be6b132bec9b8e1b78f48dddac436c5c2bb3 Mon Sep 17 00:00:00 2001
From: Paul Ferraro <paul.ferraro@redhat.com>
Date: Fri, 22 May 2020 16:17:02 -0400
Subject: [PATCH] WFLY-13132 Expose remoting connection associated with a
 ClusterTopologyListener. (#460)

---
 .../org/jboss/ejb/protocol/remote/EJBServerChannel.java  | 9 +++++++++
 .../org/jboss/ejb/server/ClusterTopologyListener.java    | 8 ++++++++
 2 files changed, 17 insertions(+)

diff --git a/src/main/java/org/jboss/ejb/protocol/remote/EJBServerChannel.java b/src/main/java/org/jboss/ejb/protocol/remote/EJBServerChannel.java
index 5602237cc..e11797a3a 100644
--- a/src/main/java/org/jboss/ejb/protocol/remote/EJBServerChannel.java
+++ b/src/main/java/org/jboss/ejb/protocol/remote/EJBServerChannel.java
@@ -1183,6 +1183,7 @@ final class ClusterTopologyWriter implements ClusterTopologyListener {
         ClusterTopologyWriter() {
         }
 
+        @Override
         public void clusterTopology(final List<ClusterInfo> clusterInfoList) {
             try (MessageOutputStream os = messageTracker.openMessageUninterruptibly()) {
                 os.writeByte(Protocol.CLUSTER_TOPOLOGY_COMPLETE);
@@ -1214,6 +1215,7 @@ public void clusterTopology(final List<ClusterInfo> clusterInfoList) {
             }
         }
 
+        @Override
         public void clusterRemoval(final List<String> clusterNames) {
             try (MessageOutputStream os = messageTracker.openMessageUninterruptibly()) {
                 os.writeByte(Protocol.CLUSTER_TOPOLOGY_REMOVAL);
@@ -1227,6 +1229,7 @@ public void clusterRemoval(final List<String> clusterNames) {
             }
         }
 
+        @Override
         public void clusterNewNodesAdded(final ClusterInfo clusterInfo) {
             try (MessageOutputStream os = messageTracker.openMessageUninterruptibly()) {
                 os.writeByte(Protocol.CLUSTER_TOPOLOGY_ADDITION);
@@ -1256,6 +1259,7 @@ public void clusterNewNodesAdded(final ClusterInfo clusterInfo) {
             }
         }
 
+        @Override
         public void clusterNodesRemoved(final List<ClusterRemovalInfo> clusterRemovalInfoList) {
             try (MessageOutputStream os = messageTracker.openMessageUninterruptibly()) {
                 os.writeByte(Protocol.CLUSTER_TOPOLOGY_NODE_REMOVAL);
@@ -1273,6 +1277,11 @@ public void clusterNodesRemoved(final List<ClusterRemovalInfo> clusterRemovalInf
                 Logs.REMOTING.trace("EJB cluster message write failed", e);
             }
         }
+
+        @Override
+        public Connection getConnection() {
+            return EJBServerChannel.this.channel.getConnection();
+        }
     }
 
     final class ModuleAvailabilityWriter implements ModuleAvailabilityListener {
diff --git a/src/main/java/org/jboss/ejb/server/ClusterTopologyListener.java b/src/main/java/org/jboss/ejb/server/ClusterTopologyListener.java
index 8950389fa..340b4ddf2 100644
--- a/src/main/java/org/jboss/ejb/server/ClusterTopologyListener.java
+++ b/src/main/java/org/jboss/ejb/server/ClusterTopologyListener.java
@@ -21,6 +21,8 @@
 import java.net.InetAddress;
 import java.util.List;
 
+import org.jboss.remoting3.Connection;
+
 /**
  * A legacy cluster topology notification client.
  *
@@ -36,6 +38,12 @@ public interface ClusterTopologyListener {
 
     void clusterNodesRemoved(List<ClusterRemovalInfo> clusterRemovalInfoList);
 
+    /**
+     * Returns the remoting connection associated with this listener
+     * @return a remoting connection
+     */
+    Connection getConnection();
+
     final class ClusterInfo {
         private final String clusterName;
         private final List<NodeInfo> nodeInfoList;