Repository navigation
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request implements client-side recovery from position-lost errors via one-response release slack and bounded one-shot replays for root plans. It also enhances instrumentation with point-in-time events and tracks in-flight gRPC calls to cut wedged streams. The review feedback highlights critical synchronization issues where class-level shared dictionaries are protected by instance-level locks instead of class-level locks, as well as a type-hint mismatch in env_float where the default parameter should support None.
dborowitz
force-pushed
the
dborowitz-resilience
branch
2 times, most recently
from
October 26, 2026 19:00
13e0624 to
568cb6d
Compare
dborowitz
force-pushed
the
dborowitz-resilience
branch
2 times, most recently
from
October 26, 2026 19:53
97db870 to
fb3ea49
Compare
dborowitz
force-pushed
the
dborowitz-resilience
branch
from
October 27, 2026 19:33
fb3ea49 to
7d3c306
Compare
dborowitz
force-pushed
the
dborowitz-resilience
branch
from
October 28, 2026 03:37
7d3c306 to
b180595
Compare
dborowitz
force-pushed
the
dborowitz-resilience
branch
2 times, most recently
from
October 28, 2026 22:29
7839943 to
a91334d
Compare
Emit point-in-time events from the local WebSocket proxy when
instrumentation is on (no-op recorder otherwise; the disabled path calls
connect_sockets exactly as before):
* proxy.eof - a forwarding direction saw its source end close:
recv returned empty (gRPC side) or the websocket
raised ConnectionClosed (adds error_type,
close_code). attrs: conn, direction, bytes_forwarded,
age_ns since the connection was accepted.
* proxy.socket_timeout - a send timed out and is retried, with timeout_s
of the gRPC-side socket. Recv timeouts are routine
polls of an idle connection and emit nothing.
* proxy.forward_error - any other exception that ends a direction
(error_type, message truncated to 200 chars),
including SendTimeout after 10 failed send retries.
* proxy.ws_connect_failed - connect_tcp_bridge failed (auth or handshake),
with http_status when the error carries one.
Also names the 10 s socket / 60 s websocket timeouts as constants, and
closes both sockets when a direction sees EOF, as the error path already
did.
get_refreshed_credentials emits auth.token_refreshed (attrs expiry_wall_ns of the new token, had_token = whether a token existed before) whenever it actually refreshed the Application Default Credentials, so a long-lived WebSocket's auth timeline can be reconstructed. No-op without instrumentation.
…t hook Make the notebook execution progress bar observable: * Events (instrumented sessions only): progress.bar_open (operation_id, total_tasks) when the handler creates a tqdm bar; progress.bar_update (operation_id, completed_tasks, total_tasks) on every later handler call for that bar that does not finish it; progress.bar_close (operation_id, completed_tasks, total_tasks, reason) when it closes and unregisters it. A query whose first progress update already says done=True gets no bar. * Every session (instrumented or not) wraps the client's _execute_and_fetch to close the query's bar in a finally, should the handler have left it open (reason=query_end); pyspark's Progress.__exit__ normally delivers the done=True that closes it (reason=done) before that. * Per-instance client wrappers now go through one installer pair, install_client_hooks / uninstall_client_hooks: the first hook on a method records the original, and uninstalling restores every original at once. stop() removes them. * ManagedSparkSession.progress_bars() returns dict[operation_id, bar]: a snapshot of the open bars, so tests and diagnostics no longer reach into the private class-level registry. * Private test hook for running the handler outside a notebook: MANAGED_SPARK_CONNECT_FORCE_PROGRESS_BAR=1 makes the handler treat the environment as interactive, and ManagedSparkSession._set_tqdm_factory(cls) substitutes a tqdm-like class (None restores the default). Events are inert for sessions without instrumentation; apart from closing a leftover bar and skipping a bar for an already finished query, the handler's behaviour in notebooks is unchanged.
The RPC liveness guard built by make_is_active_callback emits a guard.result event when instrumentation is on: mode (pre_rpc | on_error, as set by the guarded channel; direct outside it), cached (answered from the liveness cache without a GetSession RPC), active (the verdict), and error_type with active=None when the check itself raised. A cached active verdict, which every healthy guarded RPC gets, emits nothing. The existing dataproc.get_session span still covers the RPC. Without instrumentation the callback runs exactly as before.
dborowitz
force-pushed
the
dborowitz-resilience
branch
from
October 30, 2026 18:23
a91334d to
4743d18
Compare
Contributor
Author
|
/gemini review |
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces client-side recovery from INVALID_CURSOR.POSITION_NOT_AVAILABLE errors by re-running root plans, adds environment variable parsing helpers, enhances instrumentation with several new events, and implements a silence cut mechanism in the proxy to prevent hanging streams. Key feedback includes addressing a potential memory leak/OOM risk from disabling WebSocket backpressure (max_queue=None), fixing a critical protobuf bug where clearing repeated composite fields invalidates references, wrapping a potentially throwing callable in a try-except block, replacing print statements with proper logging, and resolving encapsulation violations caused by importing private members across modules.
install_client_wrappers now also wraps the client's _retrying factory so the Retrying it hands to pyspark's retry loops (including the reattachable ExecutePlan iterator, which captures the factory when built) is a subclass that emits grpc.retry after each wait that precedes a retry: attempt (1 = first retry within that loop), reason (reattach for pyspark's RetryException, error otherwise), and grpc_code and error_type of the exception being retried. The failed RPC is the grpc.* span that ended just before the event on the same thread. Per-instance, removed by uninstall_client_wrappers; nothing is touched without instrumentation.
README: an Events table (proxy.*, auth.token_refreshed, progress.*, guard.result, grpc.retry) next to the span catalogue, and a short section on ManagedSparkSession.progress_bars().
… spans The gRPC interceptor now tags ExecutePlan spans with plan_kind=root|command (request.plan's op_type oneof) and, for commands, command_kind (the command_type oneof name, e.g. sql_command, create_dataframe_view). Root plans are the executions visible server-side as SQL queries; commands (temp views, cache, conf, ...) never are, so consumers can restrict server-side joins to root plans. Only non-empty string oneof names are recorded; other request types are unaffected.
…ne-shot replay When a transport blip breaks an ExecutePlan stream and the server no longer buffers the response pyspark reattaches from, ReattachExecute fails with INVALID_CURSOR.POSITION_NOT_AVAILABLE, which pyspark's DefaultPolicy does not retry, so the whole query fails. Add client/reattach_patch.py. install_replay_on_position_lost hooks the client's _execute_and_fetch (every ManagedSparkSession installs it, innermost, under the instrumentation and progress bar hooks). When a root (read) plan fails with that error condition, it closes the abandoned operation's progress bar (reason=replay), logs a WARNING and re-runs a copy of the request once under a fresh operation_id. Command plans are never replayed. Every such failure emits a reattach.position_lost event. Opt out with MANAGED_SPARK_CONNECT_REPLAY_ON_POSITION_LOST=0, documented in a new README configuration section.
Pass max_queue=None and ping_timeout=60 to websocketclient.connect in connect_tcp_bridge. In websockets.sync.client, a bounded recv_messages queue (default 16) blocks the WebSocket reader thread when the consumer is slow, preventing the reader thread from processing incoming Pong frames within the default 20s ping_timeout and causing spurious keepalive disconnects under backpressure. ping_timeout goes from 20s to 60s because a Pong queues behind the data frames the server is already sending, so it can arrive late on a saturated link; a dead peer is still detected within ping_interval (20s) plus ping_timeout.
Until now, a GetSession error other than 503/504/429/500 (an auth RefreshError, a 403, anything unexpected) marked the session dead for good, so one auth blip failed every later RPC. Transient errors fell back to a last-known-active Session through several layers: a TTL cache, _LAST_KNOWN_ACTIVE with a 60s age limit that pre_rpc ignored, and separate meanings for unset, zero and negative TTLs. Replace all of it with one rule: a session is dead only when GetSession returns NotFound or a terminal state. get_active_s8s_session_response lets every other error propagate and leaves the liveness cache alone: - ProxiedChannel._check_active catches the error and lets the RPC proceed (after a failed RPC, it re-raises that RPC's own error). It logs a WARNING once per outage, then DEBUG until a check succeeds. - getOrCreate reuses the session already in the process and logs the same way. With nothing to reuse (a first attach), it raises a ManagedSparkConnectException that names the session. The liveness cache is now one dict of verdicts. MANAGED_SPARK_CONNECT_LIVENESS_TTL_SECONDS still means seconds when positive and no caching when 0. Unset and negative values now mean the same thing: a cached ACTIVE verdict never expires, for the pre-RPC guard and for direct checks alike. Before, unset cached for the guard only. The variable is now documented in the README and read with env_float. Tests patch an injectable clock (session._now) instead of time.monotonic.
…imeout 1. Shut down both directions of sockets (socket.SHUT_RDWR) before closing in forward_bytes on EOF or forwarding error, so the peer thread blocked in recv() wakes immediately with EOF/reset instead of waiting for its socket timeout. 2. Pass close_timeout=3 to websocketclient.connect in connect_tcp_bridge so exiting the WebSocket context manager after a broken/reset backend does not wait 10s for a Close frame handshake that will never arrive.
dborowitz
force-pushed
the
dborowitz-resilience
branch
from
October 30, 2026 20:08
4743d18 to
ea4a5dd
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Note
Followed by #217 (
dborowitz-hygiene).Summary
Hardens
google-cloud-spark-connectagainst mid-stream transport drops, half-open WebSocket blackholes, slow-consumer backpressure, and Dataproc control-plane errors in the session liveness guard, while extending the opt-in instrumentation package (#210) with fine-grained proxy, auth, retry, guard, and progress-bar events.Failure Modes Fixed
ReleaseExecute(INVALID_CURSOR.POSITION_NOT_AVAILABLE)ExecutePlanstream and the server no longer buffers the response PySpark'sExecutePlanResponseReattachableIteratorreattaches from (for example when extra frames queued in the proxy before the drop),ReattachExecutefails withINVALID_CURSOR.POSITION_NOT_AVAILABLE, which PySpark'sDefaultPolicydoes not retry.ExecutePlanreplay (install_replay_on_position_lostinclient/reattach_patch.py, configurable viaMANAGED_SPARK_CONNECT_REPLAY_ON_POSITION_LOST, default1): hooksclient._execute_and_fetchto re-run read-onlyrootplans once under a freshoperation_id(stripping priorreattach_options), closes the abandoned operation's progress bar (reason="replay"), logs aWARNING, and emitsreattach.position_lost.commandplans are never replayed.ws_drop_mid_stream(early),ws_drop_mid_stream(after_release), andtcp_reset_mid_stream(2,000,000rows,0unrecovered errors).forward_bytes,server_to_clientcaughtTimeoutErrorevery 60s (WS_RECV_TIMEOUT_S) andcontinued indefinitely whileclient_to_serverlooped every 10s (SOCKET_TIMEOUT_S), leaving gRPC streams wedged forever when an intermediate hop blackholed traffic without a TCPFIN/RST.ManagedSparkSessionProxy(inc_inflight/dec_inflightin_GuardedMultiCallable,_GuardedStream, and_GuardedFuture) and expose it tobridged_socket. Wheninflight_rpcs() > 0and no server frame arrives forMANAGED_SPARK_CONNECT_WS_SILENCE_LIMIT_S(default300s,0disables; pre-RPC idle time is excluded), log aWARNING, emitproxy.silence_cut, and shut down and close both ends so gRPC fails withUNAVAILABLEand triggersReattachExecute.ReattachExecute;long_query.silent_gap(75s)andsleep_udf(90s, 150s)completed with 0 false cuts.ping_timeoutwebsockets.sync.client.connect()defaults tomax_queue=16andping_timeout=20. When a client pauses reading for >20s (e.g. heavy Arrow/pandas conversion), the boundedrecv_messagesqueue fills and blocks the WebSocket reader thread from processing incomingPongframes, causing a self-inflicted keepalive disconnect.max_queue=None(leaving backpressure bounded by HTTP/2 and TCP flow-control windows) andping_timeout=60inconnect_tcp_bridgeso aPongqueued behind in-flight server data frames does not time out prematurely.slow_consumer(30s)andslow_consumer(60s)completed5,000,000rows with0keepalive disconnects and0reattaches.get_active_s8s_session_responsepreviously treated any error outside(503, 504, 429, 500)(such as an authRefreshErroror403) as permanent session death (_DEAD_SESSIONS), and handled transient errors through three separate caches (_LIVENESS_CACHE,_LAST_KNOWN_ACTIVEwith a 60s age limit, and_DEAD_SESSIONS) with inconsistentMANAGED_SPARK_CONNECT_LIVENESS_TTL_SECONDSsemantics._LIVENESSverdict cache and a single rule: a session is marked dead only whenGetSessionreturnsNotFoundor a terminal state (not ACTIVEandnot CREATING). All other errors propagate without touching the cache:ProxiedChannel._check_activefails open (proceeding with the RPC inpre_rpcmode or re-raising the original RPC error inon_errormode) andgetOrCreatereuses an existing in-process session unchecked, logging once per outage atWARNINGand subsequent failures atDEBUG(_OutageLog). Unifies and documentsMANAGED_SPARK_CONNECT_LIVENESS_TTL_SECONDS.503s inguard_transient_503and auth blips tolerated with 0 aborted queries, while genuinely terminated sessions still fast-fail in~200 ms.forward_bytesthread withoutshutdown(SHUT_RDWR)left the peer thread blocked inrecv()until its timeout expired, and exitingwebsocketclient.connectwaited 10s (close_timeout) for a Close handshake on dead connections.socket.shutdown(socket.SHUT_RDWR)before closing sockets in_shutdown_and_close(on EOF, silence cut, or forwarding error) and setclose_timeout=3onwebsocketclient.connect.Commits in this PR
Instrumentation Additions (Commits 1–7)
feat(instrumentation): proxy forwarding events— Emitsproxy.eof(conn,direction,bytes_forwarded,age_ns, pluserror_typeandclose_codeonConnectionClosed),proxy.socket_timeout(conn,direction,timeout_s),proxy.forward_error(conn,direction,error_type,message, includingSendTimeoutafter 10 failed send retries), andproxy.ws_connect_failed(conn,error_type,http_status); closes both sockets on EOF.feat(instrumentation): auth.token_refreshed event— Emitsauth.token_refreshed(expiry_wall_ns,had_token) whenever Application Default Credentials are refreshed inget_refreshed_credentials.feat(instrumentation): progress bar events, snapshot and headless test hook— Emitsprogress.bar_open,progress.bar_update, andprogress.bar_close(reason=done|query_end|replay); skips creating a bar when the first update already hasdone=True; addsinstall_client_hooks/uninstall_client_hooksand wraps_execute_and_fetchon every session to close any leftover bar on query completion; addsManagedSparkSession.progress_bars(),ManagedSparkSession._set_tqdm_factory(), andMANAGED_SPARK_CONNECT_FORCE_PROGRESS_BAR=1.feat(instrumentation): guard.result event for the liveness guard— Emitsguard.result(mode=pre_rpc|on_error|direct,cached,active,error_type) when the liveness guard queries the control plane, finds the session inactive, or raises (cachedactive=Trueverdicts emit nothing).feat(instrumentation): grpc.retry event from pyspark's retry loop— Wrapsclient._retryingviainstall_client_hookswithInstrumentedRetryingto emitgrpc.retry(attempt,reason=reattach|error,grpc_code,error_type) after each retry backoff.docs(instrumentation): document events and progress_bars()— Documents theEventscatalogue andManagedSparkSession.progress_bars()indocs/instrumentation.mdandREADME.md.feat(instrumentation): plan_kind and command_kind on grpc.ExecutePlan spans— Tagsgrpc.ExecutePlanspans withplan_kind=root|commandand, for commands,command_kind(e.g.sql_command,create_dataframe_view).Resilience Fixes (Commits 8–12)
fix(client): recover from INVALID_CURSOR.POSITION_NOT_AVAILABLE via one-shot replay— Addsclient/reattach_patch.py(install_replay_on_position_lost, configurable viaMANAGED_SPARK_CONNECT_REPLAY_ON_POSITION_LOST), emitsreattach.position_lost, and wires it intoManagedSparkSessionandREADME.md.fix(proxy): cut wedged streams after prolonged WebSocket silence— Tracks in-flight RPC count onManagedSparkSessionProxy/bridged_socketand cuts connections silent forMANAGED_SPARK_CONNECT_WS_SILENCE_LIMIT_S(default300s) while an RPC is in flight, emittingproxy.silence_cut.fix(proxy): prevent slow consumer from triggering keepalive ping timeout— Setsmax_queue=Noneandping_timeout=60inconnect_tcp_bridge.fix(session): fail open on control-plane errors in the liveness guard— Replaces_LIVENESS_CACHE,_LAST_KNOWN_ACTIVE, and_DEAD_SESSIONSwith a single_LIVENESSverdict cache that only marks sessions dead onNotFoundor terminal states, failing open with_OutageLograte-limited logging inProxiedChannel._check_activeand in-processgetOrCreatereuse, and documentsMANAGED_SPARK_CONNECT_LIVENESS_TTL_SECONDSinREADME.md.fix(proxy): shutdown local socket before close and shorten ws close_timeout— Callssocket.shutdown(socket.SHUT_RDWR)before closing sockets inforward_bytesand passesclose_timeout=3towebsocketclient.connect.Testing
tests/unit/test_instrumentation_events.py,tests/unit/test_reattach_patch.py,tests/unit/test_proxy.py,tests/unit/test_guard.py,tests/unit/test_core.py, andtests/unit/test_instrumentation.py.ws+guard,ws+noguard, anddirect+noguard). Like the benchmark suite from perf: eliminate per-RPC control-plane polling, cache auth/clients, and optimize proxy framing #211, this is a custom AI-generated suite that I hope to open-source, but I'm getting the improvements in first.