[SPARK-58264][CONNECT][TESTS] Use OS-assigned ports in Connect server tests

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

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,13 +40,20 @@ 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")
       .set(SQLConf.ARTIFACTS_SESSION_ISOLATION_ENABLED, true)
       .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

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 = {

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,29 +235,23 @@ 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"))
     }
   }
 
   test(
     "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"))
     }
   }
 }