diff --git a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/test/RemoteSparkSession.scala b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/test/RemoteSparkSession.scala index 8bd6c5cf01681..386aeba9011bb 100644 --- a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/test/RemoteSparkSession.scala +++ b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/test/RemoteSparkSession.scala @@ -18,7 +18,8 @@ package org.apache.spark.sql.connect.test import java.io.{File, IOException, OutputStream} import java.lang.ProcessBuilder.Redirect -import java.nio.file.Paths +import java.nio.charset.StandardCharsets.UTF_8 +import java.nio.file.{Files, Paths} import java.util.concurrent.TimeUnit import scala.concurrent.duration.FiniteDuration @@ -33,7 +34,6 @@ import org.scalatest.time.SpanSugar._ // scalastyle:ignore import org.apache.spark.SparkBuildInfo import org.apache.spark.sql.connect.SparkSession import org.apache.spark.sql.connect.client.{RetryPolicy, SparkConnectClient} -import org.apache.spark.sql.connect.common.config.ConnectCommon import org.apache.spark.sql.connect.test.IntegrationTestUtils._ import org.apache.spark.util.ArrayImplicits._ @@ -52,9 +52,22 @@ object SparkConnectServerUtils { // The equivalent command to start the connect server via command line: // bin/spark-shell --conf spark.plugins=org.apache.spark.sql.connect.SparkConnectPlugin - // Server port - val port: Int = - ConnectCommon.CONNECT_GRPC_BINDING_PORT + util.Random.nextInt(1000) + // File the server process writes its actual bound port into. + private val portFile: File = { + val f = File.createTempFile("spark-connect-server-port", ".tmp") + f.deleteOnExit() + f + } + + // The port the launched server bound to, read from `portFile` once the server reports it. + lazy val port: Int = { + start() + eventually(timeout(1.minute)) { + val reported = new String(Files.readAllBytes(portFile.toPath), UTF_8).trim + assert(reported.nonEmpty, "The Spark Connect server has not reported its port yet.") + reported.toInt + } + } @volatile private var stopped = false @@ -75,7 +88,8 @@ object SparkConnectServerUtils { command += "--driver-class-path" += connectJar command += "--class" += "org.apache.spark.sql.connect.SimpleSparkConnectService" command += "--jars" += catalystTestJar - command += "--conf" += s"spark.connect.grpc.binding.port=$port" + command += "--conf" += "spark.connect.grpc.binding.port=0" + command += "--conf" += s"spark.connect.test.portFile=${portFile.getAbsolutePath}" command ++= testConfigs command ++= log4jConfigs command += connectJar diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/SimpleSparkConnectService.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/SimpleSparkConnectService.scala index 8061e913dc0da..06c8dbcc0f6a4 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/SimpleSparkConnectService.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/SimpleSparkConnectService.scala @@ -17,6 +17,8 @@ package org.apache.spark.sql.connect +import java.nio.charset.StandardCharsets.UTF_8 +import java.nio.file.{Files, Paths} import java.util.concurrent.TimeUnit import scala.io.StdIn @@ -38,6 +40,9 @@ import org.apache.spark.sql.internal.SQLConf private[sql] object SimpleSparkConnectService { private val stopCommand = "q" + // Conf naming a file to write the actual bound port into. + private val portFileConf = "spark.connect.test.portFile" + def main(args: Array[String]): Unit = { val conf = new SparkConf() .set("spark.plugins", "org.apache.spark.sql.connect.SparkConnectPlugin") @@ -45,6 +50,10 @@ private[sql] object SimpleSparkConnectService { .set(SQLConf.ARTIFACTS_SESSION_ISOLATION_ALWAYS_APPLY_CLASSLOADER, true) val sparkSession = SparkSession.builder().config(conf).getOrCreate() val sparkContext = sparkSession.sparkContext // init spark context + // Write the actual bound port to the file, if one was configured. + sparkContext.getConf.getOption(portFileConf).foreach { path => + Files.write(Paths.get(path), SparkConnectService.localPort.toString.getBytes(UTF_8)) + } // scalastyle:off println println("Ready for client connections.") // scalastyle:on println diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SparkConnectServerTest.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SparkConnectServerTest.scala index c935370f0643f..68adb498fcce6 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SparkConnectServerTest.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SparkConnectServerTest.scala @@ -32,7 +32,6 @@ import org.apache.spark.sql.classic import org.apache.spark.sql.connect import org.apache.spark.sql.connect.client.{CustomSparkConnectBlockingStub, ExecutePlanResponseReattachableIterator, RetryPolicy, SparkConnectClient, SparkConnectStubState} import org.apache.spark.sql.connect.client.arrow.ArrowSerializer -import org.apache.spark.sql.connect.common.config.ConnectCommon import org.apache.spark.sql.connect.config.Connect import org.apache.spark.sql.connect.dsl.MockRemoteSession import org.apache.spark.sql.connect.dsl.plans._ @@ -46,9 +45,8 @@ import org.apache.spark.sql.util.CloseableIterator */ trait SparkConnectServerTest extends SharedSparkSession { - // Server port - val serverPort: Int = - ConnectCommon.CONNECT_GRPC_BINDING_PORT + util.Random.nextInt(1000) + // The port the running service is bound to, set after start. + var serverPort: Int = -1 val eventuallyTimeout = 30.seconds @@ -65,12 +63,15 @@ trait SparkConnectServerTest extends SharedSparkSession { // Other suites using mocks leave a mess in the global executionManager, // shut it down so that it's cleared before starting server. SparkConnectService.executionManager.shutdown() - // Start the real service. - withSparkEnvConfs( - (Seq( - (Connect.CONNECT_GRPC_BINDING_PORT.key, serverPort.toString)) ++ extraServerConfs): _*) { + startService() + } + + /** Starts the service on an OS-assigned free port (port 0) and records it in `serverPort`. */ + protected def startService(confs: Seq[(String, String)] = extraServerConfs): Unit = { + withSparkEnvConfs((Seq((Connect.CONNECT_GRPC_BINDING_PORT.key, "0")) ++ confs): _*) { SparkConnectService.start(spark.sparkContext) } + serverPort = SparkConnectService.localPort } override def afterAll(): Unit = { diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceKeepAliveSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceKeepAliveSuite.scala index de8d6b84e30d1..ff27b09c9206a 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceKeepAliveSuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceKeepAliveSuite.scala @@ -53,6 +53,12 @@ class SparkConnectServiceKeepAliveSuite extends SparkConnectServerTest { Connect.CONNECT_GRPC_KEEPALIVE_TIME.key -> "1s", Connect.CONNECT_GRPC_KEEPALIVE_TIMEOUT.key -> "1s") + // Stops the service and restarts it on a fresh OS-assigned port with the given confs. + private def restartService(confs: Seq[(String, String)]): Unit = { + SparkConnectService.stop(Some(30), Some(TimeUnit.SECONDS)) + startService(confs) + } + test("SPARK-58094: real SparkConnectService applies configured keepalive end-to-end") { val serverSession = SparkConnectService @@ -121,14 +127,11 @@ class SparkConnectServiceKeepAliveSuite extends SparkConnectServerTest { // behavior), this client's cadence would violate that coupled invariant; with today's fixed, // decoupled GRPC_KEEPALIVE_PERMIT_TIME_SECONDS floor (10s), it's comfortably tolerated. val clientKeepAliveMs = (SparkConnectService.GRPC_KEEPALIVE_PERMIT_TIME_SECONDS + 1) * 1000 - SparkConnectService.stop(Some(30), Some(TimeUnit.SECONDS)) - withSparkEnvConfs( - Connect.CONNECT_GRPC_BINDING_PORT.key -> serverPort.toString, - Connect.CONNECT_GRPC_KEEPALIVE_ENABLED.key -> "true", - Connect.CONNECT_GRPC_KEEPALIVE_TIME.key -> "20s", - Connect.CONNECT_GRPC_KEEPALIVE_TIMEOUT.key -> "5s") { - SparkConnectService.start(spark.sparkContext) - } + restartService( + Seq( + Connect.CONNECT_GRPC_KEEPALIVE_ENABLED.key -> "true", + Connect.CONNECT_GRPC_KEEPALIVE_TIME.key -> "20s", + Connect.CONNECT_GRPC_KEEPALIVE_TIMEOUT.key -> "5s")) try { val client = SparkConnectClient .builder() @@ -156,14 +159,11 @@ class SparkConnectServiceKeepAliveSuite extends SparkConnectServerTest { client.shutdown() } } finally { - SparkConnectService.stop(Some(30), Some(TimeUnit.SECONDS)) - withSparkEnvConfs( - Connect.CONNECT_GRPC_BINDING_PORT.key -> serverPort.toString, - Connect.CONNECT_GRPC_KEEPALIVE_ENABLED.key -> "true", - Connect.CONNECT_GRPC_KEEPALIVE_TIME.key -> "1s", - Connect.CONNECT_GRPC_KEEPALIVE_TIMEOUT.key -> "1s") { - SparkConnectService.start(spark.sparkContext) - } + restartService( + Seq( + Connect.CONNECT_GRPC_KEEPALIVE_ENABLED.key -> "true", + Connect.CONNECT_GRPC_KEEPALIVE_TIME.key -> "1s", + Connect.CONNECT_GRPC_KEEPALIVE_TIMEOUT.key -> "1s")) } } @@ -190,14 +190,11 @@ class SparkConnectServiceKeepAliveSuite extends SparkConnectServerTest { // truly idle connection -- hence one call to establish the transport, then real idle time // (no further calls) spanning several ping intervals, then a connection-count check. val clientKeepAliveMs = (SparkConnectService.GRPC_KEEPALIVE_PERMIT_TIME_SECONDS + 1) * 1000 - SparkConnectService.stop(Some(30), Some(TimeUnit.SECONDS)) - withSparkEnvConfs( - Connect.CONNECT_GRPC_BINDING_PORT.key -> serverPort.toString, - Connect.CONNECT_GRPC_KEEPALIVE_ENABLED.key -> "false", - Connect.CONNECT_GRPC_KEEPALIVE_TIME.key -> "1s", - Connect.CONNECT_GRPC_KEEPALIVE_TIMEOUT.key -> "1s") { - SparkConnectService.start(spark.sparkContext) - } + restartService( + Seq( + Connect.CONNECT_GRPC_KEEPALIVE_ENABLED.key -> "false", + Connect.CONNECT_GRPC_KEEPALIVE_TIME.key -> "1s", + Connect.CONNECT_GRPC_KEEPALIVE_TIMEOUT.key -> "1s")) val relay = new FreezableTcpRelay(serverPort) try { val client = SparkConnectClient @@ -238,14 +235,11 @@ class SparkConnectServiceKeepAliveSuite extends SparkConnectServerTest { } } finally { relay.close() - SparkConnectService.stop(Some(30), Some(TimeUnit.SECONDS)) - withSparkEnvConfs( - Connect.CONNECT_GRPC_BINDING_PORT.key -> serverPort.toString, - Connect.CONNECT_GRPC_KEEPALIVE_ENABLED.key -> "true", - Connect.CONNECT_GRPC_KEEPALIVE_TIME.key -> "1s", - Connect.CONNECT_GRPC_KEEPALIVE_TIMEOUT.key -> "1s") { - SparkConnectService.start(spark.sparkContext) - } + restartService( + Seq( + Connect.CONNECT_GRPC_KEEPALIVE_ENABLED.key -> "true", + Connect.CONNECT_GRPC_KEEPALIVE_TIME.key -> "1s", + Connect.CONNECT_GRPC_KEEPALIVE_TIMEOUT.key -> "1s")) } } @@ -253,14 +247,11 @@ class SparkConnectServiceKeepAliveSuite extends SparkConnectServerTest { "SPARK-58094: disabling spark.connect.grpc.keepAlive.enabled reverts to the pre-fix hang") { // Restart the real service with keepalive fully disabled (not just given short/aggressive // timing) to prove the flag genuinely gates the fix rather than only tuning its timing. - SparkConnectService.stop(Some(30), Some(TimeUnit.SECONDS)) - withSparkEnvConfs( - Connect.CONNECT_GRPC_BINDING_PORT.key -> serverPort.toString, - Connect.CONNECT_GRPC_KEEPALIVE_ENABLED.key -> "false", - Connect.CONNECT_GRPC_KEEPALIVE_TIME.key -> "1s", - Connect.CONNECT_GRPC_KEEPALIVE_TIMEOUT.key -> "1s") { - SparkConnectService.start(spark.sparkContext) - } + restartService( + Seq( + Connect.CONNECT_GRPC_KEEPALIVE_ENABLED.key -> "false", + Connect.CONNECT_GRPC_KEEPALIVE_TIME.key -> "1s", + Connect.CONNECT_GRPC_KEEPALIVE_TIMEOUT.key -> "1s")) try { val serverSession = SparkConnectService @@ -321,14 +312,11 @@ class SparkConnectServiceKeepAliveSuite extends SparkConnectServerTest { } } finally { // Restore the enabled server for afterAll()/subsequent tests in this suite. - SparkConnectService.stop(Some(30), Some(TimeUnit.SECONDS)) - withSparkEnvConfs( - Connect.CONNECT_GRPC_BINDING_PORT.key -> serverPort.toString, - Connect.CONNECT_GRPC_KEEPALIVE_ENABLED.key -> "true", - Connect.CONNECT_GRPC_KEEPALIVE_TIME.key -> "1s", - Connect.CONNECT_GRPC_KEEPALIVE_TIMEOUT.key -> "1s") { - SparkConnectService.start(spark.sparkContext) - } + restartService( + Seq( + Connect.CONNECT_GRPC_KEEPALIVE_ENABLED.key -> "true", + Connect.CONNECT_GRPC_KEEPALIVE_TIME.key -> "1s", + Connect.CONNECT_GRPC_KEEPALIVE_TIMEOUT.key -> "1s")) } } }