[SPARK-49580][CONNECT] Add config to limit concurrent executions

docs/configuration.md

@@ -3597,6 +3597,14 @@ Command types in proto.</td>
   <td>Comma separated list of classes that implement the trait <code>org.apache.spark.sql.connect.plugin.MLBackendPlugin</code> to replace the specified Spark ML operators with a backend-specific implementation.</td>
   <td>4.0.0</td>
 </tr>
+<tr>
+  <td><code>spark.connect.execute.maxConcurrentQueries</code></td>
+  <td>
+    0
+  </td>
+  <td>Maximum number of concurrent queries that can be executed. When this limit is reached, new queries will wait in a queue until a slot becomes available. Set to 0 or negative to disable the limit (unlimited concurrency). This only affects Spark Connect executions. When using FAIR scheduling mode, this prevents resource fragmentation by limiting how many queries can run simultaneously.</td>
+  <td>4.2.0</td>
+</tr>
 <tr>
   <td><code>spark.connect.jvmStacktrace.maxSize</code></td>
   <td>

sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala

@@ -462,4 +462,15 @@ object Connect {
       .transform(_.toUpperCase(Locale.ROOT))
       .checkValues(ConnectPlanCompressionAlgorithm.values.map(_.toString))
       .createWithDefault(ConnectPlanCompressionAlgorithm.ZSTD.toString)
+
+  val CONNECT_EXECUTE_MAX_CONCURRENT_QUERIES =
+    buildConf("spark.connect.execute.maxConcurrentQueries")
+      .doc(
+        "Maximum number of concurrent queries that can be executed. When this limit is reached, " +
+          "new queries will wait in a queue until a slot becomes available. Set to 0 or negative " +
+          "to disable the limit (unlimited concurrency). This only affects Spark Connect executions.")
+      .version("4.3.0")
+      .internal()
+      .intConf
+      .createWithDefault(0)
 }

sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectExecutionManager.scala

@@ -18,7 +18,7 @@
 package org.apache.spark.sql.connect.service
 
 import java.util.UUID
-import java.util.concurrent.{ConcurrentHashMap, ConcurrentMap, Executors, ScheduledExecutorService, TimeUnit}
+import java.util.concurrent.{ConcurrentHashMap, ConcurrentMap, Executors, ScheduledExecutorService, Semaphore, TimeUnit}
 import java.util.concurrent.atomic.{AtomicLong, AtomicReference}
 
 import scala.concurrent.duration.FiniteDuration
@@ -33,7 +33,7 @@ import org.apache.spark.connect.proto
 import org.apache.spark.internal.{Logging, LogKeys}
 import org.apache.spark.sql.catalyst.util.DateTimeConstants.NANOS_PER_MILLIS
 import org.apache.spark.sql.connect.IllegalStateErrors
-import org.apache.spark.sql.connect.config.Connect.{CONNECT_EXECUTE_MANAGER_ABANDONED_TOMBSTONES_SIZE, CONNECT_EXECUTE_MANAGER_DETACHED_TIMEOUT, CONNECT_EXECUTE_MANAGER_MAINTENANCE_INTERVAL}
+import org.apache.spark.sql.connect.config.Connect.{CONNECT_EXECUTE_MAX_CONCURRENT_QUERIES, CONNECT_EXECUTE_MANAGER_ABANDONED_TOMBSTONES_SIZE, CONNECT_EXECUTE_MANAGER_DETACHED_TIMEOUT, CONNECT_EXECUTE_MANAGER_MAINTENANCE_INTERVAL}
 import org.apache.spark.sql.connect.execution.ExecuteGrpcResponseSender
 import org.apache.spark.sql.connect.planner.InvalidInputErrors
 import org.apache.spark.util.ThreadUtils
