diff --git a/README.md b/README.md index e4bf8907..54c4e055 100644 --- a/README.md +++ b/README.md @@ -125,6 +125,16 @@ The `ManagedSparkSession.builder` provides a fluent API to configure the session | `ttl(duration)` | Sets the time-to-live (TTL) for the session using a `datetime.timedelta` object. | | `instrumentation(inst)` | Records what the library does on your behalf; see [Instrumentation (opt-in)](#instrumentation-opt-in). | +### Configuration (environment variables) + +These environment variables change how the client behaves. + +| Variable | Default | Effect | +|----------|---------|--------| +| `MANAGED_SPARK_CONNECT_REPLAY_ON_POSITION_LOST` | `1` | When a read query (`collect`, `toPandas`, `count`, ...) fails with `INVALID_CURSOR.POSITION_NOT_AVAILABLE` because the server no longer buffers the result the client has to resume from, re-run it once under a new operation id and log a `WARNING`. Commands (writes, DDL) are never re-run. Set to `0` to disable, e.g. for queries that are expensive or not deterministic. Read when the session is created. | +| `MANAGED_SPARK_CONNECT_WS_SILENCE_LIMIT_S` | `300` | Cut the proxy WebSocket connection if an RPC is in flight and the server stays silent (no frames received) for this many seconds, so a wedged stream fails instead of hanging. The check runs on each 60 s WebSocket receive poll, so the cut happens between this limit and the limit plus 60 s. Set to `0` to disable. | +| `MANAGED_SPARK_CONNECT_LIVENESS_TTL_SECONDS` | unset | Before an RPC, the client checks that the Dataproc session is still alive. Once a check finds it `ACTIVE`, that answer is reused for this many seconds: unset or negative reuses it until an RPC fails, `0` checks before every RPC. A failed RPC always triggers a fresh check. The session only counts as gone when Dataproc reports it not found or in a terminal state; when the check itself fails (for example during a control-plane outage), the RPC proceeds. | + ### Reusing Named Sessions Across Notebooks Named sessions allow you to share a single Spark session across multiple notebooks, improving efficiency by avoiding repeated session startup times and reducing costs. @@ -168,7 +178,7 @@ result conversions — as timed *spans*. It is off by default; opt in via `MANAGED_SPARK_CONNECT_INSTRUMENTATION_LOG` environment variables. See [docs/instrumentation.md](https://github.com/GoogleCloudDataproc/managed-spark-connect-python/blob/main/docs/instrumentation.md) -for usage, the API, the record schema, and the span catalogue. +for usage, the API, the record schema, and the span and event catalogues. ### Interactive Extras diff --git a/docs/instrumentation.md b/docs/instrumentation.md index 405a4fe3..16ad6027 100644 --- a/docs/instrumentation.md +++ b/docs/instrumentation.md @@ -71,12 +71,16 @@ JsonlFileSink(path, *, append=True) InMemorySink() LoggingSink(logger=None, # timed_guard(callback), install_client_wrappers(client, recorder) ``` -`wrap_client=False` leaves the pyspark client untouched (no `client.*` spans). +`wrap_client=False` installs no instrumentation wrappers on the pyspark client +(no `client.*` spans). The session's own client hooks are installed either way: +they set operation ids, display operation links, close a query's progress +bar when it ends and re-run a read query once on +`INVALID_CURSOR.POSITION_NOT_AVAILABLE` (see `reattach.position_lost` below). ## Record schema Every record is a flat JSON object with a `type`. Global attributes are -merged into `attrs`. +merged into `attrs`. The span and event catalogues are below. ```json {"type":"span","name":"grpc.ExecutePlan","span_id":"16 hex","parent_id":"16 hex | null", @@ -103,10 +107,10 @@ meta record carry them as of when they are written. | Span | Attributes | |---|---| -| `session.get_or_create` → `session.create_rpc`, `session.wait_active`, `session.connect` (create) or `session.reuse` → `dataproc.get_session`, `session.wait_endpoint`, `session.connect` (reuse) | `mode=create\|reuse`, `session_id`, `uuid`, `runtime_version`, `template`; `session.wait_active`/`session.wait_endpoint` add `poll_interval_s`, `create_time`, `state_time`, `state`, `state_history` (JSON string); `session.reuse` adds `found` | +| `session.get_or_create` → `session.create_rpc`, `session.wait_active`, `session.connect` (create) or `session.reuse` → `dataproc.get_session`, `session.wait_endpoint`, `session.connect` (reuse) | `mode=create\|reuse`, `session_id`, `uuid`, `runtime_version`, `template`; `session.wait_active`/`session.wait_endpoint` add `poll_interval_s`, `create_time`, `state_time`, `state`, `state_history` (JSON string); `session.reuse` adds `found`, and `checked=false` when `GetSession` failed without a verdict and the session already in the process is reused unchecked | | `session.stop` → `session.terminate_rpc`, `session.wait_terminating` | `terminate=true\|false`, `session_id`; `session.wait_terminating` adds `state` | | `dataproc.get_session` | `reason=rpc_guard\|reuse\|poll\|stop`, `new_client` (whether a fresh `SessionControllerClient` was built for the call) | -| `grpc.` (e.g. `grpc.ExecutePlan`), one `layer=outer` span (what the client waits for, including the liveness guard) wrapping one `layer=inner` span (pure wire time to the proxy) | `method`, `kind=unary_unary\|unary_stream\|stream_unary\|stream_stream`, `layer`, `session_id`, `operation_id`, `request_bytes`, `request_count` (client streaming), `response_count`, `response_bytes`, `grpc_code`, `guard_ns` (outer only); events `first_response`, `status`. Byte counts are protobuf `ByteSize()` of the decoded messages, not wire bytes (the proxy's websocket framing roughly doubles those); a streaming span counts the messages delivered to the consumer, freezes its `duration_ns` when the final gRPC status arrives, and emits the record when the consumer is done with the stream (iteration finished, error, `cancel()`, or the iterator being dropped — pyspark drops an `ExecutePlan` stream at `result_complete` without draining it; that last case relies on CPython's reference counting to be prompt) | +| `grpc.` (e.g. `grpc.ExecutePlan`), one `layer=outer` span (what the client waits for, including the liveness guard) wrapping one `layer=inner` span (pure wire time to the proxy) | `method`, `kind=unary_unary\|unary_stream\|stream_unary\|stream_stream`, `layer`, `session_id`, `operation_id`, `plan_kind=root\|command` and `command_kind` (the command's oneof name, e.g. `sql_command`, `create_dataframe_view`) on `ExecutePlan`, `request_bytes`, `request_count` (client streaming), `response_count`, `response_bytes`, `grpc_code`, `guard_ns` (outer only); events `first_response`, `status`. Byte counts are protobuf `ByteSize()` of the decoded messages, not wire bytes (the proxy's websocket framing roughly doubles those); a streaming span counts the messages delivered to the consumer, freezes its `duration_ns` when the final gRPC status arrives, and emits the record when the consumer is done with the stream (iteration finished, error, `cancel()`, or the iterator being dropped — pyspark drops an `ExecutePlan` stream at `result_complete` without draining it; that last case relies on CPython's reference counting to be prompt) | | `proxy.start` | `port` | | `proxy.connection` (accept to close, on the forwarding thread) → `proxy.ws_connect` → `auth.default`, `auth.refresh`, `ws.handshake` | `conn`, `bytes_in`, `bytes_out`, `frames_in`, `frames_out` (wire side, i.e. the raw TCP bytes forwarded); `proxy.ws_connect` has `conn`, `host` | | `client.plan_to_proto` | `plan_bytes` | @@ -116,3 +120,38 @@ meta record carry them as of when they are written. The outer `grpc.*` span is also the parent of the `dataproc.get_session` span the liveness guard issues, so for an RPC `outer − inner − guard_ns` is the client's own overhead and `inner` is time on the wire plus server time. + +## Events + +Point-in-time records (`{"type": "event", "name": ..., "wall_ns": ..., +"mono_ns": ..., "thread": ..., "attrs": {...}}`) for things that have no +duration of their own. Like spans they are only produced by instrumented +sessions; a session without an `Instrumentation` pays nothing for them. + +| Event | When | Attributes | +|---|---|---| +| `proxy.eof` | A forwarding direction of a proxy connection saw its source end close: the gRPC-side socket's `recv` returned empty (`client_to_server`), or the websocket raised `ConnectionClosed` (`server_to_client`); both ends are then closed | `conn`, `direction=client_to_server\|server_to_client`, `bytes_forwarded` (this direction), `age_ns` (since the connection was accepted); for the websocket case also `error_type` (`ConnectionClosedOK`/`ConnectionClosedError`) and `close_code` | +| `proxy.socket_timeout` | Sending a frame timed out and is retried (up to 10 times, then `proxy.forward_error` with `error_type=SendTimeout`). Recv timeouts are routine polls of an idle connection and emit nothing | `conn`, `direction`, `timeout_s` (of the gRPC-side TCP socket; `null` for the websocket side) | +| `proxy.forward_error` | Any other exception ended a forwarding direction (both ends are then closed) | `conn`, `direction`, `error_type`, `message` (first 200 chars) | +| `proxy.silence_cut` | With an RPC in flight, the websocket received nothing for `MANAGED_SPARK_CONNECT_WS_SILENCE_LIMIT_S` seconds, so the proxy closed both ends of the connection to unwedge the stream. The check runs on each 60 s websocket recv poll, so the cut comes between the limit and the limit plus 60 s | `conn`, `silent_s` (seconds without a received frame since the later of the last frame and the RPC being seen in flight) | +| `proxy.ws_connect_failed` | Opening the websocket for a new proxy connection failed (credentials or handshake); the `proxy.ws_connect` span records the exception too | `conn`, `error_type`, `http_status` (when the handshake was rejected with an HTTP status, else `null`) | +| `auth.token_refreshed` | The Application Default Credentials were actually refreshed while opening a websocket (inside the `auth.refresh` span) | `expiry_wall_ns` (epoch ns of the new token's expiry, `null` if unknown), `had_token` (a token existed before the refresh) | +| `progress.bar_open` | The execution progress handler created a progress bar for an operation (first progress message with tasks that is not already final) | `operation_id`, `total_tasks` | +| `progress.bar_update` | Every later progress message for an open bar, except the final one | `operation_id`, `completed_tasks`, `total_tasks` | +| `progress.bar_close` | The bar was closed and unregistered: by the handler on the final (`done=True`) message, or, if that never came, when the query's `_execute_and_fetch` returned or raised | `operation_id`, `completed_tasks`, `total_tasks`, `reason=done\|query_end\|replay` (`replay`: the query hit `INVALID_CURSOR.POSITION_NOT_AVAILABLE` and is re-run under a new operation id) | +| `guard.result` | The session liveness guard ran (before an RPC, or after a failed one) and either had to ask the control plane or found the session inactive. A cached "active" verdict, the normal case for every RPC, emits nothing | `mode=pre_rpc\|on_error` (`direct` when called outside the channel), `cached` (answered from the liveness cache, no `GetSession` RPC), `active` (the verdict; `null` when the check raised), `error_type` (only when the check raised) | +| `grpc.retry` | pyspark's retry loop is about to retry an RPC, after the policy's backoff. The failed RPC is the `grpc.*` span that ended just before, on the same `thread` | `attempt` (1 = first retry within that loop), `reason=error\|reattach` (`reattach`: the server ended the stream and pyspark reattaches), `grpc_code`, `error_type` | +| `reattach.position_lost` | A query's `_execute_and_fetch` (`collect`, `toPandas`, `count`, ...) failed with `INVALID_CURSOR.POSITION_NOT_AVAILABLE`: the server no longer buffers the response the client has to reattach from. A replay also logs a `WARNING` | `operation_id`, `replay` (whether the query is re-run once under a new operation id: only root (read) plans, and not with `MANAGED_SPARK_CONNECT_REPLAY_ON_POSITION_LOST=0`; if the re-run fails the same way, a second event has its `operation_id` and `replay=false`), `new_operation_id` (only when `replay=true`) | + +Notes: `progress.*` events require the progress handler to be active, +which it only is in interactive environments. + +## Progress bars + +The notebook progress bar is kept in a process-wide registry keyed by +operation id. `ManagedSparkSession.progress_bars()` returns a snapshot of it, +for tests and diagnostics: + +```python +ManagedSparkSession.progress_bars() -> dict[str, Any] # operation_id -> open tqdm bar (a snapshot copy) +``` diff --git a/google/cloud/managed_spark_connect/client/core.py b/google/cloud/managed_spark_connect/client/core.py index 38eea697..b06a9413 100644 --- a/google/cloud/managed_spark_connect/client/core.py +++ b/google/cloud/managed_spark_connect/client/core.py @@ -160,10 +160,22 @@ class _GuardedStream(_StreamingCall): _call: Any = None _finished = True - def __init__(self, call: Any, on_error: Callable[[BaseException], None]): + def __init__( + self, call: Any, on_error: Callable[[BaseException], None], proxy + ): self._call = call self._on_error = on_error self._span = getattr(call, "_span", None) + self._proxy = proxy + self._lock = threading.Lock() + self._dec_done = False + + def _dec(self): + with self._lock: + if self._dec_done: + return + self._dec_done = True + self._proxy.dec_inflight() def __repr__(self) -> str: return f"_GuardedStream({self._call!r})" @@ -174,12 +186,17 @@ def __iter__(self) -> Iterator[Any]: def __next__(self) -> Any: try: return next(self._call) + except StopIteration: + self._dec() + raise except BaseException as exc: - if _should_check_liveness_on_error(exc): + self._dec() + if self._on_error and _should_check_liveness_on_error(exc): self._on_error(exc) raise def _finish(self, exc: Optional[BaseException]) -> None: + self._dec() inner = self._call if isinstance(inner, _StreamingCall): inner._finish(exc) @@ -191,24 +208,37 @@ def __del__(self) -> None: pass def cancel(self): + self._dec() return self._call.cancel() class _GuardedFuture(grpc.Future): """Wraps a gRPC Future to check session liveness when the call fails.""" - def __init__(self, future: Any, on_error: Callable[[BaseException], None]): + def __init__( + self, future: Any, on_error: Callable[[BaseException], None], proxy + ): self._future = future self._on_error = on_error + self._proxy = proxy self._lock = threading.Lock() self._error_checked = False self._stored_guard_exc: Optional[BaseException] = None + self._dec_done = False + self._future.add_done_callback(self._dec) + + def _dec(self, _fut=None): + with self._lock: + if self._dec_done: + return + self._dec_done = True + self._proxy.dec_inflight() def _check_error(self, exc: BaseException) -> None: with self._lock: if not self._error_checked: self._error_checked = True - if _should_check_liveness_on_error(exc): + if self._on_error and _should_check_liveness_on_error(exc): try: self._on_error(exc) except BaseException as guard_exc: @@ -218,6 +248,7 @@ def _check_error(self, exc: BaseException) -> None: raise self._stored_guard_exc from exc def cancel(self): + self._dec() return self._future.cancel() def cancelled(self): @@ -268,11 +299,14 @@ class _GuardedMultiCallable: replacement for the channel it wraps. """ - def __init__(self, wrapped_method, guard): + def __init__(self, wrapped_method, guard, proxy): self._wrapped_method = wrapped_method self._guard = guard + self._proxy = proxy def _run_pre_guard(self) -> None: + if self._guard is None: + return token = _GUARD_MODE.set("pre_rpc") try: self._guard() @@ -280,6 +314,8 @@ def _run_pre_guard(self) -> None: _GUARD_MODE.reset(token) def _run_error_guard(self, _exc: BaseException) -> None: + if self._guard is None: + return token = _GUARD_MODE.set("on_error") try: self._guard() @@ -288,34 +324,70 @@ def _run_error_guard(self, _exc: BaseException) -> None: def __call__(self, *args, **kwargs): self._run_pre_guard() + self._proxy.inc_inflight() try: response = self._wrapped_method(*args, **kwargs) except Exception as exc: - if _should_check_liveness_on_error(exc): + self._proxy.dec_inflight() + if self._guard and _should_check_liveness_on_error(exc): self._run_error_guard(exc) raise if hasattr(response, "__next__"): - return _GuardedStream(response, self._run_error_guard) + return _GuardedStream(response, self._run_error_guard, self._proxy) + self._proxy.dec_inflight() return response def with_call(self, *args, **kwargs): self._run_pre_guard() + self._proxy.inc_inflight() try: - return self._wrapped_method.with_call(*args, **kwargs) + response = self._wrapped_method.with_call(*args, **kwargs) except Exception as exc: - if _should_check_liveness_on_error(exc): + self._proxy.dec_inflight() + if self._guard and _should_check_liveness_on_error(exc): self._run_error_guard(exc) raise + if hasattr(response[0], "__next__"): + return ( + _GuardedStream(response[0], self._run_error_guard, self._proxy), + response[1], + ) + self._proxy.dec_inflight() + return response def future(self, *args, **kwargs): self._run_pre_guard() + self._proxy.inc_inflight() try: fut = self._wrapped_method.future(*args, **kwargs) except Exception as exc: - if _should_check_liveness_on_error(exc): + self._proxy.dec_inflight() + if self._guard and _should_check_liveness_on_error(exc): self._run_error_guard(exc) raise - return _GuardedFuture(fut, self._run_error_guard) + return _GuardedFuture(fut, self._run_error_guard, self._proxy) + + +class _OutageLog: + """Logs the first failure of an outage at WARNING and the rest at DEBUG. + + ``success()`` ends the outage, so the next failure warns again. Used for + control-plane errors that would otherwise log on every RPC. Best effort: + there is no lock, so two threads failing at the same moment can both log + at WARNING. + """ + + def __init__(self, log: logging.Logger): + self._log = log + self._failing = False + + def failure(self, msg: str, *args: Any) -> None: + level = logging.DEBUG if self._failing else logging.WARNING + self._failing = True + self._log.log(level, msg, *args) + + def success(self) -> None: + self._failing = False class ProxiedChannel(grpc.Channel): @@ -337,6 +409,7 @@ def __init__(self, target_host, is_active_callback, instrumentation=None): is_active_callback ) self._is_active_callback = is_active_callback + self._guard_errors = _OutageLog(logger) self._proxy = proxy.ManagedSparkSessionProxy( 0, target_host, recorder=recorder ) @@ -364,24 +437,42 @@ def close(self): def _check_active(self): if self._is_active_callback is None: return - if _GUARD_MODE.get() == "on_error" and not self._supports_error_guard: + mode = _GUARD_MODE.get() + if mode == "on_error" and not self._supports_error_guard: + return + try: + active = self._is_active_callback() + except Exception as exc: + # Only NotFound or a terminal state is a verdict; anything else + # (a control-plane outage, an auth error) fails open. + self._guard_errors.failure( + "Error checking session liveness; %s: %s", + ( + "re-raising the original RPC error" + if mode == "on_error" + else "proceeding with the RPC" + ), + exc, + ) return - if not self._is_active_callback(): + self._guard_errors.success() + if not active: logger.warning(f"Session is no longer active") raise RuntimeError( "Session not active. Please create a new session" ) def _wrap_method(self, wrapped_method, method_name=None): + guard = self._check_active if self._is_active_callback is None: - return wrapped_method + guard = None if isinstance(method_name, bytes): method_name = method_name.decode("utf-8", errors="ignore") if isinstance(method_name, str) and method_name.endswith( _UNGUARDED_METHOD_SUFFIXES ): - return wrapped_method - return _GuardedMultiCallable(wrapped_method, self._check_active) + guard = None + return _GuardedMultiCallable(wrapped_method, guard, self._proxy) def stream_stream(self, *args, **kwargs): method_name = args[0] if args else kwargs.get("method") diff --git a/google/cloud/managed_spark_connect/client/proxy.py b/google/cloud/managed_spark_connect/client/proxy.py index 6e4f3da7..35361ae4 100755 --- a/google/cloud/managed_spark_connect/client/proxy.py +++ b/google/cloud/managed_spark_connect/client/proxy.py @@ -14,18 +14,22 @@ # See the License for the specific language governing permissions and # limitations under the License. import argparse +import calendar import contextlib +import datetime import logging import os import socket import threading import time +import websockets.exceptions import websockets.sync.client as websocketclient from google import auth as googleauth from google.auth.transport import requests as googleauthrequests +from google.cloud.managed_spark_connect import environment from google.cloud.managed_spark_connect.instrumentation import NOOP_RECORDER parser = argparse.ArgumentParser() @@ -41,6 +45,19 @@ RECV_BUFFER_SIZE = 65536 MAX_WS_MESSAGE_SIZE = 64 * 1024 * 1024 +# How long the gRPC-side socket and the websocket may block in one call. +SOCKET_TIMEOUT_S = 10 +WS_RECV_TIMEOUT_S = 60 + +CLIENT_TO_SERVER = "client_to_server" +SERVER_TO_CLIENT = "server_to_client" + +_MAX_EVENT_MESSAGE_CHARS = 200 + +# Cut a connection whose server stays silent this long while an RPC is in +# flight (see forward_bytes). 0 disables. +_SILENCE_LIMIT_ENV = "MANAGED_SPARK_CONNECT_WS_SILENCE_LIMIT_S" +_DEFAULT_SILENCE_LIMIT_S = 300.0 def clear_credentials_cache(): @@ -74,11 +91,31 @@ def get_default_credentials(scopes=None): return cached +def _expiry_wall_ns(creds): + """``creds.expiry`` (a naive UTC datetime in google-auth) as epoch ns.""" + expiry = getattr(creds, "expiry", None) + if not isinstance(expiry, datetime.datetime): + return None + # Integer arithmetic: a float timestamp loses sub-microsecond precision. + # utctimetuple() treats a naive datetime as UTC and converts an aware one. + try: + seconds = calendar.timegm(expiry.utctimetuple()) + except (OSError, OverflowError, ValueError): + return None + return seconds * 10**9 + expiry.microsecond * 1000 + + def get_refreshed_credentials(recorder=NOOP_RECORDER, scopes=None): - """Obtain and refresh Application Default Credentials for bearer-token auth.""" + """Obtain and refresh Application Default Credentials for bearer-token auth. + + With instrumentation on, an ``auth.token_refreshed`` event (attrs + ``expiry_wall_ns`` of the new token and ``had_token``, whether a token + existed before) is emitted whenever a refresh actually happened. + """ with recorder.span("auth.default"): creds, _ = get_default_credentials(scopes=scopes) with recorder.span("auth.refresh"): + refreshed_attrs = None with _AUTH_LOCK: is_valid_real_creds = ( isinstance(getattr(creds, "valid", None), bool) @@ -86,11 +123,23 @@ def get_refreshed_credentials(recorder=NOOP_RECORDER, scopes=None): and bool(getattr(creds, "token", None)) ) if not is_valid_real_creds: + had_token = bool(getattr(creds, "token", None)) try: creds.refresh(googleauthrequests.Request()) except Exception: _CREDENTIALS_CACHE.clear() raise + if recorder.enabled: + refreshed_attrs = ( + _expiry_wall_ns(creds), + had_token, + ) + if refreshed_attrs is not None: + recorder.event( + "auth.token_refreshed", + expiry_wall_ns=refreshed_attrs[0], + had_token=refreshed_attrs[1], + ) return creds @@ -98,10 +147,20 @@ class bridged_socket(object): """Socket-like object that uses a websocket-over-TCP Bridge transport. See: https://github.com/google/inverting-proxy/tree/master/utils/tcpbridge + + Args: + websocket_conn: The connected WebSocket. + inflight_rpcs: Optional zero-argument callable returning how many + RPCs are in flight through the proxy. When set, forward_bytes + cuts the connection if the server stays silent too long while + that count is positive (see + ``MANAGED_SPARK_CONNECT_WS_SILENCE_LIMIT_S``). Wrappers around + this socket that delegate attribute access keep the cut armed. """ - def __init__(self, websocket_conn): + def __init__(self, websocket_conn, inflight_rpcs=None): self._conn = websocket_conn + self.inflight_rpcs = inflight_rpcs def recv(self, buff_size): # N.B. The websockets [recv method](https://websockets.readthedocs.io/en/stable/reference/sync/client.html#websockets.sync.client.ClientConnection.recv) @@ -110,7 +169,7 @@ def recv(self, buff_size): # # We set that timeout to 60 seconds to prevent any scenarios where we wind up stuck waiting for a message from a websocket connection # that never comes. - msg = self._conn.recv(timeout=60) + msg = self._conn.recv(timeout=WS_RECV_TIMEOUT_S) return bytes.fromhex(msg) def send(self, msg_bytes): @@ -122,6 +181,42 @@ def close(self): return self._conn.close() +class SendTimeout(Exception): + """Sending one frame timed out on every retry.""" + + +def _http_status_of(exc): + """The HTTP status of a failed websocket handshake, if the error has one.""" + status = getattr(getattr(exc, "response", None), "status_code", None) + return status if isinstance(status, int) else None + + +def _shutdown_and_close(*socks): + """Shut down both directions of each socket (if supported), then close it. + + The shutdown wakes a peer thread blocked in ``recv()`` on the same socket + immediately, instead of leaving it to wait for its socket timeout. + """ + for sock in socks: + shutdown = getattr(sock, "shutdown", None) + if callable(shutdown): + try: + shutdown(socket.SHUT_RDWR) + except Exception: + pass + try: + sock.close() + except Exception: + pass + + +def _short_message(exc): + message = str(exc) + if len(message) > _MAX_EVENT_MESSAGE_CHARS: + return message[:_MAX_EVENT_MESSAGE_CHARS] + return message + + def connect_tcp_bridge(hostname, recorder=NOOP_RECORDER, conn_number=None): """Create a socket-like connection to the given hostname using websocket. @@ -130,7 +225,8 @@ def connect_tcp_bridge(hostname, recorder=NOOP_RECORDER, conn_number=None): Args: hostname: The hostname of the server running the TCP-bridge backend. - recorder: Instrumentation recorder for the ``proxy.ws_connect`` span. + recorder: Instrumentation recorder for the ``proxy.ws_connect`` span + (and a ``proxy.ws_connect_failed`` event when it fails). conn_number: The proxy connection this websocket serves, if any. Returns: @@ -138,16 +234,44 @@ def connect_tcp_bridge(hostname, recorder=NOOP_RECORDER, conn_number=None): """ path = "tcp-over-websocket-bridge/35218cb7-1201-4940-89e8-48d8f03fed96" with recorder.span("proxy.ws_connect", conn=conn_number, host=hostname): - creds = get_refreshed_credentials(recorder=recorder) - - with recorder.span("ws.handshake"): - return websocketclient.connect( - f"wss://{hostname}/{path}", - additional_headers={"Authorization": f"Bearer {creds.token}"}, - open_timeout=30, - compression=None, - max_size=MAX_WS_MESSAGE_SIZE, - ) + try: + creds = get_refreshed_credentials(recorder=recorder) + + with recorder.span("ws.handshake"): + return websocketclient.connect( + f"wss://{hostname}/{path}", + additional_headers={ + "Authorization": f"Bearer {creds.token}" + }, + open_timeout=30, + # Don't wait the default 10 s for a Close handshake from a + # backend that has already broken or reset. + close_timeout=3, + compression=None, + max_size=MAX_WS_MESSAGE_SIZE, + # Unbounded in websockets so a slow gRPC consumer does not + # pause the WS reader thread and starve keepalive PONGs; + # memory remains bounded by HTTP/2 flow-control windows. + # Trade-off: more of a result can sit between the server + # and pyspark, and the server may already have trimmed it + # from its reattach buffer if the stream then breaks + # (see client/reattach_patch.py). + max_queue=None, + # 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 (20 s) + ping_timeout. + ping_timeout=60, + ) + except Exception as exc: + if recorder.enabled: + recorder.event( + "proxy.ws_connect_failed", + conn=conn_number, + error_type=type(exc).__name__, + http_status=_http_status_of(exc), + ) + raise class _DirectionStats: @@ -164,13 +288,46 @@ def __init__(self): self.frames = 0 -def forward_bytes(name, from_sock, to_sock, stats=None): +def _emit_eof(recorder, conn, direction, stats, start_mono_ns, **extra): + if not recorder.enabled: + return + recorder.event( + "proxy.eof", + conn=conn, + direction=direction, + bytes_forwarded=stats.bytes if stats else 0, + age_ns=( + None + if start_mono_ns is None + else time.perf_counter_ns() - start_mono_ns + ), + **extra, + ) + + +def forward_bytes( + name, + from_sock, + to_sock, + stats=None, + recorder=NOOP_RECORDER, + conn=None, + direction=None, + start_mono_ns=None, +): """Continuously stream bytes from the `from_sock` to the `to_sock`. This method terminates when either the `from_sock` is closed (causing it to return a Falsy value from its `recv` method), or the first time it hits an exception. + In the ``server_to_client`` direction, if `from_sock` has a callable + ``inflight_rpcs`` attribute (see :class:`bridged_socket`), a positive + count arms the silence cut: if nothing has been received for + ``MANAGED_SPARK_CONNECT_WS_SILENCE_LIMIT_S`` seconds while an RPC is in + flight, both sockets are closed. The check runs on each receive timeout, + so the cut happens up to ``WS_RECV_TIMEOUT_S`` after the limit. + This method is intended to be run in a separate thread of execution. Args: @@ -178,12 +335,32 @@ def forward_bytes(name, from_sock, to_sock, stats=None): from_sock: A socket-like object to stream bytes from. to_sock: A socket-like object to stream bytes to. stats: Optional :class:`_DirectionStats` to count what was forwarded. + recorder: Instrumentation recorder for the ``proxy.eof``, + ``proxy.socket_timeout``, ``proxy.forward_error`` and + ``proxy.silence_cut`` events. + conn: The proxy connection number, for the events. + direction: ``client_to_server`` or ``server_to_client``. + start_mono_ns: ``perf_counter_ns`` at which the connection was + accepted, for ``age_ns`` on ``proxy.eof``. """ + silence_limit_s = environment.env_float( + _SILENCE_LIMIT_ENV, default=_DEFAULT_SILENCE_LIMIT_S, min_value=0.0 + ) + inflight_rpcs = None + if direction == SERVER_TO_CLIENT and silence_limit_s > 0: + inflight_rpcs = getattr(from_sock, "inflight_rpcs", None) + if not callable(inflight_rpcs): + inflight_rpcs = None + last_rx = time.monotonic() + inflight_since = None + while True: try: bs = from_sock.recv(RECV_BUFFER_SIZE) + last_rx = time.monotonic() if not bs: - to_sock.close() + _emit_eof(recorder, conn, direction, stats, start_mono_ns) + _shutdown_and_close(to_sock, from_sock) return frame_len = len(bs) attempt = 0 @@ -202,30 +379,80 @@ def forward_bytes(name, from_sock, to_sock, stats=None): bs = None except TimeoutError: # On timeouts during a send, we retry just the send - # to make sure we don't lose any bytes. - pass + # to make sure we don't lose any bytes. (Only the gRPC-side + # TCP socket has a send timeout; a websocket send has none.) + if recorder.enabled: + recorder.event( + "proxy.socket_timeout", + conn=conn, + direction=direction, + timeout_s=( + to_sock.gettimeout() + if isinstance(to_sock, socket.socket) + else None + ), + ) if bs: - raise Exception(f"Failed to forward bytes for {name}") + raise SendTimeout(f"Failed to forward bytes for {name}") if stats is not None: stats.bytes += frame_len stats.frames += 1 except TimeoutError: - # On timeouts during a receive, we retry the entire flow. - pass + # On timeouts during a receive, we retry the entire flow. These + # are routine polls of an idle connection, so no event, unless + # the server has been silent too long with an RPC in flight. + if inflight_rpcs is not None: + if inflight_rpcs() > 0: + now = time.monotonic() + if inflight_since is None: + inflight_since = now + # Idle time from before the RPC started doesn't count. + silent_s = now - max(last_rx, inflight_since) + if silent_s >= silence_limit_s: + logger.warning( + f"[{name}] Silence limit exceeded" + f" ({silent_s:.1f}s >= {silence_limit_s}s)" + ) + if recorder.enabled: + recorder.event( + "proxy.silence_cut", + conn=conn, + silent_s=silent_s, + ) + _shutdown_and_close(to_sock, from_sock) + return + else: + inflight_since = None except Exception as ex: logger.debug(f"[{name}] Exception forwarding bytes: {ex}") - try: - to_sock.close() - except Exception: - pass - try: - from_sock.close() - except Exception: - pass + if recorder.enabled: + if isinstance(ex, websockets.exceptions.ConnectionClosed): + # The websocket's way of saying the remote end closed. + rcvd = getattr(ex, "rcvd", None) + _emit_eof( + recorder, + conn, + direction, + stats, + start_mono_ns, + error_type=type(ex).__name__, + close_code=getattr(rcvd, "code", None), + ) + else: + recorder.event( + "proxy.forward_error", + conn=conn, + direction=direction, + error_type=type(ex).__name__, + message=_short_message(ex), + ) + _shutdown_and_close(to_sock, from_sock) return -def connect_sockets(conn_number, from_sock, to_sock, span=None): +def connect_sockets( + conn_number, from_sock, to_sock, span=None, recorder=NOOP_RECORDER +): """Create a connection between the two given ports. This method continuously streams bytes in both directions between the @@ -234,15 +461,26 @@ def connect_sockets(conn_number, from_sock, to_sock, span=None): The caller is responsible for creating and closing the supplied sockets. When ``span`` is given, ``bytes_out``/``frames_out`` (``from_sock`` to ``to_sock``) and ``bytes_in``/``frames_in`` (the other way) are set on - it once both directions have finished. + it once both directions have finished. ``recorder`` receives the + per-direction events of :func:`forward_bytes`. """ - out_stats = _DirectionStats() if span is not None else None - in_stats = _DirectionStats() if span is not None else None + track_stats = span is not None or recorder.enabled + out_stats = _DirectionStats() if track_stats else None + in_stats = _DirectionStats() if track_stats else None + start_mono_ns = getattr(span, "start_mono_ns", None) + if start_mono_ns is None and track_stats: + start_mono_ns = time.perf_counter_ns() forward_name = f"{conn_number}-forward" t1 = threading.Thread( name=forward_name, target=forward_bytes, args=[forward_name, from_sock, to_sock, out_stats], + kwargs=dict( + recorder=recorder, + conn=conn_number, + direction=CLIENT_TO_SERVER, + start_mono_ns=start_mono_ns, + ), daemon=True, ) t1.start() @@ -251,6 +489,12 @@ def connect_sockets(conn_number, from_sock, to_sock, span=None): name=backward_name, target=forward_bytes, args=[backward_name, to_sock, from_sock, in_stats], + kwargs=dict( + recorder=recorder, + conn=conn_number, + direction=SERVER_TO_CLIENT, + start_mono_ns=start_mono_ns, + ), daemon=True, ) t2.start() @@ -272,6 +516,7 @@ def forward_connection( target_host, recorder=NOOP_RECORDER, accepted_ns=None, + proxy_instance=None, ): """Create a connection to the target and forward `conn` to it. @@ -296,7 +541,12 @@ def forward_connection( accepted_ns: ``(perf_counter_ns, time_ns)`` of the accept, so the span starts when the client connected rather than when this thread got scheduled. + proxy_instance: Optional :class:`ManagedSparkSessionProxy` whose + in-flight RPC count arms the backend socket's silence cut. """ + inflight_rpcs = ( + proxy_instance.get_inflight_rpcs if proxy_instance is not None else None + ) if accepted_ns is None: span = recorder.start_span("proxy.connection", conn=conn_number) else: @@ -312,17 +562,20 @@ def forward_connection( with connect_tcp_bridge( target_host, recorder, conn_number ) as websocket_conn: - backend_socket = bridged_socket(websocket_conn) + backend_socket = bridged_socket( + websocket_conn, inflight_rpcs=inflight_rpcs + ) # Set a timeout on how long we will allow send/recv calls to block # # The code that reads and writes to this connection will retry # on timeouts, so this is a safe change. - conn.settimeout(10) + conn.settimeout(SOCKET_TIMEOUT_S) connect_sockets( conn_number, conn, backend_socket, span if recorder.enabled else None, + recorder, ) @@ -344,6 +597,11 @@ def __init__(self, port, target_host, recorder=None): self._started = False self._killed = False self._conn_number = 0 + # RPCs in flight through this proxy, maintained by core's guarded + # callables via inc_inflight()/dec_inflight(). Backend sockets read + # it through get_inflight_rpcs() to arm the silence cut. + self._inflight_rpcs = 0 + self._inflight_lock = threading.Lock() self._frontend_socket = None @property @@ -351,6 +609,20 @@ def port(self): """The local port the proxy is listening on""" return self._port + def get_inflight_rpcs(self): + """How many RPCs are in flight through this proxy.""" + return self._inflight_rpcs + + def inc_inflight(self): + """Count an RPC that has started.""" + with self._inflight_lock: + self._inflight_rpcs += 1 + + def dec_inflight(self): + """Count an RPC that has finished.""" + with self._inflight_lock: + self._inflight_rpcs -= 1 + def start(self, daemon=True): """Start the proxy. @@ -394,6 +666,7 @@ def _run(self, s): self._target_host, self._recorder, accepted_ns, + self, ], daemon=True, ).start() diff --git a/google/cloud/managed_spark_connect/client/reattach_patch.py b/google/cloud/managed_spark_connect/client/reattach_patch.py new file mode 100644 index 00000000..15eeb816 --- /dev/null +++ b/google/cloud/managed_spark_connect/client/reattach_patch.py @@ -0,0 +1,139 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. +"""Client-side recovery from ``INVALID_CURSOR.POSITION_NOT_AVAILABLE``. + +When a transport blip breaks an ``ExecutePlan`` stream, pyspark reattaches +from the last response it returned. If the server no longer buffers that +response, ``ReattachExecute`` fails with +``[INVALID_CURSOR.POSITION_NOT_AVAILABLE]``, which pyspark's +``DefaultPolicy`` does not retry, and the whole query fails. + +:func:`install_replay_on_position_lost` hooks ``client._execute_and_fetch`` +on one client instance (``collect``, ``toPandas``, ``count``, ...): when a +root (read) plan fails that way, the query is re-run once under a fresh +``operation_id``. Command plans are never re-run. Every such failure emits +``reattach.position_lost``, and a replay also logs a ``WARNING``. +``MANAGED_SPARK_CONNECT_REPLAY_ON_POSITION_LOST=0`` turns the replay off. +""" + +import logging +from typing import Any, Callable, Optional +import uuid + +from google.cloud.managed_spark_connect import environment +from google.cloud.managed_spark_connect.instrumentation import NOOP_RECORDER +from google.cloud.managed_spark_connect.instrumentation._client import ( + install_client_hooks, +) +from google.cloud.managed_spark_connect.instrumentation._grpc import _plan_kinds + +logger = logging.getLogger(__name__) + +ENV_REPLAY_ON_POSITION_LOST = "MANAGED_SPARK_CONNECT_REPLAY_ON_POSITION_LOST" +POSITION_LOST_CONDITION = "INVALID_CURSOR.POSITION_NOT_AVAILABLE" + + +def _is_position_not_available(exc: BaseException) -> bool: + cond = getattr(exc, "getCondition", None) or getattr( + exc, "getErrorClass", None + ) + return callable(cond) and cond() == POSITION_LOST_CONDITION + + +def _replay_request(req: Any) -> Any: + """A copy of ``req`` under a fresh ``operation_id``. + + pyspark's reattachable iterator appends ``reattach_options`` to the + request it is given; the copy drops them so the next attempt adds its + own. + """ + retry = type(req)() + retry.CopyFrom(req) + retry.operation_id = str(uuid.uuid4()) + keep = [ + opt + for opt in retry.request_options + if not opt.HasField("reattach_options") + ] + del retry.request_options[:] + retry.request_options.extend(keep) + return retry + + +def install_replay_on_position_lost( + client: Any, + recorder: Any = NOOP_RECORDER, + close_progress_bar: Optional[Callable[..., None]] = None, +) -> None: + """Re-run a root plan once when it fails with a lost reattach position. + + Hooks ``client._execute_and_fetch`` through + :func:`~google.cloud.managed_spark_connect.instrumentation._client.install_client_hooks` + (call it once per client; ``uninstall_client_hooks`` removes it). + ``close_progress_bar(operation_id, reason=...)``, if given, closes the + progress bar of the abandoned operation (``reason="replay"``) and, + after the replay, any bar the new operation left open. + """ + enabled = environment.env_bool(ENV_REPLAY_ON_POSITION_LOST, default=True) + + def replay_hook(inner): + def _execute_and_fetch(req, *args, **kwargs): + try: + return inner(req, *args, **kwargs) + except Exception as exc: + if not _is_position_not_available(exc): + raise + if not enabled or _plan_kinds(req)[0] != "root": + recorder.event( + "reattach.position_lost", + operation_id=req.operation_id, + replay=False, + ) + raise + retry = _replay_request(req) + recorder.event( + "reattach.position_lost", + operation_id=req.operation_id, + new_operation_id=retry.operation_id, + replay=True, + ) + logger.warning( + "Operation %s failed with %s; re-running the query once as" + " operation %s. Set %s=0 to disable.", + req.operation_id, + POSITION_LOST_CONDITION, + retry.operation_id, + ENV_REPLAY_ON_POSITION_LOST, + ) + if close_progress_bar is not None: + close_progress_bar(req.operation_id, reason="replay") + # Outside the except block, so a second failure is not chained + # to the first. + try: + return inner(retry, *args, **kwargs) + except Exception as exc: + if _is_position_not_available(exc): + recorder.event( + "reattach.position_lost", + operation_id=retry.operation_id, + replay=False, + ) + raise + finally: + if close_progress_bar is not None: + close_progress_bar(retry.operation_id) + + return _execute_and_fetch + + install_client_hooks(client, {"_execute_and_fetch": replay_hook}) diff --git a/google/cloud/managed_spark_connect/environment.py b/google/cloud/managed_spark_connect/environment.py index e19dd973..071cb493 100644 --- a/google/cloud/managed_spark_connect/environment.py +++ b/google/cloud/managed_spark_connect/environment.py @@ -12,9 +12,77 @@ # See the License for the specific language governing permissions and # limitations under the License. +import logging +import math import os import sys -from typing import Callable, Tuple, List +from typing import Callable, List, Optional, Tuple, overload + +logger = logging.getLogger(__name__) + +_TRUE_VALUES = frozenset({"1", "true", "yes", "on"}) +_FALSE_VALUES = frozenset({"0", "false", "no", "off"}) + + +def env_optional_bool(name: str) -> Optional[bool]: + """Parse an optional boolean environment variable (`1/true/yes/on` vs `0/false/no/off`).""" + raw = os.environ.get(name, "").strip().lower() + if not raw: + return None + if raw in _TRUE_VALUES: + return True + if raw in _FALSE_VALUES: + return False + return None + + +def env_bool(name: str, default: bool = False) -> bool: + """Parse a boolean environment variable. + + Returns ``default`` when it is unset or not a recognized boolean. + """ + value = env_optional_bool(name) + return default if value is None else value + + +@overload +def env_float( + name: str, default: float, *, min_value: Optional[float] = None +) -> float: + ... + + +@overload +def env_float( + name: str, default: None = None, *, min_value: Optional[float] = None +) -> Optional[float]: + ... + + +def env_float( + name: str, + default: Optional[float] = None, + *, + min_value: Optional[float] = None, +) -> Optional[float]: + """Parse a float environment variable, clamping to ``min_value`` if given. + + A malformed or non-finite value is ignored (with a warning) in favour of + ``default``. + """ + raw = os.environ.get(name, "").strip() + if not raw: + return default + try: + val = float(raw) + except ValueError: + val = math.nan + if not math.isfinite(val): + logger.warning("Ignoring invalid %s=%r", name, raw) + return default + if min_value is not None and val < min_value: + return min_value + return val def is_antigravity() -> bool: diff --git a/google/cloud/managed_spark_connect/instrumentation/__init__.py b/google/cloud/managed_spark_connect/instrumentation/__init__.py index 8a670108..771e0858 100644 --- a/google/cloud/managed_spark_connect/instrumentation/__init__.py +++ b/google/cloud/managed_spark_connect/instrumentation/__init__.py @@ -55,7 +55,8 @@ grpcio, pyarrow, Python and platform versions. Global attributes set on the :class:`Instrumentation` are merged into every record's ``attrs``. -See ``docs/instrumentation.md`` for the span catalogue. +See ``docs/instrumentation.md`` for the span and event catalogues (``proxy.*``, +``auth.token_refreshed``, ``progress.*``, ``guard.result``, ``grpc.retry``). """ from ._client import install_client_wrappers, uninstall_client_wrappers diff --git a/google/cloud/managed_spark_connect/instrumentation/_client.py b/google/cloud/managed_spark_connect/instrumentation/_client.py index bb4370d3..a109b85e 100644 --- a/google/cloud/managed_spark_connect/instrumentation/_client.py +++ b/google/cloud/managed_spark_connect/instrumentation/_client.py @@ -28,18 +28,29 @@ class level. That static method knows no client, so it records to the most recently instrumented client's recorder (a process has at most one managed session at a time). +* ``grpc.retry`` events come from a per-instance wrapper of the client's + ``_retrying`` factory: the ``Retrying`` it hands to pyspark's retry loops + (and to the reattachable ``ExecutePlan`` iterator, which captures the + factory when it is built) reports every retry it is about to make. The class-level hooks are installed once per process, on first use, and are inert while no instrumented client is registered. + +:func:`install_client_hooks` / :func:`uninstall_client_hooks` are the one +way per-instance client wrappers are installed and removed, both here and +by ``ManagedSparkSession`` for its own (always installed) hooks. """ import contextvars +import functools import threading -from typing import Any, List, Optional +from typing import Any, Callable, Dict, List, Mapping, Optional -from ._core import NOOP_RECORDER, Recorder +from ._core import NOOP_RECORDER, Recorder, error_info _RECORDER_ATTR = "_managed_spark_connect_recorder" +_ORIGINALS_ATTR = "_managed_spark_connect_hooked_methods" +_ABSENT = object() _lock = threading.Lock() _registered: List[Recorder] = [] @@ -53,6 +64,35 @@ class level. That static method knows no client, so it records to the ) +def install_client_hooks( + client: Any, hooks: Mapping[str, Callable[[Callable[..., Any]], Any]] +) -> None: + """Replace methods of one client instance with wrappers. + + ``hooks`` maps a method name to a factory that takes the current bound + method and returns its replacement. Calls stack: a later hook wraps an + earlier one. The first hook on each name records the original, and + :func:`uninstall_client_hooks` restores all of them at once. + """ + originals: Dict[str, Any] = client.__dict__.setdefault(_ORIGINALS_ATTR, {}) + for name, make_wrapper in hooks.items(): + originals.setdefault(name, client.__dict__.get(name, _ABSENT)) + inner = getattr(client, name) + wrapper = make_wrapper(inner) + wrapper.__wrapped__ = inner + setattr(client, name, wrapper) + + +def uninstall_client_hooks(client: Any) -> None: + """Restore every method :func:`install_client_hooks` replaced.""" + originals = client.__dict__.pop(_ORIGINALS_ATTR, {}) + for name, original in originals.items(): + if original is _ABSENT: + client.__dict__.pop(name, None) + else: + setattr(client, name, original) + + def _active_recorder() -> Recorder: with _lock: return _registered[-1] if _registered else NOOP_RECORDER @@ -62,6 +102,54 @@ def _recorder_of(client: Any) -> Recorder: return getattr(client, _RECORDER_ATTR, None) or NOOP_RECORDER +@functools.lru_cache(maxsize=None) +def _instrumented_retrying_class(): + """``Retrying`` subclass that emits ``grpc.retry`` before each retry. + + Built lazily (once) so importing this module does not import pyspark. + """ + from pyspark.sql.connect.client.retries import Retrying, RetryException + + class InstrumentedRetrying(Retrying): + """pyspark's retry loop, reporting each retry it decides to make. + + ``_wait`` is what pyspark calls between attempts: it either sleeps + for the policy's backoff (a retry follows) or gives up and raises + (``RetriesExceeded`` before pyspark 4.1, the last error itself from + 4.1 on). The event is emitted once the wait is over, + so ``attempt`` counts retries that actually happen (1 = first + retry). ``reason`` is ``reattach`` for pyspark's ``RetryException`` + (the server ended the stream and the iterator reattaches) and + ``error`` otherwise; ``grpc_code`` is the status of the exception + being retried. The failed RPC is the ``grpc.*`` span that ended just + before the event on the same thread. + """ + + def __init__(self, policies, *, recorder: Recorder, **kwargs): + super().__init__(policies, **kwargs) + self._recorder = recorder + self._retry_count = 0 + + def _wait(self) -> None: + exception = self._exception + super()._wait() + self._retry_count += 1 + info = error_info(exception) if exception is not None else {} + self._recorder.event( + "grpc.retry", + attempt=self._retry_count, + reason=( + "reattach" + if isinstance(exception, RetryException) + else "error" + ), + grpc_code=info.get("grpc_code"), + error_type=info.get("type"), + ) + + return InstrumentedRetrying + + def _install_class_hooks() -> None: global _hooks_installed with _lock: @@ -108,52 +196,72 @@ def install_client_wrappers(client: Any, recorder: Recorder) -> None: with _lock: _registered.append(recorder) - original_to_pandas = client.to_pandas - original_execute_and_fetch = client._execute_and_fetch - - def to_pandas(*args, **kwargs): - with recorder.span("client.to_pandas") as span: - state = [span, None] - token = _TO_PANDAS.set(state) - try: - return original_to_pandas(*args, **kwargs) - except BaseException as e: - if state[1] is not None: - state[1].record_exception(e) - raise - finally: - _TO_PANDAS.reset(token) - if state[1] is not None: - state[1].end() - - def _execute_and_fetch(*args, **kwargs): - result = original_execute_and_fetch(*args, **kwargs) - state = _TO_PANDAS.get() - if state is not None and state[1] is None: - table = result[0] if isinstance(result, tuple) else None - state[1] = recorder.start_span( - "client.arrow_to_pandas", - parent=state[0], - rows=getattr(table, "num_rows", None), - columns=getattr(table, "num_columns", None), - ) - return result + def to_pandas_hook(inner): + def to_pandas(*args, **kwargs): + with recorder.span("client.to_pandas") as span: + state = [span, None] + token = _TO_PANDAS.set(state) + try: + return inner(*args, **kwargs) + except BaseException as e: + if state[1] is not None: + state[1].record_exception(e) + raise + finally: + _TO_PANDAS.reset(token) + if state[1] is not None: + state[1].end() + + return to_pandas + + def execute_and_fetch_hook(inner): + def _execute_and_fetch(*args, **kwargs): + result = inner(*args, **kwargs) + state = _TO_PANDAS.get() + if state is not None and state[1] is None: + table = result[0] if isinstance(result, tuple) else None + state[1] = recorder.start_span( + "client.arrow_to_pandas", + parent=state[0], + rows=getattr(table, "num_rows", None), + columns=getattr(table, "num_columns", None), + ) + return result - to_pandas.__wrapped__ = original_to_pandas # type: ignore[attr-defined] - _execute_and_fetch.__wrapped__ = original_execute_and_fetch # type: ignore[attr-defined] - client.to_pandas = to_pandas - client._execute_and_fetch = _execute_and_fetch + return _execute_and_fetch + + install_client_hooks( + client, + { + "to_pandas": to_pandas_hook, + "_execute_and_fetch": execute_and_fetch_hook, + }, + ) + + policies = getattr(client, "_retry_policies", None) + original_retrying = getattr(client, "_retrying", None) + if policies is not None and callable(original_retrying): + retrying_cls = _instrumented_retrying_class() + + def retrying_hook(inner): + def _retrying(): + return retrying_cls(client._retry_policies, recorder=recorder) + + return _retrying + + install_client_hooks(client, {"_retrying": retrying_hook}) def uninstall_client_wrappers(client: Any) -> None: - """Undo :func:`install_client_wrappers` for ``client`` (if applied).""" + """Undo :func:`install_client_wrappers` for ``client`` (if applied). + + This restores every method hooked through :func:`install_client_hooks` + on ``client``, not only the instrumentation wrappers. + """ recorder = getattr(client, _RECORDER_ATTR, None) if recorder is None: return - for name in ("to_pandas", "_execute_and_fetch"): - wrapper = client.__dict__.get(name) - if wrapper is not None and hasattr(wrapper, "__wrapped__"): - del client.__dict__[name] + uninstall_client_hooks(client) delattr(client, _RECORDER_ATTR) with _lock: try: diff --git a/google/cloud/managed_spark_connect/instrumentation/_core.py b/google/cloud/managed_spark_connect/instrumentation/_core.py index 666fcac9..2b58aff8 100644 --- a/google/cloud/managed_spark_connect/instrumentation/_core.py +++ b/google/cloud/managed_spark_connect/instrumentation/_core.py @@ -485,8 +485,10 @@ class Instrumentation: client-side plan/result conversions. Records go to ``sinks``; ``attributes`` are attached to every record. - With ``wrap_client=False`` the pyspark client is left untouched: only - the session, proxy and gRPC spans are recorded. + With ``wrap_client=False`` the instrumentation installs no wrappers on + the pyspark client: only the session, proxy and gRPC spans are + recorded. The session's own client hooks (operation ids, operation + links, closing a query's progress bar) are installed either way. """ def __init__( diff --git a/google/cloud/managed_spark_connect/instrumentation/_grpc.py b/google/cloud/managed_spark_connect/instrumentation/_grpc.py index e1941788..f40a9945 100644 --- a/google/cloud/managed_spark_connect/instrumentation/_grpc.py +++ b/google/cloud/managed_spark_connect/instrumentation/_grpc.py @@ -36,7 +36,7 @@ import threading import time import weakref -from typing import Any, Callable, Dict, Iterator, Optional +from typing import Any, Callable, Dict, Iterator, Optional, Tuple import grpc @@ -83,12 +83,40 @@ def _byte_size(message: Any) -> int: return 0 +def _plan_kinds(request: Any) -> Tuple[Optional[str], Optional[str]]: + """``(plan_kind, command_kind)`` of an ``ExecutePlanRequest``. + + ``plan_kind`` is the ``plan.op_type`` oneof (``root`` or ``command``), + ``command_kind`` the ``command_type`` oneof of a command (e.g. + ``sql_command``). Either is ``None`` when unset; both are for any other + request type. + """ + from pyspark.sql.connect.proto import ExecutePlanRequest + + if not isinstance(request, ExecutePlanRequest): + return None, None + plan_kind = request.plan.WhichOneof("op_type") + if plan_kind != "command": + return plan_kind, None + return plan_kind, request.plan.command.WhichOneof("command_type") + + def _request_attrs(request: Any) -> Dict[str, Any]: attrs: Dict[str, Any] = {} for field in ("session_id", "operation_id"): value = getattr(request, field, None) if value: attrs[field] = value + # ExecutePlanRequest.plan is a oneof of root (a relation that yields + # rows, visible server-side as a SQL execution) or command (a side + # effect: temp view, cache, config, ...). Telling them apart lets a + # consumer join only root plans against server-side query listings; + # command_kind names the command (e.g. create_dataframe_view, sql_command). + plan_kind, command_kind = _plan_kinds(request) + if plan_kind is not None: + attrs["plan_kind"] = plan_kind + if command_kind is not None: + attrs["command_kind"] = command_kind return attrs diff --git a/google/cloud/managed_spark_connect/session.py b/google/cloud/managed_spark_connect/session.py index d38955cc..2adfa812 100644 --- a/google/cloud/managed_spark_connect/session.py +++ b/google/cloud/managed_spark_connect/session.py @@ -38,6 +38,7 @@ Dict, Iterable, List, + NamedTuple, Optional, Tuple, Union, @@ -48,24 +49,26 @@ from google.api_core.client_options import ClientOptions from google.api_core.exceptions import ( Aborted, - DeadlineExceeded, FailedPrecondition, - InternalServerError, InvalidArgument, NotFound, PermissionDenied, - ServiceUnavailable, - TooManyRequests, ) from google.api_core.future.polling import POLLING_PREDICATE from google.auth.exceptions import DefaultCredentialsError from google.cloud import dataproc_v1 from google.cloud.managed_spark_connect.client import ManagedSparkChannelBuilder -from google.cloud.managed_spark_connect.client.core import _GUARD_MODE +from google.cloud.managed_spark_connect.client.core import ( + _GUARD_MODE, + _OutageLog, +) from google.cloud.managed_spark_connect.client.proxy import ( clear_credentials_cache, get_default_credentials, ) +from google.cloud.managed_spark_connect.client.reattach_patch import ( + install_replay_on_position_lost, +) from google.cloud.managed_spark_connect.exceptions import ManagedSparkConnectException from google.cloud.managed_spark_connect.instrumentation import ( NOOP_RECORDER, @@ -74,6 +77,10 @@ install_client_wrappers, uninstall_client_wrappers, ) +from google.cloud.managed_spark_connect.instrumentation._client import ( + install_client_hooks, + uninstall_client_hooks, +) from google.cloud.managed_spark_connect.pypi_artifacts import PyPiArtifacts from google.cloud.dataproc_v1 import ( AuthenticationConfig, @@ -246,6 +253,14 @@ def fetch(): _CLIENT_CACHE_LOCK = threading.RLock() _CLIENT_CACHE: Dict[Tuple[Any, ...], SessionControllerClient] = {} +#: Set to a true value to run the execution progress handler outside a +#: notebook / interactive shell. A test hook, deliberately undocumented. +_FORCE_PROGRESS_BAR_ENV = "MANAGED_SPARK_CONNECT_FORCE_PROGRESS_BAR" + + +def _progress_bar_forced() -> bool: + return environment.env_bool(_FORCE_PROGRESS_BAR_ENV, default=False) + def _resolve_session_controller_client_cls(): return ( @@ -440,6 +455,9 @@ class ManagedSparkSession(SparkSession): _active_s8s_session_id: ClassVar[Optional[str]] = None _active_session_uses_custom_id: ClassVar[bool] = False _execution_progress_bar = dict() + #: Test hook: a ``tqdm``-like class used for execution progress bars + #: instead of the environment's default (see :meth:`_set_tqdm_factory`). + _tqdm_factory: ClassVar[Optional[Callable[..., Any]]] = None #: The :class:`Instrumentation` this session records to, or ``None``. instrumentation: Optional[Instrumentation] = None @@ -650,11 +668,7 @@ def __connect(self, session_response, session_name): url = f"{spark_connect_url}/;session_id={session_response.uuid};use_ssl=true" logger.debug(f"Spark Connect URL: {url}") recorder = self._recorder - mark_s8s_session_active( - session_name, - self._client_options, - session_response=session_response, - ) + mark_s8s_session_active(session_name, self._client_options) is_active = make_is_active_callback( session_name, self._client_options, recorder=recorder ) @@ -991,9 +1005,33 @@ def _get_exiting_active_session( reason="reuse", new_client=True, ): - session_response = get_active_s8s_session_response( - session_name, self._client_options - ) + try: + session_response = ( + get_active_s8s_session_response( + session_name, self._client_options + ) + ) + except Exception as exc: + # Not a verdict: fail open, like the RPC + # guard, if there is a session to reuse. + in_process = ( + ManagedSparkSession.getActiveSession() + or ManagedSparkSession._default_session + ) + if in_process is None: + raise ManagedSparkConnectException( + "Could not check Managed Spark Session" + f" {s8s_session_id}: {exc}" + ) from exc + _reuse_check_errors.failure( + "Error checking Managed Spark Session %s;" + " reusing it unchecked: %s", + s8s_session_id, + exc, + ) + reuse_span.set(found=True, checked=False) + return in_process + _reuse_check_errors.success() if session_response is not None: recorder.trace_id = session_response.uuid reuse_span.set(found=session_response is not None) @@ -1152,11 +1190,7 @@ def _handle_custom_session_id(self): ) if session_response.state == Session.State.ACTIVE: self._custom_session_response = session_response - mark_s8s_session_active( - session_name, - self._client_options, - session_response=session_response, - ) + mark_s8s_session_active(session_name, self._client_options) # Found an active session with the custom ID, set it as the active session ManagedSparkSession._active_s8s_session_id = ( self._custom_session_id @@ -1618,6 +1652,15 @@ def __init__( super().__init__(connection, user_id) self.instrumentation = instrumentation + # The _execute_and_fetch hooks stack in install order, innermost + # first: the position-lost replay (so the hooks around it see one + # call, with the caller's request, however many attempts it takes), + # then the instrumentation wrappers, then the progress bar hook below. + install_replay_on_position_lost( + self.client, + recorder=get_recorder(instrumentation), + close_progress_bar=self._close_progress_bar, + ) if instrumentation is not None and instrumentation.wrap_client: install_client_wrappers(self.client, instrumentation.recorder) @@ -1662,6 +1705,23 @@ def execute_and_fetch_as_iterator_wrapped_method( execute_and_fetch_as_iterator_wrapped_method, self.client ) + def execute_and_fetch_hook(inner): + def _execute_and_fetch(req, *args, **kwargs): + try: + return inner(req, *args, **kwargs) + finally: + # pyspark's Progress.__exit__ has delivered the done=True + # that closes the bar by now; this only closes a bar the + # handler left open. + self._close_progress_bar(req.operation_id) + + return _execute_and_fetch + + # Always installed (instrumented or not), outermost; stop() removes it. + install_client_hooks( + self.client, {"_execute_and_fetch": execute_and_fetch_hook} + ) + # Patching clearProgressHandlers method to not remove Managed Spark Progress Handler clearProgressHandlers_base_method = self.clearProgressHandlers @@ -1709,8 +1769,8 @@ def handler( return # Don't build / render progress bar for non-interactive (despite - # Ipython or non-IPython) - if not environment.is_interactive(): + # Ipython or non-IPython), unless tests forced the bar on. + if not (_progress_bar_forced() or environment.is_interactive()): return total_tasks = 0 @@ -1724,19 +1784,21 @@ def handler( if total_tasks == 0: return - # Get correct tqdm (notebook or CLI) - tqdm_pbar = self.get_tqdm_bar() + recorder = get_recorder(self.instrumentation) # Use a lock to ensure only one thread can access and modify - # the shared dictionaries at a time. + # the shared dictionaries at a time; emit events after releasing + # the lock so sink I/O never runs under the session lock. + event = None with self._lock: - if operation_id in self._execution_progress_bar: - pbar = self._execution_progress_bar[operation_id] - if pbar.total != total_tasks: - pbar.reset( - total=total_tasks - ) # This force resets the progress bar % too on next refresh - else: + pbar = self._execution_progress_bar.get(operation_id) + if pbar is None: + if done: + # Finished before its first update: nothing to show. + return + # Get correct tqdm (notebook or CLI), or the test factory + # (read off the class: a plain function would bind). + tqdm_pbar = type(self)._tqdm_factory or self.get_tqdm_bar() pbar = tqdm_pbar( total=total_tasks, leave=True, @@ -1744,6 +1806,17 @@ def handler( bar_format="{l_bar}{bar} {n_fmt}/{total_fmt} Tasks", ) self._execution_progress_bar[operation_id] = pbar + event = ( + "progress.bar_open", + { + "operation_id": operation_id, + "total_tasks": total_tasks, + }, + ) + elif pbar.total != total_tasks: + pbar.reset( + total=total_tasks + ) # This force resets the progress bar % too on next refresh # To handle skipped or failed tasks. # StageInfo proto doesn't have skipped and failed tasks information to process. @@ -1756,9 +1829,79 @@ def handler( if done: pbar.close() self._execution_progress_bar.pop(operation_id, None) + event = ( + "progress.bar_close", + { + "operation_id": operation_id, + "completed_tasks": completed_tasks, + "total_tasks": total_tasks, + "reason": "done", + }, + ) + elif event is None: + event = ( + "progress.bar_update", + { + "operation_id": operation_id, + "completed_tasks": completed_tasks, + "total_tasks": total_tasks, + }, + ) + + recorder.event(event[0], **event[1]) self.registerProgressHandler(handler) + def _close_progress_bar( + self, operation_id: Optional[str], reason: str = "query_end" + ) -> None: + """Close and unregister ``operation_id``'s progress bar, if open. + + The handler closes a bar when pyspark reports ``done=True``; this is + the fallback for a query that ends without that. A no-op when there + is no open bar. Emits ``progress.bar_close`` with ``reason``. + """ + if not operation_id: + return + with self._lock: + pbar = self._execution_progress_bar.pop(operation_id, None) + if pbar is None: + return + pbar.close() + get_recorder(self.instrumentation).event( + "progress.bar_close", + operation_id=operation_id, + completed_tasks=pbar.n, + total_tasks=pbar.total, + reason=reason, + ) + + @classmethod + def progress_bars(cls) -> Dict[str, Any]: + """A snapshot of the open execution progress bars. + + Maps ``operation_id`` to the ``tqdm`` bar the progress handler + created for it and has not closed yet. Bars are shared by every + ``ManagedSparkSession`` in the process. Intended for tests and + diagnostics; see ``docs/instrumentation.md``. + """ + with cls._lock: + return dict(cls._execution_progress_bar) + + @classmethod + def _set_tqdm_factory(cls, factory: Optional[Callable[..., Any]]) -> None: + """Use ``factory`` instead of ``tqdm`` for execution progress bars. + + A test hook. ``factory`` is called like the ``tqdm`` class + (``total=``, ``leave=``, ``dynamic_ncols=``, ``bar_format=``) and + must return an object with ``n``, ``total``, ``reset(total=)``, + ``refresh()`` and ``close()``. ``None`` restores the default. + Process-wide; use with ``MANAGED_SPARK_CONNECT_FORCE_PROGRESS_BAR=1`` + to run the progress handler headlessly. + """ + with cls._lock: + cls._tqdm_factory = factory + @staticmethod def _sql_lazy_transformation(req): # Select SQL command @@ -1946,6 +2089,7 @@ def stop(self, terminate: Optional[bool] = None) -> None: with recorder.span("session.stop") as span: self._stop(terminate, span) finally: + uninstall_client_hooks(self.client) if self.instrumentation is not None: uninstall_client_wrappers(self.client) recorder.trace_id = "" @@ -2030,9 +2174,7 @@ def terminate_s8s_session( terminate_session_request = TerminateSessionRequest() session_name = f"projects/{project_id}/locations/{region}/sessions/{active_s8s_session_id}" terminate_session_request.name = session_name - invalidate_s8s_session_active(session_name) - with _LIVENESS_CACHE_LOCK: - _DEAD_SESSIONS.add(_liveness_cache_key(session_name, client_options)) + _record_liveness(_liveness_cache_key(session_name, client_options), False) state = None try: session_client = _create_session_controller_client( @@ -2079,29 +2221,26 @@ def terminate_s8s_session( _ENV_LIVENESS_TTL_SECONDS = "MANAGED_SPARK_CONNECT_LIVENESS_TTL_SECONDS" -_MAX_TRANSIENT_FALLBACK_AGE_S = 60.0 +# The liveness cache's clock; tests patch this rather than time.monotonic. +_now = time.monotonic + + +class _LivenessEntry(NamedTuple): + """The last definitive liveness verdict for a session and when it was + reached: ACTIVE, or dead (``NotFound`` or a terminal state).""" + + active: bool + at: float + + _LIVENESS_CACHE_LOCK = threading.Lock() -_LIVENESS_CACHE: dict[tuple[str, Optional[str], Any], float] = {} -_LAST_KNOWN_ACTIVE: dict[ - tuple[str, Optional[str], Any], tuple[sessions.Session, float] -] = {} -_DEAD_SESSIONS: set[tuple[str, Optional[str], Any]] = set() -_TRANSIENT_CONTROL_PLANE_EXCEPTIONS = ( - ServiceUnavailable, - DeadlineExceeded, - TooManyRequests, - InternalServerError, -) +_LIVENESS: dict[tuple[str, Optional[str], Any], _LivenessEntry] = {} +# getOrCreate's GetSession check of the session it is about to reuse. +_reuse_check_errors = _OutageLog(logger) def _get_configured_liveness_ttl_seconds() -> Optional[float]: - raw = os.environ.get(_ENV_LIVENESS_TTL_SECONDS) - if raw is None or not raw.strip(): - return None - try: - return float(raw) - except ValueError: - return None + return environment.env_float(_ENV_LIVENESS_TTL_SECONDS, default=None) def _liveness_cache_key( @@ -2115,110 +2254,71 @@ def _liveness_cache_key( ) -def _is_known_dead( - session_name: str, client_options: Optional[ClientOptions] = None -) -> bool: - key = _liveness_cache_key(session_name, client_options) - with _LIVENESS_CACHE_LOCK: - return key in _DEAD_SESSIONS - - -def _is_cached_active( - session_name: str, client_options: Optional[ClientOptions] = None -) -> bool: - mode = _GUARD_MODE.get() - if mode == "on_error": - return False - ttl = _get_configured_liveness_ttl_seconds() - if ttl == 0: - return False - key = _liveness_cache_key(session_name, client_options) - now = time.monotonic() +def _record_liveness(key: tuple[str, Optional[str], Any], active: bool) -> None: with _LIVENESS_CACHE_LOCK: - if key in _DEAD_SESSIONS: - return False - cached_at = _LIVENESS_CACHE.get(key) - if cached_at is None: - return False - if ttl is None: - if mode == "pre_rpc": - existing = _LAST_KNOWN_ACTIVE.get(key) - if existing is not None: - _LAST_KNOWN_ACTIVE[key] = (existing[0], now) - return True - return False - if ttl < 0 or (now - cached_at) < ttl: - if mode == "pre_rpc": - existing = _LAST_KNOWN_ACTIVE.get(key) - if existing is not None: - _LAST_KNOWN_ACTIVE[key] = (existing[0], now) - return True - return False + _LIVENESS[key] = _LivenessEntry(active, _now()) def _check_cached_liveness( session_name: str, client_options: Optional[ClientOptions] = None ) -> Optional[bool]: + """The cached verdict, or ``None`` when the control plane must be asked. + + A dead verdict stands until the session is marked active again or + invalidated. An ACTIVE verdict is trusted for + ``MANAGED_SPARK_CONNECT_LIVENESS_TTL_SECONDS`` (forever when unset or + negative, never when 0). The guard that runs after a failed RPC always + asks. + """ if ( get_active_s8s_session_response is not _ORIG_GET_ACTIVE_S8S_SESSION_RESPONSE ): return None - mode = _GUARD_MODE.get() - if mode == "on_error": + if _GUARD_MODE.get() == "on_error": + return None + key = _liveness_cache_key(session_name, client_options) + with _LIVENESS_CACHE_LOCK: + entry = _LIVENESS.get(key) + if entry is None: return None - if _is_known_dead(session_name, client_options): + if not entry.active: return False - if _is_cached_active(session_name, client_options): + ttl = _get_configured_liveness_ttl_seconds() + if ttl is None or ttl < 0 or (ttl > 0 and _now() - entry.at < ttl): return True return None def mark_s8s_session_active( - session_name: str, - client_options: Optional[ClientOptions] = None, - session_response: Optional[sessions.Session] = None, + session_name: str, client_options: Optional[ClientOptions] = None ) -> None: - """Record a session as known-ACTIVE in the monotonic TTL liveness cache.""" + """Record a session as known-ACTIVE in the liveness cache.""" if not session_name: return - key = _liveness_cache_key(session_name, client_options) - ttl = _get_configured_liveness_ttl_seconds() - now = time.monotonic() - with _LIVENESS_CACHE_LOCK: - _DEAD_SESSIONS.discard(key) - existing = _LAST_KNOWN_ACTIVE.get(key) - resolved_response = session_response - if resolved_response is None and existing is not None: - resolved_response = existing[0] - if resolved_response is None: - resolved_response = Session( - name=session_name, state=Session.State.ACTIVE - ) - _LAST_KNOWN_ACTIVE[key] = (resolved_response, now) - if ttl != 0: - _LIVENESS_CACHE[key] = now + _record_liveness(_liveness_cache_key(session_name, client_options), True) def invalidate_s8s_session_active(session_name: Optional[str] = None) -> None: """Clear cached session liveness state when a session is stopped or terminated.""" with _LIVENESS_CACHE_LOCK: if session_name is None: - _LIVENESS_CACHE.clear() - _LAST_KNOWN_ACTIVE.clear() - _DEAD_SESSIONS.clear() + _LIVENESS.clear() return - for key in [k for k in _LIVENESS_CACHE if k[0] == session_name]: - _LIVENESS_CACHE.pop(key, None) - for key in [k for k in _LAST_KNOWN_ACTIVE if k[0] == session_name]: - _LAST_KNOWN_ACTIVE.pop(key, None) - for key in [k for k in _DEAD_SESSIONS if k[0] == session_name]: - _DEAD_SESSIONS.discard(key) + for key in [k for k in _LIVENESS if k[0] == session_name]: + del _LIVENESS[key] def get_active_s8s_session_response( session_name, client_options ) -> Optional[sessions.Session]: + """``GetSession``: the response while the session is ACTIVE or CREATING, + ``None`` once it is gone (``NotFound`` or any other state). + + Only those definitive answers update the liveness cache. Any other + error (a control-plane outage, an auth failure, ...) propagates and + leaves the cache alone. + """ key = _liveness_cache_key(session_name, client_options) get_session_request = GetSessionRequest() get_session_request.name = session_name @@ -2226,52 +2326,17 @@ def get_active_s8s_session_response( get_session_response = _create_session_controller_client( client_options=client_options ).get_session(get_session_request) - state = get_session_response.state - except _TRANSIENT_CONTROL_PLANE_EXCEPTIONS as e: - now = time.monotonic() - mode = _GUARD_MODE.get() - ttl = _get_configured_liveness_ttl_seconds() - with _LIVENESS_CACHE_LOCK: - entry = _LAST_KNOWN_ACTIVE.get(key) - if entry is not None: - last_active, observed_at = entry - if ( - mode == "pre_rpc" - or (mode == "on_error" and (ttl is None or ttl < 0)) - or (now - observed_at) <= _MAX_TRANSIENT_FALLBACK_AGE_S - ): - logger.warning( - "Transient error checking Managed Spark Session %s liveness; " - "assuming still active: %s", - session_name, - e, - ) - return last_active + except NotFound as e: + _record_liveness(key, False) print(f"{session_name} Managed Spark Session deleted: {e}") return None - except Exception as e: - with _LIVENESS_CACHE_LOCK: - _LIVENESS_CACHE.pop(key, None) - _LAST_KNOWN_ACTIVE.pop(key, None) - _DEAD_SESSIONS.add(key) - print(f"{session_name} Managed Spark Session deleted: {e}") - return None - if state is not None and ( - state == Session.State.ACTIVE or state == Session.State.CREATING - ): - if state == Session.State.ACTIVE: - now = time.monotonic() - ttl = _get_configured_liveness_ttl_seconds() - with _LIVENESS_CACHE_LOCK: - _DEAD_SESSIONS.discard(key) - _LAST_KNOWN_ACTIVE[key] = (get_session_response, now) - if ttl != 0: - _LIVENESS_CACHE[key] = now + state = get_session_response.state + if state == Session.State.ACTIVE: + _record_liveness(key, True) return get_session_response - with _LIVENESS_CACHE_LOCK: - _LIVENESS_CACHE.pop(key, None) - _LAST_KNOWN_ACTIVE.pop(key, None) - _DEAD_SESSIONS.add(key) + if state == Session.State.CREATING: + return get_session_response + _record_liveness(key, False) return None @@ -2295,14 +2360,24 @@ def make_is_active_callback( client_options: Any, recorder: Any = NOOP_RECORDER, ) -> Callable[[], bool]: - """Create the RPC liveness guard callback for ManagedSparkChannelBuilder.""" + """Create the RPC liveness guard callback for ManagedSparkChannelBuilder. + + With instrumentation on, an invocation emits a ``guard.result`` event + (``mode`` = ``pre_rpc``/``on_error`` as set by the channel, ``cached`` + = answered from the liveness cache without a GetSession RPC, ``active`` + = the verdict; when the check itself raised, ``active`` is ``None`` and + ``error_type`` names the exception) in addition to the + ``dataproc.get_session`` span of an actual RPC. A cached ``active`` verdict, + the common case, emits nothing. + """ mark_s8s_session_active(session_name, client_options) - def is_active() -> bool: + def check() -> Tuple[bool, bool]: + """``(active, cached)`` for one guard invocation.""" if is_s8s_session_active is _ORIG_IS_S8S_SESSION_ACTIVE: cached = _check_cached_liveness(session_name, client_options) if cached is not None: - return cached + return cached, True with recorder.span( "dataproc.get_session", reason="rpc_guard", new_client=True ): @@ -2311,11 +2386,33 @@ def is_active() -> bool: session_name, client_options ) is not None - ) + ), False with recorder.span( "dataproc.get_session", reason="rpc_guard", new_client=True ): - return is_s8s_session_active(session_name, client_options) + return is_s8s_session_active(session_name, client_options), False + + def is_active() -> bool: + if not recorder.enabled: + return check()[0] + mode = _GUARD_MODE.get() + try: + active, cached = check() + except BaseException as exc: + recorder.event( + "guard.result", + mode=mode, + cached=False, + active=None, + error_type=type(exc).__name__, + ) + raise + if not cached or not active: + # A cached "active" is the common case on every guarded RPC. + recorder.event( + "guard.result", mode=mode, cached=cached, active=active + ) + return active is_active._supports_error_guard = True # type: ignore[attr-defined] return is_active diff --git a/tests/unit/test_core.py b/tests/unit/test_core.py index f222ad56..f3c8425e 100644 --- a/tests/unit/test_core.py +++ b/tests/unit/test_core.py @@ -123,12 +123,15 @@ def test_guard_failure_raises_before_the_rpc(fake_channel): assert fake_channel.multicallables["/svc/Stream"].calls == [] -def test_no_guard_returns_the_multicallable_unchanged(fake_channel): +def test_no_guard_wraps_the_multicallable_only_to_count_rpcs(fake_channel): channel = core.ProxiedChannel("example.com", None) - assert ( - channel.unary_unary("/svc/Method") - is fake_channel.multicallables["/svc/Method"] - ) + method = channel.unary_unary("/svc/Method") + # Still wrapped so the proxy can count in-flight RPCs for its silence + # cut, but no guard runs. + assert isinstance(method, core._GuardedMultiCallable) + assert method._guard is None + assert method("a") == "response:a" + assert fake_channel.multicallables["/svc/Method"].calls == [("call", "a")] # subscribe/unsubscribe are forwarded, not wrapped. assert channel.subscribe(lambda *_: None) is None assert channel.unsubscribe(lambda *_: None) is None @@ -233,12 +236,50 @@ def guard_cb(): method.with_call("req") +def _proxy(): + """A real proxy that is never started, for its in-flight RPC counter.""" + return core.proxy.ManagedSparkSessionProxy(0, "unused") + + def test_guarded_stream_repr(): inner = iter(["a", "b"]) - stream = core._GuardedStream(inner, lambda _exc: None) + stream = core._GuardedStream(inner, lambda _exc: None, _proxy()) assert "_GuardedStream(" in repr(stream) +def test_streaming_rpc_counts_as_in_flight_across_next_calls(): + proxy = _proxy() + method = core._GuardedMultiCallable( + lambda req: iter(["a", "b"]), None, proxy + ) + stream = method("req") + assert proxy.get_inflight_rpcs() == 1 + assert next(stream) == "a" + # Still in flight between next() calls, while the caller works. + assert proxy.get_inflight_rpcs() == 1 + assert next(stream) == "b" + assert proxy.get_inflight_rpcs() == 1 + with pytest.raises(StopIteration): + next(stream) + assert proxy.get_inflight_rpcs() == 0 + # Finishing again (e.g. on garbage collection) doesn't count it twice. + stream._finish(None) + assert proxy.get_inflight_rpcs() == 0 + + +def test_unary_and_failed_rpcs_are_not_left_in_flight(): + proxy = _proxy() + assert core._GuardedMultiCallable(lambda req: "r", None, proxy)("x") == "r" + assert proxy.get_inflight_rpcs() == 0 + + def fail(req): + raise RuntimeError("boom") + + with pytest.raises(RuntimeError): + core._GuardedMultiCallable(fail, None, proxy)("x") + assert proxy.get_inflight_rpcs() == 0 + + def test_guarded_stream_forwards_grpc_call_methods_without_instrumentation(): class FakeCallIterator: @@ -287,7 +328,7 @@ def add_done_callback(self, cb): self.done_callbacks.append(cb) raw_call = FakeCallIterator() - stream = core._GuardedStream(raw_call, lambda _exc: None) + stream = core._GuardedStream(raw_call, lambda _exc: None, _proxy()) assert isinstance(stream, grpc.Call) assert next(stream) == "msg-1" assert stream.initial_metadata() == (("init-k", "init-v"),) @@ -386,6 +427,21 @@ def guard_cb(): assert guard_calls == ["pre_rpc", "on_error"] +def test_check_active_logs_warning_and_proceeds_when_guard_raises( + fake_channel, caplog +): + import logging + + def raising_guard(): + raise RuntimeError("injected control plane failure") + + channel = core.ProxiedChannel("example.com", raising_guard) + method = channel.unary_unary("/svc/Method") + with caplog.at_level(logging.WARNING): + assert method("req-1") == "response:req-1" + assert any("proceeding with the RPC" in m for m in caplog.messages) + + @pytest.mark.parametrize( "url,expected_dest", [ diff --git a/tests/unit/test_guard.py b/tests/unit/test_guard.py index 5f3dad92..80ff8b62 100644 --- a/tests/unit/test_guard.py +++ b/tests/unit/test_guard.py @@ -25,12 +25,16 @@ ServiceUnavailable, TooManyRequests, ) +from google.auth.exceptions import RefreshError from google.cloud.dataproc_v1 import Session import grpc import pytest from google.cloud.managed_spark_connect import session as session_mod from google.cloud.managed_spark_connect.client import core +from google.cloud.managed_spark_connect.exceptions import ( + ManagedSparkConnectException, +) from google.cloud.managed_spark_connect.instrumentation import ( InMemorySink, Instrumentation, @@ -62,7 +66,7 @@ def test_default_lazy_mode_skips_pre_rpc_get_session_and_span(monkeypatch): """By default (Option A), pre-RPC guard checks do 0 GetSession RPCs and 0 spans.""" inst = Instrumentation([InMemorySink()]) clock = [100.0] - monkeypatch.setattr(session_mod.time, "monotonic", lambda: clock[0]) + monkeypatch.setattr(session_mod, "_now", lambda: clock[0]) with mock.patch.object( session_mod, "_create_session_controller_client" @@ -70,7 +74,9 @@ def test_default_lazy_mode_skips_pre_rpc_get_session_and_span(monkeypatch): cb = session_mod.make_is_active_callback( SESSION_NAME, CLIENT_OPTIONS, recorder=inst.recorder ) - wrapped = core._GuardedMultiCallable(lambda req: f"ok:{req}", cb) + wrapped = core._GuardedMultiCallable( + lambda req: f"ok:{req}", cb, mock.Mock() + ) assert wrapped("r1") == "ok:r1" # Even after 60 seconds, default lazy mode does NOT poll pre-RPC. clock[0] = 160.0 @@ -282,7 +288,7 @@ def test_ttl_cache_miss_after_expiry_calls_get_session_and_reseeds(monkeypatch): monkeypatch.setenv(session_mod._ENV_LIVENESS_TTL_SECONDS, "5.0") inst = Instrumentation([InMemorySink()]) clock = [100.0] - monkeypatch.setattr(session_mod.time, "monotonic", lambda: clock[0]) + monkeypatch.setattr(session_mod, "_now", lambda: clock[0]) with mock.patch.object( session_mod, "_create_session_controller_client" @@ -347,7 +353,7 @@ def test_ttl_negative_disables_pre_rpc_polling(monkeypatch): monkeypatch.setenv(session_mod._ENV_LIVENESS_TTL_SECONDS, "-1") inst = Instrumentation([InMemorySink()]) clock = [100.0] - monkeypatch.setattr(session_mod.time, "monotonic", lambda: clock[0]) + monkeypatch.setattr(session_mod, "_now", lambda: clock[0]) with mock.patch.object( session_mod, "_create_session_controller_client" @@ -417,76 +423,51 @@ def test_invalidate_s8s_session_active_targeted_and_global(monkeypatch): @pytest.mark.parametrize( - "transient_exc", + "exc", [ ServiceUnavailable("503 Backend unavailable"), DeadlineExceeded("504 Deadline exceeded"), TooManyRequests("429 Quota exceeded"), InternalServerError("500 Internal error"), + PermissionDenied("403 Permission denied"), + RefreshError("token expired"), + RuntimeError("unexpected error"), ], ) -def test_transient_control_plane_error_falls_back_to_last_known_active( - transient_exc, caplog -): - active_session = _make_session(Session.State.ACTIVE) +def test_non_definitive_get_session_errors_propagate_and_keep_cache(exc): + session_mod.mark_s8s_session_active(SESSION_NAME, CLIENT_OPTIONS) with mock.patch.object( session_mod, "_create_session_controller_client" ) as create_client: - client = create_client.return_value - client.get_session.side_effect = [active_session, transient_exc] - - first = session_mod.get_active_s8s_session_response( - SESSION_NAME, CLIENT_OPTIONS - ) - assert first is active_session + create_client.return_value.get_session.side_effect = exc - with caplog.at_level(logging.WARNING): - second = session_mod.get_active_s8s_session_response( + with pytest.raises(type(exc)): + session_mod.get_active_s8s_session_response( SESSION_NAME, CLIENT_OPTIONS ) - assert second is active_session - assert any("Transient error" in m for m in caplog.messages) - -def test_transient_error_without_prior_active_returns_none(): - with mock.patch.object( - session_mod, "_create_session_controller_client" - ) as create_client: - client = create_client.return_value - client.get_session.side_effect = ServiceUnavailable("503") assert ( - session_mod.get_active_s8s_session_response( - SESSION_NAME, CLIENT_OPTIONS - ) - is None + session_mod.is_s8s_session_active(SESSION_NAME, CLIENT_OPTIONS) + is True ) + assert create_client.return_value.get_session.call_count == 1 -@pytest.mark.parametrize( - "terminal_exc", - [ - NotFound("404 Session not found"), - PermissionDenied("403 Permission denied"), - RuntimeError("unexpected error"), - ], -) -def test_terminal_exceptions_invalidate_cache_and_return_none(terminal_exc): - active_session = _make_session(Session.State.ACTIVE) +def test_not_found_marks_session_dead(): with mock.patch.object( session_mod, "_create_session_controller_client" ) as create_client: client = create_client.return_value client.get_session.side_effect = [ - active_session, - terminal_exc, - ServiceUnavailable("503 after terminal error"), + _make_session(Session.State.ACTIVE), + NotFound("404 Session not found"), ] assert ( session_mod.get_active_s8s_session_response( SESSION_NAME, CLIENT_OPTIONS ) - is active_session + is not None ) assert ( session_mod.get_active_s8s_session_response( @@ -494,17 +475,15 @@ def test_terminal_exceptions_invalidate_cache_and_return_none(terminal_exc): ) is None ) - # Because the terminal exception cleared _LAST_KNOWN_ACTIVE, a subsequent - # transient error also returns None. + # The dead verdict is cached: no further GetSession. assert ( - session_mod.get_active_s8s_session_response( - SESSION_NAME, CLIENT_OPTIONS - ) - is None + session_mod.is_s8s_session_active(SESSION_NAME, CLIENT_OPTIONS) + is False ) + assert client.get_session.call_count == 2 -def test_non_active_session_state_invalidates_cache(): +def test_non_active_session_state_marks_session_dead(): active_session = _make_session(Session.State.ACTIVE) terminated_session = _make_session(Session.State.TERMINATED) with mock.patch.object( @@ -525,8 +504,11 @@ def test_non_active_session_state_invalidates_cache(): ) is None ) - assert session_mod._LIVENESS_CACHE == {} - assert session_mod._LAST_KNOWN_ACTIVE == {} + assert ( + session_mod.is_s8s_session_active(SESSION_NAME, CLIENT_OPTIONS) + is False + ) + assert client.get_session.call_count == 2 def test_should_check_liveness_on_error_classifications(): @@ -638,12 +620,9 @@ def broken_stream(): assert client.get_session.call_count == 1 -def test_transient_control_plane_error_during_on_error_guard_keeps_session_alive( - monkeypatch, -): - """In default lazy mode, a newly created session survives transient 503 on error guard even after >60s.""" - clock = [100.0] - monkeypatch.setattr(session_mod.time, "monotonic", lambda: clock[0]) +def test_transient_control_plane_error_during_on_error_guard_keeps_session_alive(): + """A transient 503 from the guard after a failed RPC re-raises the RPC's + own error and does not mark the session dead.""" class TransientStreamError(grpc.RpcError): @@ -672,18 +651,20 @@ def trailing_metadata(self): "503 control plane blip" ) - # make_is_active_callback seeds _LAST_KNOWN_ACTIVE without needing a prior GetSession call. cb = session_mod.make_is_active_callback(SESSION_NAME, CLIENT_OPTIONS) channel = core.ProxiedChannel("example.com", cb) method = channel.unary_unary( "/spark.connect.SparkConnectService/Config" ) - # Advance past _MAX_TRANSIENT_FALLBACK_AGE_S (60s). - clock[0] = 250.0 with pytest.raises(TransientStreamError, match="socket closed"): method("req1") assert client.get_session.call_count == 1 + # A 503 is not a verdict: the session is still cached as active. + assert ( + session_mod.is_s8s_session_active(SESSION_NAME, CLIENT_OPTIONS) + is True + ) def test_invalidate_and_terminate_affect_pre_rpc_and_negative_ttl(monkeypatch): @@ -703,7 +684,7 @@ def check_active(): raise RuntimeError("Session not active") wrapped = core._GuardedMultiCallable( - lambda req: f"ok:{req}", check_active + lambda req: f"ok:{req}", check_active, mock.Mock() ) assert wrapped("r1") == "ok:r1" assert client.get_session.call_count == 0 @@ -747,3 +728,106 @@ def check_active(): is False ) assert client.get_session.call_count == prev_calls + + +def _proxied_method(fake_rpc): + """A ``Config`` RPC through a ``ProxiedChannel`` whose wire call is + ``fake_rpc``, guarded by the real liveness callback.""" + with ( + mock.patch.object(core.proxy, "ManagedSparkSessionProxy") as proxy_cls, + mock.patch.object(core, "DefaultChannelBuilder") as builder_cls, + ): + proxy_cls.return_value.port = 12345 + inner = builder_cls.return_value.toChannel.return_value + inner.unary_unary.return_value = fake_rpc + cb = session_mod.make_is_active_callback(SESSION_NAME, CLIENT_OPTIONS) + channel = core.ProxiedChannel("example.com", cb) + return channel.unary_unary("/spark.connect.SparkConnectService/Config") + + +def test_pre_rpc_guard_fails_open_and_warns_once_per_outage( + monkeypatch, caplog +): + # TTL 0: the guard asks the control plane before every RPC. + monkeypatch.setenv(session_mod._ENV_LIVENESS_TTL_SECONDS, "0") + with mock.patch.object( + session_mod, "_create_session_controller_client" + ) as create_client: + create_client.return_value.get_session.side_effect = [ + ServiceUnavailable("503 a"), + ServiceUnavailable("503 b"), + _make_session(Session.State.ACTIVE), + ServiceUnavailable("503 c"), + ] + method = _proxied_method( + mock.MagicMock(side_effect=lambda r: f"ok:{r}") + ) + + with caplog.at_level(logging.DEBUG, logger=core.logger.name): + results = [method(f"r{i}") for i in range(4)] + + assert results == ["ok:r0", "ok:r1", "ok:r2", "ok:r3"] + levels = [ + r.levelno + for r in caplog.records + if "Error checking session liveness" in r.getMessage() + ] + # The success in between ends the first outage. + assert levels == [logging.WARNING, logging.DEBUG, logging.WARNING] + + +def test_get_or_create_reuses_in_process_session_when_check_fails( + monkeypatch, caplog +): + ManagedSparkSession = session_mod.ManagedSparkSession + monkeypatch.setattr(ManagedSparkSession, "_active_s8s_session_id", "sess-1") + session_mod._reuse_check_errors.success() + in_process = mock.Mock() + builder = ManagedSparkSession.builder + with ( + mock.patch.object( + ManagedSparkSession, "getActiveSession", return_value=in_process + ), + mock.patch.object( + session_mod, "_create_session_controller_client" + ) as create_client, + ): + create_client.return_value.get_session.side_effect = ServiceUnavailable( + "503" + ) + with caplog.at_level(logging.WARNING): + assert builder._get_exiting_active_session() is in_process + assert builder._get_exiting_active_session() is in_process + + in_process.stop.assert_not_called() + warnings = [ + r + for r in caplog.records + if r.levelno == logging.WARNING and "sess-1" in r.getMessage() + ] + assert len(warnings) == 1 + + +def test_get_or_create_propagates_check_error_with_nothing_to_reuse( + monkeypatch, +): + ManagedSparkSession = session_mod.ManagedSparkSession + monkeypatch.setattr(ManagedSparkSession, "_active_s8s_session_id", "sess-1") + monkeypatch.setattr(ManagedSparkSession, "_default_session", None) + with ( + mock.patch.object( + ManagedSparkSession, "getActiveSession", return_value=None + ), + mock.patch.object( + session_mod, "_create_session_controller_client" + ) as create_client, + ): + create_client.return_value.get_session.side_effect = ServiceUnavailable( + "503" + ) + with pytest.raises( + ManagedSparkConnectException, + match="Could not check Managed Spark Session sess-1: 503", + ) as excinfo: + ManagedSparkSession.builder._get_exiting_active_session() + assert isinstance(excinfo.value.__cause__, ServiceUnavailable) diff --git a/tests/unit/test_instrumentation.py b/tests/unit/test_instrumentation.py index 4f7e5580..3ddd5ccb 100644 --- a/tests/unit/test_instrumentation.py +++ b/tests/unit/test_instrumentation.py @@ -882,6 +882,41 @@ def test_install_is_a_noop_when_disabled(): assert not hasattr(client, _client._RECORDER_ATTR) +def test_client_hooks_stack_and_uninstall_restores_originals(): + client = FakeClient() + calls = [] + + def own_to_pandas(plan, observations): + calls.append("own") + return "own" + + # An instance attribute that predates the hooks is restored as is. + client.to_pandas = own_to_pandas + + def tag(label): + def make(inner): + def wrapper(*args, **kwargs): + calls.append(label) + return inner(*args, **kwargs) + + return wrapper + + return make + + _client.install_client_hooks( + client, {"to_pandas": tag("a"), "_execute_and_fetch": tag("a")} + ) + _client.install_client_hooks(client, {"to_pandas": tag("b")}) + assert client.to_pandas("plan", {}) == "own" + assert calls == ["b", "a", "own"] + assert client.to_pandas.__wrapped__.__wrapped__ is own_to_pandas + + _client.uninstall_client_hooks(client) + _client.uninstall_client_hooks(client) # idempotent + assert client.__dict__["to_pandas"] is own_to_pandas + assert "_execute_and_fetch" not in client.__dict__ + + # -------------------------------------------------------------------------- # Session wiring (mocked Dataproc API, no network) # -------------------------------------------------------------------------- diff --git a/tests/unit/test_instrumentation_events.py b/tests/unit/test_instrumentation_events.py new file mode 100644 index 00000000..9d49e9c8 --- /dev/null +++ b/tests/unit/test_instrumentation_events.py @@ -0,0 +1,1685 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. +"""Tests for the instrumentation *events* the library emits (no network). + +Spans are covered by ``test_instrumentation.py``; this file covers the +point-in-time ``{"type": "event"}`` records: proxy forwarding events, auth +refreshes, progress bars, the liveness guard and pyspark retries. +""" + +import contextlib +import datetime +import os +import socket +import threading +import time +from unittest import mock + +import grpc +import pytest + +from google.cloud.managed_spark_connect.client import proxy +from google.cloud.managed_spark_connect.instrumentation import ( + NOOP_RECORDER, + InMemorySink, + Instrumentation, +) + +EVENT_KEYS = {"type", "name", "wall_ns", "mono_ns", "thread", "attrs"} + + +def events(inst, name=None): + return [ + r + for r in inst.records + if r["type"] == "event" and (name is None or r["name"] == name) + ] + + +def one_event(inst, name): + found = events(inst, name) + assert len(found) == 1, f"expected one {name!r} event, got {found}" + return found[0] + + +@pytest.fixture +def inst(): + return Instrumentation([InMemorySink()], {"run_id": "r1"}) + + +# -------------------------------------------------------------------------- +# proxy.* events +# -------------------------------------------------------------------------- + + +class _ScriptedSock: + """A socket-like object whose ``recv`` replays a script. + + Script items are bytes (returned), exceptions (raised) or callables + (called, their result returned). + """ + + def __init__(self, script, timeout=None, send=None): + self._script = list(script) + self._timeout = timeout + self._send = send + self.sent = [] + self.closed = False + + def recv(self, n): + item = self._script.pop(0) + if isinstance(item, BaseException): + raise item + if callable(item): + return item() + return item + + def send(self, data): + if self._send is not None: + return self._send(data) + self.sent.append(bytes(data)) + return len(data) + + def gettimeout(self): + return self._timeout + + def close(self): + self.closed = True + + +def test_forward_bytes_eof_event_has_direction_bytes_and_age(inst): + start = time.perf_counter_ns() - 1_000_000 + from_sock = _ScriptedSock([b"abc", b"de", b""], timeout=10) + to_sock = _ScriptedSock([]) + stats = proxy._DirectionStats() + proxy.forward_bytes( + "7-forward", + from_sock, + to_sock, + stats, + recorder=inst.recorder, + conn=7, + direction=proxy.CLIENT_TO_SERVER, + start_mono_ns=start, + ) + rec = one_event(inst, "proxy.eof") + assert set(rec) == EVENT_KEYS + assert rec["attrs"]["conn"] == 7 + assert rec["attrs"]["direction"] == "client_to_server" + assert rec["attrs"]["bytes_forwarded"] == 5 + assert rec["attrs"]["age_ns"] >= 1_000_000 + assert rec["attrs"]["run_id"] == "r1" + assert to_sock.closed and from_sock.closed + assert not events(inst, "proxy.forward_error") + + +def test_forward_bytes_eof_without_stats_or_start(inst): + proxy.forward_bytes( + "n", + _ScriptedSock([b""]), + _ScriptedSock([]), + recorder=inst.recorder, + conn=1, + direction=proxy.SERVER_TO_CLIENT, + ) + rec = one_event(inst, "proxy.eof") + assert rec["attrs"]["bytes_forwarded"] == 0 + assert rec["attrs"]["age_ns"] is None + assert rec["attrs"]["direction"] == "server_to_client" + + +def test_forward_bytes_recv_timeout_is_silent_and_continues(inst): + # Recv timeouts are routine polls of an idle connection: no event. + from_sock = _ScriptedSock([TimeoutError(), b"x", b""], timeout=10) + to_sock = _ScriptedSock([], timeout=60) + proxy.forward_bytes( + "n", + from_sock, + to_sock, + recorder=inst.recorder, + conn=2, + direction=proxy.CLIENT_TO_SERVER, + ) + assert not events(inst, "proxy.socket_timeout") + assert to_sock.sent == [b"x"] + assert one_event(inst, "proxy.eof") + + +class _FlakySendSocket(socket.socket): + """An unconnected TCP socket whose first ``timeouts`` sends time out.""" + + def __init__(self, timeouts, timeout): + super().__init__(socket.AF_INET, socket.SOCK_STREAM) + self.settimeout(timeout) + self._timeouts = timeouts + self.attempts = [] + + def send(self, data, flags=0): + self.attempts.append(bytes(data)) + if len(self.attempts) <= self._timeouts: + raise TimeoutError() + return len(data) + + +def test_forward_bytes_send_timeout_event_then_retries_send(inst): + from_sock = _ScriptedSock([b"payload", b""]) + with _FlakySendSocket(timeouts=1, timeout=10) as to_sock: + proxy.forward_bytes( + "n", + from_sock, + to_sock, + recorder=inst.recorder, + conn=3, + direction=proxy.SERVER_TO_CLIENT, + ) + rec = one_event(inst, "proxy.socket_timeout") + assert rec["attrs"] == { + "conn": 3, + "direction": "server_to_client", + "timeout_s": 10, + "run_id": "r1", + } + assert to_sock.attempts == [b"payload", b"payload"] + assert not events(inst, "proxy.forward_error") + + +def test_forward_bytes_send_timeout_on_non_socket_has_no_timeout_s(inst): + attempts = [] + + def flaky_send(data): + attempts.append(bytes(data)) + if len(attempts) == 1: + raise TimeoutError() + return len(data) + + from_sock = _ScriptedSock([b"payload", b""]) + to_sock = _ScriptedSock([], timeout=60, send=flaky_send) + proxy.forward_bytes( + "n", + from_sock, + to_sock, + recorder=inst.recorder, + conn=3, + direction=proxy.CLIENT_TO_SERVER, + ) + rec = one_event(inst, "proxy.socket_timeout") + assert rec["attrs"]["timeout_s"] is None + assert attempts == [b"payload", b"payload"] + + +def test_forward_bytes_gives_up_after_ten_send_timeouts(inst): + from_sock = _ScriptedSock([b"payload"]) + with _FlakySendSocket(timeouts=10, timeout=10) as to_sock: + proxy.forward_bytes( + "n", + from_sock, + to_sock, + recorder=inst.recorder, + conn=5, + direction=proxy.SERVER_TO_CLIENT, + ) + assert len(to_sock.attempts) == 10 + assert len(events(inst, "proxy.socket_timeout")) == 10 + rec = one_event(inst, "proxy.forward_error") + assert rec["attrs"]["conn"] == 5 + assert rec["attrs"]["error_type"] == "SendTimeout" + assert from_sock.closed + assert not events(inst, "proxy.eof") + + +def test_forward_bytes_error_event_closes_both_and_truncates_message(inst): + message = "boom " * 100 + from_sock = _ScriptedSock([ConnectionResetError(message)]) + to_sock = _ScriptedSock([]) + proxy.forward_bytes( + "n", + from_sock, + to_sock, + recorder=inst.recorder, + conn=4, + direction=proxy.CLIENT_TO_SERVER, + ) + rec = one_event(inst, "proxy.forward_error") + assert rec["attrs"]["conn"] == 4 + assert rec["attrs"]["direction"] == "client_to_server" + assert rec["attrs"]["error_type"] == "ConnectionResetError" + assert len(rec["attrs"]["message"]) == 200 + assert message.startswith(rec["attrs"]["message"]) + assert from_sock.closed and to_sock.closed + assert not events(inst, "proxy.eof") + + +def test_forward_bytes_disabled_recorder_emits_nothing_and_behaves(): + # The default (no-op) recorder path: same behaviour, no records, and + # the socket-like objects are not asked for their timeout. + from_sock = mock.MagicMock() + from_sock.recv.side_effect = [TimeoutError(), b"x", OSError("gone")] + to_sock = mock.MagicMock() + to_sock.send.return_value = 1 + proxy.forward_bytes("n", from_sock, to_sock) + from_sock.gettimeout.assert_not_called() + to_sock.gettimeout.assert_not_called() + assert to_sock.close.called and from_sock.close.called + assert NOOP_RECORDER.enabled is False + + +def test_connect_sockets_tags_directions_and_age(inst): + client, proxy_side = socket.socketpair() + bridge_side, backend = socket.socketpair() + # Like forward_connection: the gRPC-side socket polls with a timeout, + # so closing it from the other direction's thread ends this one. + proxy_side.settimeout(0.2) + span = inst.recorder.start_span("proxy.connection", conn=9) + t = threading.Thread( + target=proxy.connect_sockets, + args=[9, proxy_side, bridge_side, span, inst.recorder], + daemon=True, + ) + t.start() + client.sendall(b"hello") + assert backend.recv(1024) == b"hello" + backend.sendall(b"hi") + assert client.recv(1024) == b"hi" + # Sit idle past the 0.2 s poll so the gRPC side hits recv timeouts. + time.sleep(0.5) + # Closing the backend first: the server_to_client direction sees EOF + # and closes the client side, which ends the other direction too. + backend.close() + t.join(timeout=5) + assert not t.is_alive() + span.end() + client.close() + eofs = events(inst, "proxy.eof") + by_direction = {e["attrs"]["direction"]: e["attrs"] for e in eofs} + assert by_direction["server_to_client"]["bytes_forwarded"] == 2 + assert by_direction["server_to_client"]["conn"] == 9 + assert by_direction["server_to_client"]["age_ns"] >= 0 + for e in eofs: + assert e["thread"].startswith("9-") + # The idle recv timeouts on the gRPC side are not recorded. + assert not events(inst, "proxy.socket_timeout") + + +def test_connect_sockets_without_recorder_is_the_default_path(): + client, proxy_side = socket.socketpair() + bridge_side, backend = socket.socketpair() + t = threading.Thread( + target=proxy.connect_sockets, + args=[1, proxy_side, bridge_side], + daemon=True, + ) + t.start() + client.sendall(b"ping") + assert backend.recv(1024) == b"ping" + client.close() + backend.close() + t.join(timeout=5) + assert not t.is_alive() + + +def test_bridged_socket_recv_uses_ws_recv_timeout(): + ws = mock.MagicMock() + sock = proxy.bridged_socket(ws) + ws.recv.return_value = "6869" + assert sock.recv(10) == b"hi" + ws.recv.assert_called_once_with(timeout=proxy.WS_RECV_TIMEOUT_S) + assert proxy.WS_RECV_TIMEOUT_S == 60 + + +def test_forward_connection_sets_socket_timeout_constant(): + conn = mock.MagicMock() + ws_cm = mock.MagicMock() + ws_cm.__enter__.return_value = mock.MagicMock() + with ( + mock.patch.object(proxy, "connect_tcp_bridge", return_value=ws_cm), + mock.patch.object(proxy, "connect_sockets") as connect_sockets, + ): + proxy.forward_connection(1, conn, ("127.0.0.1", 1), "h") + conn.settimeout.assert_called_once_with(proxy.SOCKET_TIMEOUT_S) + assert proxy.SOCKET_TIMEOUT_S == 10 + args = connect_sockets.call_args + assert args.args[:4] == (1, conn, mock.ANY, None) + assert args.args[4] is NOOP_RECORDER and not args.kwargs + + +def test_connect_sockets_without_span_still_tracks_bytes_and_age_when_recorder_enabled( + inst, +): + client, proxy_side = socket.socketpair() + bridge_side, backend = socket.socketpair() + proxy_side.settimeout(0.2) + t = threading.Thread( + target=proxy.connect_sockets, + args=[12, proxy_side, bridge_side, None, inst.recorder], + daemon=True, + ) + t.start() + client.sendall(b"ping") + assert backend.recv(1024) == b"ping" + client.close() + t.join(timeout=5) + backend.close() + eofs = events(inst, "proxy.eof") + by_direction = {e["attrs"]["direction"]: e["attrs"] for e in eofs} + assert by_direction["client_to_server"]["bytes_forwarded"] == 4 + assert by_direction["client_to_server"]["age_ns"] >= 0 + + +def test_forward_connection_passes_span_and_recorder_when_enabled(inst): + conn = mock.MagicMock() + ws_cm = mock.MagicMock() + ws_cm.__enter__.return_value = mock.MagicMock() + with ( + mock.patch.object(proxy, "connect_tcp_bridge", return_value=ws_cm), + mock.patch.object(proxy, "connect_sockets") as connect_sockets, + ): + proxy.forward_connection(1, conn, ("127.0.0.1", 1), "h", inst.recorder) + args = connect_sockets.call_args.args + assert args[3].name == "proxy.connection" + assert args[4] is inst.recorder + + +def test_forward_bytes_websocket_closed_is_an_eof_with_close_code(inst): + import websockets.exceptions + import websockets.frames + + closed = websockets.exceptions.ConnectionClosedError( + websockets.frames.Close(1006, "abnormal"), None + ) + from_sock = _ScriptedSock([b"abcd", closed]) + to_sock = _ScriptedSock([]) + stats = proxy._DirectionStats() + proxy.forward_bytes( + "n", + from_sock, + to_sock, + stats, + recorder=inst.recorder, + conn=8, + direction=proxy.SERVER_TO_CLIENT, + start_mono_ns=time.perf_counter_ns(), + ) + rec = one_event(inst, "proxy.eof") + assert rec["attrs"]["direction"] == "server_to_client" + assert rec["attrs"]["bytes_forwarded"] == 4 + assert rec["attrs"]["error_type"] == "ConnectionClosedError" + assert rec["attrs"]["close_code"] == 1006 + assert rec["attrs"]["age_ns"] >= 0 + assert not events(inst, "proxy.forward_error") + assert from_sock.closed and to_sock.closed + + +class _HandshakeError(Exception): + + def __init__(self, status): + super().__init__(f"server rejected WebSocket connection: HTTP {status}") + self.response = mock.Mock(status_code=status) + + +def test_ws_connect_failed_event_with_http_status(inst): + creds = mock.Mock(token="tok") + with ( + mock.patch.object( + proxy.googleauth, "default", return_value=(creds, None) + ), + mock.patch.object( + proxy.websocketclient, "connect", side_effect=_HandshakeError(401) + ), + ): + with pytest.raises(_HandshakeError): + proxy.connect_tcp_bridge("backend.example.com", inst.recorder, 5) + rec = one_event(inst, "proxy.ws_connect_failed") + assert rec["attrs"] == { + "conn": 5, + "error_type": "_HandshakeError", + "http_status": 401, + "run_id": "r1", + } + (span,) = [ + r + for r in inst.records + if r["type"] == "span" and r["name"] == "proxy.ws_connect" + ] + assert span["status"] == "error" + + +def test_ws_connect_failed_event_on_auth_error_without_status(inst): + with mock.patch.object( + proxy.googleauth, "default", side_effect=RuntimeError("no adc") + ): + with pytest.raises(RuntimeError): + proxy.connect_tcp_bridge("backend.example.com", inst.recorder, 6) + rec = one_event(inst, "proxy.ws_connect_failed") + assert rec["attrs"]["conn"] == 6 + assert rec["attrs"]["error_type"] == "RuntimeError" + assert rec["attrs"]["http_status"] is None + assert proxy._http_status_of(OSError("y")) is None + + +def test_ws_connect_success_emits_no_failure_event(inst): + creds = mock.Mock(token="tok") + with ( + mock.patch.object( + proxy.googleauth, "default", return_value=(creds, None) + ), + mock.patch.object(proxy.websocketclient, "connect"), + ): + proxy.connect_tcp_bridge("backend.example.com", inst.recorder, 1) + assert not events(inst, "proxy.ws_connect_failed") + + +def test_ws_connect_failed_disabled_recorder_reraises_only(): + with mock.patch.object( + proxy.googleauth, "default", side_effect=RuntimeError("no adc") + ): + with pytest.raises(RuntimeError): + proxy.connect_tcp_bridge("backend.example.com") + + +# -------------------------------------------------------------------------- +# auth.token_refreshed +# -------------------------------------------------------------------------- + + +class _FakeCreds: + """Credentials whose refresh() mints a token with a known expiry.""" + + def __init__(self, token=None, valid=False, expiry=None): + self.token = token + self.valid = valid + self.expiry = expiry + self.refreshes = 0 + + def refresh(self, request): + self.refreshes += 1 + self.token = f"tok{self.refreshes}" + self.valid = True + self.expiry = datetime.datetime(2030, 1, 2, 3, 4, 5) + + +def _with_creds(creds): + return mock.patch.object( + proxy.googleauth, "default", return_value=(creds, None) + ) + + +def test_token_refreshed_event_on_first_refresh(inst): + creds = _FakeCreds() + with _with_creds(creds): + assert proxy.get_refreshed_credentials(inst.recorder) is creds + rec = one_event(inst, "auth.token_refreshed") + assert set(rec) == EVENT_KEYS + # 2030-01-02T03:04:05Z + assert rec["attrs"] == { + "expiry_wall_ns": 1_893_553_445 * 10**9, + "had_token": False, + "run_id": "r1", + } + (refresh_span,) = [ + r + for r in inst.records + if r["type"] == "span" and r["name"] == "auth.refresh" + ] + assert rec["wall_ns"] >= refresh_span["start_wall_ns"] + + +def test_token_refreshed_event_when_an_expired_token_is_renewed(inst): + creds = _FakeCreds(token="old", valid=False) + with _with_creds(creds): + proxy.get_refreshed_credentials(inst.recorder) + rec = one_event(inst, "auth.token_refreshed") + assert rec["attrs"]["had_token"] is True + assert creds.refreshes == 1 + + +def test_no_token_refreshed_event_when_credentials_are_still_valid(inst): + creds = _FakeCreds( + token="good", valid=True, expiry=datetime.datetime(2030, 1, 1) + ) + with _with_creds(creds): + proxy.get_refreshed_credentials(inst.recorder) + assert creds.refreshes == 0 + assert not events(inst, "auth.token_refreshed") + + +def test_no_token_refreshed_event_when_refresh_fails(inst): + creds = mock.Mock(valid=False, token=None) + creds.refresh.side_effect = RuntimeError("refresh failed") + with _with_creds(creds): + with pytest.raises(RuntimeError): + proxy.get_refreshed_credentials(inst.recorder) + assert not events(inst, "auth.token_refreshed") + + +def test_token_refreshed_expiry_is_none_for_non_datetime(inst): + creds = mock.Mock(valid=False, token=None) # expiry is a Mock + with _with_creds(creds): + proxy.get_refreshed_credentials(inst.recorder) + rec = one_event(inst, "auth.token_refreshed") + assert rec["attrs"]["expiry_wall_ns"] is None + aware = datetime.datetime(2030, 1, 1, tzinfo=datetime.timezone.utc) + assert proxy._expiry_wall_ns(mock.Mock(expiry=aware)) == ( + 1_893_456_000 * 10**9 + ) + + +def test_expiry_wall_ns_is_exact_and_tolerates_out_of_range(): + tz = datetime.timezone(datetime.timedelta(hours=2)) + expiry = datetime.datetime(2030, 1, 1, 2, 0, 0, 123456, tzinfo=tz) + assert proxy._expiry_wall_ns(mock.Mock(expiry=expiry)) == ( + 1_893_456_000 * 10**9 + 123_456_000 + ) + # Converting to UTC overflows past datetime.max: no expiry, no error. + minus_one = datetime.timezone(datetime.timedelta(hours=-1)) + too_late = datetime.datetime.max.replace(tzinfo=minus_one) + assert proxy._expiry_wall_ns(mock.Mock(expiry=too_late)) is None + + +def test_token_refresh_disabled_recorder_still_refreshes(): + creds = _FakeCreds() + with _with_creds(creds): + proxy.get_refreshed_credentials() + assert creds.refreshes == 1 + + +# -------------------------------------------------------------------------- +# progress.* events, progress_bars() and the forced-progress test hook +# -------------------------------------------------------------------------- + + +def _session_response(session_id, uuid): + from google.cloud.dataproc_v1 import Session + + response = Session() + response.name = ( + f"projects/test-project/locations/test-region/sessions/{session_id}" + ) + response.uuid = uuid + response.state = Session.State.ACTIVE + response.runtime_config.version = "3.0" + response.runtime_info.endpoints = { + "Spark Connect Server": "sc://spark-connect-server.example.com:443" + } + return response + + +@pytest.fixture +def dataproc_mocks(monkeypatch): + """Everything getOrCreate()/stop() touch, without a network.""" + from google.cloud.dataproc_v1 import Session + from pyspark.sql.connect.client.core import ConfigResult + from pyspark.sql.connect.proto import ConfigResponse + + from google.cloud.managed_spark_connect import ManagedSparkSession + + for key in list(os.environ): + if key.startswith("MANAGED_SPARK_CONNECT_"): + monkeypatch.delenv(key) + session_id = "sc-20240702-103952-abcdef" + uuid = "c002e4ef-fe5e-41a8-a157-160aa73e4f7f" + response = _session_response(session_id, uuid) + + with ( + mock.patch("google.auth.default") as auth_default, + mock.patch( + "google.cloud.dataproc_v1.SessionControllerClient" + ) as client_cls, + mock.patch( + "pyspark.sql.connect.client.SparkConnectClient.config", + return_value=ConfigResult.fromProto(ConfigResponse()), + ), + mock.patch.object( + ManagedSparkSession.Builder, + "generate_session_id", + return_value=session_id, + ), + mock.patch( + "google.cloud.managed_spark_connect.session.is_s8s_session_active", + return_value=True, + ), + mock.patch( + "google.cloud.managed_spark_connect.session._server_spark_version", + return_value=None, + ), + mock.patch( + "google.cloud.managed_spark_connect.environment.get_client_environment_label", + return_value="unknown", + ), + mock.patch("time.sleep"), + ): + cred = mock.MagicMock() + cred.token = "token" + auth_default.return_value = (cred, "") + client = client_cls.return_value + operation = mock.Mock() + operation.result.return_value = response + client.create_session.return_value = operation + terminating = Session() + terminating.state = Session.State.TERMINATING + client.get_session.return_value = terminating + yield {"client": client, "session_id": session_id, "uuid": uuid} + + +class FakeTqdm: + """A headless stand-in for ``tqdm`` that records what happened.""" + + def __init__(self, total=None, **kwargs): + self.total = total + self.kwargs = kwargs + self.n = 0 + self.resets = [] + self.closed = False + + def reset(self, total=None): + self.resets.append(total) + self.total = total + + def refresh(self): + pass + + def close(self): + self.closed = True + + +@pytest.fixture +def forced_progress(dataproc_mocks, monkeypatch): + """Force the progress handler on, headlessly, with FakeTqdm bars. + + Yields the list of bars created, in order. Depends on + ``dataproc_mocks`` so it runs after that fixture clears the + ``MANAGED_SPARK_CONNECT_*`` environment. + """ + from google.cloud.managed_spark_connect import ManagedSparkSession + from google.cloud.managed_spark_connect import session as session_mod + + created = [] + + def factory(**kwargs): + bar = FakeTqdm(**kwargs) + created.append(bar) + return bar + + monkeypatch.setenv(session_mod._FORCE_PROGRESS_BAR_ENV, "1") + ManagedSparkSession._set_tqdm_factory(factory) + ManagedSparkSession._execution_progress_bar.clear() + try: + yield created + finally: + ManagedSparkSession._set_tqdm_factory(None) + ManagedSparkSession._execution_progress_bar.clear() + + +@pytest.fixture +def session_factory(dataproc_mocks): + """Builds mocked ManagedSparkSessions and stops them afterwards.""" + from google.cloud.managed_spark_connect import ManagedSparkSession + + sessions = [] + + def make(inst=None): + builder = ManagedSparkSession.builder.projectId( + "test-project" + ).location("test-region") + if inst is not None: + builder = builder.instrumentation(inst) + session = builder.getOrCreate() + sessions.append(session) + return session + + try: + yield make + finally: + for session in sessions: + try: + session.stop() + except Exception: + pass + + +def _stage(num_tasks, completed, done=False): + from pyspark.sql.connect.shell.progress import StageInfo + + return StageInfo( + stage_id=1, + num_tasks=num_tasks, + num_completed_tasks=completed, + num_bytes_read=0, + done=done, + ) + + +def _handler(session): + (handler,) = session.client._progress_handlers + return handler + + +def test_progress_bar_events_and_snapshot(forced_progress, session_factory): + from google.cloud.managed_spark_connect import ManagedSparkSession + + inst = Instrumentation([InMemorySink()], {"run_id": "r1"}) + session = session_factory(inst) + handler = _handler(session) + + handler( + stages=[_stage(4, 0)], inflight_tasks=0, operation_id="op1", done=False + ) + rec = one_event(inst, "progress.bar_open") + assert rec["attrs"] == { + "operation_id": "op1", + "total_tasks": 4, + "run_id": "r1", + } + bars = ManagedSparkSession.progress_bars() + assert list(bars) == ["op1"] + (bar,) = forced_progress + assert bars["op1"] is bar + assert session.progress_bars() == bars + # A snapshot: mutating it does not touch the registry. + bars.clear() + assert list(ManagedSparkSession.progress_bars()) == ["op1"] + assert not events(inst, "progress.bar_update") + + handler( + stages=[_stage(4, 2)], inflight_tasks=2, operation_id="op1", done=False + ) + rec = one_event(inst, "progress.bar_update") + assert rec["attrs"] == { + "operation_id": "op1", + "completed_tasks": 2, + "total_tasks": 4, + "run_id": "r1", + } + assert (bar.n, bar.total, bar.closed) == (2, 4, False) + + # More stages appear: the bar is reset to the new total. + handler( + stages=[_stage(6, 3)], inflight_tasks=1, operation_id="op1", done=False + ) + assert bar.resets == [6] + assert (bar.n, bar.total) == (3, 6) + + handler( + stages=[_stage(6, 5)], inflight_tasks=0, operation_id="op1", done=True + ) + close = one_event(inst, "progress.bar_close") + assert close["attrs"] == { + "operation_id": "op1", + "completed_tasks": 6, # done with 5/6: pyspark hides skipped tasks + "total_tasks": 6, + "reason": "done", + "run_id": "r1", + } + assert bar.closed and bar.n == 6 + assert ManagedSparkSession.progress_bars() == {} + # The final done=True is reported by bar_close only, not as an update. + assert len(events(inst, "progress.bar_update")) == 2 + assert len(forced_progress) == 1 + + +def test_progress_handler_skips_a_bar_for_a_query_already_done( + forced_progress, session_factory +): + inst = Instrumentation([InMemorySink()]) + session = session_factory(inst) + _handler(session)( + stages=[_stage(3, 3)], inflight_tasks=0, operation_id="op", done=True + ) + assert forced_progress == [] + assert session.progress_bars() == {} + assert not events(inst) + + +def test_progress_handler_ignores_empty_stages_and_no_operation( + forced_progress, session_factory +): + inst = Instrumentation([InMemorySink()]) + session = session_factory(inst) + handler = _handler(session) + handler(stages=[], inflight_tasks=0, operation_id="op", done=False) + handler( + stages=[_stage(3, 1)], inflight_tasks=0, operation_id=None, done=False + ) + assert not events(inst) + assert session.progress_bars() == {} + + +def test_progress_handler_is_off_when_not_interactive_and_not_forced( + session_factory, monkeypatch +): + from google.cloud.managed_spark_connect import ManagedSparkSession + from google.cloud.managed_spark_connect import session as session_mod + + monkeypatch.delenv(session_mod._FORCE_PROGRESS_BAR_ENV, raising=False) + ManagedSparkSession._set_tqdm_factory(FakeTqdm) + try: + inst = Instrumentation([InMemorySink()]) + session = session_factory(inst) + with mock.patch.object( + session_mod.environment, "is_interactive", return_value=False + ): + _handler(session)( + stages=[_stage(3, 1)], + inflight_tasks=1, + operation_id="op", + done=False, + ) + assert session.progress_bars() == {} + assert not events(inst) + finally: + ManagedSparkSession._set_tqdm_factory(None) + + +@pytest.mark.parametrize("value", ["0", "", "no", "maybe"]) +def test_force_env_var_false_values(value, monkeypatch): + from google.cloud.managed_spark_connect import session as session_mod + + monkeypatch.setenv(session_mod._FORCE_PROGRESS_BAR_ENV, value) + assert session_mod._progress_bar_forced() is False + monkeypatch.setenv(session_mod._FORCE_PROGRESS_BAR_ENV, " TRUE ") + assert session_mod._progress_bar_forced() is True + + +@pytest.mark.parametrize("default", [True, False]) +def test_env_bool_unrecognized_value_returns_default(default, monkeypatch): + from google.cloud.managed_spark_connect import environment + + monkeypatch.setenv("MSC_TEST_BOOL", "maybe") + assert environment.env_bool("MSC_TEST_BOOL", default=default) is default + monkeypatch.setenv("MSC_TEST_BOOL", "off") + assert environment.env_bool("MSC_TEST_BOOL", default=True) is False + monkeypatch.setenv("MSC_TEST_BOOL", "on") + assert environment.env_bool("MSC_TEST_BOOL", default=False) is True + + +@pytest.mark.parametrize("raw", ["abc", "nan", "inf", "-inf"]) +def test_env_float_ignores_malformed_and_non_finite_values( + raw, monkeypatch, caplog +): + from google.cloud.managed_spark_connect import environment + + monkeypatch.setenv("MSC_TEST_FLOAT", raw) + with caplog.at_level("WARNING", logger=environment.logger.name): + assert environment.env_float("MSC_TEST_FLOAT", 2.5) == 2.5 + assert "MSC_TEST_FLOAT" in caplog.text + monkeypatch.setenv("MSC_TEST_FLOAT", "-3") + assert environment.env_float("MSC_TEST_FLOAT", 2.5, min_value=0) == 0 + + +def test_forced_progress_without_factory_uses_default_tqdm( + forced_progress, session_factory +): + import tqdm as tqdm_lib + + from google.cloud.managed_spark_connect import ManagedSparkSession + + ManagedSparkSession._set_tqdm_factory(None) + session = session_factory() # no instrumentation at all + handler = _handler(session) + handler( + stages=[_stage(2, 1)], inflight_tasks=1, operation_id="op", done=False + ) + bar = ManagedSparkSession.progress_bars()["op"] + assert isinstance(bar, tqdm_lib.tqdm) + handler( + stages=[_stage(2, 2)], inflight_tasks=0, operation_id="op", done=True + ) + assert ManagedSparkSession.progress_bars() == {} + + +def test_progress_bars_work_without_instrumentation( + forced_progress, session_factory +): + session = session_factory() + assert session.instrumentation is None + handler = _handler(session) + handler( + stages=[_stage(2, 1)], inflight_tasks=1, operation_id="op", done=False + ) + assert isinstance(session.progress_bars()["op"], FakeTqdm) + handler( + stages=[_stage(2, 2)], inflight_tasks=0, operation_id="op", done=True + ) + assert session.progress_bars() == {} + + +def _execute_request(operation_id): + from pyspark.sql.connect import proto + + req = proto.ExecutePlanRequest() + req.operation_id = operation_id + return req + + +@pytest.fixture +def fake_execute_and_fetch(): + """Replaces pyspark's _execute_and_fetch with a scripted stand-in. + + The script gets ``(client, req)`` and may drive the progress handler + like a real execution would, then return or raise. + """ + from pyspark.sql.connect.client import SparkConnectClient + + script = {} + + def fake(self, req, observations, self_destruct=False): + return script["fn"](self, req) + + with mock.patch.object(SparkConnectClient, "_execute_and_fetch", fake): + yield script + + +def _drive(client, operation_id, *, finish): + (handler,) = client._progress_handlers + handler( + stages=[_stage(3, 1)], + inflight_tasks=1, + operation_id=operation_id, + done=False, + ) + if finish: + handler( + stages=[_stage(3, 3)], + inflight_tasks=0, + operation_id=operation_id, + done=True, + ) + + +@pytest.mark.parametrize("wrap_client", [True, False]) +def test_execute_and_fetch_closes_a_bar_left_open( + forced_progress, session_factory, fake_execute_and_fetch, wrap_client +): + inst = Instrumentation( + [InMemorySink()], {"run_id": "r1"}, wrap_client=wrap_client + ) + session = session_factory(inst) + + def fn(client, req): + _drive(client, req.operation_id, finish=False) + return "result" + + fake_execute_and_fetch["fn"] = fn + assert ( + session.client._execute_and_fetch(_execute_request("op-left"), {}) + == "result" + ) + rec = one_event(inst, "progress.bar_close") + assert rec["attrs"] == { + "operation_id": "op-left", + "completed_tasks": 1, + "total_tasks": 3, + "reason": "query_end", + "run_id": "r1", + } + (bar,) = forced_progress + assert bar.closed + assert session.progress_bars() == {} + # The bar-closing hook is the session's own; wrap_client only controls + # the instrumentation wrappers. + assert ("to_pandas" in session.client.__dict__) is wrap_client + + +def test_execute_and_fetch_closes_a_bar_left_open_without_instrumentation( + forced_progress, session_factory, fake_execute_and_fetch +): + session = session_factory() + + def fn(client, req): + _drive(client, req.operation_id, finish=False) + return "result" + + fake_execute_and_fetch["fn"] = fn + session.client._execute_and_fetch(_execute_request("op-plain"), {}) + (bar,) = forced_progress + assert bar.closed + assert session.progress_bars() == {} + + +def test_execute_and_fetch_leaves_a_closed_bar_alone( + forced_progress, session_factory, fake_execute_and_fetch +): + inst = Instrumentation([InMemorySink()]) + session = session_factory(inst) + + def fn(client, req): + _drive(client, req.operation_id, finish=True) + return "result" + + fake_execute_and_fetch["fn"] = fn + session.client._execute_and_fetch(_execute_request("op-ok"), {}) + assert one_event(inst, "progress.bar_open") + rec = one_event(inst, "progress.bar_close") + assert rec["attrs"]["reason"] == "done" + + +def test_execute_and_fetch_without_a_bar_emits_nothing( + forced_progress, session_factory, fake_execute_and_fetch +): + inst = Instrumentation([InMemorySink()]) + session = session_factory(inst) + fake_execute_and_fetch["fn"] = lambda client, req: "r" + session.client._execute_and_fetch(_execute_request("op-none"), {}) + assert not events(inst) + + +@pytest.mark.parametrize("exc", [ValueError("boom"), KeyboardInterrupt()]) +def test_execute_and_fetch_closes_the_bar_on_error_and_interrupt( + forced_progress, session_factory, fake_execute_and_fetch, exc +): + inst = Instrumentation([InMemorySink()]) + session = session_factory(inst) + + def fn(client, req): + _drive(client, req.operation_id, finish=False) + raise exc + + fake_execute_and_fetch["fn"] = fn + with pytest.raises(type(exc)): + session.client._execute_and_fetch(_execute_request("op-x"), {}) + rec = one_event(inst, "progress.bar_close") + assert rec["attrs"]["operation_id"] == "op-x" + assert rec["attrs"]["reason"] == "query_end" + assert session.progress_bars() == {} + + +def test_client_hooks_are_removed_on_stop(forced_progress, session_factory): + inst = Instrumentation([InMemorySink()]) + session = session_factory(inst) + assert hasattr(session.client.__dict__["_execute_and_fetch"], "__wrapped__") + assert "to_pandas" in session.client.__dict__ + session.stop() + assert "_execute_and_fetch" not in session.client.__dict__ + assert "to_pandas" not in session.client.__dict__ + + +def test_client_hooks_are_removed_on_stop_without_instrumentation( + forced_progress, session_factory +): + session = session_factory() + assert "_execute_and_fetch" in session.client.__dict__ + session.stop() + assert "_execute_and_fetch" not in session.client.__dict__ + + +# -------------------------------------------------------------------------- +# guard.result +# -------------------------------------------------------------------------- + +SESSION_NAME = "projects/test-project/locations/test-region/sessions/sc-guard" + + +@pytest.fixture +def clean_liveness(monkeypatch): + from google.cloud.managed_spark_connect import session as session_mod + + monkeypatch.delenv(session_mod._ENV_LIVENESS_TTL_SECONDS, raising=False) + session_mod.invalidate_s8s_session_active() + yield session_mod + session_mod.invalidate_s8s_session_active() + + +@contextlib.contextmanager +def _guard_mode(mode): + from google.cloud.managed_spark_connect.client.core import _GUARD_MODE + + token = _GUARD_MODE.set(mode) + try: + yield + finally: + _GUARD_MODE.reset(token) + + +def _spans(inst, name): + return [ + r for r in inst.records if r["type"] == "span" and r["name"] == name + ] + + +def test_no_guard_result_for_a_cached_active_pre_rpc_hit(clean_liveness, inst): + session_mod = clean_liveness + is_active = session_mod.make_is_active_callback( + SESSION_NAME, None, recorder=inst.recorder + ) + with mock.patch.object( + session_mod, "_create_session_controller_client" + ) as make_client: + with _guard_mode("pre_rpc"): + assert is_active() is True + # Default liveness TTL: the pre-RPC guard trusts the cache, no RPC. The + # common case on every guarded RPC, so no event either. + make_client.assert_not_called() + assert not events(inst, "guard.result") + assert not _spans(inst, "dataproc.get_session") + + +def test_guard_result_on_error_inactive_after_rpc(clean_liveness, inst): + session_mod = clean_liveness + is_active = session_mod.make_is_active_callback( + SESSION_NAME, None, recorder=inst.recorder + ) + with mock.patch.object( + session_mod, "get_active_s8s_session_response", return_value=None + ) as get_active: + with _guard_mode("on_error"): + assert is_active() is False + get_active.assert_called_once_with(SESSION_NAME, None) + rec = one_event(inst, "guard.result") + assert set(rec) == EVENT_KEYS + assert rec["attrs"] == { + "mode": "on_error", + "cached": False, + "active": False, + "run_id": "r1", + } + (span,) = _spans(inst, "dataproc.get_session") + assert span["attrs"]["reason"] == "rpc_guard" + # The event is written after the RPC span ended. + assert rec["mono_ns"] >= span["start_mono_ns"] + span["duration_ns"] + + +def test_guard_result_active_after_rpc(clean_liveness, inst): + from google.cloud.dataproc_v1 import Session + + session_mod = clean_liveness + active_session = Session() + active_session.state = Session.State.ACTIVE + is_active = session_mod.make_is_active_callback( + SESSION_NAME, None, recorder=inst.recorder + ) + with mock.patch.object( + session_mod, + "get_active_s8s_session_response", + return_value=active_session, + ): + with _guard_mode("on_error"): + assert is_active() is True + rec = one_event(inst, "guard.result") + assert rec["attrs"]["active"] is True + assert rec["attrs"]["cached"] is False + assert "error_type" not in rec["attrs"] + + +def test_guard_result_records_error_type_when_check_raises( + clean_liveness, inst +): + from google.api_core.exceptions import ServiceUnavailable + + session_mod = clean_liveness + is_active = session_mod.make_is_active_callback( + SESSION_NAME, None, recorder=inst.recorder + ) + with mock.patch.object( + session_mod, + "get_active_s8s_session_response", + side_effect=ServiceUnavailable("control plane 503"), + ): + with _guard_mode("on_error"): + with pytest.raises(ServiceUnavailable): + is_active() + rec = one_event(inst, "guard.result") + assert rec["attrs"] == { + "mode": "on_error", + "cached": False, + "active": None, + "error_type": "ServiceUnavailable", + "run_id": "r1", + } + (span,) = _spans(inst, "dataproc.get_session") + assert span["status"] == "error" + + +def test_guard_result_records_refresh_error_from_get_session( + clean_liveness, inst +): + from google.auth.exceptions import RefreshError + + session_mod = clean_liveness + is_active = session_mod.make_is_active_callback( + SESSION_NAME, None, recorder=inst.recorder + ) + with mock.patch.object( + session_mod, "_create_session_controller_client" + ) as make_client: + make_client.return_value.get_session.side_effect = RefreshError( + "token expired" + ) + with _guard_mode("on_error"): + with pytest.raises(RefreshError): + is_active() + rec = one_event(inst, "guard.result") + assert rec["attrs"]["active"] is None + assert rec["attrs"]["error_type"] == "RefreshError" + # Not a verdict: the session is still cached as active. + with _guard_mode("pre_rpc"): + assert session_mod._check_cached_liveness(SESSION_NAME, None) is True + + +def test_guard_result_with_patched_is_s8s_session_active(clean_liveness, inst): + session_mod = clean_liveness + is_active = session_mod.make_is_active_callback( + SESSION_NAME, None, recorder=inst.recorder + ) + with mock.patch.object( + session_mod, "is_s8s_session_active", return_value=False + ): + with _guard_mode("pre_rpc"): + assert is_active() is False + rec = one_event(inst, "guard.result") + assert rec["attrs"]["mode"] == "pre_rpc" + assert rec["attrs"]["cached"] is False + assert rec["attrs"]["active"] is False + assert _spans(inst, "dataproc.get_session") + + +def test_guard_result_mode_is_direct_outside_the_channel(clean_liveness, inst): + session_mod = clean_liveness + is_active = session_mod.make_is_active_callback( + SESSION_NAME, None, recorder=inst.recorder + ) + with mock.patch.object( + session_mod, "get_active_s8s_session_response", return_value=None + ): + assert is_active() is False + assert one_event(inst, "guard.result")["attrs"]["mode"] == "direct" + + +def test_guard_result_through_guarded_multicallable(clean_liveness, inst): + from google.cloud.managed_spark_connect.client import core + + session_mod = clean_liveness + is_active = session_mod.make_is_active_callback( + SESSION_NAME, None, recorder=inst.recorder + ) + calls = [] + + def guard(): + if not is_active(): + raise RuntimeError( + "Session not active. Please create a new session" + ) + + def method(request): + calls.append(request) + return "response" + + # A guard that has to ask (the RPC path) reports through the channel. + with mock.patch.object( + session_mod, "is_s8s_session_active", return_value=True + ): + wrapped = core._GuardedMultiCallable(method, guard, mock.Mock()) + assert wrapped("req") == "response" + assert calls == ["req"] + rec = one_event(inst, "guard.result") + assert rec["attrs"]["mode"] == "pre_rpc" + assert rec["attrs"]["cached"] is False + assert rec["attrs"]["active"] is True + + +def test_guard_disabled_recorder_emits_nothing(clean_liveness): + from google.api_core.exceptions import ServiceUnavailable + + session_mod = clean_liveness + is_active = session_mod.make_is_active_callback(SESSION_NAME, None) + assert getattr(is_active, "_supports_error_guard") is True + with _guard_mode("pre_rpc"): + assert is_active() is True + with mock.patch.object( + session_mod, "get_active_s8s_session_response", return_value=None + ): + with _guard_mode("on_error"): + assert is_active() is False + with mock.patch.object( + session_mod, + "get_active_s8s_session_response", + side_effect=ServiceUnavailable("x"), + ): + with _guard_mode("on_error"): + with pytest.raises(ServiceUnavailable): + is_active() + + +# -------------------------------------------------------------------------- +# grpc.retry +# -------------------------------------------------------------------------- + + +class _RpcError(grpc.RpcError, grpc.Call): + """A grpc.RpcError with the Call surface pyspark's policies look at.""" + + def __init__(self, code, details="unavailable"): + super().__init__() + self._code = code + self._details = details + + def code(self): + return self._code + + def details(self): + return self._details + + def initial_metadata(self): + return () + + def trailing_metadata(self): + return () + + def is_active(self): + return False + + def time_remaining(self): + return None + + def cancel(self): + return False + + def add_callback(self, callback): + return False + + def __str__(self): + return f"{self._code.name}: {self._details}" + + +class _RetryClient: + """Just enough of SparkConnectClient for the retry hook.""" + + def __init__(self): + from pyspark.sql.connect.client.retries import DefaultPolicy + + self._retry_policies = [ + DefaultPolicy( + max_retries=3, + initial_backoff=1, + max_backoff=1, + backoff_multiplier=1.0, + jitter=0, + min_jitter_threshold=0, + ) + ] + + def _retrying(self): + from pyspark.sql.connect.client.retries import Retrying + + return Retrying(self._retry_policies) + + def to_pandas(self, *a, **k): + return "pandas" + + def _execute_and_fetch(self, *a, **k): + return ("table",) + + +@pytest.fixture +def retry_client(inst): + from google.cloud.managed_spark_connect.instrumentation import ( + install_client_wrappers, + uninstall_client_wrappers, + ) + + client = _RetryClient() + install_client_wrappers(client, inst.recorder) + with mock.patch("time.sleep"): + yield client + uninstall_client_wrappers(client) + + +def _run_with_retries(client, failures, exc_factory): + """pyspark's retry loop: fail ``failures`` times, then succeed.""" + calls = 0 + for attempt in client._retrying(): + with attempt: + calls += 1 + if calls <= failures: + raise exc_factory(calls) + return calls + + +def test_grpc_retry_events_count_retries_with_code_and_reason( + inst, retry_client +): + assert ( + _run_with_retries( + retry_client, + 2, + lambda n: _RpcError(grpc.StatusCode.UNAVAILABLE, f"try {n}"), + ) + == 3 + ) + retries = events(inst, "grpc.retry") + assert [r["attrs"]["attempt"] for r in retries] == [1, 2] + for rec in retries: + assert set(rec) == EVENT_KEYS + assert rec["attrs"] == { + "attempt": rec["attrs"]["attempt"], + "reason": "error", + "grpc_code": "UNAVAILABLE", + "error_type": "_RpcError", + "run_id": "r1", + } + + +def test_no_grpc_retry_event_on_success_or_non_retryable(inst, retry_client): + assert _run_with_retries(retry_client, 0, None) == 1 + with pytest.raises(_RpcError): + _run_with_retries( + retry_client, + 1, + lambda n: _RpcError(grpc.StatusCode.INVALID_ARGUMENT, "bad"), + ) + assert not events(inst, "grpc.retry") + + +def _gave_up_error(): + """What pyspark's retry loop raises when it gives up. + + pyspark < 4.1 wraps the last error in ``RetriesExceeded``; 4.1 removed + it (SPARK-53307) and re-raises the last error itself. + """ + try: + from pyspark.errors import RetriesExceeded + except ImportError: + return _RpcError + return RetriesExceeded + + +def test_grpc_retry_events_stop_when_retries_are_exceeded(inst, retry_client): + with pytest.raises(_gave_up_error()): + _run_with_retries( + retry_client, 10, lambda n: _RpcError(grpc.StatusCode.UNAVAILABLE) + ) + # max_retries=3: three retries happen, the fourth wait gives up. + assert [r["attrs"]["attempt"] for r in events(inst, "grpc.retry")] == [ + 1, + 2, + 3, + ] + + +def test_grpc_retry_immediate_retry_exception(inst, retry_client): + from pyspark.sql.connect.client.retries import RetryException + + _run_with_retries(retry_client, 1, lambda n: RetryException()) + rec = one_event(inst, "grpc.retry") + assert rec["attrs"]["reason"] == "reattach" + assert rec["attrs"]["grpc_code"] is None + assert rec["attrs"]["error_type"] == "RetryException" + + +def test_retry_wrapper_honours_policies_set_later(inst, retry_client): + from pyspark.sql.connect.client.retries import DefaultPolicy + + retry_client._retry_policies = [DefaultPolicy(max_retries=0)] + + with pytest.raises(_gave_up_error()): + _run_with_retries( + retry_client, 1, lambda n: _RpcError(grpc.StatusCode.UNAVAILABLE) + ) + assert not events(inst, "grpc.retry") + + +def test_retry_wrapper_is_installed_and_removed(inst): + from pyspark.sql.connect.client.retries import Retrying + + from google.cloud.managed_spark_connect.instrumentation import ( + install_client_wrappers, + uninstall_client_wrappers, + ) + + client = _RetryClient() + original = client._retrying + install_client_wrappers(client, inst.recorder) + assert client.__dict__["_retrying"].__wrapped__ == original + retrying = client._retrying() + assert isinstance(retrying, Retrying) + assert type(retrying).__name__ == "InstrumentedRetrying" + uninstall_client_wrappers(client) + assert "_retrying" not in client.__dict__ + assert type(client._retrying()) is Retrying + + +def test_instrumented_retrying_class_is_built_once_and_takes_recorder_by_keyword( + inst, +): + from google.cloud.managed_spark_connect.instrumentation import _client + + cls = _client._instrumented_retrying_class() + assert _client._instrumented_retrying_class() is cls + policies = _RetryClient()._retry_policies + with pytest.raises(TypeError): + cls(policies, inst.recorder) + assert cls(policies, recorder=inst.recorder)._recorder is inst.recorder + + +def test_retry_wrapper_skips_clients_without_retry_policies(inst): + from google.cloud.managed_spark_connect.instrumentation import ( + install_client_wrappers, + uninstall_client_wrappers, + ) + + class Bare: + + def to_pandas(self): + return 1 + + def _execute_and_fetch(self): + return 2 + + client = Bare() + install_client_wrappers(client, inst.recorder) + assert "_retrying" not in client.__dict__ + uninstall_client_wrappers(client) + + +def test_grpc_retry_is_wired_into_a_real_session( + forced_progress, session_factory +): + inst = Instrumentation([InMemorySink()]) + session = session_factory(inst) + assert type(session.client._retrying()).__name__ == "InstrumentedRetrying" + session.stop() + assert type(session.client._retrying()).__name__ == "Retrying" + + +# -------------------------------------------------------------------------- +# grpc.ExecutePlan span attrs: plan_kind / command_kind +# -------------------------------------------------------------------------- + + +def _plan_attrs(request): + from google.cloud.managed_spark_connect.instrumentation import _grpc + + return _grpc._request_attrs(request) + + +def test_execute_plan_root_plan_kind(): + from pyspark.sql.connect import proto + + req = proto.ExecutePlanRequest(session_id="s", operation_id="op") + req.plan.root.range.end = 10 + assert _plan_attrs(req) == { + "session_id": "s", + "operation_id": "op", + "plan_kind": "root", + } + + +def test_execute_plan_command_kind(): + from pyspark.sql.connect import proto + + req = proto.ExecutePlanRequest(session_id="s") + req.plan.command.create_dataframe_view.name = "v" + assert _plan_attrs(req) == { + "session_id": "s", + "plan_kind": "command", + "command_kind": "create_dataframe_view", + } + sql = proto.ExecutePlanRequest() + sql.plan.command.sql_command.sql = "select 1" + assert _plan_attrs(sql) == { + "plan_kind": "command", + "command_kind": "sql_command", + } + + +def test_plan_kind_absent_for_other_requests(): + from pyspark.sql.connect import proto + + assert _plan_attrs(proto.ExecutePlanRequest()) == {} + assert _plan_attrs(proto.ReleaseExecuteRequest(operation_id="op")) == { + "operation_id": "op" + } + assert _plan_attrs(proto.AnalyzePlanRequest(session_id="s")) == { + "session_id": "s" + } + assert _plan_attrs(mock.Mock(session_id="", operation_id="")) == {} + + +def test_plan_kinds_helper(): + from pyspark.sql.connect import proto + + from google.cloud.managed_spark_connect.instrumentation import _grpc + + root = proto.ExecutePlanRequest() + root.plan.root.range.end = 1 + assert _grpc._plan_kinds(root) == ("root", None) + command = proto.ExecutePlanRequest() + command.plan.command.sql_command.sql = "select 1" + assert _grpc._plan_kinds(command) == ("command", "sql_command") + assert _grpc._plan_kinds(proto.ExecutePlanRequest()) == (None, None) + assert _grpc._plan_kinds(proto.AnalyzePlanRequest()) == (None, None) + + +def test_plan_kind_on_a_real_interceptor_span(inst): + from pyspark.sql.connect import proto + + from google.cloud.managed_spark_connect.instrumentation import ( + GrpcInterceptor, + ) + + interceptor = GrpcInterceptor(inst.recorder) + details = mock.Mock(method="/spark.connect.SparkConnectService/ExecutePlan") + req = proto.ExecutePlanRequest(operation_id="op") + req.plan.command.sql_command.sql = "create temp view v as select 1" + outcome = mock.Mock() + outcome.done.return_value = True + outcome.result.return_value = proto.ExecutePlanResponse() + outcome.code.return_value = grpc.StatusCode.OK + interceptor.intercept_unary_unary(lambda d, r: outcome, details, req) + (span,) = [ + r + for r in inst.records + if r["type"] == "span" and r["name"] == "grpc.ExecutePlan" + ] + assert span["attrs"]["plan_kind"] == "command" + assert span["attrs"]["command_kind"] == "sql_command" + assert span["attrs"]["operation_id"] == "op" diff --git a/tests/unit/test_proxy.py b/tests/unit/test_proxy.py index 7235c5a3..6db96ec4 100644 --- a/tests/unit/test_proxy.py +++ b/tests/unit/test_proxy.py @@ -14,6 +14,7 @@ import socket import threading import time +from unittest import mock import pytest @@ -554,13 +555,71 @@ def test_connect_tcp_bridge_disables_compression_and_sets_max_size(): ws = proxy.connect_tcp_bridge("backend.example.com") assert ws is mock_ws_connect.return_value - mock_ws_connect.assert_called_once_with( - "wss://backend.example.com/tcp-over-websocket-bridge/35218cb7-1201-4940-89e8-48d8f03fed96", - additional_headers={"Authorization": "Bearer tok-123"}, - open_timeout=30, - compression=None, - max_size=proxy.MAX_WS_MESSAGE_SIZE, + mock_ws_connect.assert_called_once() + (url,) = mock_ws_connect.call_args.args + kwargs = mock_ws_connect.call_args.kwargs + assert url == ( + "wss://backend.example.com/tcp-over-websocket-bridge/35218cb7-1201-4940-89e8-48d8f03fed96" ) + assert kwargs["additional_headers"] == {"Authorization": "Bearer tok-123"} + assert kwargs["compression"] is None + assert kwargs["max_size"] == proxy.MAX_WS_MESSAGE_SIZE + + +def test_connect_tcp_bridge_unbounded_queue_and_ping_timeout_config(): + """A config snapshot, not a behavioural test: a slow consumer must not + block the websocket reader (``max_queue=None``), late PONGs are + tolerated for 60 s, and closing a broken backend waits at most 3 s.""" + from unittest import mock + + from google.cloud.managed_spark_connect.client import proxy + + creds = mock.Mock(token="tok-slow") + with ( + mock.patch.object( + proxy.googleauth, "default", return_value=(creds, None) + ), + mock.patch.object(proxy.websocketclient, "connect") as mock_ws_connect, + ): + proxy.connect_tcp_bridge("backend.example.com") + + kwargs = mock_ws_connect.call_args.kwargs + assert kwargs["max_queue"] is None + assert kwargs["ping_timeout"] == 60 + assert kwargs["close_timeout"] == 3 + + +def test_forward_bytes_shuts_down_before_close_on_eof_and_error(): + from google.cloud.managed_spark_connect.client import proxy + + class RecordingSock: + + def __init__(self, recv_items=()): + self.recv_items = list(recv_items) + self.events = [] + + def recv(self, _buf_size): + item = self.recv_items.pop(0) + if isinstance(item, BaseException): + raise item + return item + + def shutdown(self, how): + self.events.append(("shutdown", how)) + + def close(self): + self.events.append("close") + + shut_then_close = [("shutdown", socket.SHUT_RDWR), "close"] + + # Clean EOF and an error on from_sock both shut down, then close, each + # socket. The order across the two sockets doesn't matter. + for first_recv in (b"", ConnectionResetError("rst")): + from_sock = RecordingSock([first_recv]) + to_sock = RecordingSock() + proxy.forward_bytes("test", from_sock, to_sock) + assert to_sock.events == shut_then_close + assert from_sock.events == shut_then_close def test_managed_spark_session_proxy_context_manager_exits_cleanly(): @@ -664,7 +723,9 @@ def test_proxy_run_sets_tcp_nodelay_on_accepted_socket(): accepted_nodelay = [] done = threading.Event() - def fake_connect_sockets(conn_number, conn, backend_socket, span=None): + def fake_connect_sockets( + conn_number, conn, backend_socket, span=None, recorder=None + ): accepted_nodelay.append( conn.getsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY) ) @@ -688,3 +749,233 @@ def fake_connect_sockets(conn_number, conn, backend_socket, span=None): assert len(accepted_nodelay) == 1 assert accepted_nodelay[0] != 0 + + +class _RecordingRecorder: + enabled = True + + def __init__(self): + self.events = [] + + def event(self, name, **attrs): + self.events.append((name, attrs)) + + +def _forward_scripted_server_stream( + monkeypatch, steps, proxy_inst, limit_s, wrap=None +): + """Run forward_bytes server-to-client from a scripted WebSocket. + + The source is a real :class:`bridged_socket` carrying ``proxy_inst``'s + in-flight counter (as forward_connection builds it), optionally wrapped + by ``wrap``. Each step is ``(kind, at_s, *actions)``: the next WebSocket + ``recv()`` moves the fake monotonic clock to ``at_s``, runs ``actions`` + and then either times out (``"timeout"``, a routine poll) or returns a + frame (``"frame"``). Once the script runs out, ``recv()`` fails. + + Returns ``(silence_cut_events, steps_left)``. + """ + from google.cloud.managed_spark_connect.client import proxy + + clock = [0.0] + monkeypatch.setattr(proxy.time, "monotonic", lambda: clock[0]) + monkeypatch.setenv("MANAGED_SPARK_CONNECT_WS_SILENCE_LIMIT_S", limit_s) + remaining = list(steps) + + class ScriptedWebSocket: + + def recv(self, timeout=None): + if not remaining: + raise ConnectionError("script ended") + kind, at_s, *actions = remaining.pop(0) + clock[0] = at_s + for action in actions: + action() + if kind == "timeout": + raise TimeoutError("timed out") + return b"frame".hex() + + def close(self): + pass + + class ToSock: + + def send(self, bs): + return len(bs) + + def close(self): + pass + + from_sock = proxy.bridged_socket( + ScriptedWebSocket(), + inflight_rpcs=( + proxy_inst.get_inflight_rpcs if proxy_inst is not None else None + ), + ) + if wrap is not None: + from_sock = wrap(from_sock) + recorder = _RecordingRecorder() + proxy.forward_bytes( + "test", + from_sock, + ToSock(), + recorder=recorder, + conn=7, + direction=proxy.SERVER_TO_CLIENT, + ) + cuts = [ + attrs for name, attrs in recorder.events if name == "proxy.silence_cut" + ] + return cuts, remaining + + +def _proxy_with_inflight(n): + from google.cloud.managed_spark_connect.client import proxy + + p = proxy.ManagedSparkSessionProxy(0, "unused") + for _ in range(n): + p.inc_inflight() + return p + + +def test_silence_cut_with_rpc_in_flight(monkeypatch): + p = _proxy_with_inflight(1) + steps = [ + ("timeout", 60.0), # First poll that sees the RPC arms the timer. + ("timeout", 300.0), + ("timeout", 360.0), # 300 s of silence: cut. + ("timeout", 420.0), + ] + cuts, remaining = _forward_scripted_server_stream( + monkeypatch, steps, p, "300" + ) + assert cuts == [{"conn": 7, "silent_s": 300.0}] + assert remaining == [("timeout", 420.0)] + + +def test_silence_no_cut_without_rpc_in_flight(monkeypatch): + p = _proxy_with_inflight(0) + steps = [("timeout", 400.0), ("timeout", 800.0), ("frame", 850.0)] + cuts, remaining = _forward_scripted_server_stream( + monkeypatch, steps, p, "300" + ) + assert cuts == [] + assert remaining == [] + + +def test_silence_pre_rpc_idle_time_not_counted(monkeypatch): + p = _proxy_with_inflight(0) + steps = [ + ("timeout", 400.0), # Idle, no RPC. + ("timeout", 460.0, p.inc_inflight), # RPC starts. + ("timeout", 700.0), # 240 s since the RPC started: no cut. + ("timeout", 760.0), # 300 s: cut. + ("timeout", 820.0), + ] + cuts, remaining = _forward_scripted_server_stream( + monkeypatch, steps, p, "300" + ) + assert cuts == [{"conn": 7, "silent_s": 300.0}] + assert remaining == [("timeout", 820.0)] + + +def test_silence_received_frame_resets_timer(monkeypatch): + p = _proxy_with_inflight(1) + steps = [ + ("timeout", 60.0), + ("frame", 200.0), + ("timeout", 440.0), # 240 s since the frame: no cut. + ("timeout", 500.0), # 300 s: cut. + ("timeout", 560.0), + ] + cuts, remaining = _forward_scripted_server_stream( + monkeypatch, steps, p, "300" + ) + assert cuts == [{"conn": 7, "silent_s": 300.0}] + assert remaining == [("timeout", 560.0)] + + +def test_silence_finished_rpc_disarms_timer(monkeypatch): + p = _proxy_with_inflight(1) + steps = [ + ("timeout", 60.0), + ("timeout", 120.0, p.dec_inflight), # RPC done: timer disarmed. + ("timeout", 330.0, p.inc_inflight), # Next RPC: re-armed here. + ("timeout", 600.0), # 270 s: no cut. + ("timeout", 630.0), # 300 s: cut. + ] + cuts, remaining = _forward_scripted_server_stream( + monkeypatch, steps, p, "300" + ) + assert cuts == [{"conn": 7, "silent_s": 300.0}] + assert remaining == [] + + +def test_silence_env_0_disables(monkeypatch): + p = _proxy_with_inflight(1) + steps = [("timeout", 60.0), ("timeout", 1000.0)] + cuts, remaining = _forward_scripted_server_stream( + monkeypatch, steps, p, "0" + ) + assert cuts == [] + assert remaining == [] + + +def test_silence_cut_through_wrapper_that_delegates_attributes(monkeypatch): + # A wrapper around the backend socket (like a fault-injection bridge) + # that forwards unknown attributes keeps the cut armed. + class DelegatingWrapper: + + def __init__(self, inner): + self._inner = inner + + def recv(self, buff_size): + return self._inner.recv(buff_size) + + def close(self): + return self._inner.close() + + def __getattr__(self, name): + return getattr(self._inner, name) + + p = _proxy_with_inflight(1) + steps = [("timeout", 60.0), ("timeout", 360.0), ("timeout", 420.0)] + cuts, remaining = _forward_scripted_server_stream( + monkeypatch, steps, p, "300", wrap=DelegatingWrapper + ) + assert cuts == [{"conn": 7, "silent_s": 300.0}] + assert remaining == [("timeout", 420.0)] + + +def test_silence_no_cut_without_inflight_source(monkeypatch): + steps = [("timeout", 60.0), ("timeout", 1000.0)] + cuts, remaining = _forward_scripted_server_stream( + monkeypatch, steps, None, "300" + ) + assert cuts == [] + assert remaining == [] + + +def test_forward_connection_gives_backend_socket_the_inflight_counter(): + from google.cloud.managed_spark_connect.client import proxy + + p = _proxy_with_inflight(0) + backends = [] + + def fake_connect_sockets(conn_number, conn, backend_socket, *args): + backends.append(backend_socket) + + ws_cm = mock.MagicMock() + with ( + mock.patch.object(proxy, "connect_tcp_bridge", return_value=ws_cm), + mock.patch.object( + proxy, "connect_sockets", side_effect=fake_connect_sockets + ), + ): + proxy.forward_connection( + 1, mock.MagicMock(), ("127.0.0.1", 1), "h", proxy_instance=p + ) + (backend,) = backends + assert backend.inflight_rpcs() == 0 + p.inc_inflight() + assert backend.inflight_rpcs() == 1 diff --git a/tests/unit/test_reattach_patch.py b/tests/unit/test_reattach_patch.py new file mode 100644 index 00000000..f29e0241 --- /dev/null +++ b/tests/unit/test_reattach_patch.py @@ -0,0 +1,440 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. +"""Unit tests for the one-shot replay on INVALID_CURSOR.POSITION_NOT_AVAILABLE.""" + +import logging +import os +from unittest import mock + +import pytest +from pyspark.errors.exceptions.connect import ( + AnalysisException, + UnknownException, +) +from pyspark.sql.connect import proto + +from google.cloud.managed_spark_connect import ManagedSparkSession +from google.cloud.managed_spark_connect import session as session_mod +from google.cloud.managed_spark_connect.client.reattach_patch import ( + ENV_REPLAY_ON_POSITION_LOST, + POSITION_LOST_CONDITION, + install_replay_on_position_lost, +) +from google.cloud.managed_spark_connect.instrumentation import ( + InMemorySink, + Instrumentation, +) +from google.cloud.managed_spark_connect.instrumentation._client import ( + uninstall_client_hooks, +) + + +def _position_lost_error() -> UnknownException: + return UnknownException( + message="Cursor position is no longer available", + errorClass=POSITION_LOST_CONDITION, + ) + + +def _root_request(op_id: str = "op-1") -> proto.ExecutePlanRequest: + req = proto.ExecutePlanRequest(session_id="sess-1", operation_id=op_id) + req.plan.root.range.end = 10 + return req + + +def _command_request(op_id: str = "op-cmd") -> proto.ExecutePlanRequest: + req = proto.ExecutePlanRequest(session_id="sess-1", operation_id=op_id) + req.plan.command.sql_command.sql = "CREATE TEMP VIEW v AS SELECT 1" + return req + + +def _events(inst: Instrumentation, name: str): + return [ + r + for r in inst.records + if r.get("type") == "event" and r.get("name") == name + ] + + +class _FakeClient: + """Minimal SparkConnectClient stand-in: ``_execute_and_fetch`` records + each request it gets and delegates to ``fetch_fn``.""" + + def __init__(self, fetch_fn): + self._fetch_fn = fetch_fn + self.requests = [] + + def _execute_and_fetch(self, req, observations, self_destruct=False): + self.requests.append(req) + return self._fetch_fn(len(self.requests), req) + + +def _fail_first(attempt, req): + if attempt == 1: + raise _position_lost_error() + return ("table", None, [], [], {}) + + +def _always_fail(attempt, req): + raise _position_lost_error() + + +class _BarCloser: + + def __init__(self): + self.calls = [] + + def __call__(self, operation_id, reason="query_end"): + self.calls.append((operation_id, reason)) + + +@pytest.fixture(autouse=True) +def _no_replay_env(monkeypatch): + monkeypatch.delenv(ENV_REPLAY_ON_POSITION_LOST, raising=False) + + +def test_root_plan_replays_once_under_a_new_operation_id(caplog): + inst = Instrumentation([InMemorySink()], {"run_id": "r1"}) + bars = _BarCloser() + client = _FakeClient(_fail_first) + install_replay_on_position_lost( + client, recorder=inst.recorder, close_progress_bar=bars + ) + + req = _root_request("op-initial") + with caplog.at_level(logging.WARNING): + res = client._execute_and_fetch(req, {}) + + assert res[0] == "table" + first, second = client.requests + assert first is req + assert req.operation_id == "op-initial" # not mutated in place + assert second is not req + assert second.operation_id != "op-initial" + assert second.plan == req.plan + assert bars.calls == [ + ("op-initial", "replay"), + (second.operation_id, "query_end"), + ] + (ev,) = _events(inst, "reattach.position_lost") + assert ev["attrs"] == { + "operation_id": "op-initial", + "new_operation_id": second.operation_id, + "replay": True, + "run_id": "r1", + } + (warning,) = [r for r in caplog.records if r.levelno == logging.WARNING] + assert "op-initial" in warning.getMessage() + assert second.operation_id in warning.getMessage() + assert ENV_REPLAY_ON_POSITION_LOST in warning.getMessage() + + +def test_replay_warns_without_instrumentation(caplog): + client = _FakeClient(_fail_first) + install_replay_on_position_lost(client) + + with caplog.at_level(logging.WARNING): + client._execute_and_fetch(_root_request("op-quiet"), {}) + + assert len(client.requests) == 2 + assert any( + r.levelno == logging.WARNING and "op-quiet" in r.getMessage() + for r in caplog.records + ) + + +def test_replay_request_drops_reattach_options(): + def fetch_fn(attempt, req): + # Like pyspark's reattachable iterator: append to the request. + opt = proto.ExecutePlanRequest.RequestOption() + opt.reattach_options.reattachable = True + req.request_options.append(opt) + return _fail_first(attempt, req) + + client = _FakeClient(fetch_fn) + install_replay_on_position_lost(client) + + client._execute_and_fetch(_root_request(), {}) + + for req in client.requests: + assert ( + sum( + 1 for o in req.request_options if o.HasField("reattach_options") + ) + == 1 + ) + + +def test_command_plan_is_not_replayed(): + inst = Instrumentation([InMemorySink()]) + bars = _BarCloser() + client = _FakeClient(_always_fail) + install_replay_on_position_lost( + client, recorder=inst.recorder, close_progress_bar=bars + ) + + with pytest.raises(UnknownException): + client._execute_and_fetch(_command_request("op-cmd-1"), {}) + + assert [r.operation_id for r in client.requests] == ["op-cmd-1"] + assert bars.calls == [] + (ev,) = _events(inst, "reattach.position_lost") + assert ev["attrs"] == {"operation_id": "op-cmd-1", "replay": False} + + +def test_second_position_lost_error_propagates_unchained(): + inst = Instrumentation([InMemorySink()]) + bars = _BarCloser() + client = _FakeClient(_always_fail) + install_replay_on_position_lost( + client, recorder=inst.recorder, close_progress_bar=bars + ) + + with pytest.raises(UnknownException) as excinfo: + client._execute_and_fetch(_root_request("op-twice"), {}) + + assert excinfo.value.__context__ is None + first, second = client.requests + assert second.operation_id != "op-twice" + assert bars.calls == [ + ("op-twice", "replay"), + (second.operation_id, "query_end"), + ] + # The retry's failure is reported too, and not replayed again. + assert [ + (e["attrs"]["operation_id"], e["attrs"]["replay"]) + for e in _events(inst, "reattach.position_lost") + ] == [("op-twice", True), (second.operation_id, False)] + + +def test_other_errors_propagate_without_replay(): + inst = Instrumentation([InMemorySink()]) + + def fetch_fn(attempt, req): + raise AnalysisException( + message="Table not found", errorClass="TABLE_OR_VIEW_NOT_FOUND" + ) + + client = _FakeClient(fetch_fn) + install_replay_on_position_lost(client, recorder=inst.recorder) + + with pytest.raises(AnalysisException): + client._execute_and_fetch(_root_request(), {}) + assert len(client.requests) == 1 + assert _events(inst, "reattach.position_lost") == [] + + +def test_message_that_only_mentions_the_condition_is_not_replayed(): + def fetch_fn(attempt, req): + raise UnknownException( + message=f"[{POSITION_LOST_CONDITION}] mentioned in a message" + ) + + client = _FakeClient(fetch_fn) + install_replay_on_position_lost(client) + + with pytest.raises(UnknownException): + client._execute_and_fetch(_root_request(), {}) + assert len(client.requests) == 1 + + +def test_opt_out_env_disables_replay_but_still_emits_event(monkeypatch): + monkeypatch.setenv(ENV_REPLAY_ON_POSITION_LOST, "0") + inst = Instrumentation([InMemorySink()]) + client = _FakeClient(_fail_first) + install_replay_on_position_lost(client, recorder=inst.recorder) + + with pytest.raises(UnknownException): + client._execute_and_fetch(_root_request("op-a"), {}) + + assert len(client.requests) == 1 + (ev,) = _events(inst, "reattach.position_lost") + assert ev["attrs"] == {"operation_id": "op-a", "replay": False} + + +def test_uninstall_client_hooks_removes_replay(): + client = _FakeClient(_fail_first) + install_replay_on_position_lost(client) + assert "_execute_and_fetch" in client.__dict__ + + uninstall_client_hooks(client) + + assert "_execute_and_fetch" not in client.__dict__ + with pytest.raises(UnknownException): + client._execute_and_fetch(_root_request(), {}) + + +def test_session_builder_and_attached_session_wire_replay_and_close_bar( + monkeypatch, +): + from google.cloud.dataproc_v1 import Session + from pyspark.sql.connect.client import SparkConnectClient + from pyspark.sql.connect.client.core import ConfigResult + from pyspark.sql.connect.proto import ConfigResponse + from pyspark.sql.connect.shell.progress import StageInfo + + for key in list(os.environ): + if key.startswith("MANAGED_SPARK_CONNECT_"): + monkeypatch.delenv(key) + monkeypatch.setenv(session_mod._FORCE_PROGRESS_BAR_ENV, "1") + + closed_bars = [] + + class FakeTqdm: + + def __init__(self, total=None, **kwargs): + self.total = total + self.n = 0 + self.closed = False + + def reset(self, total=None): + self.total = total + + def refresh(self): + pass + + def close(self): + self.closed = True + closed_bars.append(self) + + ManagedSparkSession._set_tqdm_factory(FakeTqdm) + ManagedSparkSession._execution_progress_bar.clear() + + attempts = [] + + def fake_fetch(self, req, observations, self_destruct=False): + attempts.append(req.operation_id) + (handler,) = self._progress_handlers + stage = StageInfo( + stage_id=1, + num_tasks=4, + num_completed_tasks=2, + num_bytes_read=0, + done=False, + ) + handler( + stages=[stage], + inflight_tasks=2, + operation_id=req.operation_id, + done=False, + ) + if len(attempts) == 1: + raise _position_lost_error() + done_stage = StageInfo( + stage_id=1, + num_tasks=4, + num_completed_tasks=4, + num_bytes_read=0, + done=True, + ) + handler( + stages=[done_stage], + inflight_tasks=0, + operation_id=req.operation_id, + done=True, + ) + return ("ok-table", None, [], [], {}) + + response = Session() + response.name = "projects/p/locations/r/sessions/sc-test-session" + response.uuid = "11111111-2222-3333-4444-555555555555" + response.state = Session.State.ACTIVE + response.runtime_config.version = "3.0" + response.runtime_info.endpoints = { + "Spark Connect Server": "sc://spark-connect.example.com:443" + } + + try: + with ( + mock.patch( + "google.auth.default", + return_value=(mock.MagicMock(token="tok"), ""), + ), + mock.patch( + "google.cloud.dataproc_v1.SessionControllerClient" + ) as client_cls, + mock.patch( + "pyspark.sql.connect.client.SparkConnectClient.config", + return_value=ConfigResult.fromProto(ConfigResponse()), + ), + mock.patch.object( + SparkConnectClient, "_execute_and_fetch", fake_fetch + ), + mock.patch( + "google.cloud.managed_spark_connect.session.is_s8s_session_active", + return_value=True, + ), + mock.patch("time.sleep"), + ): + op = mock.Mock() + op.result.return_value = response + client_cls.return_value.create_session.return_value = op + terminating = Session() + terminating.state = Session.State.TERMINATING + client_cls.return_value.get_session.return_value = terminating + + inst = Instrumentation([InMemorySink()]) + session = ( + ManagedSparkSession.builder.projectId("p") + .location("r") + .instrumentation(inst) + .getOrCreate() + ) + try: + req = _root_request("op-first") + res = session.client._execute_and_fetch(req, {}) + assert res[0] == "ok-table" + assert len(attempts) == 2 + assert ManagedSparkSession.progress_bars() == {} + assert len(closed_bars) == 2 + (ev,) = _events(inst, "reattach.position_lost") + assert ev["attrs"]["replay"] is True + assert ev["attrs"]["operation_id"] == "op-first" + new_op_id = ev["attrs"]["new_operation_id"] + assert attempts == ["op-first", new_op_id] + bar_opens = _events(inst, "progress.bar_open") + bar_closes = _events(inst, "progress.bar_close") + assert [e["attrs"]["operation_id"] for e in bar_opens] == [ + "op-first", + new_op_id, + ] + assert [ + (e["attrs"]["operation_id"], e["attrs"]["reason"]) + for e in bar_closes + ] == [("op-first", "replay"), (new_op_id, "done")] + finally: + session.stop() + + # Attached client (direct ManagedSparkSession construction) also + # gets the replay wrapper. + attempts.clear() + attached_inst = Instrumentation([InMemorySink()]) + attached = ManagedSparkSession( + connection="sc://localhost:15002", + instrumentation=attached_inst, + ) + try: + attached._register_progress_execution_handler() + req2 = _root_request("op-attached") + res2 = attached.client._execute_and_fetch(req2, {}) + assert res2[0] == "ok-table" + assert len(attempts) == 2 + (ev2,) = _events(attached_inst, "reattach.position_lost") + assert ev2["attrs"]["operation_id"] == "op-attached" + assert ev2["attrs"]["replay"] is True + finally: + attached.stop() + finally: + ManagedSparkSession._set_tqdm_factory(None) + ManagedSparkSession._execution_progress_bar.clear()