From c9260b5edfb806f3743b27e17d53d6b91b5107d2 Mon Sep 17 00:00:00 2001 From: DreamWaking Date: Sat, 11 Jul 2026 15:48:20 +0800 Subject: [PATCH] [KYUUBI #7556] Rearm the ZooKeeper engine deregistration watcher after reconnect --- .../zookeeper/ZookeeperDiscoveryClient.scala | 211 +++++++-- .../ZookeeperDiscoveryClientSuite.scala | 410 +++++++++++++++++- 2 files changed, 589 insertions(+), 32 deletions(-) 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 bead620536a..d73a1023cf9 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 @@ -19,10 +19,11 @@ package org.apache.kyuubi.ha.client.zookeeper import java.io.IOException import java.nio.charset.StandardCharsets -import java.util.concurrent.TimeUnit +import java.util.concurrent.{ScheduledFuture, TimeUnit} import java.util.concurrent.atomic.AtomicBoolean import scala.collection.JavaConverters._ +import scala.util.control.NonFatal import com.google.common.annotations.VisibleForTesting @@ -32,7 +33,6 @@ import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_ENGINE_ID import org.apache.kyuubi.ha.HighAvailabilityConf.{HA_ENGINE_REF_ID, HA_ZK_NODE_TIMEOUT, HA_ZK_PUBLISH_CONFIGS} import org.apache.kyuubi.ha.client.{DiscoveryClient, ServiceDiscovery, ServiceNodeInfo} import org.apache.kyuubi.ha.client.zookeeper.ZookeeperClientProvider.{buildZookeeperClient, getGracefulStopThreadDelay} -import org.apache.kyuubi.ha.client.zookeeper.ZookeeperDiscoveryClient.connectionChecker import org.apache.kyuubi.shaded.curator.framework.CuratorFramework import org.apache.kyuubi.shaded.curator.framework.recipes.atomic.{AtomicValue, DistributedAtomicInteger} import org.apache.kyuubi.shaded.curator.framework.recipes.locks.InterProcessSemaphoreMutex @@ -50,7 +50,15 @@ class ZookeeperDiscoveryClient(conf: KyuubiConf) extends DiscoveryClient { private val zkClient: CuratorFramework = buildZookeeperClient(conf) @volatile private var serviceNode: PersistentNode = _ - private var watcher: DeRegisterWatcher = _ + @volatile private var watcher: DeRegisterWatcher = _ + + // Monotonic generation that tags each RECONNECTED rewatch attempt. A retry scheduled for an + // older generation cancels itself when a newer RECONNECTED arrives, so repeated reconnects + // during ZK flapping never stack independent retry chains (single-flight). + private val rewatchGeneration = new java.util.concurrent.atomic.AtomicLong(0L) + // The currently pending scheduled retry, if any. Cancelled on deregister so a stale task cannot + // rearm a watcher on a node that is already torn down. + @volatile private var rewatchRetryTask: ScheduledFuture[_] = _ override def createClient(): Unit = { zkClient.start() @@ -99,32 +107,48 @@ class ZookeeperDiscoveryClient(conf: KyuubiConf) extends DiscoveryClient { } override def monitorState(serviceDiscovery: ServiceDiscovery): Unit = { - 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 = getGracefulStopThreadDelay(conf) - connectionChecker.schedule( - new Runnable { - override def run(): Unit = if (!isConnected.get()) { - error(s"Zookeeper client connection state changed to: $newState," + - s" but failed to reconnect in ${delay / 1000} seconds." + - s" Give up retry and stop gracefully . ") - serviceDiscovery.stopGracefully(true) - } - }, - delay, - TimeUnit.MILLISECONDS) - case _ => - } + registerConnectionStateListener(serviceDiscovery) + } + + @VisibleForTesting + private[kyuubi] def registerConnectionStateListener( + serviceDiscovery: ServiceDiscovery): ConnectionStateListener = { + val listener = 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 => isConnected.set(true) + case RECONNECTED => + isConnected.set(true) + // Rearm the DeRegisterWatcher off the connection-state dispatch thread. Curator 2.12 + // dispatches ConnectionStateListener callbacks on a single-thread executor + // shared with PersistentNode's recovery listener; a foreground checkExists here could + // block that thread (and delay PersistentNode re-creating the znode and subsequent + // LOST handling) when the connection drops again mid-call. Schedule the rewatch on a + // dedicated executor so stateChanged returns immediately. + scheduleRewatch("Zookeeper client reconnected") + case LOST => + isConnected.set(false) + val delay = getGracefulStopThreadDelay(conf) + ZookeeperDiscoveryClient.connectionChecker.schedule( + new Runnable { + override def run(): Unit = if (!isConnected.get()) { + error(s"Zookeeper client connection state changed to: $newState," + + s" but failed to reconnect in ${delay / 1000} seconds." + + s" Give up retry and stop gracefully . ") + serviceDiscovery.stopGracefully(true) + } + }, + delay, + TimeUnit.MILLISECONDS) + case _ => } - }) + } + } + zkClient.getConnectionStateListenable.addListener(listener) + listener } override def tryWithLock[T](lockPath: String, timeout: Long)(f: => T): T = { @@ -252,6 +276,9 @@ class ZookeeperDiscoveryClient(conf: KyuubiConf) extends DiscoveryClient { serviceNode = null } } + // Cancel any pending rewatch retry so a stale task cannot rearm a watcher on a node that has + // just been torn down. Bump the generation first so an in-flight retry observes the change. + cancelRewatchRetry() } override def postDeregisterService(namespace: String): Boolean = { @@ -400,11 +427,117 @@ class ZookeeperDiscoveryClient(conf: KyuubiConf) extends DiscoveryClient { } private def watchNode(): Unit = { - if (zkClient.checkExists - .usingWatcher(watcher.asInstanceOf[Watcher]).forPath(serviceNode.getActualPath) == null) { + if (!watchServiceNode("service registration")) { // No node exists, throw exception throw new KyuubiException(s"Unable to create znode for this Kyuubi " + - s"instance[${watcher.instance}] on ZooKeeper.") + s"instance[${Option(watcher).map(_.instance).getOrElse("unknown")}] on ZooKeeper.") + } + } + + private def watchServiceNode(reason: String): Boolean = { + val currentServiceNode = serviceNode + val currentWatcher = watcher + if (currentServiceNode == null || currentWatcher == null) { + debug(s"Skip setting DeRegisterWatcher for $reason because service node is not registered") + return false + } + + val path = currentServiceNode.getActualPath + if (path == null) { + debug(s"Skip setting DeRegisterWatcher for $reason because service node path is not ready") + return false + } + + val serviceNodeExists = zkClient.checkExists + .usingWatcher(currentWatcher.asInstanceOf[Watcher]) + .forPath(path) != null + if (serviceNodeExists) { + info(s"Set DeRegisterWatcher on $path for $reason") + } else { + warn(s"Service node $path does not exist when setting DeRegisterWatcher for $reason") + } + serviceNodeExists + } + + /** + * Attempt to rearm the DeRegisterWatcher off the connection-state dispatch thread. Runs on a + * dedicated rewatch executor (isolated from the shared connectionChecker that the LOST graceful + * stop task may block on). Retries are generation-gated: a newer RECONNECTED bumps the + * generation, so a retry scheduled for an older generation self-cancels instead of stacking + * another independent chain. + * + * The scheduled task only invokes private methods (not field reads) so that the private + * serviceNode/watcher fields are not accessed from an anonymous inner class, which would make + * scalac mangle their names and break reflection-based test setup. + */ + private def scheduleRewatch(reason: String): Unit = { + scheduleRewatch(reason, attempt = 1, rewatchGeneration.incrementAndGet()) + } + + private def scheduleRewatch(reason: String, attempt: Int, generation: Long): Unit = { + rewatchRetryTask = ZookeeperDiscoveryClient.rewatchExecutor.schedule( + new Runnable { + override def run(): Unit = runRewatchAttempt(reason, attempt, generation) + }, + 0L, + TimeUnit.MILLISECONDS) + } + + private def runRewatchAttempt(reason: String, attempt: Int, generation: Long): Unit = { + if (!rewatchStillActive(generation)) { + debug(s"Discard stale rewatch for $reason at attempt $attempt" + + s" (generation $generation vs ${rewatchGeneration.get()})") + return + } + try { + if (watchServiceNode(s"$reason (attempt $attempt)")) { + rewatchRetryTask = null + } else if (attempt < ZookeeperDiscoveryClient.RewatchNodeMaxRetries && + rewatchStillActive(generation)) { + val nextAttempt = attempt + 1 + warn(s"Failed to rewatch service node for $reason at attempt $attempt, " + + s"will retry at attempt $nextAttempt") + rewatchRetryTask = ZookeeperDiscoveryClient.rewatchExecutor.schedule( + new Runnable { + override def run(): Unit = runRewatchAttempt(reason, nextAttempt, generation) + }, + ZookeeperDiscoveryClient.RewatchNodeRetryIntervalMs, + TimeUnit.MILLISECONDS) + } else { + warn(s"Failed to rewatch service node for $reason after $attempt attempt(s)") + rewatchRetryTask = null + } + } catch { + case NonFatal(e) => + warn(s"Failed to set DeRegisterWatcher for $reason at attempt $attempt", e) + rewatchRetryTask = null + } + } + + // A rewatch attempt is still active only if no newer RECONNECTED superseded it and the service + // node has not been torn down. Reads the volatile fields here (on the rewatch executor thread), + // never from an anonymous inner class. + private def rewatchStillActive(generation: Long): Boolean = { + generation == rewatchGeneration.get() && serviceNode != null && watcher != null + } + + private def cancelRewatchRetry(): Unit = { + rewatchGeneration.incrementAndGet() + val task = rewatchRetryTask + rewatchRetryTask = null + if (task != null) { + task.cancel(false) + } + } + + @VisibleForTesting + private[kyuubi] def rewatchServiceNodeOnce(reason: String): Boolean = { + try { + watchServiceNode(reason) + } catch { + case NonFatal(e) => + warn(s"Failed to set DeRegisterWatcher for $reason", e) + false } } @@ -416,6 +549,14 @@ class ZookeeperDiscoveryClient(conf: KyuubiConf) extends DiscoveryClient { " ZooKeeper. The server will be shut down after the last client session completes.") ZookeeperDiscoveryClient.this.deregisterService() serviceDiscovery.stopGracefully() + } else if (event.getType == Watcher.Event.EventType.NodeCreated) { + // When the service node was absent at rewatch time, exists(path, watcher) leaves a + // one-shot watch that PersistentNode's recovery createNode consumes via NodeCreated. + // The DeRegisterWatcher must rearm itself here, otherwise the next NodeDeleted would be + // missed until the following scheduled retry. Reuse the same non-throwing helper. + info(s"This Kyuubi instance $instance sees its service node recreated;" + + " rearming DeRegisterWatcher") + ZookeeperDiscoveryClient.this.rewatchServiceNodeOnce("service node recreated") } else if (event.getType == Watcher.Event.EventType.NodeDataChanged) { warn(s"This Kyuubi instance $instance now receives the NodeDataChanged event") ZookeeperDiscoveryClient.this.watchNode() @@ -425,6 +566,14 @@ class ZookeeperDiscoveryClient(conf: KyuubiConf) extends DiscoveryClient { } object ZookeeperDiscoveryClient extends Logging { + final private val RewatchNodeMaxRetries = 5 + final private val RewatchNodeRetryIntervalMs = 1000L + // Dedicated executor for RECONNECTED rewatch work, intentionally separate from the shared + // connectionChecker: a LOST graceful-stop task may block that single thread waiting for open + // sessions to drain, which would otherwise starve rewatch retries of binary/HTTP discovery + // clients sharing the static executor. + final private lazy val rewatchExecutor = + ThreadUtils.newDaemonSingleThreadScheduledExecutor("zk-rewatch") final private lazy val connectionChecker = ThreadUtils.newDaemonSingleThreadScheduledExecutor("zk-connection-checker") } 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..74948450a65 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 @@ -19,7 +19,9 @@ package org.apache.kyuubi.ha.client.zookeeper import java.io.{File, IOException} import java.net.InetAddress +import java.nio.charset.StandardCharsets import java.util +import java.util.concurrent.TimeUnit import java.util.concurrent.atomic.AtomicBoolean import javax.security.auth.login.Configuration @@ -37,8 +39,10 @@ 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.framework.state.ConnectionState import org.apache.kyuubi.shaded.curator.retry.ExponentialBackoffRetry -import org.apache.kyuubi.shaded.zookeeper.ZooDefs +import org.apache.kyuubi.shaded.zookeeper.{CreateMode, ZooDefs} import org.apache.kyuubi.shaded.zookeeper.data.ACL import org.apache.kyuubi.util.reflect.ReflectUtils._ import org.apache.kyuubi.zookeeper.EmbeddedZookeeper @@ -73,6 +77,12 @@ abstract class ZookeeperDiscoveryClientSuite extends DiscoveryClientTests def stopZk(): Unit + private def setField(target: AnyRef, fieldName: String, value: AnyRef): Unit = { + val field = target.getClass.getDeclaredField(fieldName) + field.setAccessible(true) + field.set(target, value) + } + override def beforeEach(): Unit = { startZk() conf = new KyuubiConf().set(HA_ADDRESSES, getConnectString) @@ -296,4 +306,402 @@ abstract class ZookeeperDiscoveryClientSuite extends DiscoveryClientTests zkClient.close() } + + test("rewatch service node after zookeeper reconnect") { + val namespace = "kyuubirewatch" + conf.set(HA_ZK_CONN_RETRY_POLICY, "ONE_TIME") + .set(HA_ZK_CONN_BASE_RETRY_WAIT, 1) + .set(HA_ZK_SESSION_TIMEOUT, 2000) + .set(HA_ADDRESSES, getConnectString) + .set(HA_NAMESPACE, namespace) + .set(KyuubiConf.FRONTEND_THRIFT_BINARY_BIND_PORT, 0) + + val stopped = new AtomicBoolean(false) + val service = new NoopTBinaryFrontendServer() + val discovery = new ServiceDiscovery("test discovery", service.frontendServices.head) { + override def stopGracefully(isLost: Boolean = false): Unit = { + stopped.set(true) + } + } + val zkDiscoveryClient = new ZookeeperDiscoveryClient(conf) + val basePath = s"/$namespace" + var serviceNode: PersistentNode = null + + try { + zkDiscoveryClient.createClient() + val zkClient = + getField[CuratorFramework](zkDiscoveryClient, "zkClient") + zkClient.create().creatingParentsIfNeeded().withMode(CreateMode.PERSISTENT).forPath(basePath) + + val instance = "localhost:10009" + serviceNode = new PersistentNode( + zkClient, + CreateMode.EPHEMERAL_SEQUENTIAL, + false, + s"$basePath/serviceUri=$instance;version=$KYUUBI_VERSION;sequence=", + instance.getBytes(StandardCharsets.UTF_8)) + serviceNode.start() + assert(serviceNode.waitForInitialCreate(conf.get(HA_ZK_NODE_TIMEOUT), TimeUnit.MILLISECONDS)) + + val watcher = new zkDiscoveryClient.DeRegisterWatcher(instance, discovery) + setField(zkDiscoveryClient, "serviceNode", serviceNode) + setField(zkDiscoveryClient, "watcher", watcher) + assert(zkDiscoveryClient.rewatchServiceNodeOnce("test-reconnect")) + + withDiscoveryClient(conf) { discoveryClient => + assert(discoveryClient.pathExists(serviceNode.getActualPath)) + discoveryClient.delete(serviceNode.getActualPath) + + eventually(timeout(10.seconds), interval(100.millis)) { + assert(stopped.get()) + } + } + } finally { + if (serviceNode != null) { + try { + serviceNode.close() + } catch { + case _: Exception => + } + } + zkDiscoveryClient.closeClient() + } + } + + // Drive a real RECONNECTED event through the ConnectionStateListener registered by the + // production monitorState path and prove the engine-side DeRegisterWatcher is rearmed so a + // post-reconnect delete triggers stop. Retain the registered listener instead of enumerating + // Curator's internal listener container, whose implementation differs across Curator versions. + test("real RECONNECTED rearms DeRegisterWatcher so delete triggers stop") { + val namespace = "kyuubireconnect" + val reconnectConf = new KyuubiConf() + .set(HA_ZK_CONN_RETRY_POLICY, "ONE_TIME") + .set(HA_ZK_CONN_BASE_RETRY_WAIT, 1) + .set(HA_ZK_SESSION_TIMEOUT, 2000) + .set(HA_ADDRESSES, getConnectString) + .set(HA_NAMESPACE, namespace) + .set(KyuubiConf.FRONTEND_THRIFT_BINARY_BIND_PORT, 0) + + val stopped = new AtomicBoolean(false) + val service = new NoopTBinaryFrontendServer() + val discovery = new ServiceDiscovery("test discovery", service.frontendServices.head) { + override def stopGracefully(isLost: Boolean = false): Unit = { + stopped.set(true) + } + } + val zkDiscoveryClient = new ZookeeperDiscoveryClient(reconnectConf) + val basePath = s"/$namespace" + var serviceNode: PersistentNode = null + + try { + zkDiscoveryClient.createClient() + val zkClient = getField[CuratorFramework](zkDiscoveryClient, "zkClient") + val connectionStateListener = + zkDiscoveryClient.registerConnectionStateListener(discovery) + zkClient.create().creatingParentsIfNeeded().withMode(CreateMode.PERSISTENT).forPath(basePath) + + val instance = "localhost:10010" + serviceNode = new PersistentNode( + zkClient, + CreateMode.EPHEMERAL_SEQUENTIAL, + false, + s"$basePath/serviceUri=$instance;version=$KYUUBI_VERSION;sequence=", + instance.getBytes(StandardCharsets.UTF_8)) + serviceNode.start() + assert(serviceNode.waitForInitialCreate( + reconnectConf.get(HA_ZK_NODE_TIMEOUT), + TimeUnit.MILLISECONDS)) + + val watcher = new zkDiscoveryClient.DeRegisterWatcher(instance, discovery) + setField(zkDiscoveryClient, "serviceNode", serviceNode) + setField(zkDiscoveryClient, "watcher", watcher) + // NOTE: no initial watch is armed here. This models a session reconnect where the previous + // one-shot watch from registerService is gone (ZooKeeper drops a disconnected session's + // watches without delivering an event). On the old code RECONNECTED only sets isConnected, so + // the watcher stays unarmed; on the fixed code RECONNECTED schedules the rearm. + + // Dispatch RECONNECTED through the listener implementation used by monitorState. On the + // fixed code this rearms the DeRegisterWatcher off the connection-state dispatch thread; on + // the old code it does nothing. + connectionStateListener.stateChanged(zkClient, ConnectionState.RECONNECTED) + // give the async rewatch executor a moment to rearm + Thread.sleep(1000) + + // Now delete the live service node. The rearmed watcher must observe NodeDeleted -> stop. + // On the old code (no rearm on RECONNECTED) this delete is missed and stopped stays false. + withDiscoveryClient(reconnectConf) { discoveryClient => + assert(discoveryClient.pathExists(serviceNode.getActualPath)) + discoveryClient.delete(serviceNode.getActualPath) + } + eventually(timeout(10.seconds), interval(100.millis)) { + assert(stopped.get(), "stop not triggered after RECONNECTED-delete sequence") + } + } finally { + if (serviceNode != null) { + try serviceNode.close() + catch { case _: Exception => } + } + zkDiscoveryClient.closeClient() + } + } + + // Absent path -> PersistentNode recreate (NodeCreated) -> + // DeRegisterWatcher rearms on NodeCreated -> next delete is observed. + test("absent node then NodeCreated rearms watcher so delete is observed") { + val namespace = "kyuubirearm" + conf.set(HA_ZK_CONN_RETRY_POLICY, "ONE_TIME") + .set(HA_ZK_CONN_BASE_RETRY_WAIT, 1) + .set(HA_ZK_SESSION_TIMEOUT, 2000) + .set(HA_ADDRESSES, getConnectString) + .set(HA_NAMESPACE, namespace) + .set(KyuubiConf.FRONTEND_THRIFT_BINARY_BIND_PORT, 0) + + val stopped = new AtomicBoolean(false) + val service = new NoopTBinaryFrontendServer() + val discovery = new ServiceDiscovery("test discovery", service.frontendServices.head) { + override def stopGracefully(isLost: Boolean = false): Unit = { + stopped.set(true) + } + } + val zkDiscoveryClient = new ZookeeperDiscoveryClient(conf) + val basePath = s"/$namespace" + var serviceNode: PersistentNode = null + + try { + zkDiscoveryClient.createClient() + val zkClient = getField[CuratorFramework](zkDiscoveryClient, "zkClient") + zkClient.create().creatingParentsIfNeeded().withMode(CreateMode.PERSISTENT).forPath(basePath) + + val instance = "localhost:10011" + serviceNode = new PersistentNode( + zkClient, + CreateMode.EPHEMERAL_SEQUENTIAL, + false, + s"$basePath/serviceUri=$instance;version=$KYUUBI_VERSION;sequence=", + instance.getBytes(StandardCharsets.UTF_8)) + serviceNode.start() + assert(serviceNode.waitForInitialCreate(conf.get(HA_ZK_NODE_TIMEOUT), TimeUnit.MILLISECONDS)) + + val watcher = new zkDiscoveryClient.DeRegisterWatcher(instance, discovery) + setField(zkDiscoveryClient, "serviceNode", serviceNode) + setField(zkDiscoveryClient, "watcher", watcher) + + // Capture the actual path, then close the PersistentNode. close() drops the ephemeral node + // and clears nodePath (getActualPath -> null); PersistentNode's auto-recreate stops, giving a + // deterministic absent path. Re-inject livePath into the nodePath AtomicReference so the + // production rewatch helper still resolves the path to re-arm on. + val livePath = serviceNode.getActualPath + serviceNode.close() + val nodePathField = classOf[PersistentNode].getDeclaredField("nodePath") + nodePathField.setAccessible(true) + val nodePathRef = + nodePathField.get(serviceNode).asInstanceOf[ + java.util.concurrent.atomic.AtomicReference[String]] + nodePathRef.set(livePath) + assert(serviceNode.getActualPath === livePath) + + // rewatch on the absent path leaves an exists watch that fires on NodeCreated; returns false + assert( + zkDiscoveryClient.rewatchServiceNodeOnce("absent path rewatch") === false, + "rewatch on an absent path must return false but leave an exists watch") + + // Create the node on the SAME path via an external client; ZK emits NodeCreated on the + // exists watch, and the DeRegisterWatcher's NodeCreated branch must re-arm itself. + withDiscoveryClient(conf) { discoveryClient => + discoveryClient.create(livePath, "PERSISTENT") + assert(discoveryClient.pathExists(livePath)) + } + // give the NodeCreated watch a moment to rearm via the DeRegisterWatcher NodeCreated branch + Thread.sleep(1500) + + // now delete the recreated node; the rearmed watcher must observe NodeDeleted -> stop + withDiscoveryClient(conf) { discoveryClient => + assert(discoveryClient.pathExists(livePath)) + discoveryClient.delete(livePath) + } + eventually(timeout(10.seconds), interval(100.millis)) { + assert(stopped.get(), "delete after NodeCreated rearm was not observed") + } + } finally { + if (serviceNode != null) { + try serviceNode.close() + catch { case _: Exception => } + } + zkDiscoveryClient.closeClient() + } + } + + // After deregisterService, a pending or subsequent rewatch + // must not rearm the watcher on the torn-down node. + test("deregister cancels pending rewatch and does not rearm watcher") { + val namespace = "kyuubideregister" + conf.set(HA_ZK_CONN_RETRY_POLICY, "ONE_TIME") + .set(HA_ZK_CONN_BASE_RETRY_WAIT, 1) + .set(HA_ZK_SESSION_TIMEOUT, 2000) + .set(HA_ADDRESSES, getConnectString) + .set(HA_NAMESPACE, namespace) + .set(KyuubiConf.FRONTEND_THRIFT_BINARY_BIND_PORT, 0) + + val stopped = new AtomicBoolean(false) + val service = new NoopTBinaryFrontendServer() + val discovery = new ServiceDiscovery("test discovery", service.frontendServices.head) { + override def stopGracefully(isLost: Boolean = false): Unit = { + stopped.set(true) + } + } + val zkDiscoveryClient = new ZookeeperDiscoveryClient(conf) + val basePath = s"/$namespace" + var serviceNode: PersistentNode = null + + try { + zkDiscoveryClient.createClient() + val zkClient = getField[CuratorFramework](zkDiscoveryClient, "zkClient") + zkClient.create().creatingParentsIfNeeded().withMode(CreateMode.PERSISTENT).forPath(basePath) + + val instance = "localhost:10012" + serviceNode = new PersistentNode( + zkClient, + CreateMode.EPHEMERAL_SEQUENTIAL, + false, + s"$basePath/serviceUri=$instance;version=$KYUUBI_VERSION;sequence=", + instance.getBytes(StandardCharsets.UTF_8)) + serviceNode.start() + assert(serviceNode.waitForInitialCreate(conf.get(HA_ZK_NODE_TIMEOUT), TimeUnit.MILLISECONDS)) + + val watcher = new zkDiscoveryClient.DeRegisterWatcher(instance, discovery) + setField(zkDiscoveryClient, "serviceNode", serviceNode) + setField(zkDiscoveryClient, "watcher", watcher) + // Arm a pending rewatch WITHOUT arming the actual DeRegisterWatcher yet, so deregister's + // node close does not fire NodeDeleted->stop. Dispatch RECONNECTED through the production + // monitorState listener to schedule the rearm on the dedicated rewatch executor. + val connectionStateListener = + zkDiscoveryClient.registerConnectionStateListener(discovery) + val livePath = serviceNode.getActualPath + // drop the live znode WITHOUT a watch armed, so no NodeDeleted fires + serviceNode.close() + // now dispatch RECONNECTED; the production path schedules a rearm attempt that should find + // the path absent and start retrying. deregisterService() below must cancel that retry. + connectionStateListener.stateChanged(zkClient, ConnectionState.RECONNECTED) + Thread.sleep(300) + + // Deregister tears down the node and must cancel the pending rewatch retry. + zkDiscoveryClient.deregisterService() + // serviceNode is now null; a subsequent direct rewatch attempt must be a no-op. + assert(!zkDiscoveryClient.rewatchServiceNodeOnce("after deregister")) + + // Recreate the node on the same path; the cancelled/stale rewatch must NOT have rearmed a + // watcher, so this delete must not trigger stop. + withDiscoveryClient(conf) { discoveryClient => + discoveryClient.create(livePath, "PERSISTENT") + discoveryClient.delete(livePath) + } + Thread.sleep(2000) + assert(!stopped.get(), "stale rewatch rearmed a torn-down node") + } finally { + if (serviceNode != null) { + try serviceNode.close() + catch { case _: Exception => } + } + zkDiscoveryClient.closeClient() + } + } + + // Binary and HTTP discovery each own their own + // ZookeeperDiscoveryClient; their rewatch executors and serviceNodes are independent, so a rearm + // on one client must not arm or unarm the other client's watcher. + test("binary and HTTP discovery clients rearm independently") { + val namespace = "kyuubiindependent" + conf.set(HA_ZK_CONN_RETRY_POLICY, "ONE_TIME") + .set(HA_ZK_CONN_BASE_RETRY_WAIT, 1) + .set(HA_ZK_SESSION_TIMEOUT, 2000) + .set(HA_ADDRESSES, getConnectString) + .set(HA_NAMESPACE, namespace) + .set(KyuubiConf.FRONTEND_THRIFT_BINARY_BIND_PORT, 0) + + val stoppedBinary = new AtomicBoolean(false) + val stoppedHttp = new AtomicBoolean(false) + val service = new NoopTBinaryFrontendServer() + val discoveryBinary = + new ServiceDiscovery("binary discovery", service.frontendServices.head) { + override def stopGracefully(isLost: Boolean = false): Unit = stoppedBinary.set(true) + } + val discoveryHttp = + new ServiceDiscovery("http discovery", service.frontendServices.head) { + override def stopGracefully(isLost: Boolean = false): Unit = stoppedHttp.set(true) + } + val binaryClient = new ZookeeperDiscoveryClient(conf) + val httpClient = new ZookeeperDiscoveryClient(conf) + val basePath = s"/$namespace" + var binaryNode: PersistentNode = null + var httpNode: PersistentNode = null + + try { + binaryClient.createClient() + httpClient.createClient() + val binZkClient = getField[CuratorFramework](binaryClient, "zkClient") + val httpZkClient = getField[CuratorFramework](httpClient, "zkClient") + binZkClient.create().creatingParentsIfNeeded().withMode( + CreateMode.PERSISTENT).forPath(basePath) + + val binaryInstance = "localhost:10013" + val httpInstance = "localhost:10014" + binaryNode = new PersistentNode( + binZkClient, + CreateMode.EPHEMERAL_SEQUENTIAL, + false, + s"$basePath/serviceUri=$binaryInstance;version=$KYUUBI_VERSION;sequence=", + binaryInstance.getBytes(StandardCharsets.UTF_8)) + httpNode = new PersistentNode( + httpZkClient, + CreateMode.EPHEMERAL_SEQUENTIAL, + false, + s"$basePath/serviceUri=$httpInstance;version=$KYUUBI_VERSION;sequence=", + httpInstance.getBytes(StandardCharsets.UTF_8)) + binaryNode.start() + httpNode.start() + assert(binaryNode.waitForInitialCreate(conf.get(HA_ZK_NODE_TIMEOUT), TimeUnit.MILLISECONDS)) + assert(httpNode.waitForInitialCreate(conf.get(HA_ZK_NODE_TIMEOUT), TimeUnit.MILLISECONDS)) + + setField(binaryClient, "serviceNode", binaryNode) + setField( + binaryClient, + "watcher", + new binaryClient.DeRegisterWatcher(binaryInstance, discoveryBinary)) + setField(httpClient, "serviceNode", httpNode) + setField(httpClient, "watcher", new httpClient.DeRegisterWatcher(httpInstance, discoveryHttp)) + assert(binaryClient.rewatchServiceNodeOnce("binary initial")) + assert(httpClient.rewatchServiceNodeOnce("http initial")) + + // Delete the binary node only; the HTTP watcher must remain armed and unaffected. + withDiscoveryClient(conf) { discoveryClient => + assert(discoveryClient.pathExists(binaryNode.getActualPath)) + discoveryClient.delete(binaryNode.getActualPath) + } + eventually(timeout(10.seconds), interval(100.millis)) { + assert(stoppedBinary.get()) + } + // HTTP client was not deleted and must not have stopped. + assert(!stoppedHttp.get(), "independent HTTP client stopped on binary delete") + + // Now delete the HTTP node; its own (still armed) watcher must fire. + withDiscoveryClient(conf) { discoveryClient => + assert(discoveryClient.pathExists(httpNode.getActualPath)) + discoveryClient.delete(httpNode.getActualPath) + } + eventually(timeout(10.seconds), interval(100.millis)) { + assert(stoppedHttp.get()) + } + } finally { + if (binaryNode != null) { + try binaryNode.close() + catch { case _: Exception => } + } + if (httpNode != null) { + try httpNode.close() + catch { case _: Exception => } + } + binaryClient.closeClient() + httpClient.closeClient() + } + } }