spark.connect.jvmStacktrace.maxSize |
diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala
index e2d496239d290..61a6218da92f6 100644
--- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala
+++ b/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)
}
diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectExecutionManager.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectExecutionManager.scala
index 2a936b526d96a..a8b0554fe977a 100644
--- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectExecutionManager.scala
+++ b/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
diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectExecutionManagerSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectExecutionManagerSuite.scala
index 228bb6e83a985..e4b0d6f20e250 100644
--- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectExecutionManagerSuite.scala
+++ b/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")
+ }
}
|