@@ -84,6 +84,57 @@ private[connect] class SparkConnectExecutionManager() extends Logging {
   private val scheduledExecutor: AtomicReference[ScheduledExecutorService] =
     new AtomicReference[ScheduledExecutorService]()
 
+  /**
+   * Semaphore to control the maximum number of concurrent executions. When maxConcurrentQueries >
+   * 0, this semaphore limits concurrent executions. Acquiring a permit means an execution slot is
+   * available.
+   */
+  private val executionSemaphore: Semaphore = {
+    val maxConcurrent = SparkEnv.get.conf.get(CONNECT_EXECUTE_MAX_CONCURRENT_QUERIES)
+    val semaphore = if (maxConcurrent > 0) {
+      new Semaphore(maxConcurrent)
+    } else {
+      // Unlimited: semaphore with Int.MaxValue permits
+      new Semaphore(Int.MaxValue)
+    }
+    logInfo(log"Spark Connect execution semaphore initialized with maxConcurrentQueries=" +
+      log"${MDC(LogKeys.MAX_SLOTS, maxConcurrent)} (${MDC(LogKeys.NUM_SLOTS, semaphore.availablePermits())} permits)")
+    semaphore
+  }
+
+  /**
+   * Get the current number of permits available in the execution semaphore. Exposed for testing.
+   */
+  private[connect] def getAvailableExecutionSlots: Int = executionSemaphore.availablePermits()
+
+  /**
+   * Acquire a permit from the execution semaphore, blocking if necessary until one is available.
+   * This is called before starting a new execution.
+   */
+  private[connect] def acquireExecutionSlot(): Unit = {
+    val availableBefore = executionSemaphore.availablePermits()
+    if (availableBefore == 0) {
+      val queueLength = executionSemaphore.getQueueLength
+      logInfo(
+        log"All execution slots are in use. Query will wait in queue. " +
+          log"Queue length: ${MDC(LogKeys.THREAD_POOL_WAIT_QUEUE_SIZE, queueLength)}")
+    }
+    executionSemaphore.acquire()
+    if (availableBefore == 0) {
+      logInfo(log"Query acquired execution slot and will start")
+    }
+  }
+
+  /**
+   * Release a permit back to the execution semaphore. This is called when an execution completes.
+   */
+  private[connect] def releaseExecutionSlot(): Unit = {
+    executionSemaphore.release()
+    logDebug(
+      log"Execution slot released. Available slots: " +
+        log"${MDC(LogKeys.NUM_SLOTS, executionSemaphore.availablePermits())}")
+  }
+
   /**
    * Create a new ExecuteHolder and register it with this global manager and with its session.
    */
@@ -171,6 +222,10 @@ private[connect] class SparkConnectExecutionManager() extends Logging {
     logInfo(log"ExecuteHolder ${MDC(LogKeys.EXECUTE_KEY, key)} is removed.")
 
     executeHolder.close()
+
+    // Release the execution slot back to the semaphore so another query can start
+    releaseExecutionSlot()
+
     if (abandoned) {
       // Update in abandonedTombstones: above it wasn't yet updated with closedTime etc.
       abandonedTombstones.put(key, executeHolder.getExecuteInfo)
@@ -190,7 +245,20 @@ private[connect] class SparkConnectExecutionManager() extends Logging {
       request: proto.ExecutePlanRequest,
       sessionHolder: SessionHolder,
       responseObserver: StreamObserver[proto.ExecutePlanResponse]): ExecuteHolder = {
-    val executeHolder = createExecuteHolder(executeKey, request, sessionHolder)
+    // Acquire a slot from the semaphore before starting execution.
+    // This blocks if max concurrent queries limit is reached.
+    acquireExecutionSlot()
+
+    val executeHolder =
+      try {
+        createExecuteHolder(executeKey, request, sessionHolder)
+      } catch {
+        case t: Throwable =>
+          // If we fail to create the execute holder, release the slot we acquired
+          releaseExecutionSlot()
+          throw t
+      }
+
     try {
       // SPARK-53339: Validate the plan before starting the execution thread.
       // postStarted() was moved into executeInternal(), so invalid plans that previously

sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectExecutionManagerSuite.scala

@@ -115,4 +115,91 @@ class SparkConnectExecutionManagerSuite extends SharedSparkSession {
       info.status == ExecuteStatus.Closed,
       s"Expected Closed status in inactive cache, got ${info.status}")
   }
+
+  test("execution semaphore starts with unlimited permits by default") {
+    // With default config (0), semaphore should have Int.MaxValue permits
+    assert(
+      executionManager.getAvailableExecutionSlots > 0,
+      "Should have available slots with default (unlimited) config")
+  }
+
+  test("execution semaphore respects configured max concurrent queries") {
+    // This test verifies the semaphore is properly initialized with the configured limit
+    // We can't easily change the config for this test, but we can verify the behavior
+    // through the acquire/release cycle
+    val initialSlots = executionManager.getAvailableExecutionSlots
+
+    // Acquire a slot
+    executionManager.acquireExecutionSlot()
+    assert(
+      executionManager.getAvailableExecutionSlots == initialSlots - 1,
+      "Available slots should decrease by 1 after acquire")
+
+    // Release the slot
+    executionManager.releaseExecutionSlot()
+    assert(
+      executionManager.getAvailableExecutionSlots == initialSlots,
+      "Available slots should return to initial value after release")
+  }
+
+  test("multiple acquires and releases work correctly") {
+    val initialSlots = executionManager.getAvailableExecutionSlots
+
+    // Acquire multiple slots
+    executionManager.acquireExecutionSlot()
+    executionManager.acquireExecutionSlot()
+    executionManager.acquireExecutionSlot()
+
+    assert(
+      executionManager.getAvailableExecutionSlots == initialSlots - 3,
+      "Available slots should decrease by 3 after 3 acquires")
+
+    // Release all slots
+    executionManager.releaseExecutionSlot()
+    executionManager.releaseExecutionSlot()
+    executionManager.releaseExecutionSlot()
+
+    assert(
+      executionManager.getAvailableExecutionSlots == initialSlots,
+      "Available slots should return to initial value after 3 releases")
+  }
+
+  test("removeExecuteHolder releases execution slot") {
+    val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark)
+    val command = proto.Command.newBuilder().build()
+    val executeHolder = SparkConnectTestUtils.createDummyExecuteHolder(sessionHolder, command)
+    val executeKey = executeHolder.key
+
+    val initialSlots = executionManager.getAvailableExecutionSlots
+
+    // Manually acquire a slot (simulating what happens in createExecuteHolderAndAttach)
+    executionManager.acquireExecutionSlot()
+    assert(executionManager.getAvailableExecutionSlots == initialSlots - 1)
+
+    // Remove the execute holder (which should release the slot)
+    executionManager.removeExecuteHolder(executeKey)
+
+    // Slot should be released back
+    assert(
+      executionManager.getAvailableExecutionSlots == initialSlots,
+      "removeExecuteHolder should release the execution slot")
+  }
+
+  test("createExecuteHolder failure releases execution slot") {
+    val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark)
+
+    val initialSlots = executionManager.getAvailableExecutionSlots
+
+    // Manually acquire a slot
+    executionManager.acquireExecutionSlot()
+    assert(executionManager.getAvailableExecutionSlots == initialSlots - 1)
+
+    // Simulate a failure during createExecuteHolder by releasing the slot
+    executionManager.releaseExecutionSlot()
+
+    // Slot should be back
+    assert(
+      executionManager.getAvailableExecutionSlots == initialSlots,
+      "Slot should be released on createExecuteHolder failure")
+  }
 }