From 56ca0034c90a9ec5524b398b6f49b554a564858c Mon Sep 17 00:00:00 2001 From: Kousuke Saruta Date: Mon, 27 Jul 2026 07:21:27 +0000 Subject: [PATCH 1/3] Introduce a way to wait for a port unbound to SparkConnectServiceKeepAliveSuite --- .../SparkConnectServiceKeepAliveSuite.scala | 30 ++++++++++++++----- 1 file changed, 23 insertions(+), 7 deletions(-) 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..6fcdb8e443f11 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 @@ -16,6 +16,7 @@ */ package org.apache.spark.sql.connect.service +import java.net.ServerSocket import java.util.UUID import java.util.concurrent.TimeUnit @@ -28,6 +29,7 @@ import org.apache.spark.SparkException import org.apache.spark.sql.connect.SparkConnectServerTest import org.apache.spark.sql.connect.client.{RetryPolicy, SparkConnectClient} import org.apache.spark.sql.connect.config.Connect +import org.apache.spark.util.Utils.tryWithSafeFinally /** * End-to-end test that the *real* `SparkConnectService.startGRPCService()` wiring actually @@ -53,6 +55,20 @@ class SparkConnectServiceKeepAliveSuite extends SparkConnectServerTest { Connect.CONNECT_GRPC_KEEPALIVE_TIME.key -> "1s", Connect.CONNECT_GRPC_KEEPALIVE_TIMEOUT.key -> "1s") + private def stopSparkConnectServiceAndUnbind(): Unit = { + SparkConnectService.stop(Some(30), Some(TimeUnit.SECONDS)) + Eventually.eventually(timeout(eventuallyTimeout)) { + var s: ServerSocket = null + tryWithSafeFinally { + s = new ServerSocket(serverPort) + } { + if (s != null) { + s.close() + } + } + } + } + test("SPARK-58094: real SparkConnectService applies configured keepalive end-to-end") { val serverSession = SparkConnectService @@ -121,7 +137,7 @@ 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)) + stopSparkConnectServiceAndUnbind() withSparkEnvConfs( Connect.CONNECT_GRPC_BINDING_PORT.key -> serverPort.toString, Connect.CONNECT_GRPC_KEEPALIVE_ENABLED.key -> "true", @@ -156,7 +172,7 @@ class SparkConnectServiceKeepAliveSuite extends SparkConnectServerTest { client.shutdown() } } finally { - SparkConnectService.stop(Some(30), Some(TimeUnit.SECONDS)) + stopSparkConnectServiceAndUnbind() withSparkEnvConfs( Connect.CONNECT_GRPC_BINDING_PORT.key -> serverPort.toString, Connect.CONNECT_GRPC_KEEPALIVE_ENABLED.key -> "true", @@ -190,7 +206,7 @@ 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)) + stopSparkConnectServiceAndUnbind() withSparkEnvConfs( Connect.CONNECT_GRPC_BINDING_PORT.key -> serverPort.toString, Connect.CONNECT_GRPC_KEEPALIVE_ENABLED.key -> "false", @@ -238,7 +254,7 @@ class SparkConnectServiceKeepAliveSuite extends SparkConnectServerTest { } } finally { relay.close() - SparkConnectService.stop(Some(30), Some(TimeUnit.SECONDS)) + stopSparkConnectServiceAndUnbind() withSparkEnvConfs( Connect.CONNECT_GRPC_BINDING_PORT.key -> serverPort.toString, Connect.CONNECT_GRPC_KEEPALIVE_ENABLED.key -> "true", @@ -253,7 +269,7 @@ 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)) + stopSparkConnectServiceAndUnbind() withSparkEnvConfs( Connect.CONNECT_GRPC_BINDING_PORT.key -> serverPort.toString, Connect.CONNECT_GRPC_KEEPALIVE_ENABLED.key -> "false", @@ -321,7 +337,7 @@ class SparkConnectServiceKeepAliveSuite extends SparkConnectServerTest { } } finally { // Restore the enabled server for afterAll()/subsequent tests in this suite. - SparkConnectService.stop(Some(30), Some(TimeUnit.SECONDS)) + stopSparkConnectServiceAndUnbind() withSparkEnvConfs( Connect.CONNECT_GRPC_BINDING_PORT.key -> serverPort.toString, Connect.CONNECT_GRPC_KEEPALIVE_ENABLED.key -> "true", @@ -341,7 +357,7 @@ class SparkConnectServiceKeepAliveSuite extends SparkConnectServerTest { * shared across modules since it's a small, self-contained test utility. */ private class FreezableTcpRelay(targetPort: Int) { - private val listener = new java.net.ServerSocket(0) + private val listener = new ServerSocket(0) @volatile private var frozen = false private val sockets = mutable.ListBuffer[java.net.Socket]() private val acceptedConnections = new java.util.concurrent.atomic.AtomicInteger(0) From 0c3d99b572eff304e561f0aab32f74333f9539bd Mon Sep 17 00:00:00 2001 From: Kousuke Saruta Date: Mon, 27 Jul 2026 09:06:16 +0000 Subject: [PATCH 2/3] Use tryWithResource --- .../service/SparkConnectServiceKeepAliveSuite.scala | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) 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 6fcdb8e443f11..600070cae058c 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 @@ -29,7 +29,7 @@ import org.apache.spark.SparkException import org.apache.spark.sql.connect.SparkConnectServerTest import org.apache.spark.sql.connect.client.{RetryPolicy, SparkConnectClient} import org.apache.spark.sql.connect.config.Connect -import org.apache.spark.util.Utils.tryWithSafeFinally +import org.apache.spark.util.Utils.tryWithResource /** * End-to-end test that the *real* `SparkConnectService.startGRPCService()` wiring actually @@ -59,13 +59,7 @@ class SparkConnectServiceKeepAliveSuite extends SparkConnectServerTest { SparkConnectService.stop(Some(30), Some(TimeUnit.SECONDS)) Eventually.eventually(timeout(eventuallyTimeout)) { var s: ServerSocket = null - tryWithSafeFinally { - s = new ServerSocket(serverPort) - } { - if (s != null) { - s.close() - } - } + tryWithResource(new ServerSocket(serverPort))(identity) } } From 02bcc94990a2036bf7b9ae1431ae7e56990183aa Mon Sep 17 00:00:00 2001 From: Kousuke Saruta Date: Mon, 27 Jul 2026 14:50:39 +0000 Subject: [PATCH 3/3] Remove unused code --- .../sql/connect/service/SparkConnectServiceKeepAliveSuite.scala | 1 - 1 file changed, 1 deletion(-) 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 600070cae058c..a0b249a002aba 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 @@ -58,7 +58,6 @@ class SparkConnectServiceKeepAliveSuite extends SparkConnectServerTest { private def stopSparkConnectServiceAndUnbind(): Unit = { SparkConnectService.stop(Some(30), Some(TimeUnit.SECONDS)) Eventually.eventually(timeout(eventuallyTimeout)) { - var s: ServerSocket = null tryWithResource(new ServerSocket(serverPort))(identity) } }