This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new af739ddfede9 [SPARK-58094][CONNECT] Blocking RPCs hang forever when
the connection dies mid-wait
af739ddfede9 is described below
commit af739ddfede95be42e18aa9132f256a0a9de4442
Author: Vinod KC <[email protected]>
AuthorDate: Wed Jul 15 11:37:02 2026 +0800
[SPARK-58094][CONNECT] Blocking RPCs hang forever when the connection dies
mid-wait
### What changes were proposed in this pull request?
This PR configures gRPC/HTTP2 keepalive by default on both sides of Spark
Connect, for both the JVM and Python (PySpark) clients:
- **Client (JVM)** — `SparkConnectClient.scala`'s
`Configuration.createChannel()`: sets
keepAliveTime/keepAliveTimeout/keepAliveWithoutCalls on the gRPC channel
builder. Configurable via new connection-string params grpc_keepalive_enabled /
grpc_keepalive_time_ms / grpc_keepalive_timeout_ms /
grpc_keepalive_without_calls (SparkConnectClientParser.scala) and matching
Builder.grpcKeepAliveEnabled / grpcKeepAliveTimeMs / grpcKeepAliveTimeoutMs /
grpcKeepAliveWithoutCalls methods. Default [...]
- **Client (Python)** — `core.py'`s ChannelBuilder: the same four options
(grpc_keepalive_enabled/_time_ms/_timeout_ms/_without_calls), applied via a new
_effective_channel_options() that both _insecure_channel/_secure_channel funnel
through (so any custom ChannelBuilder subclass gets the fix too, not just
DefaultChannelBuilder), same defaults as the JVM client. An explicitly-set
channelOptions/setChannelOption() value for one of these keys is never
overwritten.
- **Server** — `SparkConnectService.scala'`s startGRPCService(): sets
keepAliveTime/keepAliveTimeout/permitKeepAliveTime/permitKeepAliveWithoutCalls(true)
on the NettyServerBuilder, applied only when enabled. New static confs
spark.connect.grpc.keepAlive.enabled/.time/.timeout (Connect.scala), same
enabled/60s/20s defaults. permitKeepAliveTime is set equal to the client's
default keepAliveTime (60s) — otherwise the server's gRPC-library default
5-minute permitKeepAliveTime would rejec [...]
- Both sides can be fully disabled independently
(`spark.connect.grpc.keepAlive.enabled=false` server-side,
`grpc_keepalive_enabled=false`/`Builder.grpcKeepAliveEnabled(false)`
client-side), as an escape hatch for environments where this default-on
behavior change is undesirable (e.g. a server/client environment prone to long
GC pauses or other stalls that could otherwise trip a false-positive
disconnect).
- Documentation updated: docs/configuration.md (new server confs) and
sql/connect/docs/client-connection-string.md (new client connection-string
params).
This does not change `awaitTermination()`'s (or any other blocking call's)
semantics — blocking indefinitely while a request is genuinely still in flight
on a healthy connection is correct. The gap was purely at the transport layer:
once keepalive detects a connection that has gone silently dark, the channel
fails with `UNAVAILABLE`, and the existing bounded retry policy
(RetryPolicy.scala, maxRetries=15, capped exponential backoff) takes over and
surfaces a real exception, instead of [...]
### Why are the changes needed?
`StreamingQuery.awaitTermination()` (and, more generally, any blocking
Spark Connect RPC) can hang indefinitely and never throw if the network path
between client and server goes silently dark while the call is blocked — e.g. a
NAT gateway, load balancer, or corporate proxy silently evicting an idle
connection mapping without sending a TCP RST/FIN. This is a well-known
operational gotcha for gRPC deployments behind such middleboxes, and Spark
Connect configured no keepalive anywhere ( [...]
Root cause: with no keepalive configured, neither endpoint has any way to
detect that the underlying TCP connection has gone silently dark — both sides
simply continue believing the RPC is still legitimately in flight, forever. The
streaming query can (and did, in the reported incident) terminate independently
on the server in the meantime, but there is no live channel left to deliver
that information back to the client.
Reproduced end-to-end with two genuinely separate JVM processes (a real
Spark Connect server and a standalone client app) and a small TCP proxy in
between, frozen mid-RPC with SIGSTOP to simulate a silently-dropped connection
without either endpoint's socket seeing a close/reset — see "How was this patch
tested?" below. Pre-fix, the client's main thread parks forever in
`ArrayBlockingQueue.take()` (via ExecutePlanResponseReattachableIterator);
post-fix, keepalive detects the dead conn [...]
### Does this PR introduce _any_ user-facing change?
Yes.
- New configuration, all optional and defaulted to preserve the fixed
behavior:
- **Server** (SQLConf): spark.connect.grpc.keepAlive.enabled (default
true), spark.connect.grpc.keepAlive.time (default 60s),
spark.connect.grpc.keepAlive.timeout (default 20s).
- **Client** (connection-string params, JVM Builder methods, and Python
ChannelBuilder connection-string params): grpc_keepalive_enabled (default
true), grpc_keepalive_time_ms (default 60000), grpc_keepalive_timeout_ms
(default 20000), grpc_keepalive_without_calls (default true).
- Behavior change (default-on): both the Spark Connect client and server
now send gRPC/HTTP2 keepalive PINGs periodically. Previously, a Spark Connect
connection that went silently dark mid-RPC would leave the client blocked
forever with no exception; now, such a connection is detected within roughly
keepAliveTime + keepAliveTimeout (60s + 20s by default) and the in-flight call
fails with UNAVAILABLE: Keepalive failed. The connection is likely gone, after
which Spark Connect's existin [...]
- This is a behavior change from all previously-released Spark versions
(keepalive was never configured before); it can be fully reverted via the
.enabled/grpc_keepalive_enabled flags above for environments where it is
undesirable.
### How was this patch tested?
Unit and end-to-end tests :
- `SparkConnectClientBuilderParseTestSuite` (JVM): connection-string/CLI
parsing of the new keepalive params (including an explicit
grpc_keepalive_enabled=true case, not just the default and the disabled case),
and that defaults apply when unset — 20/20.
- `SparkConnectClientSuite` (JVM), two new tests: a FreezableTcpRelay test
helper (an in-process byte-forwarding relay that can be frozen without closing
either socket) sits between a real test client and server; once a blocking RPC
is in flight, the relay is frozen and the test asserts the client's call fails
within a bounded time with UNAVAILABLE: Keepalive failed. The connection is
likely gone instead of hanging past the test's own timeout; a second test
confirms keepalive does not [...]
- `SparkConnectServiceKeepAliveSuite` (new, server side): starts the real
SparkConnectService (not a hand-built test double) with short keepalive confs,
puts a FreezableTcpRelay in front of it, and confirms both that (a) a blocked
call fails with the keepalive UNAVAILABLE when keepalive is enabled, and (b)
disabling spark.connect.grpc.keepAlive.enabled reverts to the pre-fix hang (the
call stays genuinely blocked well past the window that would have triggered
detection) .
- Python `test_connect_channel.py`: defaults, connection-string overrides
Manual end-to-end reproduction, using two genuinely separate JVM/OS
processes (not an in-process test) plus a small Python TCP proxy in between, to
faithfully reproduce the client/server RPC boundary and a silently-dropped
connection:
1. A real Spark Connect server
2. A standalone Scala client app running a rate-source streaming query that
fails via assert_true at a chosen row value (simulating a non-retryable source
error), with a StreamingQueryListener and a watchdog thread, calling
query.awaitTermination().
3. The TCP proxy frozen with SIGSTOP shortly after awaitTermination() was
called (no TCP RST/FIN sent to either endpoint — the same observable behavior
as a NAT gateway/load balancer silently evicting an idle connection mapping).
Confirmed:
- Pre-fix: the server terminates the query independently a few seconds
later (confirming the hang is not an exception-propagation bug), but the
client's awaitTermination() never returns; a thread dump shows the main thread
permanently parked in ArrayBlockingQueue.take() via
ExecutePlanResponseReattachableIterator/GrpcRetryHandler$Retrying.retry, with
no path out.
- Post-fix: the same scenario instead surfaces a bounded exception. With
fast (1s/1s) keepalive settings, a thread dump within ~1-2 minutes shows the
client has alreadymoved off ArrayBlockingQueue.take() into
`GrpcRetryHandler$Retrying.waitAfterAttempt` (keepalive detected the dead
connection); the call eventually throws UNAVAILABLE:Keepalive failed. The
connection is likely gone (observed consistently at ~630-770s end-to-end,
governed by the existing, unmodified retry policy's own ba [...]
- The same manual reproduction was repeated against the Python client
(grpc.insecure_channel/grpc.secure_channel via ChannelBuilder), confirming both
the original gap(Python builds its own gRPC channel and initially had none of
this fix) and the fix once applied: awaitTermination() threw
SparkConnectGrpcException / UNAVAILABLE:Stream removed (ping timeout) —
grpcio's equivalent wording for the same keepalive-triggered failure.
A minimal, runnable, standalone reproduction (client + server + proxy
scripts, independent of this Spark checkout's test suite) is attached to
SPARK-58094 for reviewers who want to reproduce the issue.
### Was this patch authored or co-authored using generative AI tooling?
Yes. Generated-by: Claude Code (Sonnet 5).
Closes #57202 from vinodkc/connect-streaming-hang-repro.
Authored-by: Vinod KC <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
---
docs/configuration.md | 24 ++
python/pyspark/sql/connect/client/core.py | 96 ++++-
.../sql/tests/connect/test_connect_channel.py | 73 ++++
.../SparkConnectClientBuilderParseTestSuite.scala | 45 +++
.../connect/client/SparkConnectClientSuite.scala | 175 +++++++++
.../sql/connect/client/SparkConnectClient.scala | 69 +++-
.../connect/client/SparkConnectClientParser.scala | 16 +
.../sql/connect/common/config/ConnectCommon.scala | 9 +
sql/connect/docs/client-connection-string.md | 47 +++
.../apache/spark/sql/connect/config/Connect.scala | 43 +++
.../sql/connect/service/SparkConnectService.scala | 30 +-
.../spark/sql/connect/SparkConnectServerTest.scala | 10 +-
.../SparkConnectServiceKeepAliveSuite.scala | 412 +++++++++++++++++++++
13 files changed, 1043 insertions(+), 6 deletions(-)
diff --git a/docs/configuration.md b/docs/configuration.md
index 3eeee90d0778..4999b08908de 100644
--- a/docs/configuration.md
+++ b/docs/configuration.md
@@ -3561,6 +3561,30 @@ They are typically set via the config file and
command-line options with `--conf
<td>Sets the maximum inbound message size for the gRPC requests. Requests
with a larger payload will fail.</td>
<td>3.4.0</td>
</tr>
+<tr>
+ <td><code>spark.connect.grpc.keepAlive.enabled</code></td>
+ <td>
+ true
+ </td>
+ <td>Whether the server sends gRPC/HTTP2 keepalive PINGs to detect and
terminate silently-dead client connections. Can be turned off as an escape
hatch, e.g. if it interacts badly with a particular network path, or a server
environment is prone to stalls (long GC pauses, etc.) long enough to trip
false-positive disconnects.</td>
+ <td>4.3.0</td>
+</tr>
+<tr>
+ <td><code>spark.connect.grpc.keepAlive.time</code></td>
+ <td>
+ 60s
+ </td>
+ <td>Sets the time the server waits for the connection to be idle before
sending a gRPC/HTTP2 keepalive PING, to detect and terminate a silently-dead
connection (e.g. after a NAT gateway or load balancer drops an idle connection
mapping without closing the socket). The server separately tolerates
client-sent keepalive PINGs no more often than every 10s regardless of this
setting; a client configured with <code>grpc_keepalive_time_ms</code> below
that floor will have its connection torn [...]
+ <td>4.3.0</td>
+</tr>
+<tr>
+ <td><code>spark.connect.grpc.keepAlive.timeout</code></td>
+ <td>
+ 20s
+ </td>
+ <td>Sets how long the server waits for a keepalive PING ack before
considering the connection dead.</td>
+ <td>4.3.0</td>
+</tr>
<tr>
<td><code>spark.connect.extensions.relation.classes</code></td>
<td>
diff --git a/python/pyspark/sql/connect/client/core.py
b/python/pyspark/sql/connect/client/core.py
index db6067f25e09..1e81e780d167 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 @@ class ChannelBuilder:
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 @@ class ChannelBuilder:
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 @@ class ChannelBuilder:
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 @@ class ChannelBuilder:
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 21390d671dda..66a8f0d8abc8 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 @@ class ChannelBuilderTests(unittest.TestCase):
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 ad0f880e8f57..b9b3754252a3 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 d270b0544097..fd5be9bb555e 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 7bd7d0e4e979..8f449207b796 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 7e137a6a3e05..629b6d237e5c 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 e244fd13595b..ba2aa09435d4 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 f78a91f45f84..7d246434ddd3 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
<i>Default: </i><pre> 128 * 1024 * 1024</pre></td>
<td><pre>grpc_max_message_size=134217728</pre></td>
</tr>
+ <tr>
+ <td>grpc_keepalive_enabled</td>
+ <td>Boolean</td>
+ <td>Whether the client sends gRPC/HTTP2 keepalive PINGs to detect a
silently-dead
+ connection (e.g. a NAT gateway or load balancer dropping an idle
connection mapping
+ without closing the socket), so a blocked call fails with an error instead
of hanging
+ forever. Can be turned off as an escape hatch, e.g. in environments prone
to long stalls
+ (GC pauses, etc.) that could otherwise trip a false-positive
disconnect.<br/>
+ <i>Default: </i><pre>true</pre></td>
+ <td><pre>grpc_keepalive_enabled=false</pre></td>
+ </tr>
+ <tr>
+ <td>grpc_keepalive_time_ms</td>
+ <td>Numeric</td>
+ <td>Idle time (in milliseconds) before the client sends a keepalive PING.
A Spark Connect
+ server tolerates client PINGs no more often than every 10s; setting this
below that floor
+ will get the connection torn down as "too_many_pings".<br/>
+ <i>Default: </i><pre>60000</pre></td>
+ <td><pre>grpc_keepalive_time_ms=30000</pre></td>
+ </tr>
+ <tr>
+ <td>grpc_keepalive_timeout_ms</td>
+ <td>Numeric</td>
+ <td>Time (in milliseconds) the client waits for a keepalive PING ack
before considering
+ the connection dead.<br/>
+ <i>Default: </i><pre>20000</pre></td>
+ <td><pre>grpc_keepalive_timeout_ms=10000</pre></td>
+ </tr>
+ <tr>
+ <td>grpc_keepalive_without_calls</td>
+ <td>Boolean</td>
+ <td>Whether to keep sending keepalive PINGs when there are no in-flight
RPCs on the
+ connection.<br/>
+ <i>Default: </i><pre>true</pre></td>
+ <td><pre>grpc_keepalive_without_calls=false</pre></td>
+ </tr>
</table>
## Examples
@@ -132,6 +168,17 @@ server_url = "sc://myhost.com:443/;use_ssl=true"
server_url = "sc://myhost.com:443/;use_ssl=true;token=ABCDEFG"
```
+The next example tunes the gRPC keepalive settings, e.g. to detect a dead
connection faster
+than the 60s/20s default, or to disable it entirely.
+
+```python
+server_url =
"sc://myhost.com:443/;grpc_keepalive_time_ms=30000;grpc_keepalive_timeout_ms=10000"
+```
+
+```python
+server_url = "sc://myhost.com:443/;grpc_keepalive_enabled=false"
+```
+
### Invalid Examples
As mentioned above, Spark Connect uses a regular GRPC client and the server
path
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 e2d496239d29..f5a71a667465 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
@@ -87,6 +87,49 @@ object Connect {
.intConf
.createWithDefault(ConnectCommon.CONNECT_GRPC_MARSHALLER_RECURSION_LIMIT)
+ val CONNECT_GRPC_KEEPALIVE_ENABLED =
+ buildStaticConf("spark.connect.grpc.keepAlive.enabled")
+ .doc(
+ "Whether the server sends gRPC/HTTP2 keepalive PINGs to detect and
terminate " +
+ "silently-dead client connections (see
spark.connect.grpc.keepAlive.time / " +
+ ".timeout). Enabled by default; can be turned off as an escape
hatch, e.g. if it " +
+ "interacts badly with a particular network path, or a server
environment is prone " +
+ "to stalls (long GC pauses, etc.) long enough to trip false-positive
disconnects. " +
+ "This only controls the server's own dead-client detection; the
server's tolerance " +
+ "of client-initiated keepalive PINGs is " +
+ "independent and always in effect, so that a default-on client is
never rejected " +
+ "with too_many_pings regardless of this setting.")
+ .version("4.3.0")
+ .booleanConf
+ .createWithDefault(ConnectCommon.CONNECT_GRPC_KEEPALIVE_ENABLED)
+
+ val CONNECT_GRPC_KEEPALIVE_TIME =
+ buildStaticConf("spark.connect.grpc.keepAlive.time")
+ .doc(
+ "Sets the time the server waits for the connection to be idle before
sending a " +
+ "gRPC/HTTP2 keepalive PING, to detect and terminate a silently-dead
connection " +
+ "(e.g. after a NAT gateway or load balancer drops an idle connection
mapping " +
+ "without closing the socket). The server separately tolerates
client-sent keepalive " +
+ "PINGs no more often than every 10s regardless of this setting.
Combined with " +
+ "spark.connect.grpc.keepAlive.timeout, this bounds how long a
connection can be idle " +
+ "before either side may declare it dead; a stall on either end (e.g.
a long JVM GC " +
+ "pause) longer than that combined window can cause an
otherwise-healthy connection " +
+ "to be dropped, so raise both values in environments prone to long
GC pauses or " +
+ "other event-loop stalls.")
+ .version("4.3.0")
+ .timeConf(TimeUnit.SECONDS)
+ .createWithDefault(ConnectCommon.CONNECT_GRPC_KEEPALIVE_TIME_SECONDS)
+
+ val CONNECT_GRPC_KEEPALIVE_TIMEOUT =
+ buildStaticConf("spark.connect.grpc.keepAlive.timeout")
+ .doc(
+ "Sets how long the server waits for a keepalive PING ack before
considering the " +
+ "connection dead. See spark.connect.grpc.keepAlive.time for the
combined-window " +
+ "caveat about long GC pauses or other stalls.")
+ .version("4.3.0")
+ .timeConf(TimeUnit.SECONDS)
+ .createWithDefault(ConnectCommon.CONNECT_GRPC_KEEPALIVE_TIMEOUT_SECONDS)
+
val CONNECT_SESSION_MANAGER_DEFAULT_SESSION_TIMEOUT =
buildStaticConf("spark.connect.session.manager.defaultSessionTimeout")
.internal()
diff --git
a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala
b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala
index f2dd9c1b1cef..c76794e3b6ec 100644
---
a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala
+++
b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala
@@ -41,7 +41,7 @@ import org.apache.spark.scheduler.{LiveListenerBus,
SparkListenerEvent}
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.classic.ClassicConversions._
import org.apache.spark.sql.connect.IllegalStateErrors
-import org.apache.spark.sql.connect.config.Connect.{getAuthenticateToken,
CONNECT_GRPC_BINDING_ADDRESS, CONNECT_GRPC_BINDING_PORT,
CONNECT_GRPC_MARSHALLER_RECURSION_LIMIT, CONNECT_GRPC_MAX_INBOUND_MESSAGE_SIZE,
CONNECT_GRPC_PORT_MAX_RETRIES}
+import org.apache.spark.sql.connect.config.Connect.{getAuthenticateToken,
CONNECT_GRPC_BINDING_ADDRESS, CONNECT_GRPC_BINDING_PORT,
CONNECT_GRPC_KEEPALIVE_ENABLED, CONNECT_GRPC_KEEPALIVE_TIME,
CONNECT_GRPC_KEEPALIVE_TIMEOUT, CONNECT_GRPC_MARSHALLER_RECURSION_LIMIT,
CONNECT_GRPC_MAX_INBOUND_MESSAGE_SIZE, CONNECT_GRPC_PORT_MAX_RETRIES}
import org.apache.spark.sql.connect.execution.ConnectProgressExecutionListener
import org.apache.spark.sql.connect.ui.{SparkConnectServerAppStatusStore,
SparkConnectServerListener, SparkConnectServerTab}
import org.apache.spark.sql.connect.utils.ErrorUtils
@@ -314,6 +314,14 @@ class SparkConnectService(debug: Boolean) extends
AsyncService with BindableServ
*/
object SparkConnectService extends Logging {
+ // Floor for permitKeepAliveTime: the minimum interval the server tolerates
between a client's
+ // keepalive PINGs before striking the connection as "too_many_pings" (gRFC
A8). This is
+ // independent of the server's own keepAlive.time (its probe interval for
detecting dead
+ // clients) and must stay well below any supported client
grpc_keepalive_time_ms, not derived
+ // from it -- otherwise a faster client (or an operator raising
keepAlive.time) gets its
+ // healthy connection torn down. 10s leaves a wide margin below the 60s
client/server default.
+ private[connect] val GRPC_KEEPALIVE_PERMIT_TIME_SECONDS: Long = 10
+
private[connect] var server: Server = _
private[connect] var bindingAddress: InetSocketAddress = _
@@ -408,7 +416,25 @@ object SparkConnectService extends Logging {
case _ => NettyServerBuilder.forPort(port)
}
sb.maxInboundMessageSize(SparkEnv.get.conf.get(CONNECT_GRPC_MAX_INBOUND_MESSAGE_SIZE).toInt)
- .addService(sparkConnectService)
+ if (SparkEnv.get.conf.get(CONNECT_GRPC_KEEPALIVE_ENABLED)) {
+ val keepAliveTimeSeconds =
SparkEnv.get.conf.get(CONNECT_GRPC_KEEPALIVE_TIME)
+ // Detects a silently-dead client connection (e.g. a NAT gateway/load
balancer dropping
+ // an idle connection mapping without sending a TCP RST/FIN) via
gRPC/HTTP2 keepalive
+ // PINGs.
+ sb.keepAliveTime(keepAliveTimeSeconds, TimeUnit.SECONDS)
+ sb.keepAliveTimeout(
+ SparkEnv.get.conf.get(CONNECT_GRPC_KEEPALIVE_TIMEOUT),
+ TimeUnit.SECONDS)
+ }
+ // Independent of the server's own dead-client detection above (gRFC A8:
these are
+ // separate axes), this tolerates client-initiated keepalive PINGs so a
default-on client
+ // (which pings regardless of this server's
CONNECT_GRPC_KEEPALIVE_ENABLED setting) is
+ // never hit with a "too_many_pings" GOAWAY on an otherwise-healthy
connection. Left
+ // ungated so disabling only this server's own keepalive detection
cannot regress a
+ // healthy connection to the client into disruptive reconnects.
+ sb.permitKeepAliveTime(GRPC_KEEPALIVE_PERMIT_TIME_SECONDS,
TimeUnit.SECONDS)
+ sb.permitKeepAliveWithoutCalls(true)
+ sb.addService(sparkConnectService)
getAuthenticateToken.foreach { token =>
sb.intercept(new PreSharedKeyAuthenticationInterceptor(token))
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 d4ecbca28e9d..c935370f0643 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
@@ -54,13 +54,21 @@ trait SparkConnectServerTest extends SharedSparkSession {
val allocator = new RootAllocator()
+ // Additional server-side confs to set (via withSparkEnvConfs) before
starting the real
+ // service. Override in subclasses that need to exercise non-default Connect
server behavior
+ // (e.g. tuning gRPC keepalive) through the actual production
startGRPCService() wiring,
+ // rather than only through hand-built test doubles.
+ protected def extraServerConfs: Seq[(String, String)] = Seq.empty
+
override def beforeAll(): Unit = {
super.beforeAll()
// 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((Connect.CONNECT_GRPC_BINDING_PORT.key,
serverPort.toString)) {
+ withSparkEnvConfs(
+ (Seq(
+ (Connect.CONNECT_GRPC_BINDING_PORT.key, serverPort.toString)) ++
extraServerConfs): _*) {
SparkConnectService.start(spark.sparkContext)
}
}
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
new file mode 100644
index 000000000000..d1d273016e76
--- /dev/null
+++
b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceKeepAliveSuite.scala
@@ -0,0 +1,412 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.connect.service
+
+import java.util.UUID
+import java.util.concurrent.TimeUnit
+
+import scala.collection.mutable
+import scala.concurrent.duration.FiniteDuration
+
+import org.scalatest.concurrent.Eventually
+
+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
+
+/**
+ * End-to-end test that the *real* `SparkConnectService.startGRPCService()`
wiring actually
+ * applies the configured gRPC keepalive settings (SPARK-58094).
+ *
+ * Unlike the tests in `SparkConnectClientSuite` (sql/connect/client/jvm),
which hand-build their
+ * own `NettyServerBuilder` that only mirrors the production server
construction, this suite
+ * starts the real `SparkConnectService` (via `SparkConnectServerTest`) with
the keepalive confs
+ * tuned to short test values, puts a freezable TCP relay in front of it, and
confirms a blocked
+ * client call fails with a keepalive-triggered `UNAVAILABLE` within a bounded
time -- so a future
+ * refactor of `startGRPCService()` that silently drops the keepalive wiring
would be caught here,
+ * not just in the client-side unit tests.
+ */
+class SparkConnectServiceKeepAliveSuite extends SparkConnectServerTest {
+
+ // CONNECT_GRPC_KEEPALIVE_TIME/TIMEOUT are read via
timeConf(TimeUnit.SECONDS), so anything
+ // sub-second (e.g. "300ms") truncates to 0 whole seconds and gRPC rejects
it -- use the
+ // smallest whole-second value instead.
+ private val keepAliveMs = 1000L
+
+ override protected def extraServerConfs: Seq[(String, String)] = 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: real SparkConnectService applies configured keepalive
end-to-end") {
+ val serverSession =
+ SparkConnectService
+ .getOrCreateIsolatedSession(defaultUserId, defaultSessionId, None)
+ .session
+ serverSession.udf.register(
+ "sleep",
+ (ms: Int) => {
+ Thread.sleep(ms.toLong)
+ ms
+ })
+
+ val relay = new FreezableTcpRelay(serverPort)
+ try {
+ val client = SparkConnectClient
+ .builder()
+ .port(relay.port)
+ .sessionId(defaultSessionId)
+ .userId(defaultUserId)
+ .enableReattachableExecute()
+ .grpcKeepAliveEnabled(true)
+ .grpcKeepAliveTimeMs(keepAliveMs)
+ .grpcKeepAliveTimeoutMs(keepAliveMs)
+ .retryPolicy(RetryPolicy(maxRetries = Some(0), canRetry = _ => false,
name = "NoRetry"))
+ .build()
+ try {
+ val operationId = UUID.randomUUID().toString
+ val iter = client.execute(buildPlan("select sleep(30000) as s"),
Some(operationId))
+
+ // Wait for the query to genuinely start server-side before freezing,
mirroring the
+ // manual repro's timing lesson: freezing too early (before the
connection is actually
+ // established/in-flight) doesn't exercise the "dies while blocked"
scenario.
+ Eventually.eventually(timeout(eventuallyTimeout)) {
+
assert(SparkConnectService.executionManager.listExecuteHolders.length == 1)
+ }
+ relay.freeze()
+
+ // scalastyle:off awaitresult
+ val ex = intercept[SparkException] {
+ scala.concurrent.Await.result(
+ scala.concurrent.Future {
+ while (iter.hasNext) iter.next()
+ }(scala.concurrent.ExecutionContext.global),
+ 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 {
+ client.shutdown()
+ }
+ } finally {
+ relay.close()
+ }
+ }
+
+ test(
+ "SPARK-58094: keepalive does not disrupt a healthy long-lived connection
through the " +
+ "real server") {
+ // Exercises the scenario flagged in review: a client whose ping cadence
(11s) is faster
+ // than the server's own keepAlive.time (20s), through the real
startGRPCService() wiring.
+ // With permitKeepAliveTime derived from the server's own keepAlive.time
(the pre-review
+ // 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()
+ 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)
+ }
+ try {
+ val client = SparkConnectClient
+ .builder()
+ .port(serverPort)
+ .sessionId(defaultSessionId)
+ .userId(defaultUserId)
+ .grpcKeepAliveEnabled(true)
+ .grpcKeepAliveTimeMs(clientKeepAliveMs)
+ .grpcKeepAliveTimeoutMs(clientKeepAliveMs)
+ .retryPolicy(RetryPolicy(maxRetries = Some(0), canRetry = _ => false,
name = "NoRetry"))
+ .build()
+ try {
+ // Several real calls spaced out beyond the client's keepalive
interval, so the client's
+ // own periodic PINGs land against the real server wiring: if
permitKeepAliveTime were
+ // ever re-derived from the server's own keepAlive.time (reintroducing
the coupling
+ // flagged in review) or its wiring silently dropped, this is the
exact shape of client
+ // that would be affected.
+ for (i <- 1 to 2) {
+ val operationId = UUID.randomUUID().toString
+ val iter = client.execute(buildPlan("select 1 as s"),
Some(operationId))
+ while (iter.hasNext) iter.next()
+ Thread.sleep(clientKeepAliveMs + 2000)
+ }
+ } finally {
+ client.shutdown()
+ }
+ } finally {
+ SparkConnectService.stop()
+ 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)
+ }
+ }
+ }
+
+ test(
+ "SPARK-58094: server tolerates a default-on client's PINGs even with its
own keepalive " +
+ "detection disabled") {
+ // Exercises the scenario flagged in review:
permitKeepAliveTime/permitKeepAliveWithoutCalls
+ // used to be set only inside `if (CONNECT_GRPC_KEEPALIVE_ENABLED)`, so
disabling only the
+ // server's own dead-client detection reverted permit-tolerance to
grpc-netty's defaults
+ // (permitKeepAliveTime=5min, permitKeepAliveWithoutCalls=false). With no
outstanding calls,
+ // grpc's KeepAliveEnforcer then requires 2 hours (its
IMPLICIT_PERMIT_TIME_NANOS grace
+ // period for permitWithoutCalls=false) between "acceptable" pings instead
of the configured
+ // permit floor, so a still-default-on client's periodic without-calls
PINGs would each
+ // strike, and the 3rd strike
(io.grpc.internal.KeepAliveEnforcer.MAX_PING_STRIKES=2)
+ // forcefully closes the connection (ForcefulCloseCommand) as
too_many_pings -- even though
+ // nothing was ever actually wrong with the connection. With the permit
calls hoisted out of
+ // the `if`, the enforcer instead uses the fixed 10s permit floor
regardless of
+ // CONNECT_GRPC_KEEPALIVE_ENABLED, comfortably tolerating an 11s client
ping interval, so no
+ // strikes accrue and the connection is never closed.
+ //
+ // A test loop that issues real RPC calls between pings cannot observe
this: every server
+ // write (response headers/data) resets the enforcer's strike counter
+ // (NettyServerHandler$WriteMonitoringFrameWriter), so pings are only ever
checked against a
+ // 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()
+ 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)
+ }
+ val relay = new FreezableTcpRelay(serverPort)
+ try {
+ val client = SparkConnectClient
+ .builder()
+ .port(relay.port)
+ .sessionId(defaultSessionId)
+ .userId(defaultUserId)
+ .grpcKeepAliveEnabled(true)
+ .grpcKeepAliveTimeMs(clientKeepAliveMs)
+ .grpcKeepAliveTimeoutMs(clientKeepAliveMs)
+ .retryPolicy(RetryPolicy(maxRetries = Some(0), canRetry = _ => false,
name = "NoRetry"))
+ .build()
+ try {
+ def runOneCall(): Unit = {
+ val operationId = UUID.randomUUID().toString
+ val iter = client.execute(buildPlan("select 1 as s"),
Some(operationId))
+ while (iter.hasNext) iter.next()
+ }
+
+ // Establishes the transport; this write resets the enforcer's strike
counter, so the
+ // idle window below starts clean.
+ runOneCall()
+ val connectionsAfterFirstCall = relay.connectionCount
+
+ // Stay genuinely idle (no calls) for several client ping intervals,
so at least 3
+ // without-calls PINGs land on the server with nothing resetting the
strike counter in
+ // between -- enough to trip too_many_pings if the permit calls were
re-gated.
+ Thread.sleep(clientKeepAliveMs * 4)
+
+ runOneCall()
+ assert(
+ relay.connectionCount == connectionsAfterFirstCall,
+ "server forcefully closed the connection under a default-on client's
idle keepalive " +
+ "PINGs (too_many_pings) even though its own keepAlive.enabled was
false -- the " +
+ "client had to open a new TCP connection to complete the second
call")
+ } finally {
+ client.shutdown()
+ }
+ } finally {
+ relay.close()
+ SparkConnectService.stop()
+ 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)
+ }
+ }
+ }
+
+ 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()
+ 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)
+ }
+ try {
+ val serverSession =
+ SparkConnectService
+ .getOrCreateIsolatedSession(defaultUserId, defaultSessionId, None)
+ .session
+ serverSession.udf.register(
+ "sleep",
+ (ms: Int) => {
+ Thread.sleep(ms.toLong)
+ ms
+ })
+
+ val relay = new FreezableTcpRelay(serverPort)
+ try {
+ val client = SparkConnectClient
+ .builder()
+ .port(relay.port)
+ .sessionId(defaultSessionId)
+ .userId(defaultUserId)
+ .enableReattachableExecute()
+ .grpcKeepAliveEnabled(false)
+ .grpcKeepAliveTimeMs(keepAliveMs)
+ .grpcKeepAliveTimeoutMs(keepAliveMs)
+ .retryPolicy(RetryPolicy(maxRetries = Some(0), canRetry = _ =>
false, name = "NoRetry"))
+ .build()
+ try {
+ val operationId = UUID.randomUUID().toString
+ val iter = client.execute(buildPlan("select sleep(30000) as s"),
Some(operationId))
+
+ Eventually.eventually(timeout(eventuallyTimeout)) {
+
assert(SparkConnectService.executionManager.listExecuteHolders.length == 1)
+ }
+ relay.freeze()
+
+ // With keepalive disabled on both sides, the call must NOT resolve
within a window
+ // well past the 1s/1s that would have triggered detection if it
were enabled --
+ // confirming the flag genuinely disables the fix, matching the
original bug's
+ // "blocks forever, no exception" symptom, rather than the call just
being slow.
+ // scalastyle:off awaitresult
+ val stillBlocked =
+ try {
+ scala.concurrent.Await.result(
+ scala.concurrent.Future {
+ while (iter.hasNext) iter.next()
+ }(scala.concurrent.ExecutionContext.global),
+ FiniteDuration(5, TimeUnit.SECONDS))
+ false
+ } catch {
+ case _: java.util.concurrent.TimeoutException => true
+ }
+ // scalastyle:on awaitresult
+ assert(stillBlocked, "expected the call to remain hung with
keepalive disabled")
+ } finally {
+ client.shutdown()
+ }
+ } finally {
+ relay.close()
+ }
+ } finally {
+ // Restore the enabled server for afterAll()/subsequent tests in this
suite.
+ SparkConnectService.stop()
+ 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)
+ }
+ }
+ }
+}
+
+/**
+ * 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). Mirrors
the equivalent
+ * helper in `SparkConnectClientSuite` (sql/connect/client/jvm); duplicated
here rather than
+ * 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)
+ @volatile private var frozen = false
+ private val sockets = mutable.ListBuffer[java.net.Socket]()
+ private val acceptedConnections = new
java.util.concurrent.atomic.AtomicInteger(0)
+
+ private val acceptThread = new Thread(() => {
+ try {
+ while (true) {
+ val clientSocket = listener.accept()
+ acceptedConnections.incrementAndGet()
+ 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
+
+ // Number of distinct client TCP connections accepted so far -- useful to
detect a
+ // client-transparent reconnect (e.g. after the server forcefully closes the
connection)
+ // that wouldn't otherwise surface as a visible failure to the application.
+ def connectionCount: Int = acceptedConnections.get()
+
+ // 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 =>
+ }
+ })
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
