[KYUUBI #7529][DATA-AGENT] Support cross-server session routing

externals/kyuubi-data-agent-engine/src/main/scala/org/apache/kyuubi/engine/dataagent/session/DataAgentSessionImpl.scala

@@ -16,9 +16,11 @@
  */
 package org.apache.kyuubi.engine.dataagent.session
 
+import scala.util.control.NonFatal
+
 import org.apache.kyuubi.{KYUUBI_VERSION, KyuubiSQLException}
 import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_SESSION_HANDLE_KEY
-import org.apache.kyuubi.session.{AbstractSession, SessionHandle, SessionManager}
+import org.apache.kyuubi.session.{AbstractSession, SessionHandle}
 import org.apache.kyuubi.shaded.hive.service.rpc.thrift.{TGetInfoType, TGetInfoValue, TProtocolVersion}
 
 class DataAgentSessionImpl(
@@ -27,20 +29,20 @@ class DataAgentSessionImpl(
     password: String,
     ipAddress: String,
     conf: Map[String, String],
-    sessionManager: SessionManager)
-  extends AbstractSession(protocol, user, password, ipAddress, conf, sessionManager) {
+    dataAgentSessionManager: DataAgentSessionManager)
+  extends AbstractSession(protocol, user, password, ipAddress, conf, dataAgentSessionManager) {
 
   override val handle: SessionHandle =
     conf.get(KYUUBI_SESSION_HANDLE_KEY).map(SessionHandle.fromUUID).getOrElse(SessionHandle())
 
-  private val dataAgentProvider =
-    sessionManager.asInstanceOf[DataAgentSessionManager].dataAgentProvider
+  private val dataAgentProvider = dataAgentSessionManager.dataAgentProvider
 
   override def open(): Unit = {
     info(s"Starting to open data agent session.")
     dataAgentProvider.open(handle.identifier.toString, user)
     try {
       super.open()
+      dataAgentSessionManager.registerRoute(handle.identifier.toString, user)
     } catch {
       case e: Throwable =>
         try {
@@ -72,6 +74,11 @@ class DataAgentSessionImpl(
   }
 
   override def close(): Unit = {
+    try {
+      dataAgentSessionManager.unregisterRoute(handle.identifier.toString)
+    } catch {
+      case NonFatal(e) => warn("Failed to unregister data agent session route", e)
+    }
     try {
       dataAgentProvider.close(handle.identifier.toString)
     } finally {

externals/kyuubi-data-agent-engine/src/main/scala/org/apache/kyuubi/engine/dataagent/session/DataAgentSessionManager.scala

@@ -18,11 +18,13 @@ package org.apache.kyuubi.engine.dataagent.session
 
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.config.KyuubiConf.ENGINE_SHARE_LEVEL
-import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_SESSION_HANDLE_KEY
+import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_SERVER_HA_NAMESPACE_KEY
 import org.apache.kyuubi.engine.ShareLevel
-import org.apache.kyuubi.engine.dataagent.DataAgentEngine
+import org.apache.kyuubi.engine.dataagent.{DataAgentEngine, DataAgentTBinaryFrontendService}
 import org.apache.kyuubi.engine.dataagent.operation.DataAgentOperationManager
 import org.apache.kyuubi.engine.dataagent.provider.DataAgentProvider
+import org.apache.kyuubi.ha.HighAvailabilityConf.{HA_ENGINE_REF_ID, HA_NAMESPACE}
+import org.apache.kyuubi.ha.client.{DataAgentSessionRoute, EngineServiceDiscovery}
 import org.apache.kyuubi.operation.OperationManager
 import org.apache.kyuubi.session.{Session, SessionHandle, SessionManager}
 import org.apache.kyuubi.shaded.hive.service.rpc.thrift.TProtocolVersion
@@ -50,10 +52,34 @@ class DataAgentSessionManager(name: String)
       password: String,
       ipAddress: String,
       conf: Map[String, String]): Session = {
-    conf.get(KYUUBI_SESSION_HANDLE_KEY).map(SessionHandle.fromUUID)
-      .flatMap(getSessionOption).getOrElse {
-        new DataAgentSessionImpl(protocol, user, password, ipAddress, conf, this)
-      }
+    new DataAgentSessionImpl(protocol, user, password, ipAddress, conf, this)
+  }
+
+  private def discoveryClient = DataAgentEngine.currentEngine
+    .flatMap(_.frontendServices.collectFirst {
+      case frontend: DataAgentTBinaryFrontendService => frontend
+    }).flatMap(_.discoveryService.collect {
+      case service: EngineServiceDiscovery => service.discoveryClient
+    })
+
+  private[dataagent] def registerRoute(sessionId: String, user: String): Unit = {
+    for {
+      serverSpace <- conf.getOption(KYUUBI_SERVER_HA_NAMESPACE_KEY)
+      engineRefId <- conf.get(HA_ENGINE_REF_ID)
+      discovery <- discoveryClient
+    } {
+      val route = DataAgentSessionRoute(conf.get(HA_NAMESPACE), engineRefId, user)
+      DataAgentSessionRoute.register(discovery, serverSpace, sessionId, route)
+    }
+  }
+
+  private[dataagent] def unregisterRoute(sessionId: String): Unit = {
+    for {
+      serverSpace <- conf.getOption(KYUUBI_SERVER_HA_NAMESPACE_KEY)
+      discovery <- discoveryClient
+    } {
+      DataAgentSessionRoute.unregister(discovery, serverSpace, sessionId)
+    }
   }
 
   override def stop(): Unit = {

kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala

@@ -21,6 +21,7 @@ object KyuubiReservedKeys {
   final val KYUUBI_CLIENT_IP_KEY = "kyuubi.client.ipAddress"
   final val KYUUBI_CLIENT_VERSION_KEY = "kyuubi.client.version"
   final val KYUUBI_SERVER_IP_KEY = "kyuubi.server.ipAddress"
+  final val KYUUBI_SERVER_HA_NAMESPACE_KEY = "kyuubi.server.ha.namespace"
   final val KYUUBI_SESSION_USER_KEY = "kyuubi.session.user"
   final val KYUUBI_SESSION_SIGN_PUBLICKEY = "kyuubi.session.sign.publickey"
   final val KYUUBI_SESSION_USER_SIGN = "kyuubi.session.user.sign"

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

@@ -0,0 +1,103 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.ha.client
+
+import java.nio.charset.StandardCharsets
+import java.util.Base64
+
+import scala.util.control.NonFatal
+
+case class DataAgentSessionRoute(engineSpace: String, engineRefId: String, user: String)
+
+object DataAgentSessionRoute {
+  private val ROOT_SUFFIX = "DATA_AGENT_sessions"
+
+  private def root(serverSpace: String): String = s"${serverSpace}_$ROOT_SUFFIX"
+
+  private def path(serverSpace: String, sessionId: String): String =
+    DiscoveryPaths.makePath(root(serverSpace), sessionId)
+
+  def register(
+      discovery: DiscoveryClient,
+      serverSpace: String,
+      sessionId: String,
+      route: DataAgentSessionRoute): Unit = {
+    val routePath = path(serverSpace, sessionId)
+    val entry = encode(route)
+    children(discovery, routePath) match {
+      case Nil =>
+        val entryPath = DiscoveryPaths.makePath(routePath, entry)
+        try discovery.create(entryPath, "PERSISTENT")
+        catch {
+          case NonFatal(_) if discovery.pathExists(entryPath) =>
+        }
+      case existing :: Nil if existing == entry =>
+      case _ =>
+        throw new IllegalStateException(s"A different route exists for session $sessionId")
+    }
+  }
+
+  def find(
+      discovery: DiscoveryClient,
+      serverSpace: String,
+      sessionId: String): Option[DataAgentSessionRoute] = {
+    val routePath = path(serverSpace, sessionId)
+    children(discovery, routePath) match {
+      case Nil => None
+      case entry :: Nil => Some(decode(entry))
+      case _ => throw new IllegalStateException(s"Multiple routes found for session $sessionId")
+    }
+  }
+
+  def unregister(
+      discovery: DiscoveryClient,
+      serverSpace: String,
+      sessionId: String): Unit = {
+    val routePath = path(serverSpace, sessionId)
+    try {
+      discovery.delete(routePath, deleteChildren = true)
+    } catch {
+      case NonFatal(_) if discovery.pathNonExists(routePath, isPrefix = true) =>
+    }
+  }
+
+  private def encode(route: DataAgentSessionRoute): String = {
+    Seq(route.engineSpace, route.engineRefId, route.user)
+      .map(encodeValue)
+      .mkString(".")
+  }
+
+  private def encodeValue(value: String): String =
+    Base64.getUrlEncoder.withoutPadding()
+      .encodeToString(value.getBytes(StandardCharsets.UTF_8))
+
+  private def children(discovery: DiscoveryClient, routePath: String): List[String] = {
+    try discovery.getChildren(routePath)
+    catch {
+      case NonFatal(_) if discovery.pathNonExists(routePath, isPrefix = true) => Nil
+    }
+  }
+
+  private def decode(entry: String): DataAgentSessionRoute = {
+    val values = entry.split("\\.", -1)
+    require(values.length == 3, "Invalid Data Agent session route")
+    def decodeValue(value: String): String =
+      new String(Base64.getUrlDecoder.decode(value), StandardCharsets.UTF_8)
+    DataAgentSessionRoute(decodeValue(values(0)), decodeValue(values(1)), decodeValue(values(2)))
+  }
+}

kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/DiscoveryClientTests.scala

@@ -170,4 +170,38 @@ trait DiscoveryClientTests extends KyuubiFunSuite {
       assert(data == dataFromGet)
     }
   }
+
+  test("data agent session route lifecycle") {
+    testDataAgentSessionRouteLifecycle()
+  }
+
+  protected def testDataAgentSessionRouteLifecycle(): Unit = {
+    withDiscoveryClient(conf) { discoveryClient =>
+      val serverSpace = "server-space"
+      val sessionId = "session-id"
+      val route = DataAgentSessionRoute(
+        "engine space/with.delimiters",
+        "engine-ref:id/with.delimiters",
+        "alice+route@example.com")
+
+      assert(DataAgentSessionRoute.find(discoveryClient, serverSpace, sessionId).isEmpty)
+      DataAgentSessionRoute.register(discoveryClient, serverSpace, sessionId, route)
+      DataAgentSessionRoute.register(discoveryClient, serverSpace, sessionId, route)
+      assert(DataAgentSessionRoute.find(discoveryClient, serverSpace, sessionId).contains(route))
+
+      val conflictingRoute = route.copy(engineRefId = "another-engine-ref-id")
+      intercept[IllegalStateException] {
+        DataAgentSessionRoute.register(
+          discoveryClient,
+          serverSpace,
+          sessionId,
+          conflictingRoute)
+      }
+      assert(DataAgentSessionRoute.find(discoveryClient, serverSpace, sessionId).contains(route))
+
+      DataAgentSessionRoute.unregister(discoveryClient, serverSpace, sessionId)
+      DataAgentSessionRoute.unregister(discoveryClient, serverSpace, sessionId)
+      assert(DataAgentSessionRoute.find(discoveryClient, serverSpace, sessionId).isEmpty)
+    }
+  }
 }

kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/etcd/EtcdDiscoveryClientSuite.scala

@@ -110,4 +110,8 @@ class EtcdDiscoveryClientSuite extends DiscoveryClientTests {
       assert(!discoveryClient.pathExists(path))
     }
   }
+
+  test("etcd data agent session route lifecycle") {
+    testDataAgentSessionRouteLifecycle()
+  }
 }

kyuubi-server/src/main/scala/org/apache/kyuubi/client/KyuubiSyncThriftClient.scala

@@ -62,6 +62,18 @@ class KyuubiSyncThriftClient private (
   // Visible for testing.
   private[kyuubi] def remoteSessionHandle: TSessionHandle = _remoteSessionHandle
 
+  private[kyuubi] def useExistingSessionHandle(sessionHandle: SessionHandle): Unit = {
+    require(_remoteSessionHandle == null, "A session handle is already set")
+    _remoteSessionHandle = sessionHandle.toTSessionHandle
+  }
+
+  private[kyuubi] def closeTransport(): Unit = {
+    Seq(protocol).union(engineAliveProbeProtocol.toSeq).foreach { tProtocol =>
+      if (tProtocol.getTransport.isOpen) tProtocol.getTransport.close()
+    }
+    shutdownAsyncRequestExecutor()
+  }
+
   @volatile private var _aliveProbeSessionHandle: TSessionHandle = _
   @volatile private var _remoteEngineBroken: Boolean = false
   private[kyuubi] def remoteEngineBroken: Boolean = _remoteEngineBroken
@@ -514,4 +526,16 @@ private[kyuubi] object KyuubiSyncThriftClient extends Logging {
       aliveProbeInterval,
       aliveTimeout)
   }
+
+  def createClientWithExistingSessionHandle(
+      user: String,
+      password: String,
+      host: String,
+      port: Int,
+      conf: KyuubiConf,
+      sessionHandle: SessionHandle): KyuubiSyncThriftClient = {
+    val client = createClient(user, password, host, port, conf)
+    client.useExistingSessionHandle(sessionHandle)
+    client
+  }
 }

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/EngineRef.scala

@@ -31,7 +31,7 @@ import org.apache.commons.lang3.StringUtils
 import org.apache.kyuubi.{KYUUBI_VERSION, KyuubiSQLException, Logging, Utils}
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.config.KyuubiConf._
-import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_ENGINE_SUBMIT_TIME_KEY
+import org.apache.kyuubi.config.KyuubiReservedKeys.{KYUUBI_ENGINE_SUBMIT_TIME_KEY, KYUUBI_SERVER_HA_NAMESPACE_KEY}
 import org.apache.kyuubi.engine.EngineType._
 import org.apache.kyuubi.engine.ShareLevel.{CONNECTION, GROUP, SERVER, SERVER_LOCAL, ShareLevel}
 import org.apache.kyuubi.engine.dataagent.DataAgentProcessBuilder
@@ -228,6 +228,7 @@ private[kyuubi] class EngineRef(
     var engineRef = discoveryClient.getServerHost(engineSpace)
     if (engineRef.nonEmpty) return engineRef.get
 
+    conf.set(KYUUBI_SERVER_HA_NAMESPACE_KEY, serverSpace)
     conf.set(HA_NAMESPACE, engineSpace)
     conf.set(HA_ENGINE_REF_ID, engineRefId)
     val started = System.currentTimeMillis()