(incubator-livy) branch master updated: [LIVY-1043] Track ZK connection transitions in Livy

[gg...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

gyogal pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-livy.git


The following commit(s) were added to refs/heads/master by this push:
     new 717bf747 [LIVY-1043] Track ZK connection transitions in Livy
717bf747 is described below

commit 717bf7478233fd0f92f55bfe18d6816723697e78
Author: Arnav Balyan <60...@users.noreply.github.com>
AuthorDate: Mon Apr 27 19:24:32 2026 +0530

    [LIVY-1043] Track ZK connection transitions in Livy
    
    ## What changes were proposed in this pull request?
    
     - Currently ZKManager does not observe connection transitions. When underlying connection changes, Livy is not aware of it, only ZK side logs show up.
     - Additionally register to ConnectionStateListener, which adds/listens to Livy owned logs for the connection changes.
     - Same events can be used for recovery in the future.
    
    ## How was this patch tested?
    
     - Unit tests
---
 .../livy/server/recovery/ZooKeeperManager.scala      | 17 +++++++++++++++++
 .../server/recovery/ZooKeeperStateStoreSpec.scala    | 20 ++++++++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/server/src/main/scala/org/apache/livy/server/recovery/ZooKeeperManager.scala b/server/src/main/scala/org/apache/livy/server/recovery/ZooKeeperManager.scala
index 60d98cc0..eb10f631 100644
--- a/server/src/main/scala/org/apache/livy/server/recovery/ZooKeeperManager.scala
+++ b/server/src/main/scala/org/apache/livy/server/recovery/ZooKeeperManager.scala
@@ -23,6 +23,7 @@ import scala.reflect.ClassTag
 import org.apache.curator.framework.api.UnhandledErrorListener
 import org.apache.curator.framework.CuratorFramework
 import org.apache.curator.framework.CuratorFrameworkFactory
+import org.apache.curator.framework.state.{ConnectionState, ConnectionStateListener}
 import org.apache.curator.retry.RetryNTimes
 import org.apache.zookeeper.KeeperException.NoNodeException
 
@@ -85,6 +86,22 @@ class ZooKeeperManager(
     }
   })
 
+  curatorClient.getConnectionStateListenable.addListener(new ConnectionStateListener {
+    override def stateChanged(client: CuratorFramework, newState: ConnectionState): Unit = {
+      newState match {
+        case ConnectionState.SUSPENDED =>
+          warn("ZooKeeper connection suspended; recovery state operations will block " +
+            "until the connection is re-established.")
+        case ConnectionState.RECONNECTED =>
+          info("ZooKeeper connection re-established within the session timeout.")
+        case ConnectionState.LOST =>
+          error("ZooKeeper session lost; ephemeral nodes and watches held by this client " +
+            "have been discarded. Subsequent operations will run against a new session.")
+        case _ =>
+      }
+    }
+  })
+
   def start(): Unit = {
     curatorClient.start()
   }
diff --git a/server/src/test/scala/org/apache/livy/server/recovery/ZooKeeperStateStoreSpec.scala b/server/src/test/scala/org/apache/livy/server/recovery/ZooKeeperStateStoreSpec.scala
index 0cd35db2..f3ffd9b6 100644
--- a/server/src/test/scala/org/apache/livy/server/recovery/ZooKeeperStateStoreSpec.scala
+++ b/server/src/test/scala/org/apache/livy/server/recovery/ZooKeeperStateStoreSpec.scala
@@ -22,7 +22,9 @@ import scala.collection.JavaConverters._
 import org.apache.curator.framework.CuratorFramework
 import org.apache.curator.framework.api._
 import org.apache.curator.framework.listen.Listenable
+import org.apache.curator.framework.state.{ConnectionState, ConnectionStateListener}
 import org.apache.zookeeper.data.Stat
+import org.mockito.ArgumentCaptor
 import org.mockito.Mockito._
 import org.scalatest.FunSpec
 import org.scalatest.Matchers._
@@ -43,6 +45,8 @@ class ZooKeeperStateStoreSpec extends FunSpec with LivyBaseUnitTestSuite {
       val curatorClient = mock[CuratorFramework]
       when(curatorClient.getUnhandledErrorListenable())
         .thenReturn(mock[Listenable[UnhandledErrorListener]])
+      when(curatorClient.getConnectionStateListenable())
+        .thenReturn(mock[Listenable[ConnectionStateListener]])
       val zkManager = new ZooKeeperManager(conf, Some(curatorClient))
       zkManager.start()
       val stateStore = new ZooKeeperStateStore(conf, zkManager)
@@ -186,6 +190,22 @@ class ZooKeeperStateStoreSpec extends FunSpec with LivyBaseUnitTestSuite {
       System.getProperty("zookeeper.sasl.clientconfig") shouldBe "Client"
     }
 
+    it("should register a ConnectionStateListener that handles all connection states") {
+      val curatorClient = mock[CuratorFramework]
+      when(curatorClient.getUnhandledErrorListenable())
+        .thenReturn(mock[Listenable[UnhandledErrorListener]])
+      val listenable = mock[Listenable[ConnectionStateListener]]
+      when(curatorClient.getConnectionStateListenable()).thenReturn(listenable)
+
+      new ZooKeeperManager(conf, Some(curatorClient))
+
+      val captor = ArgumentCaptor.forClass(classOf[ConnectionStateListener])
+      verify(listenable).addListener(captor.capture())
+      ConnectionState.values.foreach { state =>
+        captor.getValue.stateChanged(curatorClient, state)
+      }
+    }
+
     it("should not set SASL system properties when ZK SASL is disabled") {
       System.clearProperty("zookeeper.sasl.client")
       System.clearProperty("zookeeper.sasl.clientconfig")