[zhangxinyao88]

Why are the changes needed?

Related to #7838. In the Spark 4.3 CI run, explicit shutdown raced with ZooKeeper's deletion watcher. The watcher hit a null pointer exception in deregisterService() before signaling graceful shutdown, leaving the server waiting.

Take and clear the service node under a lock, then close it outside the lock. This lets the watcher continue graceful shutdown if another caller already took the node. A queued data-change event skips rearming when the node is gone.

How was this patch tested?

On Linux/x86 with JDK 25, the new regression test observed two close calls on the baseline and one with the fix. All 12 tests in EmbeddedZookeeperDiscoveryClientSuite passed with the fix. Spotless and Scalastyle passed.

Was this patch assisted by generative AI tooling?

Assisted-by: Codex (GPT-6)

[wangzhigang1999]

Thanks for adding these diagnostics! I checked the JDK25/Spark4.3 artifacts and found a useful clue.

After DescribeSessionSuite passes, the ZooKeeper NodeDeleted watcher throws an NPE in ZookeeperDiscoveryClient.deregisterService():248. It appears that the shutdown thread and watcher race on serviceNode: one clears it between the other’s null check and close() call.

The exception prevents the watcher from reaching stopGracefully(), leaving gracefulShutdownLatch unreleased. Both thread dumps show the ScalaTest main thread waiting in KyuubiServiceDiscovery.stop():35, with Maven waiting for the test JVM to exit. This explains the hang in this run.

Would you be interested in following up with a fix and a regression test for this race? Thanks again for helping capture the evidence!

[Copilot]

🟡 Changes recommended

A queued data-change callback can drop the deletion watch during deregistration, leaving graceful shutdown blocked indefinitely in an ordering the current test does not cover.

1 open issue

• Rearm the node watch after deregistration to avoid shutdown hangs · New

What changed in this PR

This PR targets #7838’s shutdown hang by coordinating concurrent ZooKeeper service deregistration in Kyuubi’s service-discovery layer.

Changes:

• Atomically take and clear the service node, then close it outside the lock.
• Skip watcher rearming when the node or its path is unavailable.
• Add a deterministic regression test for explicit deregistration racing a deletion callback.

File	Description
kyuubi-ha/​src/​test/​scala/​org/​apache/​kyuubi/​ha/​client/​zookeeper/​ZookeeperDiscoveryClientSuite.scala	Adds concurrent deregistration coverage, checking single-close behavior and graceful-shutdown signaling.
kyuubi-ha/​src/​main/​scala/​org/​apache/​kyuubi/​ha/​client/​zookeeper/​ZookeeperDiscoveryClient.scala	Synchronizes service-node ownership and guards watcher rearming after deregistration.

_{Review effort: Max}

[zhangxinyao88]

Fixed in 1c7124f. Explicit shutdown now waits for active sessions without relying on the ZooKeeper deletion callback. The queued-callback regression passed in the JDK 25 / Spark 4.3 fork run.

[zhangxinyao88]

Would you be interested in following up with a fix and a regression test for this race?

Yes, both are in this PR. I also added a regression for the case where the ZooKeeper callback is still queued during shutdown. On 1c7124f, the new regression passed and DescribeSessionSuite completed shutdown in the JDK 25 / Spark 4.3 fork run.

Original comment

[wangzhigang1999]

LGTM. The fix addresses the concurrent deregistration race and removes shutdown’s dependency on the deletion callback. Regression tests and ACK fault-injection checks passed.

[wangzhigang1999]

As a non-blocking follow-up, session draining still blocks the ZooKeeper event thread or the shared connection checker, and concurrent shutdown paths could use better coordination. Would you be interested in opening a separate issue and working on these lifecycle improvements?

[ruanwenjun]

It seems like a simpler approach might be to make registerService and deregisterService mutually exclusive; otherwise, a situation might still arise where a node has already been deleted but a watch is still creating.

[zhangxinyao88]

I reproduced this by calling registerService() and deregisterService() concurrently. Deregistration returned before registration stored the node, so the node stayed registered. A lock would cover that case. I haven’t reproduced it through normal service startup and shutdown. Would it make sense to track this as a separate follow-up?

[ruanwenjun]

I think holding a lock for the entire method might be the simplest approach. It could fix many concurrency issues like this while keeping the code simple.

[wangzhigang1999]

My initial concern was holding the lock during the network I/O in node.close(). I tested this yesterday with network delays and connection loss and didn’t observe any additional shutdown issues, so I’m fine with synchronizing both methods.

[zhangxinyao88]

Thanks for testing this with delays and connection loss. Do you see locking both methods as a blocker for merging #7840, or would it be okay to track it as a separate follow-up?