[KYUUBI #7655][SERVER] Support virtual threads in the binary frontend

docs/configuration/settings.md

@@ -277,6 +277,7 @@ You can configure the Kyuubi properties in `$KYUUBI_HOME/conf/kyuubi-defaults.co
 | kyuubi.frontend.thrift.binary.ssl.disallowed.protocols     | SSLv2,SSLv3        | SSL versions to disallow for Kyuubi thrift binary frontend.                                                                                                                                                                                                                                                                                                                                                                                                                                                                 | set      | 1.7.0  |
 | kyuubi.frontend.thrift.binary.ssl.enabled                  | false              | Set this to true for using SSL encryption in thrift binary frontend server.                                                                                                                                                                                                                                                                                                                                                                                                                                                 | boolean  | 1.7.0  |
 | kyuubi.frontend.thrift.binary.ssl.include.ciphersuites                         || A comma-separated list of include SSL cipher suite names for thrift binary frontend.                                                                                                                                                                                                                                                                                                                                                                                                                                        | seq      | 1.7.0  |
+| kyuubi.frontend.thrift.binary.virtualThreads.enabled       | false              | Whether to use virtual threads for the Kyuubi server thrift binary frontend workers. This requires Java 21 or later. The maximum number of concurrent workers remains limited by kyuubi.frontend.thrift.max.worker.threads. The minimum worker threads and worker keepalive configurations do not apply in this mode.                                                                                                                                                                                                       | boolean  | 1.13.0 |
 | kyuubi.frontend.thrift.client.max.message.size             | 1073741824         | Maximum message size in bytes a thrift client will receive.                                                                                                                                                                                                                                                                                                                                                                                                                                                                 | int      | 1.9.3  |
 | kyuubi.frontend.thrift.http.bind.host                      | <undefined>  | Hostname or IP of the machine on which to run the thrift frontend service via http protocol.                                                                                                                                                                                                                                                                                                                                                                                                                                | string   | 1.6.0  |
 | kyuubi.frontend.thrift.http.bind.port                      | 10010              | Port of the machine on which to run the thrift frontend service via http protocol.                                                                                                                                                                                                                                                                                                                                                                                                                                          | int      | 1.6.0  |

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

@@ -635,6 +635,18 @@ object KyuubiConf {
       .immutable
       .fallbackConf(FRONTEND_BIND_PORT)
 
+  val FRONTEND_THRIFT_BINARY_VIRTUAL_THREADS_ENABLED: ConfigEntry[Boolean] =
+    buildConf("kyuubi.frontend.thrift.binary.virtualThreads.enabled")
+      .doc("Whether to use virtual threads for the Kyuubi server thrift binary frontend " +
+        "workers. This requires Java 21 or later. The maximum number of concurrent workers " +
+        "remains limited by kyuubi.frontend.thrift.max.worker.threads. The minimum worker " +
+        "threads and worker keepalive configurations do not apply in this mode.")
+      .version("1.13.0")
+      .audience(SERVER)
+      .immutable
+      .booleanConf
+      .createWithDefault(false)
+
   val FRONTEND_THRIFT_HTTP_BIND_HOST: ConfigEntry[Option[String]] =
     buildConf("kyuubi.frontend.thrift.http.bind.host")
       .doc("Hostname or IP of the machine on which to run the thrift frontend service " +

kyuubi-common/src/main/scala/org/apache/kyuubi/service/TBinaryFrontendService.scala

@@ -29,7 +29,7 @@ import org.apache.kyuubi.shaded.hive.service.rpc.thrift._
 import org.apache.kyuubi.shaded.thrift.protocol.TBinaryProtocol
 import org.apache.kyuubi.shaded.thrift.server.{TServer, TThreadPoolServer}
 import org.apache.kyuubi.shaded.thrift.transport.{TServerSocket, TSSLTransportFactory}
-import org.apache.kyuubi.util.NamedThreadFactory
+import org.apache.kyuubi.util.{NamedThreadFactory, ThreadUtils}
 
 /**
  * Apache Thrift based hive service rpc
@@ -65,13 +65,19 @@ abstract class TBinaryFrontendService(name: String)
       val minThreads = conf.get(FRONTEND_THRIFT_MIN_WORKER_THREADS)
       val maxThreads = conf.get(FRONTEND_THRIFT_MAX_WORKER_THREADS)
       val keepAliveTime = conf.get(FRONTEND_THRIFT_WORKER_KEEPALIVE_TIME)
-      val executor = new ThreadPoolExecutor(
-        minThreads,
-        maxThreads,
-        keepAliveTime,
-        TimeUnit.MILLISECONDS,
-        new SynchronousQueue[Runnable](),
-        new NamedThreadFactory(name + "Handler-Pool", false))
+      val useVirtualThreads = isServer() &&
+        conf.get(FRONTEND_THRIFT_BINARY_VIRTUAL_THREADS_ENABLED)
+      val executor = if (useVirtualThreads) {
+        ThreadUtils.newBoundedVirtualThreadPerTaskExecutor(maxThreads, name + "Handler")
+      } else {
+        new ThreadPoolExecutor(
+          minThreads,
+          maxThreads,
+          keepAliveTime,
+          TimeUnit.MILLISECONDS,
+          new SynchronousQueue[Runnable](),
+          new NamedThreadFactory(name + "Handler-Pool", false))
+      }
       val transFactory = authFactory.getTTransportFactory
       val tProcFactory = authFactory.getTProcessorFactory(this)
       val tServerSocket =
@@ -116,8 +122,13 @@ abstract class TBinaryFrontendService(name: String)
       // TCP Server
       server = Some(new TThreadPoolServer(args))
       server.foreach(_.setServerEventHandler(new FeTServerEventHandler))
+      val workerThreadDescription = if (useVirtualThreads) {
+        s"at most $maxThreads virtual worker threads"
+      } else {
+        s"[$minThreads, $maxThreads] platform worker threads"
+      }
       info(s"Initializing $name on ${serverAddr.getHostName}:${_actualPort} with" +
-        s" [$minThreads, $maxThreads] worker threads")
+        s" $workerThreadDescription")
     } catch {
       case e: Throwable =>
         error(e)

kyuubi-common/src/main/scala/org/apache/kyuubi/util/ThreadUtils.scala

@@ -26,6 +26,35 @@ import org.apache.kyuubi.{KyuubiException, Logging}
 
 object ThreadUtils extends Logging {
 
+  def newBoundedVirtualThreadPerTaskExecutor(
+      maxConcurrentTasks: Int,
+      threadNamePrefix: String): ExecutorService = {
+    require(maxConcurrentTasks > 0, "maxConcurrentTasks must be positive")
+    new BoundedExecutorService(
+      newVirtualThreadPerTaskExecutor(threadNamePrefix),
+      maxConcurrentTasks)
+  }
+
+  private def newVirtualThreadPerTaskExecutor(threadNamePrefix: String): ExecutorService =
+    try {
+      val builder = classOf[Thread].getMethod("ofVirtual").invoke(null)
+      val builderClass = Class.forName("java.lang.Thread$Builder")
+      val namedBuilder = builderClass
+        .getMethod("name", classOf[String], java.lang.Long.TYPE)
+        .invoke(builder, s"$threadNamePrefix-", Long.box(0L))
+      val threadFactory = builderClass
+        .getMethod("factory")
+        .invoke(namedBuilder)
+        .asInstanceOf[ThreadFactory]
+      classOf[Executors]
+        .getMethod("newThreadPerTaskExecutor", classOf[ThreadFactory])
+        .invoke(null, threadFactory)
+        .asInstanceOf[ExecutorService]
+    } catch {
+      case e: ReflectiveOperationException =>
+        throw new IllegalStateException("Virtual threads require Java 21 or later", e)
+    }
+
   def newDaemonSingleThreadScheduledExecutor(
       threadName: String,
       executeExistingDelayedTasksAfterShutdown: Boolean = true): ScheduledExecutorService = {
@@ -150,4 +179,45 @@ object ThreadUtils extends Logging {
       delay,
       timeUnit)
   }
+
+  private class BoundedExecutorService(
+      delegate: ExecutorService,
+      maxConcurrentTasks: Int) extends AbstractExecutorService {
+
+    private val permits = new Semaphore(maxConcurrentTasks)
+
+    override def shutdown(): Unit = delegate.shutdown()
+
+    override def shutdownNow(): java.util.List[Runnable] = delegate.shutdownNow()
+
+    override def isShutdown: Boolean = delegate.isShutdown
+
+    override def isTerminated: Boolean = delegate.isTerminated
+
+    override def awaitTermination(timeout: Long, unit: TimeUnit): Boolean =
+      delegate.awaitTermination(timeout, unit)
+
+    override def execute(command: Runnable): Unit = {
+      if (!permits.tryAcquire()) {
+        throw new RejectedExecutionException(
+          s"Maximum concurrent task limit $maxConcurrentTasks reached")
+      }
+
+      try {
+        delegate.execute(new Runnable {
+          override def run(): Unit = {
+            try {
+              command.run()
+            } finally {
+              permits.release()
+            }
+          }
+        })
+      } catch {
+        case t: Throwable =>
+          permits.release()
+          throw t
+      }
+    }
+  }
 }

kyuubi-common/src/test/scala/org/apache/kyuubi/config/KyuubiConfSuite.scala

@@ -352,6 +352,16 @@ class KyuubiConfSuite extends KyuubiFunSuite {
     }
   }
 
+  test("getEngineConf excludes the server thrift binary virtual thread config") {
+    val kyuubiConf = KyuubiConf(false)
+    kyuubiConf.set(FRONTEND_THRIFT_BINARY_VIRTUAL_THREADS_ENABLED, true)
+
+    EngineType.values.foreach { engineType =>
+      assert(!kyuubiConf.getEngineConf(engineType)
+        .contains(FRONTEND_THRIFT_BINARY_VIRTUAL_THREADS_ENABLED.key))
+    }
+  }
+
   test("getEngineConf passes through reserved keys") {
     val kyuubiConf = KyuubiConf(false)
     kyuubiConf.set(KyuubiReservedKeys.KYUUBI_SERVER_IP_KEY, "10.0.0.1")

kyuubi-common/src/test/scala/org/apache/kyuubi/util/ThreadUtilsSuite.scala

@@ -17,7 +17,7 @@
 
 package org.apache.kyuubi.util
 
-import java.util.concurrent.TimeUnit
+import java.util.concurrent.{ConcurrentLinkedQueue, CountDownLatch, RejectedExecutionException, TimeUnit}
 
 import org.apache.kyuubi.KyuubiFunSuite
 
@@ -76,4 +76,62 @@ class ThreadUtilsSuite extends KyuubiFunSuite {
     ThreadUtils.shutdown(pool)
     assert(pool.isShutdown)
   }
+
+  test("New bounded virtual thread per task executor") {
+    val virtualThreadsSupported =
+      try {
+        classOf[Thread].getMethod("isVirtual")
+        true
+      } catch {
+        case _: NoSuchMethodException => false
+      }
+
+    if (!virtualThreadsSupported) {
+      val error = intercept[IllegalStateException] {
+        ThreadUtils.newBoundedVirtualThreadPerTaskExecutor(2, "ThreadUtilsVirtualTest")
+      }
+      assert(error.getMessage.contains("Java 21"))
+    } else {
+      val executor =
+        ThreadUtils.newBoundedVirtualThreadPerTaskExecutor(2, "ThreadUtilsVirtualTest")
+      val ready = new CountDownLatch(2)
+      val release = new CountDownLatch(1)
+      val threadNames = new ConcurrentLinkedQueue[String]()
+      val tasksAreVirtual = new ConcurrentLinkedQueue[Boolean]()
+      val isVirtual = classOf[Thread].getMethod("isVirtual")
+
+      def blockingTask: Runnable = new Runnable {
+        override def run(): Unit = {
+          threadNames.add(Thread.currentThread().getName)
+          tasksAreVirtual.add(isVirtual.invoke(Thread.currentThread()).asInstanceOf[Boolean])
+          ready.countDown()
+          release.await()
+        }
+      }
+
+      try {
+        val first = executor.submit(blockingTask)
+        val second = executor.submit(blockingTask)
+        assert(ready.await(10, TimeUnit.SECONDS))
+        intercept[RejectedExecutionException](executor.submit(blockingTask))
+        release.countDown()
+        first.get(10, TimeUnit.SECONDS)
+        second.get(10, TimeUnit.SECONDS)
+
+        val last = executor.submit(new Runnable {
+          override def run(): Unit = {
+            threadNames.add(Thread.currentThread().getName)
+            tasksAreVirtual.add(isVirtual.invoke(Thread.currentThread()).asInstanceOf[Boolean])
+          }
+        })
+        last.get(10, TimeUnit.SECONDS)
+        assert(threadNames.size() === 3)
+        assert(threadNames.toArray.forall(_.toString.startsWith("ThreadUtilsVirtualTest-")))
+        assert(tasksAreVirtual.toArray.forall(_.asInstanceOf[Boolean]))
+      } finally {
+        release.countDown()
+        ThreadUtils.shutdown(executor)
+      }
+    }
+  }
 }