[KYUUBI #7838] Fix concurrent ZooKeeper service deregistration

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

@@ -32,7 +32,8 @@ class KyuubiServiceDiscovery(
     if (!isServerLost.get()) {
       discoveryClient.deregisterService()
       discoveryClient.closeClient()
-      gracefulShutdownLatch.await() // wait for graceful shutdown triggered by watcher
+      // Closing the client can prevent delivery of the deletion notification.
+      waitUntilNoActiveSessions()
     } else {
       warn(s"The Zookeeper ensemble is LOST")
     }

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

@@ -17,7 +17,7 @@
 
 package org.apache.kyuubi.ha.client
 
-import java.util.concurrent.{CountDownLatch, TimeUnit}
+import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicBoolean
 
 import org.apache.kyuubi.Logging
@@ -35,7 +35,6 @@ abstract class ServiceDiscovery(
     name: String,
     val fe: FrontendService) extends AbstractService(name) {
 
-  protected val gracefulShutdownLatch = new CountDownLatch(1)
   protected val isServerLost = new AtomicBoolean(false)
 
   /**
@@ -64,16 +63,18 @@ abstract class ServiceDiscovery(
     super.start()
   }
 
-  // stop the server genteelly
-  def stopGracefully(isLost: Boolean = false): Unit = {
+  final protected def waitUntilNoActiveSessions(): Unit = {
     var activeSessionCount = fe.be.sessionManager.getActiveUserSessionCount
     while (activeSessionCount > 0) {
       info(s"$activeSessionCount connection(s) are active, delay shutdown")
       Thread.sleep(TimeUnit.SECONDS.toMillis(10))
       activeSessionCount = fe.be.sessionManager.getActiveUserSessionCount
     }
+  }
+
+  def stopGracefully(isLost: Boolean = false): Unit = {
+    waitUntilNoActiveSessions()
     isServerLost.set(isLost)
-    gracefulShutdownLatch.countDown()
     fe.serverable.stop()
   }
 }

kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/zookeeper/ZookeeperDiscoveryClient.scala

@@ -230,27 +230,33 @@ class ZookeeperDiscoveryClient(conf: KyuubiConf) extends DiscoveryClient {
       external: Boolean = false): Unit = {
     val instance = serviceDiscovery.fe.connectionUrl
     watcher = new DeRegisterWatcher(instance, serviceDiscovery)
-    serviceNode = createPersistentNode(
+    val node = createPersistentNode(
       conf,
       namespace,
       instance,
       version,
       external,
       serviceDiscovery.fe.attributes)
-    // Set a watch on the serviceNode
+    synchronized {
+      serviceNode = node
+    }
     watchNode()
   }
 
   override def deregisterService(): Unit = {
-    // close the EPHEMERAL_SEQUENTIAL node in zk
-    if (serviceNode != null) {
+    // Explicit shutdown and the deletion watcher can deregister concurrently.
+    val node = synchronized {
+      val current = serviceNode
+      serviceNode = null
+      current
+    }
+    if (node != null) {
+      val path = node.getActualPath
       try {
-        serviceNode.close()
+        node.close()
       } catch {
         case e @ (_: IOException | _: KeeperException) =>
-          error("Failed to close the persistent ephemeral znode" + serviceNode.getActualPath, e)
-      } finally {
-        serviceNode = null
+          error(s"Failed to close the persistent ephemeral znode $path", e)
       }
     }
   }
@@ -401,11 +407,12 @@ class ZookeeperDiscoveryClient(conf: KyuubiConf) extends DiscoveryClient {
   }
 
   private def watchNode(): Unit = {
-    if (zkClient.checkExists
-        .usingWatcher(watcher.asInstanceOf[Watcher]).forPath(serviceNode.getActualPath) == null) {
-      // No node exists, throw exception
-      throw new KyuubiException(s"Unable to create znode for this Kyuubi " +
-        s"instance[${watcher.instance}] on ZooKeeper.")
+    // A queued NodeDataChanged event can arrive after deregistration.
+    Option(serviceNode).flatMap(node => Option(node.getActualPath)).foreach { path =>
+      if (zkClient.checkExists.usingWatcher(watcher.asInstanceOf[Watcher]).forPath(path) == null) {
+        throw new KyuubiException(s"Unable to create znode for this Kyuubi " +
+          s"instance[${watcher.instance}] on ZooKeeper.")
+      }
     }
   }
 

kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/zookeeper/ZookeeperDiscoveryClientSuite.scala

@@ -20,7 +20,8 @@ package org.apache.kyuubi.ha.client.zookeeper
 import java.io.{File, IOException}
 import java.net.InetAddress
 import java.util
-import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.{CountDownLatch, Executors, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger}
 import javax.security.auth.login.Configuration
 
 import scala.collection.JavaConverters._
@@ -37,8 +38,9 @@ import org.apache.kyuubi.ha.client.zookeeper.ZookeeperClientProvider._
 import org.apache.kyuubi.jdbc.hive.strategy.{ServerSelectStrategy, ServerSelectStrategyFactory}
 import org.apache.kyuubi.service._
 import org.apache.kyuubi.shaded.curator.framework.{CuratorFramework, CuratorFrameworkFactory}
+import org.apache.kyuubi.shaded.curator.framework.recipes.nodes.PersistentNode
 import org.apache.kyuubi.shaded.curator.retry.ExponentialBackoffRetry
-import org.apache.kyuubi.shaded.zookeeper.ZooDefs
+import org.apache.kyuubi.shaded.zookeeper.{CreateMode, WatchedEvent, Watcher, ZooDefs}
 import org.apache.kyuubi.shaded.zookeeper.data.ACL
 import org.apache.kyuubi.util.reflect.ReflectUtils._
 import org.apache.kyuubi.zookeeper.EmbeddedZookeeper
@@ -199,6 +201,104 @@ abstract class ZookeeperDiscoveryClientSuite extends DiscoveryClientTests
     }
   }
 
+  test("node deletion can complete graceful shutdown during explicit deregistration") {
+    val client = new ZookeeperDiscoveryClient(conf)
+    val zkClient = getField[CuratorFramework](client, "zkClient")
+    val path = "/concurrent-deregistration"
+    val closeStarted = new CountDownLatch(1)
+    val finishClose = new CountDownLatch(1)
+    val gracefulShutdown = new CountDownLatch(1)
+    val closeCalls = new AtomicInteger()
+    val serviceNode =
+      new PersistentNode(zkClient, CreateMode.EPHEMERAL, false, path, Array.empty[Byte]) {
+        override def close(): Unit = {
+          if (closeCalls.incrementAndGet() == 1) {
+            closeStarted.countDown()
+            assert(finishClose.await(10, TimeUnit.SECONDS))
+          }
+        }
+      }
+    val serviceNodeField = classOf[ZookeeperDiscoveryClient].getDeclaredField("serviceNode")
+    serviceNodeField.setAccessible(true)
+    serviceNodeField.set(client, serviceNode)
+
+    val discovery = new ServiceDiscovery("test", null) {
+      override def stopGracefully(isLost: Boolean): Unit = gracefulShutdown.countDown()
+    }
+    val watcher = new client.DeRegisterWatcher("test", discovery)
+    val executor = Executors.newSingleThreadExecutor()
+    try {
+      val deregistration = executor.submit(new Runnable {
+        override def run(): Unit = client.deregisterService()
+      })
+      assert(closeStarted.await(10, TimeUnit.SECONDS))
+      watcher.process(new WatchedEvent(
+        Watcher.Event.EventType.NodeDeleted,
+        Watcher.Event.KeeperState.SyncConnected,
+        path))
+      assert(gracefulShutdown.getCount == 0)
+      finishClose.countDown()
+      deregistration.get(10, TimeUnit.SECONDS)
+      assert(closeCalls.get() == 1)
+    } finally {
+      finishClose.countDown()
+      executor.shutdownNow()
+      client.closeClient()
+    }
+  }
+
+  test("explicit server stop completes while ZooKeeper watch callbacks are queued") {
+    val namespace = "queued-watch-shutdown"
+    conf.set(HA_NAMESPACE, namespace).set(KyuubiConf.FRONTEND_THRIFT_BINARY_BIND_PORT, 0)
+    var discovery: ServiceDiscovery = null
+    val server = new NoopTBinaryFrontendServer() {
+      override val frontendServices: Seq[NoopTBinaryFrontendService] = Seq(
+        new NoopTBinaryFrontendService(this) {
+          override val discoveryService: Option[Service] = {
+            discovery = new KyuubiServiceDiscovery(this)
+            Some(discovery)
+          }
+        })
+    }
+    server.initialize(conf)
+    server.start()
+
+    val callbackStarted = new CountDownLatch(1)
+    val releaseCallback = new CountDownLatch(1)
+    val executor = Executors.newSingleThreadExecutor()
+    try {
+      withDiscoveryClient(conf) { observer =>
+        val basePath = s"/$namespace"
+        val servicePath = s"$basePath/${observer.getChildren(basePath).head}"
+        val barrierPath = s"$basePath/barrier"
+        observer.create(barrierPath, CreateMode.PERSISTENT.name())
+        val zkClient = getField[CuratorFramework](discovery.discoveryClient, "zkClient")
+        zkClient.checkExists().usingWatcher(new Watcher {
+          override def process(event: WatchedEvent): Unit = {
+            if (event.getType == Watcher.Event.EventType.NodeDataChanged) {
+              callbackStarted.countDown()
+              releaseCallback.await()
+            }
+          }
+        }).forPath(barrierPath)
+        observer.setData(barrierPath, Array[Byte](1))
+        assert(callbackStarted.await(10, TimeUnit.SECONDS))
+        observer.setData(servicePath, observer.getData(servicePath))
+
+        val shutdown = executor.submit(new Runnable {
+          override def run(): Unit = server.stop()
+        })
+        shutdown.get(10, TimeUnit.SECONDS)
+        assert(discovery.getServiceState == ServiceState.STOPPED)
+        assert(observer.pathNonExists(servicePath))
+      }
+    } finally {
+      releaseCallback.countDown()
+      executor.shutdownNow()
+      server.stop()
+    }
+  }
+
   test("watcher for zookeeper") {
     val namespace = "kyuubiwatcher"
     var discovery: ServiceDiscovery = null