diff --git a/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/KyuubiServiceDiscovery.scala b/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/KyuubiServiceDiscovery.scala index 31618e46596..895db7b0464 100644 --- a/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/KyuubiServiceDiscovery.scala +++ b/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") } diff --git a/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/ServiceDiscovery.scala b/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/ServiceDiscovery.scala index 2b05a8a7b8c..96fab6a1ad2 100644 --- a/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/ServiceDiscovery.scala +++ b/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() } } diff --git a/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/zookeeper/ZookeeperDiscoveryClient.scala b/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/zookeeper/ZookeeperDiscoveryClient.scala index a559258eecb..9358aaf64ed 100644 --- a/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/zookeeper/ZookeeperDiscoveryClient.scala +++ b/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.") + } } } diff --git a/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/zookeeper/ZookeeperDiscoveryClientSuite.scala b/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/zookeeper/ZookeeperDiscoveryClientSuite.scala index eef26dc00e5..149d2cf160c 100644 --- a/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/zookeeper/ZookeeperDiscoveryClientSuite.scala +++ b/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