spark.connect.extensions.relation.classes |
diff --git a/python/pyspark/sql/connect/client/core.py b/python/pyspark/sql/connect/client/core.py
index db6067f25e096..1e81e780d167e 100644
--- a/python/pyspark/sql/connect/client/core.py
+++ b/python/pyspark/sql/connect/client/core.py
@@ -212,9 +212,23 @@ class ChannelBuilder:
PARAM_USER_ID = "user_id"
PARAM_USER_AGENT = "user_agent"
PARAM_SESSION_ID = "session_id"
+ PARAM_GRPC_KEEPALIVE_ENABLED = "grpc_keepalive_enabled"
+ PARAM_GRPC_KEEPALIVE_TIME_MS = "grpc_keepalive_time_ms"
+ PARAM_GRPC_KEEPALIVE_TIMEOUT_MS = "grpc_keepalive_timeout_ms"
+ PARAM_GRPC_KEEPALIVE_WITHOUT_CALLS = "grpc_keepalive_without_calls"
GRPC_MAX_MESSAGE_LENGTH_DEFAULT = 128 * 1024 * 1024
+ # Detects a silently-dead connection (e.g. a NAT gateway/load balancer dropping an idle
+ # connection mapping without sending a TCP RST/FIN) via gRPC/HTTP2 keepalive PINGs, so a
+ # blocked RPC (such as streaming query awaitTermination()) surfaces as UNAVAILABLE instead
+ # of hanging forever. Mirrors the JVM client's defaults (SparkConnectClient.scala). See
+ # SPARK-58094.
+ GRPC_DEFAULT_KEEPALIVE_ENABLED = True
+ GRPC_DEFAULT_KEEPALIVE_TIME_MS = 60 * 1000
+ GRPC_DEFAULT_KEEPALIVE_TIMEOUT_MS = 20 * 1000
+ GRPC_DEFAULT_KEEPALIVE_WITHOUT_CALLS = True
+
GRPC_DEFAULT_OPTIONS = [
("grpc.max_send_message_length", GRPC_MAX_MESSAGE_LENGTH_DEFAULT),
("grpc.max_receive_message_length", GRPC_MAX_MESSAGE_LENGTH_DEFAULT),
@@ -281,7 +295,7 @@ def host(self) -> str:
raise PySparkNotImplementedError
def _insecure_channel(self, target: Any, **kwargs: Any) -> grpc.Channel:
- channel = grpc.insecure_channel(target, options=self._channel_options, **kwargs)
+ channel = grpc.insecure_channel(target, options=self._effective_channel_options(), **kwargs)
if len(self._interceptors) > 0:
logger.debug(f"Applying interceptors ({self._interceptors})")
@@ -289,7 +303,9 @@ def _insecure_channel(self, target: Any, **kwargs: Any) -> grpc.Channel:
return channel
def _secure_channel(self, target: Any, credentials: Any, **kwargs: Any) -> grpc.Channel:
- channel = grpc.secure_channel(target, credentials, options=self._channel_options, **kwargs)
+ channel = grpc.secure_channel(
+ target, credentials, options=self._effective_channel_options(), **kwargs
+ )
if len(self._interceptors) > 0:
logger.debug(f"Applying interceptors ({self._interceptors})")
@@ -311,6 +327,74 @@ def token(self) -> Optional[str]:
ChannelBuilder.PARAM_TOKEN, os.environ.get("SPARK_CONNECT_AUTHENTICATE_TOKEN")
)
+ @property
+ def keepalive_enabled(self) -> bool:
+ """
+ Whether the client sends gRPC/HTTP2 keepalive PINGs to detect a silently-dead
+ connection. Enabled by default; can be turned off as an escape hatch, e.g. if it
+ interacts badly with a particular network path, or a client environment is prone to
+ stalls long enough to trip false-positive disconnects.
+ """
+ return (
+ self.getDefault(
+ ChannelBuilder.PARAM_GRPC_KEEPALIVE_ENABLED,
+ str(ChannelBuilder.GRPC_DEFAULT_KEEPALIVE_ENABLED),
+ ).lower()
+ == "true"
+ )
+
+ @property
+ def keepalive_time_ms(self) -> int:
+ """Idle time (in milliseconds) before sending a gRPC/HTTP2 keepalive PING."""
+ return int(
+ self.getDefault(
+ ChannelBuilder.PARAM_GRPC_KEEPALIVE_TIME_MS,
+ ChannelBuilder.GRPC_DEFAULT_KEEPALIVE_TIME_MS,
+ )
+ )
+
+ @property
+ def keepalive_timeout_ms(self) -> int:
+ """Time (in milliseconds) to wait for a keepalive PING ack before failing."""
+ return int(
+ self.getDefault(
+ ChannelBuilder.PARAM_GRPC_KEEPALIVE_TIMEOUT_MS,
+ ChannelBuilder.GRPC_DEFAULT_KEEPALIVE_TIMEOUT_MS,
+ )
+ )
+
+ @property
+ def keepalive_without_calls(self) -> bool:
+ """Whether to keep sending keepalive PINGs when there are no in-flight RPCs."""
+ return (
+ self.getDefault(
+ ChannelBuilder.PARAM_GRPC_KEEPALIVE_WITHOUT_CALLS,
+ str(ChannelBuilder.GRPC_DEFAULT_KEEPALIVE_WITHOUT_CALLS),
+ ).lower()
+ == "true"
+ )
+
+ def _effective_channel_options(self) -> List[Tuple[str, Any]]:
+ """
+ Returns ``self._channel_options`` with the keepalive options
+ (:attr:`keepalive_enabled`/:attr:`keepalive_time_ms`/:attr:`keepalive_timeout_ms`/
+ :attr:`keepalive_without_calls`) applied, unless a caller already set one of those
+ specific keys explicitly via ``channelOptions``/:meth:`setChannelOption`, in which case
+ the explicit value wins.
+ """
+ options = list(self._channel_options)
+ if not self.keepalive_enabled:
+ return options
+ existing_keys = {k for k, _ in options}
+ for key, value in (
+ ("grpc.keepalive_time_ms", self.keepalive_time_ms),
+ ("grpc.keepalive_timeout_ms", self.keepalive_timeout_ms),
+ ("grpc.keepalive_permit_without_calls", 1 if self.keepalive_without_calls else 0),
+ ):
+ if key not in existing_keys:
+ options.append((key, value))
+ return options
+
def metadata(self) -> Iterable[Tuple[str, str]]:
"""
Builds the GRPC specific metadata list to be injected into the request. All
@@ -330,6 +414,10 @@ def metadata(self) -> Iterable[Tuple[str, str]]:
ChannelBuilder.PARAM_USER_ID,
ChannelBuilder.PARAM_USER_AGENT,
ChannelBuilder.PARAM_SESSION_ID,
+ ChannelBuilder.PARAM_GRPC_KEEPALIVE_ENABLED,
+ ChannelBuilder.PARAM_GRPC_KEEPALIVE_TIME_MS,
+ ChannelBuilder.PARAM_GRPC_KEEPALIVE_TIMEOUT_MS,
+ ChannelBuilder.PARAM_GRPC_KEEPALIVE_WITHOUT_CALLS,
]
]
@@ -399,6 +487,10 @@ class DefaultChannelBuilder(ChannelBuilder):
>>> cb = DefaultChannelBuilder("sc://localhost/;use_ssl=true;token=aaa")
... cb.secure
True
+
+ >>> cb = DefaultChannelBuilder("sc://localhost/;grpc_keepalive_time_ms=30000")
+ ... cb.keepalive_time_ms
+ 30000
"""
@staticmethod
diff --git a/python/pyspark/sql/tests/connect/test_connect_channel.py b/python/pyspark/sql/tests/connect/test_connect_channel.py
index 21390d671ddad..66a8f0d8abc87 100644
--- a/python/pyspark/sql/tests/connect/test_connect_channel.py
+++ b/python/pyspark/sql/tests/connect/test_connect_channel.py
@@ -129,6 +129,79 @@ def test_metadata_with_session_id(self):
chan = DefaultChannelBuilder("sc://host/")
self.assertIsNone(chan.session_id)
+ def test_keepalive_defaults(self):
+ # SPARK-58094
+ chan = DefaultChannelBuilder("sc://host")
+ self.assertTrue(chan.keepalive_enabled)
+ self.assertEqual(60 * 1000, chan.keepalive_time_ms)
+ self.assertEqual(20 * 1000, chan.keepalive_timeout_ms)
+ self.assertTrue(chan.keepalive_without_calls)
+
+ options = dict(chan._effective_channel_options())
+ self.assertEqual(60 * 1000, options["grpc.keepalive_time_ms"])
+ self.assertEqual(20 * 1000, options["grpc.keepalive_timeout_ms"])
+ self.assertEqual(1, options["grpc.keepalive_permit_without_calls"])
+
+ def test_keepalive_connection_string_overrides(self):
+ # SPARK-58094
+ chan = DefaultChannelBuilder(
+ "sc://host/;grpc_keepalive_enabled=false;grpc_keepalive_time_ms=12345;"
+ "grpc_keepalive_timeout_ms=6789;grpc_keepalive_without_calls=false"
+ )
+ self.assertFalse(chan.keepalive_enabled)
+ self.assertEqual(12345, chan.keepalive_time_ms)
+ self.assertEqual(6789, chan.keepalive_timeout_ms)
+ self.assertFalse(chan.keepalive_without_calls)
+
+ # Disabled entirely: none of the keepalive options are applied to the channel.
+ options = dict(chan._effective_channel_options())
+ self.assertNotIn("grpc.keepalive_time_ms", options)
+ self.assertNotIn("grpc.keepalive_timeout_ms", options)
+ self.assertNotIn("grpc.keepalive_permit_without_calls", options)
+
+ def test_keepalive_enabled_with_custom_values(self):
+ # SPARK-58094
+ chan = DefaultChannelBuilder(
+ "sc://host/;grpc_keepalive_enabled=true;grpc_keepalive_time_ms=1000;"
+ "grpc_keepalive_timeout_ms=500;grpc_keepalive_without_calls=false"
+ )
+ self.assertTrue(chan.keepalive_enabled)
+ options = dict(chan._effective_channel_options())
+ self.assertEqual(1000, options["grpc.keepalive_time_ms"])
+ self.assertEqual(500, options["grpc.keepalive_timeout_ms"])
+ self.assertEqual(0, options["grpc.keepalive_permit_without_calls"])
+
+ def test_keepalive_explicit_channel_option_wins(self):
+ # SPARK-58094: an explicitly-provided channel option for one of the keepalive keys is
+ # not overwritten/duplicated by the default-derived value.
+ chan = DefaultChannelBuilder("sc://host", [("grpc.keepalive_time_ms", 999)])
+ options = chan._effective_channel_options()
+ self.assertEqual(
+ [k for k, _ in options].count("grpc.keepalive_time_ms"),
+ 1,
+ "only one occurrence, no duplicate",
+ )
+ self.assertEqual(
+ next(v for k, v in options if k == "grpc.keepalive_time_ms"),
+ 999,
+ "explicit value is not overwritten by the default",
+ )
+ # The other keepalive keys are still applied normally.
+ self.assertIn("grpc.keepalive_timeout_ms", dict(options))
+
+ def test_keepalive_params_excluded_from_metadata(self):
+ # SPARK-58094
+ chan = DefaultChannelBuilder(
+ "sc://host/;grpc_keepalive_enabled=false;grpc_keepalive_time_ms=1000;"
+ "grpc_keepalive_timeout_ms=500;grpc_keepalive_without_calls=false;param1=abc"
+ )
+ md = dict(chan.metadata())
+ self.assertNotIn(ChannelBuilder.PARAM_GRPC_KEEPALIVE_ENABLED, md)
+ self.assertNotIn(ChannelBuilder.PARAM_GRPC_KEEPALIVE_TIME_MS, md)
+ self.assertNotIn(ChannelBuilder.PARAM_GRPC_KEEPALIVE_TIMEOUT_MS, md)
+ self.assertNotIn(ChannelBuilder.PARAM_GRPC_KEEPALIVE_WITHOUT_CALLS, md)
+ self.assertIn("param1", md)
+
def test_channel_options(self):
# SPARK-47694
chan = DefaultChannelBuilder(
diff --git a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientBuilderParseTestSuite.scala b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientBuilderParseTestSuite.scala
index ad0f880e8f57a..b9b3754252a33 100644
--- a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientBuilderParseTestSuite.scala
+++ b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientBuilderParseTestSuite.scala
@@ -18,6 +18,7 @@ package org.apache.spark.sql.connect.client
import java.util.UUID
+import org.apache.spark.sql.connect.common.config.ConnectCommon
import org.apache.spark.sql.connect.test.ConnectFunSuite
/**
@@ -49,6 +50,15 @@ class SparkConnectClientBuilderParseTestSuite extends ConnectFunSuite {
argumentTest("user_name", "alice", _.userName.get)
argumentTest("user_agent", "robert", _.userAgent.split(" ")(0))
argumentTest("session_id", UUID.randomUUID().toString, _.sessionId.get)
+ argumentTest("grpc_keepalive_enabled", "false", _.grpcKeepAliveEnabled.toString)
+ argumentTest("grpc_keepalive_time_ms", "12345", _.grpcKeepAliveTimeMs.toString)
+ argumentTest("grpc_keepalive_timeout_ms", "6789", _.grpcKeepAliveTimeoutMs.toString)
+ argumentTest("grpc_keepalive_without_calls", "false", _.grpcKeepAliveWithoutCalls.toString)
+
+ test("Argument - grpc_keepalive_enabled explicit true") {
+ val builder = build("--grpc_keepalive_enabled", "true")
+ assert(builder.grpcKeepAliveEnabled)
+ }
test("Argument - remote") {
val builder =
@@ -61,6 +71,41 @@ class SparkConnectClientBuilderParseTestSuite extends ConnectFunSuite {
assert(builder.sessionId.isEmpty)
}
+ test("Argument - remote with keepalive params (disabled)") {
+ val builder = build(
+ "--remote",
+ "sc://srv.apache.org/;grpc_keepalive_enabled=false;grpc_keepalive_time_ms=15000;" +
+ "grpc_keepalive_timeout_ms=5000;grpc_keepalive_without_calls=false")
+ assert(!builder.grpcKeepAliveEnabled)
+ assert(builder.grpcKeepAliveTimeMs === 15000L)
+ assert(builder.grpcKeepAliveTimeoutMs === 5000L)
+ assert(!builder.grpcKeepAliveWithoutCalls)
+ }
+
+ test("Argument - remote with keepalive params (explicitly enabled)") {
+ val builder = build(
+ "--remote",
+ "sc://srv.apache.org/;grpc_keepalive_enabled=true;grpc_keepalive_time_ms=15000;" +
+ "grpc_keepalive_timeout_ms=5000;grpc_keepalive_without_calls=true")
+ assert(builder.grpcKeepAliveEnabled)
+ assert(builder.grpcKeepAliveTimeMs === 15000L)
+ assert(builder.grpcKeepAliveTimeoutMs === 5000L)
+ assert(builder.grpcKeepAliveWithoutCalls)
+ }
+
+ test("keepalive defaults apply when unset") {
+ val builder = build("--host", "localhost")
+ assert(
+ builder.configuration.grpcKeepAliveEnabled === ConnectCommon.CONNECT_GRPC_KEEPALIVE_ENABLED)
+ assert(
+ builder.configuration.grpcKeepAliveTimeMs ===
+ ConnectCommon.CONNECT_GRPC_KEEPALIVE_TIME_SECONDS * 1000)
+ assert(
+ builder.configuration.grpcKeepAliveTimeoutMs ===
+ ConnectCommon.CONNECT_GRPC_KEEPALIVE_TIMEOUT_SECONDS * 1000)
+ assert(builder.configuration.grpcKeepAliveWithoutCalls)
+ }
+
test("Argument - use_ssl") {
val builder = build("--use_ssl")
assert(builder.sslEnabled)
diff --git a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientSuite.scala b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientSuite.scala
index d270b0544097e..fd5be9bb555ee 100644
--- a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientSuite.scala
+++ b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientSuite.scala
@@ -898,6 +898,110 @@ class SparkConnectClientSuite extends ConnectFunSuite {
}
}
+ test(
+ "SPARK-58094: gRPC keepalive surfaces a bounded failure on a silently dropped " +
+ "connection instead of hanging forever") {
+ val requestReceived = new CountDownLatch(1)
+ val serverLatch = new CountDownLatch(1)
+ val slowService = new DummySparkConnectService {
+ override def analyzePlan(
+ request: AnalyzePlanRequest,
+ responseObserver: StreamObserver[AnalyzePlanResponse]): Unit = {
+ requestReceived.countDown()
+ // Blocks far longer than the test's keepalive settings, standing in for a query that
+ // is still genuinely running server-side while the client's connection goes dark.
+ serverLatch.await(30, TimeUnit.SECONDS)
+ super.analyzePlan(request, responseObserver)
+ }
+ }
+ val keepAliveMs = 300L
+ server = NettyServerBuilder
+ .forPort(0)
+ .addService(slowService)
+ // The server must permit client PINGs at least as frequently as the client's own
+ // keepAliveTime, or it will tear down the connection as "too_many_pings" even when
+ // healthy. In production (SparkConnectService.scala) this floor is a fixed constant
+ // decoupled from any per-connection setting; here it's simply set to match the test
+ // client's cadence.
+ .permitKeepAliveTime(keepAliveMs, TimeUnit.MILLISECONDS)
+ .permitKeepAliveWithoutCalls(true)
+ .build()
+ .start()
+ service = slowService
+ val relay = new FreezableTcpRelay(server.getPort)
+ try {
+ client = SparkConnectClient
+ .builder()
+ .connectionString(s"sc://localhost:${relay.port}")
+ .grpcKeepAliveEnabled(true)
+ .grpcKeepAliveTimeMs(keepAliveMs)
+ .grpcKeepAliveTimeoutMs(keepAliveMs)
+ .retryPolicy(RetryPolicy(maxRetries = Some(0), canRetry = _ => false, name = "NoRetry"))
+ .build()
+
+ val resultPromise = scala.concurrent.Promise[Unit]()
+ val callThread = new Thread(() => {
+ try {
+ client.analyze(proto.AnalyzePlanRequest.newBuilder().setSessionId("abc123").build())
+ resultPromise.success(())
+ } catch {
+ case t: Throwable => resultPromise.failure(t)
+ }
+ })
+ callThread.setDaemon(true)
+ callThread.start()
+
+ assert(requestReceived.await(10, TimeUnit.SECONDS), "server never received the request")
+ // Freeze the relay without closing either socket, simulating a NAT gateway/load balancer
+ // silently dropping an idle connection mapping (no TCP RST/FIN to either endpoint).
+ relay.freeze()
+
+ // Bounded by Await's timeout: without the keepalive fix this never completes and the
+ // test would fail with a TimeoutException instead of surfacing UNAVAILABLE.
+ // scalastyle:off awaitresult
+ val ex = intercept[SparkException] {
+ scala.concurrent.Await.result(resultPromise.future, FiniteDuration(15, TimeUnit.SECONDS))
+ }
+ // scalastyle:on awaitresult
+ // A keepalive-triggered UNAVAILABLE carries no wrapped cause (same as DEADLINE_EXCEEDED,
+ // see GrpcExceptionConverter.toThrowable), so the status code/description is only in the
+ // message.
+ assert(ex.getMessage.contains("UNAVAILABLE"))
+ assert(ex.getMessage.contains("Keepalive failed"))
+ } finally {
+ serverLatch.countDown()
+ relay.close()
+ }
+ }
+
+ test("SPARK-58094: keepalive does not disrupt a healthy long-lived connection") {
+ val keepAliveMs = 300L
+ server = NettyServerBuilder
+ .forPort(0)
+ .addService(new DummySparkConnectService)
+ .permitKeepAliveTime(keepAliveMs, TimeUnit.MILLISECONDS)
+ .permitKeepAliveWithoutCalls(true)
+ .build()
+ .start()
+ client = SparkConnectClient
+ .builder()
+ .connectionString(s"sc://localhost:${server.getPort}")
+ .grpcKeepAliveEnabled(true)
+ .grpcKeepAliveTimeMs(keepAliveMs)
+ .grpcKeepAliveTimeoutMs(keepAliveMs)
+ .build()
+
+ // Several calls spaced out well beyond the keepalive interval: if permitKeepAliveTime were
+ // mistuned (e.g. left at the gRPC default of 5 minutes while the client pings every 300ms),
+ // the server would tear the connection down as "too_many_pings" and these would fail.
+ for (i <- 1 to 3) {
+ val response =
+ client.analyze(proto.AnalyzePlanRequest.newBuilder().setSessionId(s"abc$i").build())
+ assert(response.getSessionId === s"abc$i")
+ Thread.sleep(keepAliveMs * 2)
+ }
+ }
+
test("analyzePlan with short deadline fires, then succeeds after disabling deadlines") {
val latch = new CountDownLatch(1)
val slowService = new DummySparkConnectService {
@@ -1067,6 +1171,77 @@ class SparkConnectClientSuite extends ConnectFunSuite {
}
}
+/**
+ * A minimal bidirectional TCP relay that can be "frozen" mid-connection without closing either
+ * side's socket -- simulating a NAT gateway/load balancer/corporate proxy that silently drops an
+ * idle connection's mapping (no TCP RST/FIN sent to either endpoint). While frozen, no bytes are
+ * read from or written to either socket, so gRPC/HTTP2 keepalive PING frames also stop flowing,
+ * exactly as they would over a genuinely dead network path.
+ */
+private class FreezableTcpRelay(targetPort: Int) {
+ private val listener = new java.net.ServerSocket(0)
+ @volatile private var frozen = false
+ private val sockets = mutable.ListBuffer[java.net.Socket]()
+
+ private val acceptThread = new Thread(() => {
+ try {
+ while (true) {
+ val clientSocket = listener.accept()
+ val serverSocket = new java.net.Socket("127.0.0.1", targetPort)
+ sockets.synchronized {
+ sockets += clientSocket
+ sockets += serverSocket
+ }
+ pump(clientSocket, serverSocket)
+ pump(serverSocket, clientSocket)
+ }
+ } catch {
+ case _: java.io.IOException => // listener closed, relay shutting down
+ }
+ })
+ acceptThread.setDaemon(true)
+ acceptThread.start()
+
+ def port: Int = listener.getLocalPort
+
+ // Once frozen, bytes read from either socket are dropped rather than forwarded: no data
+ // (including keepalive PING/PONG frames) reaches the other side, but neither socket is
+ // closed -- the same observable behavior as a middlebox silently black-holing the connection.
+ def freeze(): Unit = frozen = true
+
+ private def pump(src: java.net.Socket, dst: java.net.Socket): Unit = {
+ val t = new Thread(() => {
+ try {
+ val in = src.getInputStream
+ val out = dst.getOutputStream
+ val buf = new Array[Byte](8192)
+ var n = 0
+ while ({ n = in.read(buf); n != -1 }) {
+ if (!frozen) {
+ out.write(buf, 0, n)
+ out.flush()
+ }
+ }
+ } catch {
+ case _: java.io.IOException => // socket closed, nothing left to pump
+ }
+ })
+ t.setDaemon(true)
+ t.start()
+ }
+
+ def close(): Unit = {
+ listener.close()
+ sockets.synchronized(sockets.foreach { s =>
+ try {
+ s.close()
+ } catch {
+ case _: java.io.IOException =>
+ }
+ })
+ }
+}
+
class DummySparkConnectService() extends SparkConnectServiceGrpc.SparkConnectServiceImplBase {
private var inputPlan: proto.Plan = _
diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala
index 7bd7d0e4e9790..8f449207b7968 100644
--- a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala
+++ b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala
@@ -20,7 +20,7 @@ package org.apache.spark.sql.connect.client
import java.net.URI
import java.nio.charset.StandardCharsets.UTF_8
import java.util.{Base64, Locale, UUID}
-import java.util.concurrent.Executor
+import java.util.concurrent.{Executor, TimeUnit}
import scala.collection.mutable
import scala.jdk.CollectionConverters._
@@ -816,6 +816,10 @@ object SparkConnectClient {
val PARAM_USER_AGENT = "user_agent"
val PARAM_SESSION_ID = "session_id"
val PARAM_GRPC_MAX_MESSAGE_SIZE = "grpc_max_message_size"
+ val PARAM_GRPC_KEEPALIVE_ENABLED = "grpc_keepalive_enabled"
+ val PARAM_GRPC_KEEPALIVE_TIME_MS = "grpc_keepalive_time_ms"
+ val PARAM_GRPC_KEEPALIVE_TIMEOUT_MS = "grpc_keepalive_timeout_ms"
+ val PARAM_GRPC_KEEPALIVE_WITHOUT_CALLS = "grpc_keepalive_without_calls"
}
private def verifyURI(uri: URI): Unit = {
@@ -878,6 +882,48 @@ object SparkConnectClient {
def grpcMaxRecursionLimit: Int = _configuration.grpcMaxRecursionLimit
+ /**
+ * Whether the client sends gRPC/HTTP2 keepalive PINGs to detect a silently-dead connection
+ * (see `grpcKeepAliveTimeMs`/`grpcKeepAliveTimeoutMs`). Enabled by default; can be turned off
+ * as an escape hatch, e.g. if it interacts badly with a particular network path, or a client
+ * environment is prone to stalls long enough to trip false-positive disconnects.
+ */
+ def grpcKeepAliveEnabled(enabled: Boolean): Builder = {
+ _configuration = _configuration.copy(grpcKeepAliveEnabled = enabled)
+ this
+ }
+
+ def grpcKeepAliveEnabled: Boolean = _configuration.grpcKeepAliveEnabled
+
+ def grpcKeepAliveTimeMs(timeMs: Long): Builder = {
+ _configuration = _configuration.copy(grpcKeepAliveTimeMs = timeMs)
+ this
+ }
+
+ def grpcKeepAliveTimeMs: Long = _configuration.grpcKeepAliveTimeMs
+
+ def grpcKeepAliveTimeoutMs(timeoutMs: Long): Builder = {
+ _configuration = _configuration.copy(grpcKeepAliveTimeoutMs = timeoutMs)
+ this
+ }
+
+ def grpcKeepAliveTimeoutMs: Long = _configuration.grpcKeepAliveTimeoutMs
+
+ /**
+ * Whether to keep sending keepalive PINGs when there are no in-flight RPCs on the channel.
+ * Defaults to true, so a long-idle connection (e.g. a parked/unused session) is still
+ * monitored and not just calls that are actively blocked. Set to false to avoid the resulting
+ * steady background PING traffic for deployments with very large numbers of mostly-idle
+ * connections, at the cost of only detecting a dead connection once the next RPC is attempted
+ * on it.
+ */
+ def grpcKeepAliveWithoutCalls(enabled: Boolean): Builder = {
+ _configuration = _configuration.copy(grpcKeepAliveWithoutCalls = enabled)
+ this
+ }
+
+ def grpcKeepAliveWithoutCalls: Boolean = _configuration.grpcKeepAliveWithoutCalls
+
def option(key: String, value: String): Builder = {
_configuration = _configuration.copy(metadata = _configuration.metadata + ((key, value)))
this
@@ -905,6 +951,11 @@ object SparkConnectClient {
if (java.lang.Boolean.valueOf(value)) enableSsl() else disableSsl()
case URIParams.PARAM_SESSION_ID => sessionId(value)
case URIParams.PARAM_GRPC_MAX_MESSAGE_SIZE => grpcMaxMessageSize(value.toInt)
+ case URIParams.PARAM_GRPC_KEEPALIVE_ENABLED => grpcKeepAliveEnabled(value.toBoolean)
+ case URIParams.PARAM_GRPC_KEEPALIVE_TIME_MS => grpcKeepAliveTimeMs(value.toLong)
+ case URIParams.PARAM_GRPC_KEEPALIVE_TIMEOUT_MS => grpcKeepAliveTimeoutMs(value.toLong)
+ case URIParams.PARAM_GRPC_KEEPALIVE_WITHOUT_CALLS =>
+ grpcKeepAliveWithoutCalls(value.toBoolean)
case _ => option(key, value)
}
}
@@ -1048,6 +1099,10 @@ object SparkConnectClient {
sessionId: Option[String] = None,
grpcMaxMessageSize: Int = ConnectCommon.CONNECT_GRPC_MAX_MESSAGE_SIZE,
grpcMaxRecursionLimit: Int = ConnectCommon.CONNECT_GRPC_MARSHALLER_RECURSION_LIMIT,
+ grpcKeepAliveEnabled: Boolean = ConnectCommon.CONNECT_GRPC_KEEPALIVE_ENABLED,
+ grpcKeepAliveTimeMs: Long = ConnectCommon.CONNECT_GRPC_KEEPALIVE_TIME_SECONDS * 1000,
+ grpcKeepAliveTimeoutMs: Long = ConnectCommon.CONNECT_GRPC_KEEPALIVE_TIMEOUT_SECONDS * 1000,
+ grpcKeepAliveWithoutCalls: Boolean = ConnectCommon.CONNECT_GRPC_KEEPALIVE_WITHOUT_CALLS,
allowArrowBatchChunking: Boolean = true,
preferredArrowChunkSize: Option[Int] = None) {
@@ -1098,6 +1153,18 @@ object SparkConnectClient {
interceptors.foreach(channelBuilder.intercept(_))
channelBuilder.maxInboundMessageSize(grpcMaxMessageSize)
+ if (grpcKeepAliveEnabled) {
+ // Detects a silently-dead connection (e.g. a NAT gateway/load balancer dropping an idle
+ // connection mapping without sending a TCP RST/FIN) so a blocked RPC (such as streaming
+ // query awaitTermination()) surfaces as UNAVAILABLE instead of hanging forever. Note: a
+ // stall on either end longer than keepAliveTime + keepAliveTimeout (e.g. a long JVM GC
+ // pause) can also cause an otherwise-healthy connection to be dropped -- raise both
+ // values (or disable via grpcKeepAliveEnabled) in environments prone to long GC pauses
+ // or other event-loop stalls.
+ channelBuilder.keepAliveTime(grpcKeepAliveTimeMs, TimeUnit.MILLISECONDS)
+ channelBuilder.keepAliveTimeout(grpcKeepAliveTimeoutMs, TimeUnit.MILLISECONDS)
+ channelBuilder.keepAliveWithoutCalls(grpcKeepAliveWithoutCalls)
+ }
channelBuilder.build()
}
diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClientParser.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClientParser.scala
index 7e137a6a3e051..629b6d237e5c4 100644
--- a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClientParser.scala
+++ b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClientParser.scala
@@ -42,6 +42,10 @@ private[sql] object SparkConnectClientParser {
| --user_agent USER_AGENT The User-Agent Client information (only intended for logging purposes by the server).
| --session_id SESSION_ID Session Id of the user connecting.
| --grpc_max_message_size SIZE Maximum message size allowed for gRPC messages in bytes.
+ | --grpc_keepalive_enabled BOOL Whether to enable gRPC/HTTP2 keepalive PINGs.
+ | --grpc_keepalive_time_ms MS Idle time before sending a gRPC/HTTP2 keepalive PING.
+ | --grpc_keepalive_timeout_ms MS Time to wait for a keepalive PING ack before failing.
+ | --grpc_keepalive_without_calls BOOL Whether to keep sending keepalive PINGs when idle.
| --option KEY=VALUE Key-value pair that is used to further configure the session.
""".stripMargin
// scalastyle:on line.size.limit
@@ -92,6 +96,18 @@ private[sql] object SparkConnectClientParser {
case "--grpc_max_message_size" :: tail =>
val (value, remainder) = extract("--grpc_max_message_size", tail)
parse(remainder, builder.grpcMaxMessageSize(value.toInt))
+ case "--grpc_keepalive_enabled" :: tail =>
+ val (value, remainder) = extract("--grpc_keepalive_enabled", tail)
+ parse(remainder, builder.grpcKeepAliveEnabled(value.toBoolean))
+ case "--grpc_keepalive_time_ms" :: tail =>
+ val (value, remainder) = extract("--grpc_keepalive_time_ms", tail)
+ parse(remainder, builder.grpcKeepAliveTimeMs(value.toLong))
+ case "--grpc_keepalive_timeout_ms" :: tail =>
+ val (value, remainder) = extract("--grpc_keepalive_timeout_ms", tail)
+ parse(remainder, builder.grpcKeepAliveTimeoutMs(value.toLong))
+ case "--grpc_keepalive_without_calls" :: tail =>
+ val (value, remainder) = extract("--grpc_keepalive_without_calls", tail)
+ parse(remainder, builder.grpcKeepAliveWithoutCalls(value.toBoolean))
case unsupported :: _ =>
throw new IllegalArgumentException(s"$unsupported is an unsupported argument.")
}
diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/common/config/ConnectCommon.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/common/config/ConnectCommon.scala
index e244fd13595b2..ba2aa09435d42 100644
--- a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/common/config/ConnectCommon.scala
+++ b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/common/config/ConnectCommon.scala
@@ -21,4 +21,13 @@ private[sql] object ConnectCommon {
val CONNECT_GRPC_PORT_MAX_RETRIES: Int = 0
val CONNECT_GRPC_MAX_MESSAGE_SIZE: Int = 128 * 1024 * 1024
val CONNECT_GRPC_MARSHALLER_RECURSION_LIMIT: Int = 1024
+ // Detects a silently-dead connection (e.g. a NAT gateway/load balancer dropping an idle
+ // connection mapping without sending a TCP RST/FIN) via gRPC/HTTP2 keepalive PINGs, so a
+ // blocked RPC surfaces as UNAVAILABLE instead of hanging forever. Guarded by
+ // CONNECT_GRPC_KEEPALIVE_ENABLED so it can be turned off entirely as an escape hatch (e.g. if
+ // it interacts badly with a particular network path, or with long GC pauses).
+ val CONNECT_GRPC_KEEPALIVE_ENABLED: Boolean = true
+ val CONNECT_GRPC_KEEPALIVE_TIME_SECONDS: Long = 60
+ val CONNECT_GRPC_KEEPALIVE_TIMEOUT_SECONDS: Long = 20
+ val CONNECT_GRPC_KEEPALIVE_WITHOUT_CALLS: Boolean = true
}
diff --git a/sql/connect/docs/client-connection-string.md b/sql/connect/docs/client-connection-string.md
index f78a91f45f841..7d246434ddd37 100644
--- a/sql/connect/docs/client-connection-string.md
+++ b/sql/connect/docs/client-connection-string.md
@@ -109,6 +109,42 @@ sc://host:port/;param1=value;param2=value
Default: 128 * 1024 * 1024 |
grpc_max_message_size=134217728 |
+