HBASE-28563 Closing ZooKeeper in ZKMainServer (#5869)

Signed-off-by: Duo Zhang <zhangduo@apache.org>
Reviewed-by: Andor Molnár <andor@apache.org>

hbase-zookeeper/src/main/java/org/apache/hadoop/hbase/zookeeper/ZKMainServer.java

@@ -17,10 +17,12 @@
  */
 package org.apache.hadoop.hbase.zookeeper;
 
+import java.io.Closeable;
 import java.io.IOException;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.HBaseConfiguration;
 import org.apache.hadoop.hbase.HBaseInterfaceAudience;
+import org.apache.hadoop.hbase.ZooKeeperConnectionException;
 import org.apache.yetus.audience.InterfaceAudience;
 import org.apache.zookeeper.ZooKeeperMain;
 import org.apache.zookeeper.cli.CliException;
@@ -38,18 +40,23 @@ public String parse(final Configuration c) {
   }
 
   /**
-   * ZooKeeper 3.4.6 broke being able to pass commands on command line. See ZOOKEEPER-1897. This
-   * class is a hack to restore this faclity.
+   * ZooKeeper 3.4.6 broke being able to pass commands on command line. See ZOOKEEPER-1897,
+   * ZOOKEEPER-4804. This class is a hack to restore this faclity.
    */
-  private static class HACK_UNTIL_ZOOKEEPER_1897_ZooKeeperMain extends ZooKeeperMain {
+  private static class HACK_UNTIL_ZOOKEEPER_1897_ZooKeeperMain extends ZooKeeperMain
+    implements Closeable {
     public HACK_UNTIL_ZOOKEEPER_1897_ZooKeeperMain(String[] args)
       throws IOException, InterruptedException {
       super(args);
       // Make sure we are connected before we proceed. Can take a while on some systems. If we
       // run the command without being connected, we get ConnectionLoss KeeperErrorConnection...
       // Make it 30seconds. We dont' have a config in this context and zk doesn't have
       // a timeout until after connection. 30000ms is default for zk.
-      ZooKeeperHelper.ensureConnectedZooKeeper(this.zk, 30000);
+      try {
+        ZooKeeperHelper.ensureConnectedZooKeeper(this.zk, 30000);
+      } catch (ZooKeeperConnectionException e) {
+        this.zk.close();
+      }
     }
 
     /**
@@ -62,6 +69,15 @@ void runCmdLine() throws IOException, InterruptedException, CliException {
       processCmd(this.cl);
       System.exit(0);
     }
+
+    @Override
+    public void close() throws IOException {
+      try {
+        this.zk.close();
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+      }
+    }
   }
 
   /**
@@ -109,9 +125,10 @@ public static void main(String[] args) throws Exception {
     // ZOOKEEPER-1897 was committed to zookeeper-3.4.6 but elsewhere in this class we say
     // 3.4.6 breaks command-processing; TODO.
     if (hasCommandLineArguments(args)) {
-      HACK_UNTIL_ZOOKEEPER_1897_ZooKeeperMain zkm =
-        new HACK_UNTIL_ZOOKEEPER_1897_ZooKeeperMain(newArgs);
-      zkm.runCmdLine();
+      try (HACK_UNTIL_ZOOKEEPER_1897_ZooKeeperMain zkm =
+        new HACK_UNTIL_ZOOKEEPER_1897_ZooKeeperMain(newArgs)) {
+        zkm.runCmdLine();
+      }
     } else {
       ZooKeeperMain.main(newArgs);
     }