[KYUUBI #1446] Decouple zookeeper from abstract ServiceDiscovery

<!--
Thanks for sending a pull request!

Here are some tips for you:
  1. If this is your first time, please read our contributor guidelines: https://kyuubi.readthedocs.io/en/latest/community/contributions.html
  2. If the PR is related to an issue in https://github.com/apache/incubator-kyuubi/issues, add '[KYUUBI #XXXX]' in your PR title, e.g., '[KYUUBI #XXXX] Your PR title ...'.
  3. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP][KYUUBI #XXXX] Your PR title ...'.
-->

### _Why are the changes needed?_
<!--
Please clarify why the changes are needed. For instance,
  1. If you add a feature, you can talk about the use case of it.
  2. If you fix a bug, you can clarify why it is a bug.
-->
Decouple zookeeper from abstract ServiceDiscovery

### _How was this patch tested?_
- [X] Add some test cases that check the changes thoroughly including negative and positive cases if possible

- [ ] Add screenshots for manual tests if appropriate

- [X] [Run test](https://kyuubi.readthedocs.io/en/latest/develop_tools/testing.html#running-tests) locally before make a pull request

Closes #1448 from hddong/decouple-zeekeeper.

Closes #1446

0294c3c [hongdongdong] remove
5da2f94 [hongdongdong] fix
be066f6 [hongdongdong] fix
0303855 [hongdongdong] [KYUUBI #1446] Decouple zookeeper from abstract ServiceDiscovery

Authored-by: hongdongdong <hongdongdong@cmss.chinamobile.com>
Signed-off-by: Kent Yao <yao@apache.org>

Author: hddong

kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/EngineServiceDiscovery.scala

@@ -31,13 +31,14 @@ class EngineServiceDiscovery(
     fe: FrontendService) extends ServiceDiscovery("EngineServiceDiscovery", fe) {
 
   override def stop(): Unit = synchronized {
-    closeServiceNode()
+    discoveryClient.deregisterService()
     conf.get(ENGINE_SHARE_LEVEL) match {
       // For connection level, we should clean up the namespace in zk in case the disk stress.
-      case "CONNECTION" if namespace != null =>
+      case "CONNECTION" =>
         try {
-          zkClient.delete().deletingChildrenIfNeeded().forPath(namespace)
-          info("Clean up discovery service due to this is connection share level.")
+          if (discoveryClient.postDeregisterService) {
+            info("Clean up discovery service due to this is connection share level.")
+          }
         } catch {
           case NonFatal(e) =>
             warn("Failed to clean up Spark engine before stop.", e)

kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/ServiceDiscovery.scala

@@ -17,28 +17,20 @@
 
 package org.apache.kyuubi.ha.client
 
-import java.io.IOException
 import java.nio.charset.StandardCharsets
-import java.util.concurrent.TimeUnit
-import java.util.concurrent.atomic.AtomicBoolean
 
 import scala.collection.JavaConverters._
 
 import com.google.common.annotations.VisibleForTesting
 import org.apache.curator.framework.CuratorFramework
-import org.apache.curator.framework.recipes.nodes.PersistentNode
-import org.apache.curator.framework.state.{ConnectionState, ConnectionStateListener}
-import org.apache.curator.framework.state.ConnectionState.{CONNECTED, LOST, RECONNECTED}
 import org.apache.curator.utils.ZKPaths
-import org.apache.zookeeper.{CreateMode, KeeperException, WatchedEvent, Watcher}
-import org.apache.zookeeper.CreateMode.PERSISTENT
-import org.apache.zookeeper.KeeperException.NodeExistsException
 
-import org.apache.kyuubi.{KYUUBI_VERSION, KyuubiException, Logging}
+import org.apache.kyuubi.Logging
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.ha.HighAvailabilityConf._
+import org.apache.kyuubi.ha.client.zookeeper.ServiceDiscoveryClient
+import org.apache.kyuubi.ha.client.zookeeper.ServiceDiscoveryClient.createServiceNode
 import org.apache.kyuubi.service.{AbstractService, FrontendService}
-import org.apache.kyuubi.util.{KyuubiHadoopUtils, ThreadUtils}
 
 /**
  * A abstract service for service discovery
@@ -48,92 +40,31 @@ import org.apache.kyuubi.util.{KyuubiHadoopUtils, ThreadUtils}
  */
 abstract class ServiceDiscovery(
     name: String,
-    fe: FrontendService) extends AbstractService(name) {
+    val fe: FrontendService) extends AbstractService(name) {
 
-  import ServiceDiscovery._
-  import ZooKeeperClientProvider._
+  private var _discoveryClient: ServiceDiscoveryClient = _
 
-  private var _zkClient: CuratorFramework = _
-  private var _serviceNode: PersistentNode = _
-
-  /**
-   * a pre-defined namespace used to publish the instance of the associate service
-   */
-  private var _namespace: String = _
-
-  def zkClient: CuratorFramework = _zkClient
-
-  def serviceNode: PersistentNode = _serviceNode
-
-  def namespace: String = _namespace
+  def discoveryClient: ServiceDiscoveryClient = _discoveryClient
 
   override def initialize(conf: KyuubiConf): Unit = {
     this.conf = conf
-    _namespace = conf.get(HA_ZK_NAMESPACE)
-    val maxSleepTime = conf.get(HA_ZK_CONN_MAX_RETRY_WAIT)
-    val maxRetries = conf.get(HA_ZK_CONN_MAX_RETRIES)
-    _zkClient = buildZookeeperClient(conf)
-    zkClient.getConnectionStateListenable.addListener(new ConnectionStateListener {
-      private val isConnected = new AtomicBoolean(false)
 
-      override def stateChanged(client: CuratorFramework, newState: ConnectionState): Unit = {
-        info(s"Zookeeper client connection state changed to: $newState")
-        newState match {
-          case CONNECTED | RECONNECTED => isConnected.set(true)
-          case LOST =>
-            isConnected.set(false)
-            val delay = maxRetries.toLong * maxSleepTime
-            connectionChecker.schedule(
-              new Runnable {
-                override def run(): Unit = if (!isConnected.get()) {
-                  error(s"Zookeeper client connection state changed to: $newState, but failed to" +
-                    s" reconnect in ${delay / 1000} seconds. Give up retry. ")
-                  stopGracefully()
-                }
-              },
-              delay,
-              TimeUnit.MILLISECONDS)
-          case _ =>
-        }
-      }
-    })
-    zkClient.start()
+    _discoveryClient = new ServiceDiscoveryClient(this)
+    discoveryClient.createClient(conf)
+
     super.initialize(conf)
   }
 
   override def start(): Unit = {
-    val instance = fe.connectionUrl
-    _serviceNode = createServiceNode(conf, zkClient, namespace, instance)
-    // Set a watch on the serviceNode
-    val watcher = new DeRegisterWatcher
-    if (zkClient.checkExists.usingWatcher(watcher).forPath(serviceNode.getActualPath) == null) {
-      // No node exists, throw exception
-      throw new KyuubiException(s"Unable to create znode for this Kyuubi " +
-        s"instance[${fe.connectionUrl}] on ZooKeeper.")
-    }
+    discoveryClient.registerService(conf)
     super.start()
   }
 
   override def stop(): Unit = {
-    closeServiceNode()
-    if (zkClient != null) zkClient.close()
+    discoveryClient.closeClient()
     super.stop()
   }
 
-  // close the EPHEMERAL_SEQUENTIAL node in zk
-  protected def closeServiceNode(): Unit = {
-    if (_serviceNode != null) {
-      try {
-        _serviceNode.close()
-      } catch {
-        case e: IOException =>
-          error("Failed to close the persistent ephemeral znode" + serviceNode.getActualPath, e)
-      } finally {
-        _serviceNode = null
-      }
-    }
-  }
-
   // stop the server genteelly
   def stopGracefully(): Unit = {
     stop()
@@ -143,23 +74,10 @@ abstract class ServiceDiscovery(
     fe.serverable.stop()
   }
 
-  class DeRegisterWatcher extends Watcher {
-    override def process(event: WatchedEvent): Unit = {
-      if (event.getType == Watcher.Event.EventType.NodeDeleted) {
-        warn(s"This Kyuubi instance ${fe.connectionUrl} is now de-registered from" +
-          s" ZooKeeper. The server will be shut down after the last client session completes.")
-        stopGracefully()
-      }
-    }
-  }
-
 }
 
 object ServiceDiscovery extends Logging {
 
-  final private lazy val connectionChecker =
-    ThreadUtils.newDaemonSingleThreadScheduledExecutor("zk-connection-checker")
-
   def supportServiceDiscovery(conf: KyuubiConf): Boolean = {
     val zkEnsemble = conf.get(HA_ZK_QUORUM)
     zkEnsemble != null && zkEnsemble.nonEmpty
@@ -234,97 +152,6 @@ object ServiceDiscovery extends Logging {
       external: Boolean = false): String = {
     createServiceNode(conf, zkClient, namespace, instance, version, external).getActualPath
   }
-
-  private def createServiceNode(
-      conf: KyuubiConf,
-      zkClient: CuratorFramework,
-      namespace: String,
-      instance: String,
-      version: Option[String] = None,
-      external: Boolean = false): PersistentNode = {
-    val ns = ZKPaths.makePath(null, namespace)
-    try {
-      zkClient
-        .create()
-        .creatingParentsIfNeeded()
-        .withMode(PERSISTENT)
-        .forPath(ns)
-    } catch {
-      case _: NodeExistsException => // do nothing
-      case e: KeeperException =>
-        throw new KyuubiException(s"Failed to create namespace '$ns'", e)
-    }
-
-    val session = conf.get(HA_ZK_ENGINE_REF_ID)
-      .map(refId => s"refId=$refId;").getOrElse("")
-    val pathPrefix = ZKPaths.makePath(
-      namespace,
-      s"serviceUri=$instance;version=${version.getOrElse(KYUUBI_VERSION)};${session}sequence=")
-    var serviceNode: PersistentNode = null
-    val createMode =
-      if (external) CreateMode.PERSISTENT_SEQUENTIAL
-      else CreateMode.EPHEMERAL_SEQUENTIAL
-    val znodeData =
-      if (conf.get(HA_ZK_PUBLIST_CONFIGS) && session.isEmpty) {
-        addConfsToPublish(conf, instance)
-      } else {
-        instance
-      }
-    try {
-      serviceNode = new PersistentNode(
-        zkClient,
-        createMode,
-        false,
-        pathPrefix,
-        znodeData.getBytes(StandardCharsets.UTF_8))
-      serviceNode.start()
-      val znodeTimeout = conf.get(HA_ZK_NODE_TIMEOUT)
-      if (!serviceNode.waitForInitialCreate(znodeTimeout, TimeUnit.MILLISECONDS)) {
-        throw new KyuubiException(s"Max znode creation wait time $znodeTimeout s exhausted")
-      }
-      info(s"Created a ${serviceNode.getActualPath} on ZooKeeper for KyuubiServer uri: " + instance)
-    } catch {
-      case e: Exception =>
-        if (serviceNode != null) {
-          serviceNode.close()
-        }
-        throw new KyuubiException(
-          s"Unable to create a znode for this server instance: $instance",
-          e)
-    }
-    serviceNode
-  }
-
-  /**
-   * Refer to the implementation of HIVE-11581 to simplify user connection parameters.
-   * https://issues.apache.org/jira/browse/HIVE-11581
-   * HiveServer2 should store connection params in ZK
-   * when using dynamic service discovery for simpler client connection string.
-   */
-  private def addConfsToPublish(conf: KyuubiConf, instance: String): String = {
-    if (!instance.contains(":")) {
-      return instance
-    }
-    val hostPort = instance.split(":", 2)
-    val confsToPublish = collection.mutable.Map[String, String]()
-
-    // Hostname
-    confsToPublish += ("hive.server2.thrift.bind.host" -> hostPort(0))
-    // Transport mode
-    confsToPublish += ("hive.server2.transport.mode" -> "binary")
-    // Transport specific confs
-    confsToPublish += ("hive.server2.thrift.port" -> hostPort(1))
-    confsToPublish += ("hive.server2.thrift.sasl.qop" -> conf.get(KyuubiConf.SASL_QOP))
-    // Auth specific confs
-    val authenticationMethod = conf.get(KyuubiConf.AUTHENTICATION_METHOD).mkString(",")
-    confsToPublish += ("hive.server2.authentication" -> authenticationMethod)
-    if (authenticationMethod.equalsIgnoreCase("KERBEROS")) {
-      confsToPublish += ("hive.server2.authentication.kerberos.principal" ->
-        conf.get(KyuubiConf.SERVER_PRINCIPAL).map(KyuubiHadoopUtils.getServerPrincipal)
-          .getOrElse(""))
-    }
-    confsToPublish.map { case (k, v) => k + "=" + v }.mkString(";")
-  }
 }
 
 case class ServiceNodeInfo(