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 b9a3de5bc231 [SPARK-58162][CONNECT] Bound the unbounded RetryException
retry loop in reattach path
b9a3de5bc231 is described below
commit b9a3de5bc23190d21517e7111a7bfaa7e09d7bdd
Author: Vinod KC <[email protected]>
AuthorDate: Mon Jul 20 22:36:10 2026 +0800
[SPARK-58162][CONNECT] Bound the unbounded RetryException retry loop in
reattach path
### What changes were proposed in this pull request?
When a Spark Connect client runs a query, it stays connected to the server
through a "reattach" mechanism that can reconnect if the stream drops. If
certain errors happen during that reconnect, the client throws an internal
RetryException and tries again.
The problem: this particular retry has no limit — no max attempts, no
waiting between tries, no time cap. If the error keeps happening, the client
retries forever and the query just hangs.
This can happen in two situations:
1. The server says it can't find the operation or session (for example, a
load balancer sent the request to the wrong server that never got it).
2. The reconnect keeps timing out (the connection is dead but nobody
noticed).
This PR puts a time limit on that retry loop:
- It adds a setting, `maxRetryExceptionElapsedTime`, that defaults to 1
hour. The client keeps retrying as before, but once it has spent a total of 1
hour retrying this specific error, it gives up.
- When it gives up, it reports the real underlying error (like a timeout)
instead of a confusing internal marker, so the user knows what actually went
wrong.
- The same fix is applied to both the Scala client and the Python (PySpark)
client so they behave identically.
- The 1-hour limit can be changed by advanced users through the client
builder (Scala) or a constructor argument (Python). There's no new config file
setting or command-line flag.
- The 1-hour limit only counts time spent stuck in the RetryException retry
loop — not the total runtime of your query.
- The clock starts on the first RetryException and only accumulates while
the client is repeatedly failing to reattach (connection dead, operation not
found, reconnect timing out).
- As long as the server is making progress and sending results back, no
RetryException is being thrown, so the clock never starts. A query can run for
6 hours, 12 hours, or longer with zero impact.
### Why are the changes needed?
Today, the only reason this loop ever stops is that a separate feature —
gRPC keepalive, which detects dead connections — happens to be on and working.
That is not something the retry logic should depend on:
- The "operation/session not found" case doesn't benefit from keepalive at
all; it can spin forever regardless.
- If keepalive is turned off or misconfigured on the client, the loop can
hang even against a perfectly healthy server.
Under the default setup the hang is usually avoided, so this is best seen
as closing a hidden gap and hardening the client — the retry loop now stops on
its own instead of depending on another feature being configured correctly
everywhere.
### Does this PR introduce _any_ user-facing change?
Yes. Behavior is unchanged until the 1-hour default limit is reached — at
which point a previously infinite hang now surfaces the real underlying error
instead. Advanced users can adjust the limit via the new Builder method /
constructor argument.
Scope of the **limit**: the timer is a field on GrpcRetryHandler.Retrying,
so it spans exactly one retry() invocation and is discarded when that call
returns — it does not accumulate across the query's or session's lifetime.
Because the existing reattach iterator issues a fresh retry { ... } per
next()/hasNext(), every call that makes successful progress resets the budget
to zero. It only advances on consecutive DEADLINE_EXCEEDED-driven reattaches
within a single stuck call; idle str [...]
### How was this patch tested?
Added tests in `SparkConnectClientRetriesSuite`, `test_client_retries.py`
covering: the limit is enforced, below-limit retries still work
(non-regression), the fallback when no underlying cause is attached, the timer
starting only at the first `RetryException`, and the default value.
One `SparkConnectClientSuite` test exercises the whole path end-to-end
against a test server that never responds.
Each test was checked to hang without the fix reproducing the original bug
and confirming they are real regression guards.
### Was this patch authored or co-authored using generative AI tooling?
Yes, Generated-by: Claude Code (sonnet-5)
Closes #57320 from vinodkc/connect-retry-exception-bound.
Authored-by: Vinod KC <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
---
python/pyspark/sql/connect/client/core.py | 26 ++++-
python/pyspark/sql/connect/client/reattach.py | 2 +-
python/pyspark/sql/connect/client/retries.py | 34 +++++-
.../tests/connect/client/test_client_retries.py | 112 ++++++++++++++++++
.../client/SparkConnectClientRetriesSuite.scala | 128 +++++++++++++++++++++
.../connect/client/SparkConnectClientSuite.scala | 42 +++++++
.../sql/connect/client/GrpcRetryHandler.scala | 57 ++++++++-
.../spark/sql/connect/client/RetryPolicy.scala | 4 +-
.../sql/connect/client/SparkConnectClient.scala | 17 +++
.../sql/connect/client/SparkConnectStubState.scala | 4 +-
10 files changed, 413 insertions(+), 13 deletions(-)
diff --git a/python/pyspark/sql/connect/client/core.py
b/python/pyspark/sql/connect/client/core.py
index 1e81e780d167..f0f1a23789dc 100644
--- a/python/pyspark/sql/connect/client/core.py
+++ b/python/pyspark/sql/connect/client/core.py
@@ -75,7 +75,12 @@ from pyspark.sql.connect.client.artifact import
ArtifactManager
from pyspark.sql.connect.logging import logger
from pyspark.sql.connect.profiler import ConnectProfilerCollector
from pyspark.sql.connect.client.reattach import
ExecutePlanResponseReattachableIterator
-from pyspark.sql.connect.client.retries import RetryPolicy, Retrying,
DefaultPolicy
+from pyspark.sql.connect.client.retries import (
+ RetryPolicy,
+ Retrying,
+ DefaultPolicy,
+ DEFAULT_MAX_RETRY_EXCEPTION_ELAPSED_TIME,
+)
from pyspark.sql.connect.conversion import (
storage_level_to_proto,
proto_to_storage_level,
@@ -814,6 +819,7 @@ class SparkConnectClient(object):
allow_arrow_batch_chunking: bool = True,
preferred_arrow_chunk_size: Optional[int] = None,
rpc_deadlines: Optional[RpcDeadlines] = None,
+ max_retry_exception_elapsed_time: Optional[float] = None,
):
"""
Creates a new SparkSession for the Spark Connect interface.
@@ -865,6 +871,13 @@ class SparkConnectClient(object):
Per-RPC gRPC call timeouts in seconds (10 min for most RPCs,
1 hour for analyze/addArtifacts, none for non-reattachable
execute).
Use :meth:`RpcDeadlines.disabled` to turn off all deadlines.
+ max_retry_exception_elapsed_time : float, optional
+ Maximum cumulative elapsed time in seconds the client will keep
retrying a
+ RetryException (raised internally when a reattach attempt keeps
hitting
+ DEADLINE_EXCEEDED, or when the initial ExecutePlan never reached
the server) before
+ giving up and raising the underlying error. Defaults to
+
:data:`~pyspark.sql.connect.client.retries.DEFAULT_MAX_RETRY_EXCEPTION_ELAPSED_TIME`
+ (1 hour).
"""
self.thread_local = threading.local()
@@ -881,6 +894,12 @@ class SparkConnectClient(object):
default_policy = DefaultPolicy(**retry_policy_args)
self.set_retry_policies([default_policy])
+ self._max_retry_exception_elapsed_time = (
+ max_retry_exception_elapsed_time
+ if max_retry_exception_elapsed_time is not None
+ else DEFAULT_MAX_RETRY_EXCEPTION_ELAPSED_TIME
+ )
+
if self._builder.session_id is None:
# Generate a unique session ID for this client. This UUID must be
unique to allow
# concurrent Spark sessions of the same user. If the channel is
closed, creating
@@ -984,7 +1003,10 @@ class SparkConnectClient(object):
self._progress_handlers.remove(handler)
def _retrying(self) -> "Retrying":
- return Retrying(self._retry_policies)
+ return Retrying(
+ self._retry_policies,
+
max_retry_exception_elapsed_time=self._max_retry_exception_elapsed_time,
+ )
def disable_reattachable_execute(self) -> "SparkConnectClient":
self._use_reattachable_execute = False
diff --git a/python/pyspark/sql/connect/client/reattach.py
b/python/pyspark/sql/connect/client/reattach.py
index 6babe7a8e428..f1d06320866e 100644
--- a/python/pyspark/sql/connect/client/reattach.py
+++ b/python/pyspark/sql/connect/client/reattach.py
@@ -299,7 +299,7 @@ class ExecutePlanResponseReattachableIterator(Generator):
timeout=self._reattachable_execute_plan_timeout,
)
)
- raise RetryException()
+ raise RetryException() from e
elif e.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
# The per-RPC deadline fired. The server-side operation is
still alive; clear
# the iterator and raise RetryException so the retry loop
opens a fresh
diff --git a/python/pyspark/sql/connect/client/retries.py
b/python/pyspark/sql/connect/client/retries.py
index 86ed004c40fe..f16968f6391f 100644
--- a/python/pyspark/sql/connect/client/retries.py
+++ b/python/pyspark/sql/connect/client/retries.py
@@ -198,6 +198,12 @@ class AttemptManager:
return None
+# Default cumulative elapsed-time bound (in seconds) for RetryException-driven
retries. Please
+# synchronize changes here with the Scala side's
GrpcRetryHandler.maxRetryExceptionElapsedTime
+# default in org.apache.spark.sql.connect.client.GrpcRetryHandler.
+DEFAULT_MAX_RETRY_EXCEPTION_ELAPSED_TIME = 60 * 60 # 1 hour
+
+
class Retrying:
"""
This class is a point of entry into the retry logic.
@@ -214,20 +220,32 @@ class Retrying:
it will be raised if the retries limit would exceed.
Exceptions not considered retriable will be passed through transparently.
+
+ RetryException is a special case: it bypasses policies entirely and is
always retried
+ immediately, but only up to a cumulative elapsed-time bound
(max_retry_exception_elapsed_time,
+ measured from the first RetryException seen by this instance) rather than
being retried
+ unconditionally forever.
"""
def __init__(
self,
policies: typing.Union[RetryPolicy, typing.Iterable[RetryPolicy]],
sleep: Callable[[float], None] = time.sleep,
+ max_retry_exception_elapsed_time: float =
DEFAULT_MAX_RETRY_EXCEPTION_ELAPSED_TIME,
+ now: Callable[[], float] = time.monotonic,
) -> None:
if isinstance(policies, RetryPolicy):
policies = [policies]
self._policies: List[RetryPolicyState] = [policy.to_state() for policy
in policies]
self._sleep = sleep
+ self._max_retry_exception_elapsed_time =
max_retry_exception_elapsed_time
+ self._now = now
self._exception: Optional[BaseException] = None
self._done = False
+ # Set to self._now() when the first RetryException is observed; used
to bound the
+ # total time spent in RetryException-driven retries (see _wait below).
+ self._retry_exception_start_time: Optional[float] = None
def can_retry(self, exception: BaseException) -> bool:
if isinstance(exception, RetryException):
@@ -256,7 +274,19 @@ class Retrying:
exception = self._last_exception()
if isinstance(exception, RetryException):
- # Considered immediately retriable
+ if self._retry_exception_start_time is None:
+ self._retry_exception_start_time = self._now()
+ elapsed = self._now() - self._retry_exception_start_time
+ if elapsed >= self._max_retry_exception_elapsed_time:
+ logger.debug(f"Given up on retrying. error: {repr(exception)}")
+ warnings.warn("[RETRIES_EXCEEDED] The maximum number of
retries has been exceeded.")
+ # Unwrap the underlying cause (both raise sites in
+ # ExecutePlanResponseReattachableIterator use `from e`) so the
caller sees a
+ # real, actionable error instead of the bare RetryException
marker.
+ cause = exception.__cause__ if exception.__cause__ is not None
else exception
+ raise cause
+ # Considered immediately retriable, as long as the cumulative
elapsed time above
+ # has not been exceeded.
logger.debug(f"Got error: {repr(exception)}. Retrying.")
return
@@ -301,7 +331,7 @@ class Retrying:
class RetryException(Exception):
"""
An exception that can be thrown upstream when inside retry and which is
always retryable
- even without policies
+ even without policies, bounded by Retrying's
max_retry_exception_elapsed_time.
"""
diff --git a/python/pyspark/sql/tests/connect/client/test_client_retries.py
b/python/pyspark/sql/tests/connect/client/test_client_retries.py
index aa8612e1e232..ebfac7411406 100644
--- a/python/pyspark/sql/tests/connect/client/test_client_retries.py
+++ b/python/pyspark/sql/tests/connect/client/test_client_retries.py
@@ -31,6 +31,8 @@ if should_test_connect:
from pyspark.sql.connect.client.retries import (
Retrying,
DefaultPolicy,
+ RetryException,
+ DEFAULT_MAX_RETRY_EXCEPTION_ELAPSED_TIME,
)
from pyspark.sql.tests.connect.client.test_client import (
TestPolicy,
@@ -50,6 +52,18 @@ if should_test_connect:
def times(self):
return list(self._times)
+ class FakeClock:
+ """Injectable monotonic clock for testing elapsed-time bounds without
real waits."""
+
+ def __init__(self):
+ self._t = 0.0
+
+ def advance(self, seconds: float):
+ self._t += seconds
+
+ def now(self) -> float:
+ return self._t
+
def create_test_exception_with_details(
msg: str,
code: grpc.StatusCode = grpc.StatusCode.INTERNAL,
@@ -256,6 +270,104 @@ class
SparkConnectClientRetriesTestCase(unittest.TestCase):
self.assertEqual(tries, 1)
self.assertEqual(len(sleep_tracker.times), 0)
+ def test_retry_exception_bounded_by_elapsed_time(self):
+ # Each attempt simulates a full 10-min reattach deadline elapsing
before RetryException
+ # is raised. The elapsed clock starts on the first RetryException,
after that attempt
+ # already advanced the clock, so with a 1-hour bound and a
10-min-per-attempt advance the
+ # loop gives up after bound // increment + 1 = 7 attempts: attempts
1-6 leave elapsed
+ # < 1h, and attempt 7 leaves elapsed = 6 * 10min = 1h >= bound,
tripping the throw.
+ client = SparkConnectClient("sc://foo/;token=bar")
+ clock = FakeClock()
+ cause = TestException("deadline exceeded",
code=grpc.StatusCode.DEADLINE_EXCEEDED)
+ tries = 0
+ with self.assertRaises(TestException) as cm:
+ for attempt in Retrying(
+ client._retry_policies,
+ sleep=lambda t: None,
+ now=clock.now,
+
max_retry_exception_elapsed_time=DEFAULT_MAX_RETRY_EXCEPTION_ELAPSED_TIME,
+ ):
+ with attempt:
+ tries += 1
+ clock.advance(10 * 60)
+ raise RetryException() from cause
+ self.assertIs(cm.exception, cause)
+ self.assertEqual(tries, 7)
+
+ def test_retry_exception_below_bound_keeps_retrying(self):
+ # Non-regression: well under the elapsed-time budget, RetryException
retries should
+ # keep behaving as before (immediate retry, no policy consulted) until
success.
+ client = SparkConnectClient("sc://foo/;token=bar")
+ clock = FakeClock()
+ tries = 0
+ for attempt in Retrying(
+ client._retry_policies,
+ sleep=lambda t: None,
+ now=clock.now,
+
max_retry_exception_elapsed_time=DEFAULT_MAX_RETRY_EXCEPTION_ELAPSED_TIME,
+ ):
+ with attempt:
+ tries += 1
+ if tries <= 3:
+ clock.advance(10)
+ raise RetryException()
+ self.assertEqual(tries, 4)
+
+ def
test_retry_exception_bare_no_cause_raises_self_when_bound_exceeded(self):
+ # If a RetryException has no attached cause (e.g. the
OPERATION_NOT_FOUND/
+ # SESSION_NOT_FOUND path, which today raises a bare RetryException),
fall back to
+ # raising the bare RetryException itself once the bound is exceeded.
+ client = SparkConnectClient("sc://foo/;token=bar")
+ clock = FakeClock()
+ with self.assertRaises(RetryException):
+ for attempt in Retrying(
+ client._retry_policies,
+ sleep=lambda t: None,
+ now=clock.now,
+ max_retry_exception_elapsed_time=100,
+ ):
+ with attempt:
+ clock.advance(30)
+ raise RetryException()
+
+ def
test_retry_exception_elapsed_clock_starts_at_first_retry_exception(self):
+ # The elapsed-time clock should start only when the first
RetryException is observed,
+ # not at Retrying construction, and should be unaffected by
interleaved ordinary
+ # (policy-governed) retries that happen first.
+ client = SparkConnectClient("sc://foo/;token=bar")
+ clock = FakeClock()
+ cause = TestException("deadline exceeded",
code=grpc.StatusCode.DEADLINE_EXCEEDED)
+ sleep_tracker = SleepTimeTracker()
+ tries = 0
+ with self.assertRaises(TestException) as cm:
+ for attempt in Retrying(
+ client._retry_policies,
+ sleep=sleep_tracker.sleep,
+ now=clock.now,
+
max_retry_exception_elapsed_time=DEFAULT_MAX_RETRY_EXCEPTION_ELAPSED_TIME,
+ ):
+ with attempt:
+ tries += 1
+ if tries <= 2:
+ # Ordinary, policy-governed retries -- must not
consume any of the
+ # RetryException elapsed-time budget.
+ raise TestException("Retryable error",
grpc.StatusCode.UNAVAILABLE)
+ clock.advance(10 * 60)
+ raise RetryException() from cause
+ self.assertIs(cm.exception, cause)
+ # 2 ordinary retries + 7 RetryException retries (same count as the
isolated test above).
+ self.assertEqual(tries, 2 + 7)
+
+ def test_default_max_retry_exception_elapsed_time(self):
+ client = SparkConnectClient("sc://foo/;token=bar")
+ self.assertEqual(
+ client._max_retry_exception_elapsed_time,
DEFAULT_MAX_RETRY_EXCEPTION_ELAPSED_TIME
+ )
+
+ client2 = SparkConnectClient("sc://foo/;token=bar",
max_retry_exception_elapsed_time=42)
+ self.assertEqual(client2._max_retry_exception_elapsed_time, 42)
+
self.assertEqual(client2._retrying()._max_retry_exception_elapsed_time, 42)
+
def test_deadline_exceeded_is_not_retried_for_non_retryable_codes(self):
# Sanity check: ABORTED is not retried by DefaultPolicy (unless
matching cluster message)
policy = DefaultPolicy()
diff --git
a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientRetriesSuite.scala
b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientRetriesSuite.scala
index 595f48b9f548..db352c922322 100644
---
a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientRetriesSuite.scala
+++
b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientRetriesSuite.scala
@@ -48,6 +48,13 @@ class SparkConnectClientRetriesSuite extends ConnectFunSuite
with Eventually {
def totalSleep: Long = data.sum
}
+ /** Injectable clock for testing elapsed-time bounds without real waits. */
+ private class FakeClock {
+ private var current: Long = 0L
+ def advanceMillis(ms: Long): Unit = current += ms * 1000000L
+ def nowNanos(): Long = current
+ }
+
/** Helper function for creating a test exception with retry_delay */
private def createTestExceptionWithDetails(
msg: String,
@@ -295,6 +302,127 @@ class SparkConnectClientRetriesSuite extends
ConnectFunSuite with Eventually {
assert(dummyFn.counter == 1)
}
+ test("RetryException retries are bounded by elapsed time") {
+ // Each attempt simulates a full 10-min reattach deadline elapsing before
RetryException is
+ // raised. The elapsed clock starts on the first RetryException, after
that attempt already
+ // advanced the clock, so with a 1-hour bound and a 10-min-per-attempt
advance the loop gives
+ // up after bound / increment + 1 = 7 attempts: attempts 1-6 leave elapsed
< 1h, and attempt
+ // 7 leaves elapsed = 6 * 10min = 1h >= bound, tripping the throw.
+ val clock = new FakeClock()
+ val cause = new StatusRuntimeException(Status.DEADLINE_EXCEEDED)
+ var tries = 0
+ val retryHandler = new GrpcRetryHandler(
+ RetryPolicy.defaultPolicies(),
+ sleep = _ => {},
+ maxRetryExceptionElapsedTime = FiniteDuration(1, "hour"),
+ nowNanos = clock.nowNanos)
+
+ val thrown = intercept[StatusRuntimeException] {
+ retryHandler.retry {
+ tries += 1
+ clock.advanceMillis(10 * 60 * 1000)
+ val error = new GrpcRetryHandler.RetryException()
+ error.addSuppressed(cause)
+ throw error
+ }
+ }
+ assert(thrown eq cause)
+ assert(tries == 7)
+ }
+
+ test("RetryException below the elapsed bound keeps retrying") {
+ // Non-regression: well under the elapsed-time budget, RetryException
retries should keep
+ // behaving as before (immediate retry, no policy consulted) until success.
+ val clock = new FakeClock()
+ var tries = 0
+ val retryHandler = new GrpcRetryHandler(
+ RetryPolicy.defaultPolicies(),
+ sleep = _ => {},
+ maxRetryExceptionElapsedTime = FiniteDuration(1, "hour"),
+ nowNanos = clock.nowNanos)
+
+ val result = retryHandler.retry {
+ tries += 1
+ if (tries <= 3) {
+ clock.advanceMillis(10 * 1000)
+ throw new GrpcRetryHandler.RetryException()
+ }
+ 42
+ }
+ assert(result == 42)
+ assert(tries == 4)
+ }
+
+ test("RetryException with no suppressed cause throws itself once bound
exceeded") {
+ val clock = new FakeClock()
+ val retryHandler = new GrpcRetryHandler(
+ RetryPolicy.defaultPolicies(),
+ sleep = _ => {},
+ maxRetryExceptionElapsedTime = FiniteDuration(100, "s"),
+ nowNanos = clock.nowNanos)
+
+ intercept[GrpcRetryHandler.RetryException] {
+ retryHandler.retry {
+ clock.advanceMillis(30 * 1000)
+ throw new GrpcRetryHandler.RetryException()
+ }
+ }
+ }
+
+ test("elapsed-time clock for RetryException starts only at first
occurrence") {
+ // The elapsed-time clock should start only when the first RetryException
is observed, not
+ // at Retrying construction, and should be unaffected by interleaved
ordinary
+ // (policy-governed) retries that happen first.
+ val clock = new FakeClock()
+ val cause = new StatusRuntimeException(Status.DEADLINE_EXCEEDED)
+ val st = new SleepTimeTracker()
+ var tries = 0
+ val retryHandler = new GrpcRetryHandler(
+ RetryPolicy.defaultPolicies(),
+ sleep = st.sleep,
+ maxRetryExceptionElapsedTime = FiniteDuration(1, "hour"),
+ nowNanos = clock.nowNanos)
+
+ val thrown = intercept[StatusRuntimeException] {
+ retryHandler.retry {
+ tries += 1
+ if (tries <= 2) {
+ // Ordinary, policy-governed retries -- must not consume any of the
RetryException
+ // elapsed-time budget.
+ throw new StatusRuntimeException(Status.UNAVAILABLE)
+ }
+ clock.advanceMillis(10 * 60 * 1000)
+ val error = new GrpcRetryHandler.RetryException()
+ error.addSuppressed(cause)
+ throw error
+ }
+ }
+ assert(thrown eq cause)
+ // 2 ordinary retries + 7 RetryException retries (same count as the
isolated test above).
+ assert(tries == 2 + 7)
+ }
+
+ test("default maxRetryExceptionElapsedTime is 1 hour") {
+ // Uses the default maxRetryExceptionElapsedTime (not passed) -- if this
default ever
+ // changes, update this assertion together with the doc comment on
GrpcRetryHandler's
+ // constructor.
+ val clock = new FakeClock()
+ var tries = 0
+ val retryHandler = new GrpcRetryHandler(
+ RetryPolicy.defaultPolicies(),
+ sleep = _ => {},
+ nowNanos = clock.nowNanos)
+
+ intercept[GrpcRetryHandler.RetryException] {
+ retryHandler.retry {
+ tries += 1
+ clock.advanceMillis(10 * 60 * 1000)
+ throw new GrpcRetryHandler.RetryException()
+ }
+ }
+ assert(tries == 7)
+ }
+
test("RpcDeadlines.disabled creates an instance with all deadlines set to
None") {
val disabled = RpcDeadlines.disabled
assert(disabled.reattachableExecutePlan.isEmpty)
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 fd5be9bb555e..24e637a1f215 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
@@ -1116,6 +1116,48 @@ class SparkConnectClientSuite extends ConnectFunSuite {
}
}
+ test(
+ "SPARK-58162: maxRetryExceptionElapsedTime set via Builder bounds the " +
+ "RetryException retry loop end-to-end") {
+ // Every attempt (the initial ExecutePlan and every subsequent
ReattachExecute) never
+ // responds at all, so the client's reattach deadline fires every time --
DEADLINE_EXCEEDED
+ // -> RetryException, on every attempt, forever -- unless the elapsed-time
bound kicks in.
+ // The point of this test is that the bound is configured only through the
real
+ // Builder -> Configuration -> SparkConnectStubState -> GrpcRetryHandler
wiring (unlike
+ // SparkConnectClientRetriesSuite, which constructs GrpcRetryHandler
directly), so it guards
+ // against a regression that silently drops the
maxRetryExceptionElapsedTime passthrough in
+ // SparkConnectStubState.
+ val stallingService = new DummySparkConnectService {
+ override def executePlan(
+ request: ExecutePlanRequest,
+ responseObserver: StreamObserver[ExecutePlanResponse]): Unit = {}
+
+ override def reattachExecute(
+ request: proto.ReattachExecuteRequest,
+ responseObserver: StreamObserver[ExecutePlanResponse]): Unit = {}
+ }
+ server =
NettyServerBuilder.forPort(0).addService(stallingService).build().start()
+ service = stallingService
+ client = SparkConnectClient
+ .builder()
+ .connectionString(s"sc://localhost:${server.getPort}")
+ .rpcDeadlines(
+ RpcDeadlines(
+ reattachableExecutePlan = Some(FiniteDuration(50,
TimeUnit.MILLISECONDS)),
+ reattachExecute = Some(FiniteDuration(50, TimeUnit.MILLISECONDS))))
+ .maxRetryExceptionElapsedTime(FiniteDuration(200, TimeUnit.MILLISECONDS))
+ .enableReattachableExecute()
+ .build()
+
+ val ex = intercept[SparkException] {
+ val iter = client.execute(buildPlan("select 1"))
+ iter.foreach(_ => ())
+ }
+ assert(
+ ex.getCause.asInstanceOf[StatusRuntimeException].getStatus.getCode ==
+ Status.Code.DEADLINE_EXCEEDED)
+ }
+
test("non-reattachable executePlan has no client-side deadline") {
val latch = new CountDownLatch(1)
val slowService = new DummySparkConnectService {
diff --git
a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/GrpcRetryHandler.scala
b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/GrpcRetryHandler.scala
index d92dc902fedc..b62cf7ac58ed 100644
---
a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/GrpcRetryHandler.scala
+++
b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/GrpcRetryHandler.scala
@@ -17,17 +17,21 @@
package org.apache.spark.sql.connect.client
+import scala.concurrent.duration.FiniteDuration
import scala.util.control.NonFatal
import io.grpc.stub.StreamObserver
import org.apache.spark.internal.Logging
-import org.apache.spark.internal.LogKeys.{ERROR, NUM_RETRY, POLICY,
RETRY_WAIT_TIME}
+import org.apache.spark.internal.LogKeys.{ELAPSED_TIME, ERROR, NUM_RETRY,
POLICY, RETRY_WAIT_TIME}
import org.apache.spark.sql.util.{CloseableIterator, WrappedCloseableIterator}
private[sql] class GrpcRetryHandler(
private val policies: Seq[RetryPolicy],
- private val sleep: Long => Unit = Thread.sleep) {
+ private val sleep: Long => Unit = Thread.sleep,
+ private val maxRetryExceptionElapsedTime: FiniteDuration =
+ GrpcRetryHandler.DEFAULT_MAX_RETRY_EXCEPTION_ELAPSED_TIME,
+ private val nowNanos: () => Long = () => System.nanoTime()) {
def this(policy: RetryPolicy, sleep: Long => Unit) = this(List(policy),
sleep)
def this(policy: RetryPolicy) = this(policy, Thread.sleep)
@@ -35,7 +39,9 @@ private[sql] class GrpcRetryHandler(
/**
* Retries the given function with exponential backoff according to the
client's retryPolicy.
*/
- def retry[T](fn: => T): T = new GrpcRetryHandler.Retrying(policies, sleep,
fn).retry()
+ def retry[T](fn: => T): T =
+ new GrpcRetryHandler.Retrying(policies, sleep, fn,
maxRetryExceptionElapsedTime, nowNanos)
+ .retry()
/**
* Generalizes the retry logic for RPC calls that return an iterator.
@@ -151,6 +157,11 @@ private[sql] class GrpcRetryHandler(
private[sql] object GrpcRetryHandler extends Logging {
+ // Default cumulative elapsed-time bound for RetryException-driven retries
(see
+ // Retrying.waitAfterAttempt below). Single source of truth for this
default; please keep the
+ // Python side (retries.py's DEFAULT_MAX_RETRY_EXCEPTION_ELAPSED_TIME) in
sync with this value.
+ val DEFAULT_MAX_RETRY_EXCEPTION_ELAPSED_TIME: FiniteDuration =
FiniteDuration(1, "hour")
+
/**
* Class managing the state of the retrying logic during a single retryable
block.
* @param retryPolicies
@@ -159,14 +170,29 @@ private[sql] object GrpcRetryHandler extends Logging {
* typically Thread.sleep
* @param fn
* the function to compute
+ * @param maxRetryExceptionElapsedTime
+ * maximum cumulative wall-clock time to keep retrying a
[[RetryException]] (which is
+ * otherwise retried unconditionally, without consulting any
[[RetryPolicy]]) before giving
+ * up. Measured from the first [[RetryException]] seen by this instance.
+ * @param nowNanos
+ * typically System.nanoTime, used to measure maxRetryExceptionElapsedTime
* @tparam T
* result of function fn
*/
- class Retrying[T](retryPolicies: Seq[RetryPolicy], sleep: Long => Unit, fn:
=> T) {
+ class Retrying[T](
+ retryPolicies: Seq[RetryPolicy],
+ sleep: Long => Unit,
+ fn: => T,
+ maxRetryExceptionElapsedTime: FiniteDuration =
DEFAULT_MAX_RETRY_EXCEPTION_ELAPSED_TIME,
+ nowNanos: () => Long = () => System.nanoTime()) {
private var currentRetryNum: Int = 0
private var exceptionList: Seq[Throwable] = Seq.empty
private val policies: Seq[RetryPolicy.RetryPolicyState] =
retryPolicies.map(_.toState)
+ // Set to the result of nowNanos() when the first RetryException is
observed; used to bound
+ // the total time spent in RetryException-driven retries (see
waitAfterAttempt below).
+ private var retryExceptionStartNanos: Option[Long] = None
+
def canRetry(throwable: Throwable): Boolean = {
throwable.isInstanceOf[RetryException] || policies.exists(p =>
p.canRetry(throwable))
}
@@ -187,7 +213,26 @@ private[sql] object GrpcRetryHandler extends Logging {
val lastException = exceptionList.head
if (lastException.isInstanceOf[RetryException]) {
- // retry exception is considered immediately retriable without any
policies.
+ val startNanos = retryExceptionStartNanos.getOrElse {
+ val now = nowNanos()
+ retryExceptionStartNanos = Some(now)
+ now
+ }
+ val elapsedNanos = nowNanos() - startNanos
+ if (elapsedNanos >= maxRetryExceptionElapsedTime.toNanos) {
+ logWarning(
+ log"Non-Fatal error during RPC execution: ${MDC(ERROR,
lastException)}, " +
+ log"exceeded maxRetryExceptionElapsedTime " +
+ log"(elapsed=${MDC(ELAPSED_TIME, elapsedNanos / 1000000)} ms, " +
+ log"currentRetryNum=${MDC(NUM_RETRY, currentRetryNum)})")
+ logWarning(log"[RETRIES_EXCEEDED] The maximum number of retries has
been exceeded.")
+ // Unwrap the underlying cause (both throw sites in
+ // ExecutePlanResponseReattachableIterator attach it via
addSuppressed) so the
+ // caller sees a real, actionable error instead of the bare
RetryException marker.
+ throw lastException.getSuppressed.headOption.getOrElse(lastException)
+ }
+ // retry exception is considered immediately retriable without any
policies, as long
+ // as the cumulative elapsed time above has not been exceeded.
logWarning(
log"Non-Fatal error during RPC execution: ${MDC(ERROR,
lastException)}, " +
log"retrying (currentRetryNum=${MDC(NUM_RETRY, currentRetryNum)})")
@@ -232,7 +277,7 @@ private[sql] object GrpcRetryHandler extends Logging {
/**
* An exception that can be thrown upstream when inside retry and which will
be always retryable
- * without any policies.
+ * without any policies, bounded by Retrying's maxRetryExceptionElapsedTime.
*/
class RetryException extends Throwable
}
diff --git
a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/RetryPolicy.scala
b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/RetryPolicy.scala
index b396ee96a4e5..1e2a86485846 100644
---
a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/RetryPolicy.scala
+++
b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/RetryPolicy.scala
@@ -173,7 +173,9 @@ object RetryPolicy extends Logging {
// ExecutePlanResponseReattachableIterator, which converts it to
RetryException so the
// server-side operation continues and a fresh ReattachExecute is
issued. We do not retry
// other RPCs on deadline: those are non-idempotent or the deadline
signals a genuine
- // timeout that retrying won't fix.
+ // timeout that retrying won't fix. RetryException retries bypass
RetryPolicy entirely
+ // (see GrpcRetryHandler.Retrying.waitAfterAttempt), but are bounded
by a separate
+ // cumulative elapsed-time cap (GrpcRetryHandler's
maxRetryExceptionElapsedTime).
false
case _ => false
}
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 8f449207b796..eeb14fcfc2c7 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
@@ -23,6 +23,7 @@ import java.util.{Base64, Locale, UUID}
import java.util.concurrent.{Executor, TimeUnit}
import scala.collection.mutable
+import scala.concurrent.duration.FiniteDuration
import scala.jdk.CollectionConverters._
import scala.util.Properties
import scala.util.control.NonFatal
@@ -809,6 +810,20 @@ object SparkConnectClient {
this
}
+ /**
+ * Sets the maximum cumulative wall-clock time the client will keep
retrying a
+ * [[GrpcRetryHandler.RetryException]] (raised internally when a reattach
attempt keeps
+ * hitting DEADLINE_EXCEEDED, or when the initial ExecutePlan never
reached the server) before
+ * giving up and surfacing the underlying error. Defaults to 1 hour.
+ *
+ * @return
+ * this builder.
+ */
+ def maxRetryExceptionElapsedTime(duration: FiniteDuration): Builder = {
+ _configuration = _configuration.copy(maxRetryExceptionElapsedTime =
duration)
+ this
+ }
+
private object URIParams {
val PARAM_USER_ID = "user_id"
val PARAM_USE_SSL = "use_ssl"
@@ -1094,6 +1109,8 @@ object SparkConnectClient {
sys.env.getOrElse("SPARK_CONNECT_USER_AGENT", DEFAULT_USER_AGENT)),
retryPolicies: Seq[RetryPolicy] = RetryPolicy.defaultPolicies(),
rpcDeadlines: RpcDeadlines = RpcDeadlines(),
+ maxRetryExceptionElapsedTime: FiniteDuration =
+ GrpcRetryHandler.DEFAULT_MAX_RETRY_EXCEPTION_ELAPSED_TIME,
useReattachableExecute: Boolean = true,
interceptors: List[ClientInterceptor] = List.empty,
sessionId: Option[String] = None,
diff --git
a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectStubState.scala
b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectStubState.scala
index 4d2d15073f22..ad1e21a6441b 100644
---
a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectStubState.scala
+++
b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectStubState.scala
@@ -38,7 +38,9 @@ class SparkConnectStubState(
}
// Manages the retry handler logic used by the stubs.
- lazy val retryHandler = new GrpcRetryHandler(configuration.retryPolicies)
+ lazy val retryHandler = new GrpcRetryHandler(
+ configuration.retryPolicies,
+ maxRetryExceptionElapsedTime = configuration.maxRetryExceptionElapsedTime)
// Responsible to convert the GRPC Status exceptions into Spark exceptions.
lazy val exceptionConverter: GrpcExceptionConverter =
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
