[SPARK-58094][ Connect] Blocking RPCs hang forever when the connection dies mid-wait

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 down as "too_many_pings".</td>
+  <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>

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,15 +295,17 @@ 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})")
             channel = grpc.intercept_channel(channel, *self._interceptors)
         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

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(

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)