From f727ff84bdb8d36a20484d5e75e07c8efe9e5dd8 Mon Sep 17 00:00:00 2001
From: brycezhongqing <zhonchen+1@linkedin.com>
Date: Mon, 7 Oct 2024 18:47:21 -0700
Subject: [PATCH] resolve comments

---
 .../balancer/servers/ConnectionManager.java   | 27 ++--------------
 .../servers/ZooKeeperConnectionManager.java   | 31 ++++++++++++++-----
 2 files changed, 26 insertions(+), 32 deletions(-)

diff --git a/d2/src/main/java/com/linkedin/d2/balancer/servers/ConnectionManager.java b/d2/src/main/java/com/linkedin/d2/balancer/servers/ConnectionManager.java
index 86e5f97d8..050555501 100644
--- a/d2/src/main/java/com/linkedin/d2/balancer/servers/ConnectionManager.java
+++ b/d2/src/main/java/com/linkedin/d2/balancer/servers/ConnectionManager.java
@@ -9,24 +9,20 @@
 public abstract class ConnectionManager
 {
   private final ZooKeeperAnnouncer[] _servers;
-  private final String _zkConnectString;
-  private final int _zkSessionTimeout;
-  private final String _zkBasePath;
 
   private static final Logger LOG = LoggerFactory.getLogger(ConnectionManager.class);
 
-  protected ConnectionManager(ZooKeeperAnnouncer[] servers, String zkConnectString, int zkSessionTimeout, String zkBasePath)
+  protected ConnectionManager(ZooKeeperAnnouncer[] servers)
   {
     _servers = servers;
-    _zkConnectString = zkConnectString;
-    _zkSessionTimeout = zkSessionTimeout;
-    _zkBasePath = zkBasePath;
   }
 
   abstract public void start(Callback<None> callback);
 
   abstract public void shutdown(final Callback<None> callback);
 
+  abstract public String getAnnouncementTargetIdentifier();
+
   public void markDownAllServers(final Callback<None> callback)
   {
     Callback<None> markDownCallback;
@@ -94,21 +90,4 @@ public ZooKeeperAnnouncer[] getAnnouncers()
   {
     return _servers;
   }
-
-  abstract public boolean isSessionEstablished();
-
-  public String getZooKeeperConnectString()
-  {
-    return _zkConnectString;
-  }
-
-  public int getZooKeeperSessionTimeout()
-  {
-    return _zkSessionTimeout;
-  }
-
-  public String getZooKeeperBasePath()
-  {
-    return _zkBasePath;
-  }
 }
diff --git a/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperConnectionManager.java b/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperConnectionManager.java
index f1d607305..7d36bc218 100644
--- a/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperConnectionManager.java
+++ b/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperConnectionManager.java
@@ -79,10 +79,7 @@ public ZooKeeperConnectionManager(ZKPersistentConnection zkConnection,
                                     ZKStoreFactory<UriProperties,ZooKeeperEphemeralStore<UriProperties>> factory,
                                     ZooKeeperAnnouncer... servers)
   {
-    super(servers,
-          zkConnection.getZKConnection().getConnectString(),
-          zkConnection.getZKConnection().getTimeout(),
-          zkBasePath);
+    super(servers);
     _zkBasePath = zkBasePath;
     _zkConnection = zkConnection;
     _factory = factory;
@@ -98,10 +95,7 @@ public ZooKeeperConnectionManager(String zkConnectString, int zkSessionTimeout,
                                     ZKStoreFactory<UriProperties,ZooKeeperEphemeralStore<UriProperties>> factory,
                                     ZooKeeperAnnouncer... servers)
   {
-    super(servers,
-          zkConnectString,
-          zkSessionTimeout,
-          zkBasePath);
+    super(servers);
     _zkConnectString = zkConnectString;
     _zkSessionTimeout = zkSessionTimeout;
     _zkBasePath = zkBasePath;
@@ -189,6 +183,7 @@ protected None convertResponse(None none) throws Exception
       zkCloseCallback.onSuccess(None.none());
     }
   }
+
   private class Listener implements ZKPersistentConnection.EventListener
   {
     @Override
@@ -301,8 +296,28 @@ public interface ZKStoreFactory<P, Z extends ZooKeeperStore<P>>
   }
 
   @Override
+  public String getAnnouncementTargetIdentifier()
+  {
+    return getZooKeeperConnectString();
+  }
+
   public boolean isSessionEstablished()
   {
     return _sessionEstablished;
   }
+
+  public String getZooKeeperConnectString()
+  {
+    return _zkConnectString;
+  }
+
+  public int getZooKeeperSessionTimeout()
+  {
+    return _zkSessionTimeout;
+  }
+
+  public String getZooKeeperBasePath()
+  {
+    return _zkBasePath;
+  }
 }