From d001e62c18e3f79c453fd35631304815e6b5c151 Mon Sep 17 00:00:00 2001 From: Dave Borowitz Date: Thu, 15 Oct 2026 12:55:54 -0700 Subject: [PATCH 1/5] chore: ignore .venv/ --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index a7c34230..d1d12304 100644 --- a/.gitignore +++ b/.gitignore @@ -9,3 +9,4 @@ __pycache__ build/ dist/ +.venv/ From fe48fce5cc2db0e2c5428be27bbab0205bebfddf Mon Sep 17 00:00:00 2001 From: Dave Borowitz Date: Thu, 15 Oct 2026 12:55:54 -0700 Subject: [PATCH 2/5] fix: keep with_call()/future() on guarded ProxiedChannel methods ProxiedChannel._wrap_method replaced each gRPC multi-callable with a plain function, dropping the with_call() and future() entry points. Anything that layers grpc.intercept_channel() on top of the channel (upstream ChannelBuilder.add_interceptor, tracing/metrics interceptors) calls with_call() for unary-unary and stream-unary RPCs and failed with an AttributeError. Wrap the multi-callable in a small object that runs the liveness guard before __call__, with_call and future alike. subscribe()/unsubscribe() return None and are now forwarded instead of being wrapped. --- .../managed_spark_connect/client/core.py | 56 +++++-- tests/unit/test_core.py | 156 ++++++++++++++++++ 2 files changed, 197 insertions(+), 15 deletions(-) create mode 100644 tests/unit/test_core.py diff --git a/google/cloud/managed_spark_connect/client/core.py b/google/cloud/managed_spark_connect/client/core.py index 843e741c..298f6682 100644 --- a/google/cloud/managed_spark_connect/client/core.py +++ b/google/cloud/managed_spark_connect/client/core.py @@ -81,6 +81,34 @@ def _direct_channel(self) -> grpc.Channel: ) +class _GuardedMultiCallable: + """A gRPC multi-callable that runs a guard before every invocation. + + ``grpc.Channel.unary_unary()`` and friends return multi-callables that + are invoked either directly (``stub.Method(request)``), through + ``with_call()`` (as ``grpc.intercept_channel`` does for unary-unary and + stream-unary RPCs) or through ``future()``. All three entry points are + forwarded here so a guarded :class:`ProxiedChannel` stays a drop-in + replacement for the channel it wraps. + """ + + def __init__(self, wrapped_method, guard): + self._wrapped_method = wrapped_method + self._guard = guard + + def __call__(self, *args, **kwargs): + self._guard() + return self._wrapped_method(*args, **kwargs) + + def with_call(self, *args, **kwargs): + self._guard() + return self._wrapped_method.with_call(*args, **kwargs) + + def future(self, *args, **kwargs): + self._guard() + return self._wrapped_method.future(*args, **kwargs) + + class ProxiedChannel(grpc.Channel): def __init__(self, target_host, is_active_callback): @@ -105,22 +133,20 @@ def close(self): self._proxy.stop() return ret + def _check_active(self): + if ( + self._is_active_callback is not None + and not self._is_active_callback() + ): + logger.warning(f"Session is no longer active") + raise RuntimeError( + "Session not active. Please create a new session" + ) + def _wrap_method(self, wrapped_method): if self._is_active_callback is None: return wrapped_method - - def checked_method(*margs, **mkwargs): - if ( - self._is_active_callback is not None - and not self._is_active_callback() - ): - logger.warning(f"Session is no longer active") - raise RuntimeError( - "Session not active. Please create a new session" - ) - return wrapped_method(*margs, **mkwargs) - - return checked_method + return _GuardedMultiCallable(wrapped_method, self._check_active) def stream_stream(self, *args, **kwargs): return self._wrap_method(self._wrapped.stream_stream(*args, **kwargs)) @@ -129,7 +155,7 @@ def stream_unary(self, *args, **kwargs): return self._wrap_method(self._wrapped.stream_unary(*args, **kwargs)) def subscribe(self, *args, **kwargs): - return self._wrap_method(self._wrapped.subscribe(*args, **kwargs)) + return self._wrapped.subscribe(*args, **kwargs) def unary_stream(self, *args, **kwargs): return self._wrap_method(self._wrapped.unary_stream(*args, **kwargs)) @@ -138,4 +164,4 @@ def unary_unary(self, *args, **kwargs): return self._wrap_method(self._wrapped.unary_unary(*args, **kwargs)) def unsubscribe(self, *args, **kwargs): - return self._wrap_method(self._wrapped.unsubscribe(*args, **kwargs)) + return self._wrapped.unsubscribe(*args, **kwargs) diff --git a/tests/unit/test_core.py b/tests/unit/test_core.py new file mode 100644 index 00000000..6201df98 --- /dev/null +++ b/tests/unit/test_core.py @@ -0,0 +1,156 @@ +# Copyright 2024 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 channel builder and the guarded proxied channel. + +No network: the local proxy and the real gRPC channel are replaced by fakes. +""" + +from unittest import mock + +import grpc +import pytest + +from google.cloud.managed_spark_connect.client import core + + +class FakeMultiCallable: + """Stands in for the object returned by ``grpc.Channel.unary_unary``.""" + + def __init__(self): + self.calls = [] + + def __call__(self, request, **kwargs): + self.calls.append(("call", request)) + return f"response:{request}" + + def with_call(self, request, **kwargs): + self.calls.append(("with_call", request)) + return f"response:{request}", mock.Mock(spec=grpc.Call) + + def future(self, request, **kwargs): + self.calls.append(("future", request)) + return mock.Mock(spec=grpc.Future) + + +class FakeChannel(grpc.Channel): + """A grpc.Channel whose multi-callables are :class:`FakeMultiCallable`.""" + + def __init__(self): + self.multicallables = {} + self.closed = False + + def _get(self, method): + return self.multicallables.setdefault(method, FakeMultiCallable()) + + def unary_unary(self, method, *args, **kwargs): + return self._get(method) + + def unary_stream(self, method, *args, **kwargs): + return self._get(method) + + def stream_unary(self, method, *args, **kwargs): + return self._get(method) + + def stream_stream(self, method, *args, **kwargs): + return self._get(method) + + def subscribe(self, callback, try_to_connect=False): + return None + + def unsubscribe(self, callback): + return None + + def close(self): + self.closed = True + + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() + return False + + +@pytest.fixture +def fake_channel(): + channel = FakeChannel() + with ( + mock.patch.object(core.proxy, "ManagedSparkSessionProxy") as proxy_cls, + mock.patch.object(core, "DefaultChannelBuilder") as builder_cls, + ): + proxy_cls.return_value.port = 12345 + builder_cls.return_value.toChannel.return_value = channel + yield channel + + +def test_guarded_multicallable_exposes_with_call_and_future(fake_channel): + guard = mock.Mock(return_value=True) + channel = core.ProxiedChannel("example.com", guard) + + method = channel.unary_unary("/svc/Method") + assert method("a") == "response:a" + response, call = method.with_call("b") + assert response == "response:b" + assert isinstance(call, grpc.Call) + assert method.future("c") is not None + + assert guard.call_count == 3 + assert fake_channel.multicallables["/svc/Method"].calls == [ + ("call", "a"), + ("with_call", "b"), + ("future", "c"), + ] + + +def test_guard_failure_raises_before_the_rpc(fake_channel): + channel = core.ProxiedChannel("example.com", lambda: False) + method = channel.unary_stream("/svc/Stream") + with pytest.raises(RuntimeError, match="Session not active"): + method("a") + with pytest.raises(RuntimeError, match="Session not active"): + method.with_call("a") + assert fake_channel.multicallables["/svc/Stream"].calls == [] + + +def test_no_guard_returns_the_multicallable_unchanged(fake_channel): + channel = core.ProxiedChannel("example.com", None) + assert ( + channel.unary_unary("/svc/Method") + is fake_channel.multicallables["/svc/Method"] + ) + # subscribe/unsubscribe are forwarded, not wrapped. + assert channel.subscribe(lambda *_: None) is None + assert channel.unsubscribe(lambda *_: None) is None + + +def test_proxied_channel_works_under_intercept_channel(fake_channel): + """grpc.intercept_channel calls with_call() on unary-unary methods.""" + + class Recording(grpc.UnaryUnaryClientInterceptor): + + def __init__(self): + self.methods = [] + + def intercept_unary_unary(self, continuation, details, request): + self.methods.append(details.method) + return continuation(details, request) + + guard = mock.Mock(return_value=True) + interceptor = Recording() + channel = grpc.intercept_channel( + core.ProxiedChannel("example.com", guard), interceptor + ) + assert channel.unary_unary("/svc/Method")("req") == "response:req" + assert interceptor.methods == ["/svc/Method"] + assert guard.call_count == 1 From 1024e145a61c05065a2e0cf007bc62f9dd71b433 Mon Sep 17 00:00:00 2001 From: Dave Borowitz Date: Thu, 15 Oct 2026 13:27:16 -0700 Subject: [PATCH 3/5] fix: honor add_interceptor() on ManagedSparkChannelBuilder ManagedSparkChannelBuilder inherits add_interceptor() from the upstream ChannelBuilder, but toChannel() returned the ProxiedChannel directly, so registered interceptors were silently ignored (the inner channel is built by a fresh DefaultChannelBuilder with an empty interceptor list). Apply the builder's interceptors around the ProxiedChannel, matching what ChannelBuilder._insecure_channel/_secure_channel do upstream. --- .../managed_spark_connect/client/core.py | 6 +++- tests/unit/test_core.py | 29 +++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/google/cloud/managed_spark_connect/client/core.py b/google/cloud/managed_spark_connect/client/core.py index 298f6682..ecf88b3c 100644 --- a/google/cloud/managed_spark_connect/client/core.py +++ b/google/cloud/managed_spark_connect/client/core.py @@ -59,7 +59,11 @@ def toChannel(self) -> grpc.Channel: return self._proxied_channel() def _proxied_channel(self) -> grpc.Channel: - return ProxiedChannel(self.host, self._is_active_callback) + channel = ProxiedChannel(self.host, self._is_active_callback) + if self._interceptors: + logger.debug(f"Applying interceptors ({self._interceptors})") + channel = grpc.intercept_channel(channel, *self._interceptors) + return channel def _direct_channel(self) -> grpc.Channel: destination = f"{self.host}:{self.port}" diff --git a/tests/unit/test_core.py b/tests/unit/test_core.py index 6201df98..7d3869e7 100644 --- a/tests/unit/test_core.py +++ b/tests/unit/test_core.py @@ -154,3 +154,32 @@ def intercept_unary_unary(self, continuation, details, request): assert channel.unary_unary("/svc/Method")("req") == "response:req" assert interceptor.methods == ["/svc/Method"] assert guard.call_count == 1 + + +def test_channel_builder_honors_add_interceptor(fake_channel): + + class Recording(grpc.UnaryUnaryClientInterceptor): + + def __init__(self): + self.methods = [] + + def intercept_unary_unary(self, continuation, details, request): + self.methods.append(details.method) + return continuation(details, request) + + interceptor = Recording() + builder = core.ManagedSparkChannelBuilder( + "sc://example.com:443/;use_ssl=true", is_active_callback=lambda: True + ) + builder.add_interceptor(interceptor) + channel = builder.toChannel() + + assert channel.unary_unary("/svc/Method")("req") == "response:req" + assert interceptor.methods == ["/svc/Method"] + + +def test_channel_builder_without_interceptors_returns_proxied_channel( + fake_channel, +): + builder = core.ManagedSparkChannelBuilder("sc://example.com:443") + assert isinstance(builder.toChannel(), core.ProxiedChannel) From bbde1630a8225fb157898878cf661521bb2855e2 Mon Sep 17 00:00:00 2001 From: Dave Borowitz Date: Thu, 15 Oct 2026 13:45:45 -0700 Subject: [PATCH 4/5] feat: opt-in instrumentation of sessions, RPCs, proxy and client Add google.cloud.managed_spark_connect.instrumentation: a small, dependency-free tracing layer that records timed spans (plus events, counters and one meta record) to pluggable sinks (JSONL file, in-memory, logging, callback). Off by default: without an Instrumentation the library goes through a shared no-op recorder, installs no gRPC interceptor and leaves the pyspark client untouched. Opt in with ManagedSparkSession.builder.instrumentation(inst), or via MANAGED_SPARK_CONNECT_INSTRUMENTATION_FILE / _LOG. The session exposes `session.instrumentation` and closes it on stop(). What is recorded: - session lifecycle: get_or_create (create_rpc, wait_active with Dataproc state history, connect; or reuse with get_session, wait_endpoint, connect), stop (terminate_rpc, wait_terminating), and every GetSession (reason=rpc_guard|reuse|poll|stop). - every Spark Connect RPC at two layers of the channel stack: an outer span around the ProxiedChannel (what the client waits for, including the liveness guard, reported as guard_ns) and an inner span around the plain channel to the local proxy. Streaming calls go through an iterator proxy that keeps the grpc.Call surface, counts responses, ends the span at the final status timestamp, and handles abandoned streams via __del__ and weakref status callbacks. - the local proxy: start, each connection on the forwarding thread with byte/frame counters, and each websocket connect (ADC lookup, token refresh, handshake). - client-side conversions: plan building, Arrow to rows, Arrow to pandas. Spans nest through a contextvar and snapshot global attributes when they start, so a trace reads top-down from the user call to the wire. See docs/instrumentation.md for the API, the record schema and the span catalogue. --- README.md | 14 + docs/instrumentation.md | 116 ++ .../managed_spark_connect/client/core.py | 40 +- .../managed_spark_connect/client/proxy.py | 149 +- .../instrumentation/__init__.py | 100 ++ .../instrumentation/_client.py | 162 ++ .../instrumentation/_core.py | 566 +++++++ .../instrumentation/_grpc.py | 421 +++++ .../instrumentation/_sinks.py | 154 ++ google/cloud/managed_spark_connect/session.py | 335 +++- tests/unit/test_instrumentation.py | 1489 +++++++++++++++++ tests/unit/test_proxy.py | 188 +++ 12 files changed, 3622 insertions(+), 112 deletions(-) create mode 100644 docs/instrumentation.md create mode 100644 google/cloud/managed_spark_connect/instrumentation/__init__.py create mode 100644 google/cloud/managed_spark_connect/instrumentation/_client.py create mode 100644 google/cloud/managed_spark_connect/instrumentation/_core.py create mode 100644 google/cloud/managed_spark_connect/instrumentation/_grpc.py create mode 100644 google/cloud/managed_spark_connect/instrumentation/_sinks.py create mode 100644 tests/unit/test_instrumentation.py diff --git a/README.md b/README.md index 32803d27..97508091 100644 --- a/README.md +++ b/README.md @@ -116,6 +116,7 @@ The `ManagedSparkSession.builder` provides a fluent API to configure the session | `sessionTemplate(profile)` | Sets the Session Template to use. | | `subnetwork(subnet)` | Sets the subnetwork URI for 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). | ### Reusing Named Sessions Across Notebooks @@ -149,6 +150,19 @@ To create or connect to a named session: 5. If a session with the specified ID is already in a `TERMINATED` or `FAILED` state (or is `TERMINATING`), the client emits a warning, deletes the old session, and creates a new session with the same ID. Deleting the old session removes its previous metadata and Spark UI history from Dataproc, though logs remain available in Cloud Logging. +### Instrumentation (opt-in) + +The library can record what it does on your behalf — Dataproc session +lifecycle steps, every Spark Connect gRPC call (with and without the local +proxy's overhead), the proxy's own connections, and the client-side plan and +result conversions — as timed *spans*. It is off by default; opt in via +`ManagedSparkSession.builder.instrumentation(...)` or the +`MANAGED_SPARK_CONNECT_INSTRUMENTATION_FILE` and +`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. + ### Interactive Extras For notebooks and other interactive use, include the `interactive` extra (along diff --git a/docs/instrumentation.md b/docs/instrumentation.md new file mode 100644 index 00000000..fa91a66b --- /dev/null +++ b/docs/instrumentation.md @@ -0,0 +1,116 @@ +# Instrumentation (opt-in) + +The library can record what it does on your behalf — Dataproc session +lifecycle steps, every Spark Connect gRPC call (with and without the local +proxy's overhead), the proxy's own connections and the client-side plan and +result conversions — as timed *spans*. It is off by default and adds no work +to a session that has not opted in. + +Opt in from code: + +```python +from google.cloud.managed_spark_connect import ManagedSparkSession +from google.cloud.managed_spark_connect.instrumentation import ( + Instrumentation, JsonlFileSink, InMemorySink, +) + +inst = Instrumentation( + sinks=[JsonlFileSink("spark-connect-events.jsonl"), InMemorySink()], + attributes={"run_id": "nightly-42"}, # attached to every record +) +spark = ( + ManagedSparkSession.builder + .projectId(project).location(region) + .instrumentation(inst) + .getOrCreate() +) +spark.range(10).collect() + +spark.instrumentation is inst # True +inst.records # records of the InMemorySink(s) +inst.set_attribute("step", "load") # tag subsequent records +spark.stop() # flushes and closes the sinks +``` + +Or from the environment, without code changes (the builder wins when both +are used): + +| Variable | Effect | +|---|---| +| `MANAGED_SPARK_CONNECT_INSTRUMENTATION_FILE=/path/events.jsonl` | Append one JSON record per line to the file. | +| `MANAGED_SPARK_CONNECT_INSTRUMENTATION_LOG=1` | Log every record at `DEBUG` level through the `google.cloud.managed_spark_connect.instrumentation` logger. | + +## API + +```python +class Instrumentation: + def __init__(self, sinks=(), attributes=None, *, wrap_client=True) + @classmethod from_env() -> Instrumentation | None + recorder: Recorder; sinks: list[Sink] + set_attribute(key, value); remove_attribute(key) # global attrs on every subsequent record + records -> list[dict] # concatenated InMemorySink records + wrap_channel(channel, layer="outer") -> grpc.Channel + flush(); close() # idempotent; file sinks reopen on the next write + +class Recorder: # thread-safe + span(name, **attrs) -> ContextManager[Span] # nests under the current span (contextvars) + start_span(name, parent=None, **attrs) -> Span # for threads that own their spans; call end() + event(name, **attrs); counter(name, delta=1, **attrs) + trace_id: str; enabled: bool + +class Span: + span_id, parent_id, trace_id, name + set(**attrs); add_event(name, **attrs); record_exception(exc); end(end_mono_ns=None) + +# Sinks: write(record) / flush() / close() +JsonlFileSink(path, *, append=True) InMemorySink() LoggingSink(logger=None, level=DEBUG) CallbackSink(fn) + +# Lower level: GrpcInterceptor(recorder, layer), wrap_channel(channel, inst_or_recorder, layer), +# timed_guard(callback), install_client_wrappers(client, recorder) +``` + +`wrap_client=False` leaves the pyspark client untouched (no `client.*` spans). + +## Record schema + +Every record is a flat JSON object with a `type`. Global attributes are +merged into `attrs`. + +```json +{"type":"span","name":"grpc.ExecutePlan","span_id":"16 hex","parent_id":"16 hex | null", + "trace_id":"", + "start_wall_ns":0,"start_mono_ns":0,"duration_ns":0, + "status":"ok | error","error":{"type":"","message":"","grpc_code":"UNAVAILABLE"} | null, + "thread":"MainThread","attrs":{},"events":[{"name":"first_response","offset_ns":0}]} +{"type":"event","name":"...","wall_ns":0,"mono_ns":0,"thread":"...","attrs":{}} +{"type":"counter","name":"...","delta":1,"wall_ns":0,"attrs":{}} +{"type":"meta","wall_ns":0,"attrs":{"library_version":"","pyspark_version":"","grpcio_version":"", + "pyarrow_version":"","python":"","platform":"","pid":0}} +``` + +`start_wall_ns` is `time.time_ns()`; `start_mono_ns`, `duration_ns` and +event offsets come from `time.perf_counter_ns()`. One `meta` record is +written when an `Instrumentation` is created. Spans started on threads the +library owns (proxy connections, reattach pool threads) have +`parent_id: null` and the thread's name in `thread`. A span's `attrs` carry +the global attributes as they were when the span *started* (so a background +RPC stays attributed to the step that issued it); events, counters and the +meta record carry them as of when they are written. + +## Spans + +| 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.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) | +| `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` | +| `client.arrow_to_rows` | `rows` (recorded to the most recently instrumented session when several share a process) | +| `client.to_pandas` → `client.arrow_to_pandas` | `rows`, `columns` on the conversion span | + +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. diff --git a/google/cloud/managed_spark_connect/client/core.py b/google/cloud/managed_spark_connect/client/core.py index ecf88b3c..159de2a9 100644 --- a/google/cloud/managed_spark_connect/client/core.py +++ b/google/cloud/managed_spark_connect/client/core.py @@ -18,6 +18,7 @@ from pyspark.sql.connect.client import DefaultChannelBuilder from . import proxy +from .. import instrumentation as instrumentation_lib logger = logging.getLogger(__name__) @@ -40,8 +41,9 @@ class ManagedSparkChannelBuilder(DefaultChannelBuilder): True """ - def __init__(self, url, is_active_callback=None): + def __init__(self, url, is_active_callback=None, instrumentation=None): self._is_active_callback = is_active_callback + self._instrumentation = instrumentation super().__init__(url) def toChannel(self) -> grpc.Channel: @@ -59,11 +61,17 @@ def toChannel(self) -> grpc.Channel: return self._proxied_channel() def _proxied_channel(self) -> grpc.Channel: - channel = ProxiedChannel(self.host, self._is_active_callback) + channel = ProxiedChannel( + self.host, self._is_active_callback, self._instrumentation + ) if self._interceptors: logger.debug(f"Applying interceptors ({self._interceptors})") channel = grpc.intercept_channel(channel, *self._interceptors) - return channel + # Outermost, so the span covers the user's interceptors and the + # liveness guard as well as the RPC itself. + return instrumentation_lib.wrap_channel( + channel, self._instrumentation, instrumentation_lib.OUTER + ) def _direct_channel(self) -> grpc.Channel: destination = f"{self.host}:{self.port}" @@ -114,15 +122,31 @@ def future(self, *args, **kwargs): class ProxiedChannel(grpc.Channel): + """A gRPC channel to a Spark Connect server behind the local proxy. + + Every RPC first runs ``is_active_callback`` (when given) and is rejected + with a ``RuntimeError`` if the session is gone. With ``instrumentation`` + the callback is timed and the inner, pure-wire channel to the proxy is + traced as well (see the ``instrumentation`` package). + """ - def __init__(self, target_host, is_active_callback): + def __init__(self, target_host, is_active_callback, instrumentation=None): + recorder = instrumentation_lib.get_recorder(instrumentation) + if recorder.enabled: + is_active_callback = instrumentation_lib.timed_guard( + is_active_callback + ) self._is_active_callback = is_active_callback - self._proxy = proxy.ManagedSparkSessionProxy(0, target_host) + self._proxy = proxy.ManagedSparkSessionProxy( + 0, target_host, recorder=recorder + ) self._proxy.start() self._proxied_connect_url = f"sc://localhost:{self._proxy.port}" - self._wrapped = DefaultChannelBuilder( - self._proxied_connect_url - ).toChannel() + self._wrapped = instrumentation_lib.wrap_channel( + DefaultChannelBuilder(self._proxied_connect_url).toChannel(), + recorder, + instrumentation_lib.INNER, + ) def __enter__(self): return self diff --git a/google/cloud/managed_spark_connect/client/proxy.py b/google/cloud/managed_spark_connect/client/proxy.py index cf680439..1d26852c 100755 --- a/google/cloud/managed_spark_connect/client/proxy.py +++ b/google/cloud/managed_spark_connect/client/proxy.py @@ -18,12 +18,15 @@ import logging import socket import threading +import time 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.instrumentation import NOOP_RECORDER + parser = argparse.ArgumentParser() parser.add_argument("port") parser.add_argument("target_host") @@ -59,7 +62,7 @@ def close(self): return self._conn.close() -def connect_tcp_bridge(hostname): +def connect_tcp_bridge(hostname, recorder=NOOP_RECORDER, conn_number=None): """Create a socket-like connection to the given hostname using websocket. The backend server connected to over the websocket connection must be @@ -67,24 +70,44 @@ def connect_tcp_bridge(hostname): Args: hostname: The hostname of the server running the TCP-bridge backend. + recorder: Instrumentation recorder for the ``proxy.ws_connect`` span. + conn_number: The proxy connection this websocket serves, if any. Returns: A socket-like object with `recv` and `send` methods. """ path = "tcp-over-websocket-bridge/35218cb7-1201-4940-89e8-48d8f03fed96" - creds, _ = googleauth.default( - scopes=["https://www.googleapis.com/auth/cloud-platform"] - ) - creds.refresh(googleauthrequests.Request()) + with recorder.span("proxy.ws_connect", conn=conn_number, host=hostname): + with recorder.span("auth.default"): + creds, _ = googleauth.default( + scopes=["https://www.googleapis.com/auth/cloud-platform"] + ) + with recorder.span("auth.refresh"): + creds.refresh(googleauthrequests.Request()) + + with recorder.span("ws.handshake"): + return websocketclient.connect( + f"wss://{hostname}/{path}", + additional_headers={"Authorization": f"Bearer {creds.token}"}, + open_timeout=30, + ) + + +class _DirectionStats: + """Byte and frame counters for one forwarding direction. + + Each direction is driven by exactly one thread, so plain attributes are + enough; the totals are read only after both threads have finished. + """ - return websocketclient.connect( - f"wss://{hostname}/{path}", - additional_headers={"Authorization": f"Bearer {creds.token}"}, - open_timeout=30, - ) + __slots__ = ("bytes", "frames") + + def __init__(self): + self.bytes = 0 + self.frames = 0 -def forward_bytes(name, from_sock, to_sock): +def forward_bytes(name, from_sock, to_sock, stats=None): """Continuously stream bytes from the `from_sock` to the `to_sock`. This method terminates when either the `from_sock` is closed (causing @@ -97,6 +120,7 @@ def forward_bytes(name, from_sock, to_sock): name: forwarding thread name 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. """ while True: try: @@ -104,6 +128,9 @@ def forward_bytes(name, from_sock, to_sock): if not bs: to_sock.close() return + if stats is not None: + stats.bytes += len(bs) + stats.frames += 1 attempt = 0 while bs and (attempt < 10): attempt += 1 @@ -125,19 +152,24 @@ def forward_bytes(name, from_sock, to_sock): return -def connect_sockets(conn_number, from_sock, to_sock): +def connect_sockets(conn_number, from_sock, to_sock, span=None): """Create a connection between the two given ports. This method continuously streams bytes in both directions between the given `from_sock` and `to_sock` socket-like objects. 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. """ + out_stats = _DirectionStats() if span is not None else None + in_stats = _DirectionStats() if span is not None else None forward_name = f"{conn_number}-forward" t1 = threading.Thread( name=forward_name, target=forward_bytes, - args=[forward_name, from_sock, to_sock], + args=[forward_name, from_sock, to_sock, out_stats], daemon=True, ) t1.start() @@ -145,15 +177,29 @@ def connect_sockets(conn_number, from_sock, to_sock): t2 = threading.Thread( name=backward_name, target=forward_bytes, - args=[backward_name, to_sock, from_sock], + args=[backward_name, to_sock, from_sock, in_stats], daemon=True, ) t2.start() t1.join() t2.join() - - -def forward_connection(conn_number, conn, addr, target_host): + if span is not None: + span.set( + bytes_in=in_stats.bytes, + bytes_out=out_stats.bytes, + frames_in=in_stats.frames, + frames_out=out_stats.frames, + ) + + +def forward_connection( + conn_number, + conn, + addr, + target_host, + recorder=NOOP_RECORDER, + accepted_ns=None, +): """Create a connection to the target and forward `conn` to it. This method creates a socket-like object holding a connection to the given @@ -165,16 +211,42 @@ def forward_connection(conn_number, conn, addr, target_host): This method should be run inside a daemon thread so that it will not block program termination. + + Args: + conn_number: Sequence number of this connection within the proxy. + conn: The accepted client socket. + addr: The client address (for logging). + target_host: The backend host to connect to. + recorder: Instrumentation recorder. A ``proxy.connection`` span + covers this call (its children: ``proxy.ws_connect`` and the + auth/handshake steps) and carries the traffic counters. + 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. """ - with conn: - with connect_tcp_bridge(target_host) as websocket_conn: - backend_socket = bridged_socket(websocket_conn) - # 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) - connect_sockets(conn_number, conn, backend_socket) + if accepted_ns is None: + span = recorder.start_span("proxy.connection", conn=conn_number) + else: + span = recorder.start_span_at( + "proxy.connection", *accepted_ns, conn=conn_number + ) + with span: + with conn: + with connect_tcp_bridge( + target_host, recorder, conn_number + ) as websocket_conn: + backend_socket = bridged_socket(websocket_conn) + # 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) + connect_sockets( + conn_number, + conn, + backend_socket, + span if recorder.enabled else None, + ) class ManagedSparkSessionProxy(object): @@ -188,9 +260,10 @@ class ManagedSparkSessionProxy(object): Default Credentials. """ - def __init__(self, port, target_host): + def __init__(self, port, target_host, recorder=None): self._port = port self._target_host = target_host + self._recorder = recorder if recorder is not None else NOOP_RECORDER self._started = False self._killed = False self._conn_number = 0 @@ -209,10 +282,12 @@ def start(self, daemon=True): if self._started: raise Exception("Managed Spark session proxy already started") self._started = True - s = threading.Semaphore(value=0) - t = threading.Thread(target=self._run, args=[s], daemon=daemon) - t.start() - s.acquire() + with self._recorder.span("proxy.start") as span: + s = threading.Semaphore(value=0) + t = threading.Thread(target=self._run, args=[s], daemon=daemon) + t.start() + s.acquire() + span.set(port=self._port) def _run(self, s): with socket.create_server(("127.0.0.1", self._port)) as frontend_socket: @@ -221,11 +296,21 @@ def _run(self, s): s.release() while not self._killed: conn, addr = frontend_socket.accept() + accepted_ns = None + if self._recorder.enabled: + accepted_ns = (time.perf_counter_ns(), time.time_ns()) logger.debug(f"Accepted a connection from {addr}...") self._conn_number += 1 threading.Thread( target=forward_connection, - args=[self._conn_number, conn, addr, self._target_host], + args=[ + self._conn_number, + conn, + addr, + self._target_host, + self._recorder, + accepted_ns, + ], daemon=True, ).start() diff --git a/google/cloud/managed_spark_connect/instrumentation/__init__.py b/google/cloud/managed_spark_connect/instrumentation/__init__.py new file mode 100644 index 00000000..8a670108 --- /dev/null +++ b/google/cloud/managed_spark_connect/instrumentation/__init__.py @@ -0,0 +1,100 @@ +# Copyright 2024 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. +"""Opt-in instrumentation for Managed Spark Connect sessions. + +Off by default. Turn it on by passing an :class:`Instrumentation` to the +session builder, or through the environment:: + + from google.cloud.managed_spark_connect import ManagedSparkSession + from google.cloud.managed_spark_connect.instrumentation import ( + Instrumentation, JsonlFileSink, InMemorySink, + ) + + inst = Instrumentation( + sinks=[JsonlFileSink("events.jsonl"), InMemorySink()], + attributes={"run_id": "2024-01-01-a"}, + ) + spark = ( + ManagedSparkSession.builder.projectId(p).location(r) + .instrumentation(inst) + .getOrCreate() + ) + spark.range(10).collect() + spark.instrumentation.records # the InMemorySink's records + + # or, without code changes: + # MANAGED_SPARK_CONNECT_INSTRUMENTATION_FILE=/tmp/events.jsonl + # MANAGED_SPARK_CONNECT_INSTRUMENTATION_LOG=1 + +Records are dicts with a ``type`` of ``span``, ``event``, ``counter`` or +``meta``. A span record looks like:: + + {"type": "span", "name": "grpc.ExecutePlan", "span_id": "16 hex", + "parent_id": "16 hex or null", "trace_id": "", + "start_wall_ns": ..., "start_mono_ns": ..., "duration_ns": ..., + "status": "ok|error", + "error": {"type": "", "message": "", "grpc_code": "UNAVAILABLE"} | null, + "thread": "MainThread", "attrs": {...}, + "events": [{"name": "first_response", "offset_ns": ...}]} + +``start_wall_ns`` is ``time.time_ns()``; ``start_mono_ns`` and every +duration/offset come from ``time.perf_counter_ns()``. Events and counters +carry ``wall_ns`` (and events ``mono_ns``) instead of a start/duration. The +single ``meta`` record per :class:`Instrumentation` lists library, pyspark, +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. +""" + +from ._client import install_client_wrappers, uninstall_client_wrappers +from ._core import ( + ENV_FILE, + ENV_LOG, + NOOP_RECORDER, + Instrumentation, + Recorder, + Span, + get_recorder, +) +from ._grpc import INNER, OUTER, GrpcInterceptor, timed_guard, wrap_channel +from ._sinks import ( + CallbackSink, + InMemorySink, + JsonlFileSink, + LoggingSink, + Sink, +) + +__all__ = [ + "CallbackSink", + "ENV_FILE", + "ENV_LOG", + "GrpcInterceptor", + "INNER", + "InMemorySink", + "Instrumentation", + "JsonlFileSink", + "LoggingSink", + "NOOP_RECORDER", + "OUTER", + "Recorder", + "Sink", + "Span", + "get_recorder", + "install_client_wrappers", + "timed_guard", + "uninstall_client_wrappers", + "wrap_channel", +] diff --git a/google/cloud/managed_spark_connect/instrumentation/_client.py b/google/cloud/managed_spark_connect/instrumentation/_client.py new file mode 100644 index 00000000..bb4370d3 --- /dev/null +++ b/google/cloud/managed_spark_connect/instrumentation/_client.py @@ -0,0 +1,162 @@ +# Copyright 2024 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. +"""Spans for the client-side work around an RPC: plan building and result +conversion. + +Nothing here is touched unless an instrumented session is created. Then: + +* ``client.to_pandas`` and ``client.arrow_to_pandas`` come from wrappers + installed on that one ``SparkConnectClient`` instance (``to_pandas`` and + ``_execute_and_fetch``): the conversion span starts when the Arrow table + has been fetched and ends when ``to_pandas`` returns. +* ``client.plan_to_proto`` wraps ``LogicalPlan.to_proto`` at class level, + because pyspark builds the plan on the DataFrame before handing it to the + client. The wrapper dispatches on the client it is given: only clients + that carry a recorder get a span, others pay one attribute lookup. +* ``client.arrow_to_rows`` wraps ``ArrowTableToRowsConversion.convert`` at + 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). + +The class-level hooks are installed once per process, on first use, and +are inert while no instrumented client is registered. +""" + +import contextvars +import threading +from typing import Any, List, Optional + +from ._core import NOOP_RECORDER, Recorder + +_RECORDER_ATTR = "_managed_spark_connect_recorder" + +_lock = threading.Lock() +_registered: List[Recorder] = [] +_hooks_installed = False + +# Set while an instrumented ``to_pandas`` runs: [to_pandas span, conversion +# span or None]. Lets the ``_execute_and_fetch`` wrapper start the +# conversion span exactly when the fetch finishes. +_TO_PANDAS: contextvars.ContextVar[Optional[list]] = contextvars.ContextVar( + "managed_spark_connect_to_pandas", default=None +) + + +def _active_recorder() -> Recorder: + with _lock: + return _registered[-1] if _registered else NOOP_RECORDER + + +def _recorder_of(client: Any) -> Recorder: + return getattr(client, _RECORDER_ATTR, None) or NOOP_RECORDER + + +def _install_class_hooks() -> None: + global _hooks_installed + with _lock: + if _hooks_installed: + return + _hooks_installed = True + + from pyspark.sql.connect.plan import LogicalPlan + from pyspark.sql.conversion import ArrowTableToRowsConversion + + original_to_proto = LogicalPlan.to_proto + + def to_proto(self, session, *args, **kwargs): + recorder = _recorder_of(session) + if not recorder.enabled: + return original_to_proto(self, session, *args, **kwargs) + with recorder.span("client.plan_to_proto") as span: + plan = original_to_proto(self, session, *args, **kwargs) + span.set(plan_bytes=plan.ByteSize()) + return plan + + to_proto.__wrapped__ = original_to_proto # type: ignore[attr-defined] + LogicalPlan.to_proto = to_proto + + original_convert = ArrowTableToRowsConversion.convert + + def convert(table, schema, *args, **kwargs): + recorder = _active_recorder() + if not recorder.enabled: + return original_convert(table, schema, *args, **kwargs) + with recorder.span("client.arrow_to_rows", rows=table.num_rows): + return original_convert(table, schema, *args, **kwargs) + + convert.__wrapped__ = original_convert # type: ignore[attr-defined] + ArrowTableToRowsConversion.convert = staticmethod(convert) + + +def install_client_wrappers(client: Any, recorder: Recorder) -> None: + """Instrument one ``SparkConnectClient`` instance with ``recorder``.""" + if not recorder.enabled or getattr(client, _RECORDER_ATTR, None): + return + _install_class_hooks() + setattr(client, _RECORDER_ATTR, recorder) + 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 + + 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 + + +def uninstall_client_wrappers(client: Any) -> None: + """Undo :func:`install_client_wrappers` for ``client`` (if applied).""" + 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] + delattr(client, _RECORDER_ATTR) + with _lock: + try: + _registered.remove(recorder) + except ValueError: + pass diff --git a/google/cloud/managed_spark_connect/instrumentation/_core.py b/google/cloud/managed_spark_connect/instrumentation/_core.py new file mode 100644 index 00000000..06c1a1f8 --- /dev/null +++ b/google/cloud/managed_spark_connect/instrumentation/_core.py @@ -0,0 +1,566 @@ +# Copyright 2024 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. +"""Spans, the recorder that creates them and the ``Instrumentation`` handle. + +Everything here is deliberately small and dependency free: a span is a +plain object that turns into a ``dict`` when it ends, the recorder tracks +the current span with a :class:`contextvars.ContextVar` so nesting works +across ``with`` blocks (and across threads whenever the context is copied), +and the no-op recorder used while instrumentation is off never allocates. +""" + +import contextvars +import importlib.metadata +import logging +import os +import platform +import random +import sys +import threading +import time +from typing import ( + Any, + Callable, + ContextManager, + Dict, + List, + Mapping, + Optional, + Sequence, +) + +from ._sinks import ( + CallbackSink, + InMemorySink, + JsonlFileSink, + LoggingSink, + Sink, +) + +logger = logging.getLogger(__name__) + +ENV_FILE = "MANAGED_SPARK_CONNECT_INSTRUMENTATION_FILE" +ENV_LOG = "MANAGED_SPARK_CONNECT_INSTRUMENTATION_LOG" + +_CURRENT_SPAN: contextvars.ContextVar[Optional["Span"]] = ( + contextvars.ContextVar("managed_spark_connect_current_span", default=None) +) + +_TRUE_VALUES = {"1", "true", "yes", "on"} + + +def _new_id() -> str: + return f"{random.getrandbits(64):016x}" + + +def _grpc_code_name(exc: BaseException) -> Optional[str]: + """Best-effort gRPC status code name for ``exc``. + + Handles ``grpc.RpcError`` (which carries ``code()``) and + ``google.api_core`` exceptions (which carry ``grpc_status_code``). + """ + code = None + code_fn = getattr(exc, "code", None) + if callable(code_fn): + try: + code = code_fn() + except Exception: # pragma: no cover - defensive + code = None + if code is None: + code = getattr(exc, "grpc_status_code", None) + if code is None: + return None + return getattr(code, "name", str(code)) + + +def error_info(exc: BaseException) -> Dict[str, Any]: + """The ``error`` object recorded for an exception.""" + message = None + details = getattr(exc, "details", None) + if callable(details): + # grpc.RpcError: the status message, not the multi-line repr. + try: + message = details() + except Exception: # pragma: no cover - defensive + message = None + if not isinstance(message, str): + message = str(exc) + return { + "type": type(exc).__name__, + "message": message, + "grpc_code": _grpc_code_name(exc), + } + + +class Span: + """A timed, named unit of work. + + Use it as a context manager (``with recorder.span(...) as span:``) to + have it become the parent of spans opened inside the block, or call + :meth:`end` yourself for spans that outlive a single frame (proxy + connections, streaming RPCs). + """ + + __slots__ = ( + "name", + "span_id", + "parent_id", + "trace_id", + "attrs", + "events", + "start_wall_ns", + "start_mono_ns", + "thread", + "status", + "error", + "_recorder", + "_ended", + "_lock", + "_token", + ) + + def __init__( + self, + recorder: "Recorder", + name: str, + parent_id: Optional[str], + attrs: Dict[str, Any], + *, + start_mono_ns: Optional[int] = None, + start_wall_ns: Optional[int] = None, + ): + self._recorder = recorder + self.name = name + self.span_id = _new_id() + self.parent_id = parent_id + self.trace_id = recorder.trace_id + # Global attributes are snapshotted now, not when the span ends, so a + # span is attributed to the trial/step that started it even if it + # outlives it (background ReleaseExecute RPCs, proxy connections). + self.attrs = recorder._attributes() + self.attrs.update(attrs) + self.events: List[Dict[str, Any]] = [] + self.start_mono_ns = ( + time.perf_counter_ns() if start_mono_ns is None else start_mono_ns + ) + self.start_wall_ns = ( + time.time_ns() if start_wall_ns is None else start_wall_ns + ) + self.thread = threading.current_thread().name + self.status = "ok" + self.error: Optional[Dict[str, Any]] = None + self._ended = False + self._lock = threading.Lock() + self._token = None + + def set(self, **attrs: Any) -> None: + """Add or overwrite attributes.""" + with self._lock: + if not self._ended: + self.attrs.update(attrs) + + def add_event(self, name: str, **attrs: Any) -> None: + """Record a point in time inside the span (``offset_ns`` is added).""" + event = { + "name": name, + "offset_ns": time.perf_counter_ns() - self.start_mono_ns, + } + if attrs: + event.update(attrs) + with self._lock: + if not self._ended: + self.events.append(event) + + def record_exception(self, exc: BaseException) -> None: + """Mark the span failed; the error keeps type, message, grpc code.""" + info = error_info(exc) + with self._lock: + if self._ended: + return + self.status = "error" + self.error = info + + @property + def ended(self) -> bool: + return self._ended + + def end(self, end_mono_ns: Optional[int] = None) -> None: + """Finish the span and hand its record to the sinks. Idempotent. + + ``end_mono_ns`` (a ``perf_counter_ns`` value) freezes the end at an + earlier instant than now, for spans whose end was observed before + they could be emitted (a streaming RPC whose status arrived while + the consumer still held the iterator). + """ + if end_mono_ns is None: + end_mono_ns = time.perf_counter_ns() + with self._lock: + if self._ended: + return + self._ended = True + self._recorder._emit( + { + "type": "span", + "name": self.name, + "span_id": self.span_id, + "parent_id": self.parent_id, + "trace_id": self.trace_id or self._recorder.trace_id, + "start_wall_ns": self.start_wall_ns, + "start_mono_ns": self.start_mono_ns, + "duration_ns": max(0, end_mono_ns - self.start_mono_ns), + "status": self.status, + "error": self.error, + "thread": self.thread, + "attrs": self.attrs, + "events": self.events, + } + ) + + def __enter__(self) -> "Span": + self._token = _CURRENT_SPAN.set(self) + return self + + def __exit__(self, exc_type, exc, tb) -> bool: + if exc is not None: + self.record_exception(exc) + token, self._token = self._token, None + if token is not None: + try: + _CURRENT_SPAN.reset(token) + except ValueError: + # Exited in a different context than it was entered in. + _CURRENT_SPAN.set(None) + self.end() + return False + + def __repr__(self) -> str: + return ( + f"Span({self.name!r}, id={self.span_id}, parent={self.parent_id})" + ) + + +class _NoopSpan: + """The span handed out while instrumentation is disabled. + + One shared instance; every method is a no-op, so ``with + recorder.span(...)`` costs a method call and nothing else. + """ + + __slots__ = () + + name = "" + span_id = "" + parent_id = None + trace_id = "" + attrs: Dict[str, Any] = {} + events: List[Dict[str, Any]] = [] + ended = True + + def set(self, **attrs: Any) -> None: + pass + + def add_event(self, name: str, **attrs: Any) -> None: + pass + + def record_exception(self, exc: BaseException) -> None: + pass + + def end(self, end_mono_ns: Optional[int] = None) -> None: + pass + + def __enter__(self) -> "_NoopSpan": + return self + + def __exit__(self, exc_type, exc, tb) -> bool: + return False + + +_NOOP_SPAN = _NoopSpan() + + +class Recorder: + """Creates spans, events and counters. Thread-safe. + + ``span()`` nests under whatever span is current in the calling context; + ``start_span()`` does not look at the context (pass ``parent`` + explicitly), which is what threads that own their own spans want. + """ + + enabled = True + + def __init__( + self, + emit: Callable[[Dict[str, Any]], None], + attributes: Callable[[], Dict[str, Any]] = dict, + ): + self._emit = emit + self._attributes = attributes + self.trace_id = "" + + def current_span(self) -> Optional[Span]: + """The span that ``span()`` would nest under, if any.""" + return _CURRENT_SPAN.get() + + def span(self, name: str, **attrs: Any) -> ContextManager[Span]: + parent = _CURRENT_SPAN.get() + return Span(self, name, parent.span_id if parent else None, attrs) + + def start_span( + self, name: str, parent: Optional[Span] = None, **attrs: Any + ) -> Span: + return Span(self, name, parent.span_id if parent else None, attrs) + + def start_span_at( + self, + name: str, + start_mono_ns: int, + start_wall_ns: int, + parent: Optional[Span] = None, + **attrs: Any, + ) -> Span: + """Like :meth:`start_span`, for work that began before the span could + be created (``start_*_ns`` as ``perf_counter_ns``/``time_ns``).""" + return Span( + self, + name, + parent.span_id if parent else None, + attrs, + start_mono_ns=start_mono_ns, + start_wall_ns=start_wall_ns, + ) + + def event(self, name: str, **attrs: Any) -> None: + self._emit( + { + "type": "event", + "name": name, + "wall_ns": time.time_ns(), + "mono_ns": time.perf_counter_ns(), + "thread": threading.current_thread().name, + "attrs": attrs, + } + ) + + def counter(self, name: str, delta: float = 1, **attrs: Any) -> None: + self._emit( + { + "type": "counter", + "name": name, + "delta": delta, + "wall_ns": time.time_ns(), + "attrs": attrs, + } + ) + + +class _NoopRecorder: + """Drop-in for :class:`Recorder` that records nothing.""" + + enabled = False + + @property + def trace_id(self) -> str: + return "" + + @trace_id.setter + def trace_id(self, value: str) -> None: + pass + + def current_span(self) -> None: + return None + + def span(self, name: str, **attrs: Any) -> ContextManager[_NoopSpan]: + return _NOOP_SPAN + + def start_span(self, name: str, parent=None, **attrs: Any) -> _NoopSpan: + return _NOOP_SPAN + + def start_span_at( + self, name, start_mono_ns, start_wall_ns, parent=None, **attrs + ): + return _NOOP_SPAN + + def event(self, name: str, **attrs: Any) -> None: + pass + + def counter(self, name: str, delta: float = 1, **attrs: Any) -> None: + pass + + +NOOP_RECORDER = _NoopRecorder() +# The name library code reaches for when it has no Instrumentation. +_NOOP_RECORDER = NOOP_RECORDER + + +def _version_of(distribution: str) -> Optional[str]: + try: + return importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + return None + + +def _pyspark_version() -> Optional[str]: + try: + from pyspark.version import __version__ + + return __version__ + except Exception: + return None + + +def _meta_attrs() -> Dict[str, Any]: + return { + "library_version": _version_of("google-cloud-spark-connect"), + "pyspark_version": _pyspark_version(), + "grpcio_version": _version_of("grpcio"), + "pyarrow_version": _version_of("pyarrow"), + "python": platform.python_version(), + "platform": platform.platform(), + "pid": os.getpid(), + } + + +class Instrumentation: + """Opt-in tracing of what the library does on the user's behalf. + + Pass an instance to ``ManagedSparkSession.builder.instrumentation()`` + (or set the environment variables read by :meth:`from_env`) and the + library records spans for session lifecycle steps, every gRPC call at + both sides of the local proxy, the proxy's own connections and the + 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. + """ + + def __init__( + self, + sinks: Sequence[Sink] = (), + attributes: Optional[Mapping[str, Any]] = None, + *, + wrap_client: bool = True, + ): + self.sinks: List[Sink] = list(sinks) + self.wrap_client = wrap_client + self._attributes: Dict[str, Any] = dict(attributes or {}) + self._lock = threading.Lock() + self._failed_sinks: set = set() + self.recorder = Recorder(self._emit, self._attribute_snapshot) + self._emit( + {"type": "meta", "wall_ns": time.time_ns(), "attrs": _meta_attrs()} + ) + + @classmethod + def from_env(cls) -> Optional["Instrumentation"]: + """An instance configured from the environment, or ``None``. + + ``MANAGED_SPARK_CONNECT_INSTRUMENTATION_FILE=`` appends JSONL + records to ````; ``MANAGED_SPARK_CONNECT_INSTRUMENTATION_LOG`` + set to a true value logs records at DEBUG level. Both may be set. + """ + sinks: List[Sink] = [] + path = os.environ.get(ENV_FILE) + if path: + sinks.append(JsonlFileSink(path, append=True)) + if os.environ.get(ENV_LOG, "").strip().lower() in _TRUE_VALUES: + sinks.append(LoggingSink()) + if not sinks: + return None + return cls(sinks) + + @property + def enabled(self) -> bool: + return True + + @property + def attributes(self) -> Dict[str, Any]: + """A copy of the global attributes.""" + return self._attribute_snapshot() + + def _attribute_snapshot(self) -> Dict[str, Any]: + with self._lock: + return dict(self._attributes) + + def set_attribute(self, key: str, value: Any) -> None: + """Attach ``key=value`` to every record emitted from now on.""" + with self._lock: + self._attributes[key] = value + + def remove_attribute(self, key: str) -> None: + with self._lock: + self._attributes.pop(key, None) + + @property + def records(self) -> List[Dict[str, Any]]: + """The records of all :class:`InMemorySink` sinks, concatenated.""" + records: List[Dict[str, Any]] = [] + for sink in self.sinks: + if isinstance(sink, InMemorySink): + records.extend(sink.records) + return records + + def wrap_channel(self, channel, layer: str = "outer"): + """``channel`` with this instrumentation's gRPC interceptor added.""" + from ._grpc import wrap_channel + + return wrap_channel(channel, self.recorder, layer) + + def flush(self) -> None: + for sink in self.sinks: + self._call_sink(sink, "flush") + + def close(self) -> None: + """Flush and close every sink. Safe to call more than once.""" + for sink in self.sinks: + self._call_sink(sink, "flush") + self._call_sink(sink, "close") + + def _emit(self, record: Dict[str, Any]) -> None: + with self._lock: + if self._attributes and record["type"] != "span": + # Spans carry the globals they started with (see Span); + # events, counters and meta get them as of now. + attrs = record["attrs"] + for key, value in self._attributes.items(): + attrs.setdefault(key, value) + for sink in self.sinks: + try: + sink.write(record) + except Exception as e: + self._warn_sink(sink, "write", e) + + def _call_sink(self, sink: Sink, method: str) -> None: + try: + getattr(sink, method)() + except Exception as e: + self._warn_sink(sink, method, e) + + def _warn_sink(self, sink: Sink, method: str, exc: Exception) -> None: + # Telemetry must never take the user's job down; warn once per sink. + if id(sink) not in self._failed_sinks: + self._failed_sinks.add(id(sink)) + logger.warning( + "Instrumentation sink %r failed in %s(): %s", sink, method, exc + ) + + def __repr__(self) -> str: + return f"Instrumentation(sinks={self.sinks!r})" + + +def get_recorder(instrumentation: Optional[Instrumentation]): + """The recorder of ``instrumentation``, or the no-op recorder.""" + if instrumentation is None: + return NOOP_RECORDER + return instrumentation.recorder diff --git a/google/cloud/managed_spark_connect/instrumentation/_grpc.py b/google/cloud/managed_spark_connect/instrumentation/_grpc.py new file mode 100644 index 00000000..e1941788 --- /dev/null +++ b/google/cloud/managed_spark_connect/instrumentation/_grpc.py @@ -0,0 +1,421 @@ +# Copyright 2024 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. +"""A gRPC client interceptor that records one span per RPC. + +The library installs it at two layers of the channel stack: + +* ``outer``: around the whole :class:`ProxiedChannel`, so the span covers + everything the pyspark client waits for, including the session liveness + guard that runs before each RPC (its share is reported as ``guard_ns``). +* ``inner``: around the plain gRPC channel to the local proxy, so the span + is the pure wire time of the RPC. + +Inner spans are children of outer spans; both are children of whatever span +is current on the calling thread (e.g. ``client.to_pandas``). + +Interceptors receive requests before serialization, so ``request_bytes``, +``session_id`` and ``operation_id`` come straight from the request proto. +Streaming responses are counted by a thin iterator proxy that forwards the +``grpc.Call`` surface (``code()``, ``trailing_metadata()``, ``cancel()``, +...) to the real call, because pyspark reads it on errors. +""" + +import contextvars +import functools +import threading +import time +import weakref +from typing import Any, Callable, Dict, Iterator, Optional + +import grpc + +from ._core import _CURRENT_SPAN, NOOP_RECORDER, Recorder, Span + +OUTER = "outer" +INNER = "inner" + +_GUARD_NS: contextvars.ContextVar[int] = contextvars.ContextVar( + "managed_spark_connect_guard_ns", default=0 +) + + +def timed_guard(callback: Optional[Callable[[], bool]]): + """Wrap a session liveness callback so its duration is attributed. + + The time spent in ``callback`` is added to a context variable that the + outer interceptor reads after the RPC and records as ``guard_ns``. + ``None`` stays ``None``; a callback that is already wrapped is returned + as is, so wrapping twice does not double count. + """ + if callback is None or getattr(callback, "_timed_guard", None) is True: + return callback + + @functools.wraps(callback) + def guarded() -> bool: + start = time.perf_counter_ns() + try: + return callback() + finally: + _GUARD_NS.set(_GUARD_NS.get() + time.perf_counter_ns() - start) + + guarded._timed_guard = True # type: ignore[attr-defined] + return guarded + + +def _byte_size(message: Any) -> int: + byte_size = getattr(message, "ByteSize", None) + if byte_size is None: + return 0 + try: + return byte_size() + except Exception: # pragma: no cover - defensive + return 0 + + +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 + return attrs + + +def _code_name(call: Any) -> Optional[str]: + code_fn = getattr(call, "code", None) + if code_fn is None: + return None + try: + code = code_fn() + except Exception: + return None + return getattr(code, "name", str(code)) if code is not None else None + + +def _is_pending_future(outcome: Any) -> bool: + done = getattr(outcome, "done", None) + if done is None: + return False + try: + return not done() + except Exception: # pragma: no cover - defensive + return False + + +def _weak_status_callback(proxy: "_StreamingCall") -> Callable[[], None]: + """A status callback that does not keep ``proxy`` alive. + + gRPC holds registered callbacks until the RPC terminates; a strong + reference would form a cycle (call -> callback -> proxy -> call) and + delay the proxy's finalizer, which is what ends the span when the + consumer drops an undrained stream. + """ + ref = weakref.ref(proxy) + + def on_status() -> None: + live = ref() + if live is not None: + live._record_status() + + return on_status + + +class _CountingRequestIterator: + """Counts the request messages of a client-streaming RPC.""" + + def __init__(self, requests: Any, span: Span): + self._requests = iter(requests) + self._span = span + self.count = 0 + self.bytes = 0 + + def __iter__(self) -> "_CountingRequestIterator": + return self + + def __next__(self) -> Any: + message = next(self._requests) + if self.count == 0: + self._span.set(**_request_attrs(message)) + self.count += 1 + self.bytes += _byte_size(message) + self._span.set(request_count=self.count, request_bytes=self.bytes) + return message + + +class _StreamingCall(grpc.Call): + """Iterator proxy around a server-streaming call. + + Records ``first_response`` when the first message arrives and ``status`` + when gRPC reports the final status (from the channel's own thread, via + ``add_callback``, independent of how fast the consumer iterates). The + span ends when the consumer is done with the stream: at the end of the + iteration, on an error, on ``cancel()``, or when the consumer drops the + iterator without draining it (pyspark stops reading an ``ExecutePlan`` + stream as soon as it sees ``result_complete``), which is caught by + ``__del__``. ``response_count``/``response_bytes`` therefore count the + messages delivered to the consumer. + """ + + # Class-level defaults keep __del__/__getattr__ safe on a half-built + # instance (e.g. when __init__ raised). + _call: Any = None + _finished = True + + def __init__(self, call: Any, span: Span): + self._call = call + self._span = span + self._count = 0 + self._bytes = 0 + self._lock = threading.Lock() + self._finished = False + self._status_recorded = False + self._status_mono_ns: Optional[int] = None + add_callback = getattr(call, "add_callback", None) + if add_callback is None or not add_callback( + _weak_status_callback(self) + ): + # No callback support, or the RPC is already over. + self._record_status() + + # -- iteration ----------------------------------------------------- + + def __iter__(self) -> Iterator[Any]: + return self + + def __next__(self) -> Any: + try: + message = next(self._call) + except StopIteration: + self._finish(None) + raise + except BaseException as e: + self._finish(e) + raise + self._count += 1 + self._bytes += _byte_size(message) + if self._count == 1: + self._span.add_event("first_response") + return message + + def _record_status(self) -> None: + # Runs on the gRPC thread (status callback) or the consumer thread + # (_finish); the lock makes the check-and-set atomic so the status + # event is recorded exactly once. + with self._lock: + if self._status_recorded or self._finished: + return + is_active = getattr(self._call, "is_active", None) + if is_active is not None and is_active(): + # code() would block until the RPC completes; the callback + # that fires on completion records the status then. + return + self._status_recorded = True + self._status_mono_ns = time.perf_counter_ns() + code = _code_name(self._call) + if code is not None: + self._span.set(grpc_code=code) + self._span.add_event("status") + + def _finish(self, exc: Optional[BaseException]) -> None: + self._record_status() + with self._lock: + if self._finished: + return + self._finished = True + inner = self._call + if isinstance(inner, _StreamingCall): + # Nested layers: the inner span closes before the outer one. + inner._finish(exc) + self._span.set(response_count=self._count, response_bytes=self._bytes) + if exc is not None: + self._span.record_exception(exc) + # The RPC is over when its status arrived, even if the consumer + # held on to the iterator for a while (pyspark's reattach iterator + # is released only after its ReleaseExecute has been queued). + self._span.end(end_mono_ns=self._status_mono_ns) + + def __del__(self) -> None: + try: + self._finish(None) + except Exception: # pragma: no cover - never raise from a finalizer + pass + + # -- grpc.Call / grpc.RpcContext ----------------------------------- + + def initial_metadata(self): + return self._call.initial_metadata() + + def trailing_metadata(self): + return self._call.trailing_metadata() + + def code(self): + return self._call.code() + + def details(self): + return self._call.details() + + def debug_error_string(self): + return self._call.debug_error_string() + + def is_active(self) -> bool: + return self._call.is_active() + + def time_remaining(self): + return self._call.time_remaining() + + def cancel(self): + result = self._call.cancel() + self._span.set(cancelled=True) + self._finish(None) + return result + + def add_callback(self, callback) -> bool: + return self._call.add_callback(callback) + + def __getattr__(self, name: str) -> Any: + # Anything else (result(), exception(), running(), ...) belongs to + # the real call. + return getattr(self._call, name) + + def __repr__(self) -> str: + return f"_StreamingCall({self._span!r})" + + +class GrpcInterceptor( + grpc.UnaryUnaryClientInterceptor, + grpc.UnaryStreamClientInterceptor, + grpc.StreamUnaryClientInterceptor, + grpc.StreamStreamClientInterceptor, +): + """Records a ``grpc.`` span for every RPC on a channel.""" + + def __init__(self, recorder: Recorder, layer: str = OUTER): + if layer not in (OUTER, INNER): + raise ValueError(f"layer must be 'outer' or 'inner', got {layer!r}") + self._recorder = recorder + self._layer = layer + + # -- helpers ------------------------------------------------------- + + def _start(self, details: grpc.ClientCallDetails, kind: str) -> Span: + method = details.method + if isinstance(method, bytes): + method = method.decode() + return self._recorder.span( + "grpc." + method.rsplit("/", 1)[-1], + method=method, + kind=kind, + layer=self._layer, + ) + + def _run(self, span: Span, continuation: Callable, details, request): + """Call ``continuation`` with ``span`` current, attributing guard time. + + Exceptions raised synchronously (e.g. a failed liveness guard) end + the span as an error and propagate. + """ + span_token = _CURRENT_SPAN.set(span) + guard_token = _GUARD_NS.set(0) if self._layer == OUTER else None + try: + return continuation(details, request) + except BaseException as e: + if guard_token is not None: + span.set(guard_ns=_GUARD_NS.get()) + span.record_exception(e) + span.end() + raise + finally: + if guard_token is not None: + if not span.ended: + span.set(guard_ns=_GUARD_NS.get()) + _GUARD_NS.reset(guard_token) + try: + _CURRENT_SPAN.reset(span_token) + except ValueError: # pragma: no cover - different context + _CURRENT_SPAN.set(None) + + def _finish_unary(self, span: Span, outcome: Any) -> None: + code = _code_name(outcome) + if code is not None: + span.set(grpc_code=code) + span.add_event("status") + try: + response = outcome.result() + except BaseException as e: + span.record_exception(e) + else: + span.set(response_count=1, response_bytes=_byte_size(response)) + span.end() + + def _unary_response(self, span: Span, outcome: Any) -> Any: + if _is_pending_future(outcome): + outcome.add_done_callback( + lambda done: self._finish_unary(span, done) + ) + else: + self._finish_unary(span, outcome) + return outcome + + # -- grpc interceptor interface ----------------------------------- + + def intercept_unary_unary(self, continuation, client_call_details, request): + span = self._start(client_call_details, "unary_unary") + span.set(request_bytes=_byte_size(request), **_request_attrs(request)) + outcome = self._run(span, continuation, client_call_details, request) + return self._unary_response(span, outcome) + + def intercept_unary_stream( + self, continuation, client_call_details, request + ): + span = self._start(client_call_details, "unary_stream") + span.set(request_bytes=_byte_size(request), **_request_attrs(request)) + call = self._run(span, continuation, client_call_details, request) + return _StreamingCall(call, span) + + def intercept_stream_unary( + self, continuation, client_call_details, request_iterator + ): + span = self._start(client_call_details, "stream_unary") + requests = _CountingRequestIterator(request_iterator, span) + outcome = self._run(span, continuation, client_call_details, requests) + return self._unary_response(span, outcome) + + def intercept_stream_stream( + self, continuation, client_call_details, request_iterator + ): + span = self._start(client_call_details, "stream_stream") + requests = _CountingRequestIterator(request_iterator, span) + call = self._run(span, continuation, client_call_details, requests) + return _StreamingCall(call, span) + + +def _as_recorder(instrumentation_or_recorder: Any): + if instrumentation_or_recorder is None: + return NOOP_RECORDER + recorder = getattr(instrumentation_or_recorder, "recorder", None) + return recorder if recorder is not None else instrumentation_or_recorder + + +def wrap_channel( + channel: grpc.Channel, instrumentation_or_recorder: Any, layer: str = OUTER +) -> grpc.Channel: + """``channel`` with a :class:`GrpcInterceptor` at ``layer`` in front. + + Accepts an :class:`Instrumentation`, a :class:`Recorder` or ``None``; + returns ``channel`` itself when instrumentation is disabled. + """ + recorder = _as_recorder(instrumentation_or_recorder) + if not recorder.enabled: + return channel + return grpc.intercept_channel(channel, GrpcInterceptor(recorder, layer)) diff --git a/google/cloud/managed_spark_connect/instrumentation/_sinks.py b/google/cloud/managed_spark_connect/instrumentation/_sinks.py new file mode 100644 index 00000000..89f59696 --- /dev/null +++ b/google/cloud/managed_spark_connect/instrumentation/_sinks.py @@ -0,0 +1,154 @@ +# Copyright 2024 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. +"""Sinks: where instrumentation records go. + +A sink receives one ``dict`` per record (see the module docstring of +:mod:`google.cloud.managed_spark_connect.instrumentation` for the schema). +Sinks must tolerate being called from several threads at once; the +:class:`Instrumentation` serializes ``write`` calls, but ``flush`` and +``close`` may race with them. +""" + +import json +import logging +import os +import threading +from typing import Any, Callable, Dict, List, Optional, Protocol + + +class Sink(Protocol): + """The sink interface. Any object with these three methods will do.""" + + def write(self, record: Dict[str, Any]) -> None: + ... + + def flush(self) -> None: + ... + + def close(self) -> None: + ... + + +def _to_json(record: Dict[str, Any]) -> str: + # default=str: attribute values are usually scalars, but keep odd ones + # (enums, timestamps) from taking the whole record down. + return json.dumps( + record, separators=(",", ":"), ensure_ascii=False, default=str + ) + + +class JsonlFileSink: + """Appends one JSON object per line to a file. + + The file is line buffered, so every record reaches the OS as soon as it + is written and a crash loses at most the record being written. A closed + sink reopens its file (in append mode) on the next ``write``, so one sink + can outlive the session that closed it. + """ + + def __init__(self, path: "str | os.PathLike[str]", *, append: bool = True): + self.path = os.fspath(path) + self._append = append + self._lock = threading.Lock() + self._file = None + self._open(truncate=not append) + + def _open(self, truncate: bool) -> None: + self._file = open( + self.path, + "w" if truncate else "a", + buffering=1, + encoding="utf-8", + newline="\n", + ) + + def write(self, record: Dict[str, Any]) -> None: + line = _to_json(record) + "\n" + with self._lock: + if self._file is None: + self._open(truncate=False) + self._file.write(line) + + def flush(self) -> None: + with self._lock: + if self._file is not None: + self._file.flush() + + def close(self) -> None: + with self._lock: + if self._file is not None: + self._file.close() + self._file = None + + def __repr__(self) -> str: + return f"JsonlFileSink({self.path!r})" + + +class InMemorySink: + """Keeps records in a list, for tests and short interactive sessions.""" + + def __init__(self): + self.records: List[Dict[str, Any]] = [] + + def write(self, record: Dict[str, Any]) -> None: + self.records.append(record) + + def flush(self) -> None: + pass + + def close(self) -> None: + pass + + def clear(self) -> None: + self.records.clear() + + +class LoggingSink: + """Logs every record as one JSON line through :mod:`logging`.""" + + def __init__( + self, + logger: Optional[logging.Logger] = None, + level: int = logging.DEBUG, + ): + self.logger = logger or logging.getLogger( + "google.cloud.managed_spark_connect.instrumentation" + ) + self.level = level + + def write(self, record: Dict[str, Any]) -> None: + if self.logger.isEnabledFor(self.level): + self.logger.log(self.level, "%s", _to_json(record)) + + def flush(self) -> None: + pass + + def close(self) -> None: + pass + + +class CallbackSink: + """Hands every record to a callable, for custom exporters.""" + + def __init__(self, fn: Callable[[Dict[str, Any]], None]): + self.fn = fn + + def write(self, record: Dict[str, Any]) -> None: + self.fn(record) + + def flush(self) -> None: + pass + + def close(self) -> None: + pass diff --git a/google/cloud/managed_spark_connect/session.py b/google/cloud/managed_spark_connect/session.py index d6470acd..6b70a953 100644 --- a/google/cloud/managed_spark_connect/session.py +++ b/google/cloud/managed_spark_connect/session.py @@ -45,6 +45,13 @@ from google.auth.exceptions import DefaultCredentialsError from google.cloud.managed_spark_connect.client import ManagedSparkChannelBuilder from google.cloud.managed_spark_connect.exceptions import ManagedSparkConnectException +from google.cloud.managed_spark_connect.instrumentation import ( + NOOP_RECORDER, + Instrumentation, + get_recorder, + install_client_wrappers, + uninstall_client_wrappers, +) from google.cloud.managed_spark_connect.pypi_artifacts import PyPiArtifacts from google.cloud.dataproc_v1 import ( AuthenticationConfig, @@ -267,6 +274,39 @@ def _is_valid_session_id(session_id: str) -> bool: return bool(re.match(pattern, session_id)) +def _timestamp_str(value: Any) -> Optional[str]: + if value is None: + return None + isoformat = getattr(value, "isoformat", None) + return isoformat() if callable(isoformat) else str(value) + + +def _session_state_attrs(session: Any) -> Dict[str, Any]: + """Instrumentation attributes describing a Dataproc session's lifecycle. + + Best effort: anything missing (e.g. on a mocked session) is left out. + """ + attrs: Dict[str, Any] = {} + try: + attrs["create_time"] = _timestamp_str(session.create_time) + attrs["state_time"] = _timestamp_str(session.state_time) + state = getattr(session, "state", None) + attrs["state"] = getattr(state, "name", str(state)) + attrs["state_history"] = json.dumps( + [ + { + "state": getattr(h.state, "name", str(h.state)), + "state_message": h.state_message, + "state_start_time": _timestamp_str(h.state_start_time), + } + for h in session.state_history + ] + ) + except Exception: # pragma: no cover - mocks, partial protos + pass + return attrs + + class ManagedSparkSession(SparkSession): """The entry point to programming Spark with the Dataset and DataFrame API. @@ -297,6 +337,9 @@ class ManagedSparkSession(SparkSession): _active_session_uses_custom_id: ClassVar[bool] = False _execution_progress_bar = dict() + #: The :class:`Instrumentation` this session records to, or ``None``. + instrumentation: Optional[Instrumentation] = None + class Builder(SparkSession.Builder): def __init__(self): @@ -316,6 +359,18 @@ def __init__(self): SessionControllerClient ] = None self._session_template_controller_client = None + self._instrumentation: Optional[Instrumentation] = None + self._last_session_response: Optional[Session] = None + + @property + def _recorder(self): + return get_recorder(self._instrumentation) + + def instrumentation(self, instrumentation: Optional[Instrumentation]): + """Record what the library does (see the ``instrumentation`` + package). Takes precedence over the environment variables.""" + self._instrumentation = instrumentation + return self @property def session_controller_client(self) -> SessionControllerClient: @@ -471,6 +526,13 @@ def create(self) -> "ManagedSparkSession": def __create_spark_connect_session_from_s8s( self, session_response, session_name ) -> "ManagedSparkSession": + recorder = self._recorder + recorder.trace_id = session_response.uuid + self._last_session_response = session_response + with recorder.span("session.connect"): + return self.__connect(session_response, session_name) + + def __connect(self, session_response, session_name): ManagedSparkSession._active_s8s_session_uuid = session_response.uuid ManagedSparkSession._project_id = self._project_id ManagedSparkSession._region = self._region @@ -480,15 +542,27 @@ def __create_spark_connect_session_from_s8s( ) url = f"{spark_connect_url}/;session_id={session_response.uuid};use_ssl=true" logger.debug(f"Spark Connect URL: {url}") + recorder = self._recorder + + def is_active(): + with recorder.span( + "dataproc.get_session", reason="rpc_guard", new_client=True + ): + return is_s8s_session_active( + session_name, self._client_options + ) + self._channel_builder = ManagedSparkChannelBuilder( url, - is_active_callback=lambda: is_s8s_session_active( - session_name, self._client_options - ), + is_active_callback=is_active, + instrumentation=self._instrumentation, ) assert self._channel_builder is not None - session = ManagedSparkSession(connection=self._channel_builder) + session = ManagedSparkSession( + connection=self._channel_builder, + instrumentation=self._instrumentation, + ) # Register handler for Cell Execution Progress bar session._register_progress_execution_handler() @@ -586,21 +660,26 @@ def create_session_pbar(): self._client_options, ) ) - operation = SessionControllerClient( - client_options=self._client_options - ).create_session(session_request) + with self._recorder.span("session.create_rpc"): + operation = SessionControllerClient( + client_options=self._client_options + ).create_session(session_request) self._display_session_link_on_creation(session_id) self._display_view_session_details_button(session_id) create_session_pbar_thread.start() - session_response: Session = operation.result( - polling=retry.Retry( - predicate=POLLING_PREDICATE, - initial=5.0, # seconds - maximum=5.0, # seconds - multiplier=1.0, - timeout=1200, # seconds + with self._recorder.span( + "session.wait_active", poll_interval_s=5.0 + ) as wait_span: + session_response: Session = operation.result( + polling=retry.Retry( + predicate=POLLING_PREDICATE, + initial=5.0, # seconds + maximum=5.0, # seconds + multiplier=1.0, + timeout=1200, # seconds + ) ) - ) + wait_span.set(**_session_state_attrs(session_response)) stop_create_session_pbar_event.set() create_session_pbar_thread.join() file_path = ( @@ -666,13 +745,26 @@ def create_session_pbar(): def _wait_for_session_available( self, session_name: str, timeout: int = 1200 + ) -> Session: + with self._recorder.span( + "session.wait_endpoint", poll_interval_s=5.0 + ) as span: + session = self._poll_for_session_endpoint(session_name, timeout) + span.set(**_session_state_attrs(session)) + return session + + def _poll_for_session_endpoint( + self, session_name: str, timeout: int ) -> Session: start_time = time.time() while time.time() - start_time < timeout: try: - session = self.session_controller_client.get_session( - name=session_name - ) + with self._recorder.span( + "dataproc.get_session", reason="poll", new_client=False + ): + session = self.session_controller_client.get_session( + name=session_name + ) if "Spark Connect Server" in session.runtime_info.endpoints: return session if session.state in ( @@ -770,45 +862,63 @@ def _get_exiting_active_session( session_name = f"projects/{self._project_id}/locations/{self._region}/sessions/{s8s_session_id}" session_response = None session = None + recorder = self._recorder if s8s_session_id is not None: - session_response = get_active_s8s_session_response( - session_name, self._client_options + reuse_span = recorder.span( + "session.reuse", session_id=s8s_session_id ) - session = ManagedSparkSession.getActiveSession() - - if session is None: - session = ManagedSparkSession._default_session + else: + reuse_span = NOOP_RECORDER.span("session.reuse") + with reuse_span: + if s8s_session_id is not None: + with recorder.span( + "dataproc.get_session", reason="reuse", new_client=True + ): + session_response = get_active_s8s_session_response( + session_name, self._client_options + ) + if session_response is not None: + recorder.trace_id = session_response.uuid + reuse_span.set(found=session_response is not None) + session = ManagedSparkSession.getActiveSession() - if session_response is not None: if session is None: - session_response = self._wait_for_session_available( - session_name - ) - session = self.__create_spark_connect_session_from_s8s( - session_response, session_name + session = ManagedSparkSession._default_session + + if session_response is not None: + self._last_session_response = session_response + if session is None: + session_response = self._wait_for_session_available( + session_name + ) + session = self.__create_spark_connect_session_from_s8s( + session_response, session_name + ) + links = _session_links( + self._region, + self._project_id, + s8s_session_id, + html=False, ) - links = _session_links( - self._region, self._project_id, s8s_session_id, html=False - ) - versions = self._session_versions(session) - subject = ( - f"Using existing Managed Spark Session {s8s_session_id} " - f"{versions}" - ).rstrip() - print( - f"{subject}; configuration changes may not be applied: " - f"{links.url}" - ) - self._display_view_session_details_button(s8s_session_id) - return session - else: - if session is not None: + versions = self._session_versions(session) + subject = ( + f"Using existing Managed Spark Session {s8s_session_id} " + f"{versions}" + ).rstrip() print( - f"{s8s_session_id} Managed Spark Session is not active, stopping and creating a new one" + f"{subject}; configuration changes may not be applied: " + f"{links.url}" ) - session.stop() + self._display_view_session_details_button(s8s_session_id) + return session + else: + if session is not None: + print( + f"{s8s_session_id} Managed Spark Session is not active, stopping and creating a new one" + ) + session.stop() - return None + return None def getOrCreate(self) -> "ManagedSparkSession": with ManagedSparkSession._lock: @@ -829,13 +939,35 @@ def getOrCreate(self) -> "ManagedSparkSession": f"Error while creating Managed Spark Session: location is not set" ) - # Handle custom session ID by setting it early and letting existing logic handle it - if self._custom_session_id: - self._handle_custom_session_id() - - session = self._get_exiting_active_session() - if session is None: - session = self.__create() + if self._instrumentation is None: + existing_session = ( + ManagedSparkSession.getActiveSession() + or ManagedSparkSession._default_session + ) + if ( + existing_session is not None + and existing_session.instrumentation is not None + and ( + self._custom_session_id is None + or self._custom_session_id + == ManagedSparkSession._active_s8s_session_id + ) + ): + self._instrumentation = existing_session.instrumentation + else: + self._instrumentation = Instrumentation.from_env() + + with self._recorder.span("session.get_or_create") as span: + # Handle custom session ID by setting it early and letting existing logic handle it + if self._custom_session_id: + self._handle_custom_session_id() + + session = self._get_exiting_active_session() + mode = "reuse" + if session is None: + mode = "create" + session = self.__create() + span.set(mode=mode, **self._session_span_attrs()) # Register this session as the instantiated SparkSession for compatibility # with tools and libraries that expect SparkSession._instantiatedSession @@ -845,6 +977,21 @@ def getOrCreate(self) -> "ManagedSparkSession": return session + def _session_span_attrs(self) -> Dict[str, Any]: + """Identity attributes of the session for ``session.*`` spans.""" + attrs: Dict[str, Any] = { + "session_id": ManagedSparkSession._active_s8s_session_id, + "uuid": ManagedSparkSession._active_s8s_session_uuid, + } + response = self._last_session_response + if response is not None: + try: + attrs["runtime_version"] = response.runtime_config.version + attrs["template"] = response.session_template + except Exception: # pragma: no cover - mocks + pass + return attrs + def _handle_custom_session_id(self): """Handle custom session ID by checking if it exists and setting _active_s8s_session_id.""" session_response = self._get_session_by_id(self._custom_session_id) @@ -1160,9 +1307,12 @@ def _get_session_by_id(self, session_id: str) -> Optional[Session]: try: get_request = GetSessionRequest(name=session_name) - session = self.session_controller_client.get_session( - get_request - ) + with self._recorder.span( + "dataproc.get_session", reason="reuse", new_client=False + ): + session = self.session_controller_client.get_session( + get_request + ) logger.debug( f"Found existing session {session_id} in state: {session.state}" @@ -1217,9 +1367,12 @@ def _wait_for_termination(self, session_name: str, timeout: int = 180): while time.time() - start_time < timeout: try: get_request = GetSessionRequest(name=session_name) - session = self.session_controller_client.get_session( - get_request - ) + with self._recorder.span( + "dataproc.get_session", reason="stop", new_client=False + ): + session = self.session_controller_client.get_session( + get_request + ) if session.state in [ Session.State.TERMINATED, @@ -1257,6 +1410,7 @@ def __init__( self, connection: Union[str, ManagedSparkChannelBuilder], user_id: Optional[str] = None, + instrumentation: Optional[Instrumentation] = None, ): """ Creates a new ManagedSparkSession for the Spark Connect interface. @@ -1270,10 +1424,17 @@ def __init__( user_id : str, optional If not set, will default to the $USER environment. Defining the user ID as part of the connection string takes precedence. + instrumentation : :class:`Instrumentation`, optional + Records spans for the client-side work of this session (see the + ``instrumentation`` package). Exposed as ``self.instrumentation``. """ super().__init__(connection, user_id) + self.instrumentation = instrumentation + if instrumentation is not None and instrumentation.wrap_client: + install_client_wrappers(self.client, instrumentation.recorder) + execute_plan_request_base_method = ( self.client._execute_plan_request_with_metadata ) @@ -1594,6 +1755,17 @@ def stop(self, terminate: Optional[bool] = None) -> None: >>> spark.stop(terminate=False) """ + recorder = get_recorder(self.instrumentation) + try: + with recorder.span("session.stop") as span: + self._stop(terminate, span) + finally: + if self.instrumentation is not None: + uninstall_client_wrappers(self.client) + recorder.trace_id = "" + self.instrumentation.close() + + def _stop(self, terminate: Optional[bool], span) -> None: with ManagedSparkSession._lock: if ManagedSparkSession._active_s8s_session_id is not None: # Determine if we should terminate the server-side session @@ -1605,6 +1777,10 @@ def stop(self, terminate: Optional[bool] = None) -> None: else: should_terminate = terminate + span.set( + terminate=should_terminate, + session_id=ManagedSparkSession._active_s8s_session_id, + ) if should_terminate: # Terminate the server-side session logger.debug( @@ -1615,6 +1791,7 @@ def stop(self, terminate: Optional[bool] = None) -> None: ManagedSparkSession._region, ManagedSparkSession._active_s8s_session_id, self._client_options, + recorder=get_recorder(self.instrumentation), ) else: # Client-side cleanup only @@ -1654,10 +1831,16 @@ def stop(self, terminate: Optional[bool] = None) -> None: def terminate_s8s_session( - project_id, region, active_s8s_session_id, client_options=None + project_id, + region, + active_s8s_session_id, + client_options=None, + recorder=None, ): from google.cloud.dataproc_v1 import SessionControllerClient + if recorder is None: + recorder = NOOP_RECORDER logger.debug(f"Terminating Managed Spark Session: {active_s8s_session_id}") terminate_session_request = TerminateSessionRequest() session_name = f"projects/{project_id}/locations/{region}/sessions/{active_s8s_session_id}" @@ -1665,18 +1848,26 @@ def terminate_s8s_session( state = None try: session_client = SessionControllerClient(client_options=client_options) - session_client.terminate_session(terminate_session_request) + with recorder.span("session.terminate_rpc"): + session_client.terminate_session(terminate_session_request) get_session_request = GetSessionRequest() get_session_request.name = session_name state = Session.State.ACTIVE - while ( - state != Session.State.TERMINATING - and state != Session.State.TERMINATED - and state != Session.State.FAILED - ): - session = session_client.get_session(get_session_request) - state = session.state - time.sleep(1) + with recorder.span( + "session.wait_terminating", poll_interval_s=1.0 + ) as wait_span: + while ( + state != Session.State.TERMINATING + and state != Session.State.TERMINATED + and state != Session.State.FAILED + ): + with recorder.span( + "dataproc.get_session", reason="stop", new_client=False + ): + session = session_client.get_session(get_session_request) + state = session.state + time.sleep(1) + wait_span.set(state=getattr(state, "name", str(state))) except NotFound: logger.debug( f"{active_s8s_session_id} Managed Spark Session already deleted" diff --git a/tests/unit/test_instrumentation.py b/tests/unit/test_instrumentation.py new file mode 100644 index 00000000..05729a3d --- /dev/null +++ b/tests/unit/test_instrumentation.py @@ -0,0 +1,1489 @@ +# Copyright 2024 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 google.cloud.managed_spark_connect.instrumentation (no network).""" + +import json +import logging +import os +import threading +from unittest import mock + +import grpc +import pytest +from pyspark.sql.connect import proto + +from google.cloud.managed_spark_connect import instrumentation as inst_lib +from google.cloud.managed_spark_connect.client import core +from google.cloud.managed_spark_connect.instrumentation import ( + CallbackSink, + GrpcInterceptor, + InMemorySink, + Instrumentation, + JsonlFileSink, + LoggingSink, + NOOP_RECORDER, + Recorder, + get_recorder, + install_client_wrappers, + timed_guard, + uninstall_client_wrappers, + wrap_channel, +) +from google.cloud.managed_spark_connect.instrumentation import _client + +SPAN_KEYS = { + "type", + "name", + "span_id", + "parent_id", + "trace_id", + "start_wall_ns", + "start_mono_ns", + "duration_ns", + "status", + "error", + "thread", + "attrs", + "events", +} + + +def spans(inst, name=None): + return [ + r + for r in inst.records + if r["type"] == "span" and (name is None or r["name"] == name) + ] + + +def one(inst, name): + found = spans(inst, name) + assert len(found) == 1, f"expected one {name!r} span, got {found}" + return found[0] + + +@pytest.fixture +def inst(): + return Instrumentation([InMemorySink()], {"run_id": "r1"}) + + +# -------------------------------------------------------------------------- +# Recorder / Span / Instrumentation +# -------------------------------------------------------------------------- + + +def test_meta_record_is_first_and_carries_versions(inst): + meta = inst.records[0] + assert meta["type"] == "meta" + assert meta["attrs"]["run_id"] == "r1" + for key in ("library_version", "pyspark_version", "grpcio_version"): + assert key in meta["attrs"] + assert meta["attrs"]["python"] + + +def test_span_schema_nesting_and_events(inst): + with inst.recorder.span("outer", a=1) as outer: + with inst.recorder.span("inner") as inner: + inner.add_event("first_response") + inner.set(b=2) + o, i = one(inst, "outer"), one(inst, "inner") + assert set(o) == SPAN_KEYS + assert o["parent_id"] is None + assert i["parent_id"] == o["span_id"] + assert len(o["span_id"]) == 16 and int(o["span_id"], 16) >= 0 + assert o["attrs"] == {"run_id": "r1", "a": 1} + assert i["attrs"] == {"run_id": "r1", "b": 2} + assert i["events"][0]["name"] == "first_response" + assert 0 <= i["events"][0]["offset_ns"] <= i["duration_ns"] + assert o["status"] == "ok" and o["error"] is None + assert o["thread"] == threading.current_thread().name + assert i["start_mono_ns"] >= o["start_mono_ns"] + assert o["duration_ns"] >= i["duration_ns"] + # Records are JSON serializable as emitted. + json.dumps(inst.records) + + +def test_span_records_exception_with_grpc_code(inst): + class FakeRpcError(grpc.RpcError): + + def code(self): + return grpc.StatusCode.UNAVAILABLE + + def __str__(self): + return "gone" + + with pytest.raises(FakeRpcError): + with inst.recorder.span("rpc"): + raise FakeRpcError() + rec = one(inst, "rpc") + assert rec["status"] == "error" + assert rec["error"] == { + "type": "FakeRpcError", + "message": "gone", + "grpc_code": "UNAVAILABLE", + } + + +def test_span_records_api_core_exception_code(inst): + from google.api_core.exceptions import NotFound + + with pytest.raises(NotFound): + with inst.recorder.span("x"): + raise NotFound("nope") + assert one(inst, "x")["error"]["grpc_code"] == "NOT_FOUND" + + +def test_global_attributes_apply_to_subsequent_records(inst): + inst.set_attribute("trial_id", "t1") + inst.recorder.event("e") + inst.remove_attribute("trial_id") + inst.recorder.counter("c", 5, unit="bytes") + events = [r for r in inst.records if r["type"] == "event"] + counters = [r for r in inst.records if r["type"] == "counter"] + assert events[0]["attrs"] == {"run_id": "r1", "trial_id": "t1"} + assert {"name", "wall_ns", "mono_ns", "thread"} <= set(events[0]) + assert counters[0]["delta"] == 5 + assert counters[0]["attrs"] == {"run_id": "r1", "unit": "bytes"} + + +def test_trace_id_from_recorder(inst): + inst.recorder.trace_id = "uuid-1" + with inst.recorder.span("a"): + pass + assert one(inst, "a")["trace_id"] == "uuid-1" + # A span started before the trace id was known picks it up at end. + inst.recorder.trace_id = "" + span = inst.recorder.start_span("b") + inst.recorder.trace_id = "uuid-2" + span.end() + assert one(inst, "b")["trace_id"] == "uuid-2" + + +def test_start_span_on_a_thread_has_no_parent(inst): + def worker(): + span = inst.recorder.start_span("thread-owned", conn=1) + span.end() + + with inst.recorder.span("main"): + t = threading.Thread(target=worker, name="proxy-1") + t.start() + t.join() + explicit = inst.recorder.start_span( + "explicit", parent=inst.recorder.current_span() + ) + explicit.end() + owned = one(inst, "thread-owned") + assert owned["parent_id"] is None + assert owned["thread"] == "proxy-1" + assert owned["attrs"]["conn"] == 1 + assert one(inst, "explicit")["parent_id"] == one(inst, "main")["span_id"] + + +def test_end_is_idempotent(inst): + span = inst.recorder.start_span("once") + span.end() + span.add_event("late") + span.end() + assert len(spans(inst, "once")) == 1 + assert one(inst, "once")["events"] == [] + + +def test_noop_recorder_is_allocation_free_and_disabled(): + assert NOOP_RECORDER.enabled is False + a = NOOP_RECORDER.span("x", k=1) + b = NOOP_RECORDER.span("y") + assert a is b + with a as s: + s.set(z=1) + s.add_event("e") + s.record_exception(ValueError()) + assert NOOP_RECORDER.start_span("z") is a + NOOP_RECORDER.event("e") + NOOP_RECORDER.counter("c") + NOOP_RECORDER.trace_id = "ignored" + assert NOOP_RECORDER.trace_id == "" + assert get_recorder(None) is NOOP_RECORDER + inst = Instrumentation() + assert get_recorder(inst) is inst.recorder + assert inst.records == [] + + +def test_sink_failures_are_contained(inst, caplog): + class Broken: + + def write(self, record): + raise IOError("disk full") + + def flush(self): + pass + + def close(self): + pass + + inst.sinks.append(Broken()) + with caplog.at_level(logging.WARNING): + inst.recorder.event("a") + inst.recorder.event("b") + assert [r["name"] for r in inst.records if r["type"] == "event"] == [ + "a", + "b", + ] + assert sum("disk full" in m for m in caplog.messages) == 1 + + +# -------------------------------------------------------------------------- +# Sinks and from_env +# -------------------------------------------------------------------------- + + +def test_jsonl_file_sink_appends_and_reopens(tmp_path): + path = tmp_path / "events.jsonl" + sink = JsonlFileSink(path) + sink.write({"type": "event", "name": "a", "attrs": {}}) + sink.close() + sink.write({"type": "event", "name": "b", "attrs": {"v": b"x"}}) + sink.close() + sink.close() # idempotent + lines = path.read_text(encoding="utf-8").splitlines() + assert [json.loads(l)["name"] for l in lines] == ["a", "b"] + assert json.loads(lines[1])["attrs"]["v"] == "b'x'" + + JsonlFileSink(path, append=False).close() + assert path.read_text() == "" + + +def test_logging_and_callback_sinks(caplog): + seen = [] + log = logging.getLogger("test.instrumentation") + inst = Instrumentation( + [LoggingSink(log, logging.INFO), CallbackSink(seen.append)] + ) + with caplog.at_level(logging.INFO, logger="test.instrumentation"): + inst.recorder.event("hello") + assert any('"name":"hello"' in m for m in caplog.messages) + assert [r["type"] for r in seen] == ["meta", "event"] + + +def test_from_env(tmp_path, monkeypatch): + monkeypatch.delenv(inst_lib.ENV_FILE, raising=False) + monkeypatch.delenv(inst_lib.ENV_LOG, raising=False) + assert Instrumentation.from_env() is None + + monkeypatch.setenv(inst_lib.ENV_LOG, "1") + inst = Instrumentation.from_env() + assert [type(s) for s in inst.sinks] == [LoggingSink] + + path = tmp_path / "out.jsonl" + monkeypatch.setenv(inst_lib.ENV_FILE, str(path)) + monkeypatch.setenv(inst_lib.ENV_LOG, "false") + inst = Instrumentation.from_env() + assert [type(s) for s in inst.sinks] == [JsonlFileSink] + inst.close() + assert json.loads(path.read_text().splitlines()[0])["type"] == "meta" + + +# -------------------------------------------------------------------------- +# gRPC interceptor +# -------------------------------------------------------------------------- + + +class FakeCall(grpc.Call): + """A finished unary call.""" + + def __init__(self, code=grpc.StatusCode.OK): + self._code = code + + def initial_metadata(self): + return (("i", "1"),) + + def trailing_metadata(self): + return (("t", "1"),) + + def code(self): + return self._code + + def details(self): + return "details" + + def is_active(self): + return False + + def time_remaining(self): + return None + + def cancel(self): + return False + + def add_callback(self, callback): + return False + + +class FakeStreamCall(FakeCall): + """A server-streaming call that yields ``messages`` then ends. + + ``error`` (an RpcError instance) is raised after the messages. Status + callbacks fire when the stream is exhausted, like gRPC's spin thread. + """ + + def __init__(self, messages, error=None): + super().__init__(error.code() if error else grpc.StatusCode.OK) + self._messages = iter(messages) + self._error = error + self._callbacks = [] + self._done = False + self.cancelled = False + + def __iter__(self): + return self + + def __next__(self): + try: + return next(self._messages) + except StopIteration: + self._finish() + if self._error is not None: + raise self._error + raise + + def _finish(self): + self._done = True + callbacks, self._callbacks = self._callbacks, [] + for cb in callbacks: + cb() + + def is_active(self): + return not self._done + + def cancel(self): + self.cancelled = True + self._done = True + return True + + def add_callback(self, callback): + if self._done: + return False + self._callbacks.append(callback) + return True + + def debug_error_string(self): + return "debug" + + +class FakeRpcError(grpc.RpcError, grpc.Call): + + def __init__(self, code): + self._code = code + + def code(self): + return self._code + + def details(self): + return "boom" + + 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 result(self, timeout=None): + # Like grpc._channel._InactiveRpcError: the error is its own future. + raise self + + +class FakeMultiCallable: + + def __init__(self, channel, method): + self.channel = channel + self.method = method + + def _respond(self, request): + handler = self.channel.handlers.get(self.method) + return handler(request) if handler else proto.ConfigResponse() + + def __call__(self, request, **kwargs): + result = self._respond(request) + if isinstance(result, Exception): + raise result + return result + + def with_call(self, request, **kwargs): + result = self._respond(request) + if isinstance(result, Exception): + raise result + return result, FakeCall() + + def future(self, request, **kwargs): + raise NotImplementedError + + +class FakeChannel(grpc.Channel): + """A channel whose RPCs are answered by ``handlers[method](request)``.""" + + def __init__(self): + self.handlers = {} + self.closed = False + + def unary_unary(self, method, *args, **kwargs): + return FakeMultiCallable(self, method) + + unary_stream = stream_unary = stream_stream = unary_unary + + def subscribe(self, callback, try_to_connect=False): + pass + + def unsubscribe(self, callback): + pass + + def close(self): + self.closed = True + + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() + return False + + +EXECUTE = "/spark.connect.SparkConnectService/ExecutePlan" +CONFIG = "/spark.connect.SparkConnectService/Config" +ARTIFACTS = "/spark.connect.SparkConnectService/AddArtifacts" + + +def execute_request(): + return proto.ExecutePlanRequest(session_id="sess-1", operation_id="op-1") + + +def responses(n): + return [ + proto.ExecutePlanResponse(session_id="sess-1", operation_id="op-1") + for _ in range(n) + ] + + +def test_unary_unary_span(inst): + fake = FakeChannel() + fake.handlers[CONFIG] = lambda req: proto.ConfigResponse(session_id="s") + channel = wrap_channel(fake, inst, "inner") + request = proto.ConfigRequest(session_id="sess-1") + response = channel.unary_unary(CONFIG)(request) + assert response.session_id == "s" + rec = one(inst, "grpc.Config") + assert rec["attrs"] == { + "run_id": "r1", + "method": CONFIG, + "kind": "unary_unary", + "layer": "inner", + "request_bytes": request.ByteSize(), + "session_id": "sess-1", + "grpc_code": "OK", + "response_count": 1, + "response_bytes": response.ByteSize(), + } + assert [e["name"] for e in rec["events"]] == ["status"] + assert rec["status"] == "ok" + + +def test_unary_unary_error(inst): + fake = FakeChannel() + fake.handlers[CONFIG] = lambda req: FakeRpcError( + grpc.StatusCode.UNAVAILABLE + ) + channel = wrap_channel(fake, inst.recorder, "outer") + with pytest.raises(grpc.RpcError): + channel.unary_unary(CONFIG)(proto.ConfigRequest()) + rec = one(inst, "grpc.Config") + assert rec["status"] == "error" + assert rec["error"]["grpc_code"] == "UNAVAILABLE" + assert rec["attrs"]["grpc_code"] == "UNAVAILABLE" + assert rec["attrs"]["guard_ns"] == 0 + + +def test_unary_stream_span_and_iterator_proxy(inst): + fake = FakeChannel() + msgs = responses(3) + fake.handlers[EXECUTE] = lambda req: FakeStreamCall(msgs) + channel = wrap_channel(fake, inst, "outer") + call = channel.unary_stream(EXECUTE)(execute_request()) + + assert isinstance(call, grpc.Call) + assert call.initial_metadata() == (("i", "1"),) + assert call.trailing_metadata() == (("t", "1"),) + assert call.details() == "details" + assert call.debug_error_string() == "debug" + assert call.time_remaining() is None + assert call.is_active() is True + assert spans(inst, "grpc.ExecutePlan") == [] # not ended yet + + got = list(call) + assert got == msgs + assert call.code() == grpc.StatusCode.OK + rec = one(inst, "grpc.ExecutePlan") + attrs = rec["attrs"] + assert attrs["kind"] == "unary_stream" + assert attrs["layer"] == "outer" + assert attrs["operation_id"] == "op-1" + assert attrs["session_id"] == "sess-1" + assert attrs["request_bytes"] == execute_request().ByteSize() + assert attrs["response_count"] == 3 + assert attrs["response_bytes"] == sum(m.ByteSize() for m in msgs) + assert attrs["grpc_code"] == "OK" + assert [e["name"] for e in rec["events"]] == ["first_response", "status"] + assert rec["status"] == "ok" + + +def test_unary_stream_error_mid_stream(inst): + fake = FakeChannel() + error = FakeRpcError(grpc.StatusCode.INTERNAL) + fake.handlers[EXECUTE] = lambda req: FakeStreamCall(responses(1), error) + channel = wrap_channel(fake, inst) + call = channel.unary_stream(EXECUTE)(execute_request()) + next(call) + with pytest.raises(FakeRpcError) as info: + next(call) + assert info.value is error # the real call is what pyspark inspects + rec = one(inst, "grpc.ExecutePlan") + assert rec["status"] == "error" + assert rec["error"]["grpc_code"] == "INTERNAL" + assert rec["attrs"]["response_count"] == 1 + + +def test_unary_stream_cancel_ends_span(inst): + fake = FakeChannel() + fake.handlers[EXECUTE] = lambda req: FakeStreamCall(responses(5)) + channel = wrap_channel(fake, inst) + call = channel.unary_stream(EXECUTE)(execute_request()) + next(call) + assert call.cancel() is True + rec = one(inst, "grpc.ExecutePlan") + assert rec["attrs"]["cancelled"] is True + assert rec["attrs"]["response_count"] == 1 + + +def test_unary_stream_status_callback_from_grpc_thread(inst): + """The status event comes from add_callback, before the consumer + observes the end of the stream.""" + fake = FakeChannel() + stream = FakeStreamCall(responses(1)) + fake.handlers[EXECUTE] = lambda req: stream + channel = wrap_channel(fake, inst) + call = channel.unary_stream(EXECUTE)(execute_request()) + next(call) + stream._finish() # gRPC reports completion + assert spans(inst, "grpc.ExecutePlan") == [] + with pytest.raises(StopIteration): + next(call) + rec = one(inst, "grpc.ExecutePlan") + assert [e["name"] for e in rec["events"]] == ["first_response", "status"] + + +def test_stream_unary_counts_requests(inst): + fake = FakeChannel() + seen = [] + + def handler(request_iterator): + seen.extend(request_iterator) + return proto.AddArtifactsResponse() + + fake.handlers[ARTIFACTS] = handler + channel = wrap_channel(fake, inst) + chunks = [proto.AddArtifactsRequest(session_id="sess-1") for _ in range(4)] + channel.stream_unary(ARTIFACTS)(iter(chunks)) + assert seen == chunks + rec = one(inst, "grpc.AddArtifacts") + assert rec["attrs"]["kind"] == "stream_unary" + assert rec["attrs"]["request_count"] == 4 + assert rec["attrs"]["request_bytes"] == sum(c.ByteSize() for c in chunks) + assert rec["attrs"]["session_id"] == "sess-1" + assert rec["attrs"]["response_count"] == 1 + + +def test_stream_stream_span(inst): + fake = FakeChannel() + fake.handlers[EXECUTE] = lambda reqs: FakeStreamCall( + responses(len(list(reqs))) + ) + channel = wrap_channel(fake, inst) + call = channel.stream_stream(EXECUTE)(iter([execute_request()] * 2)) + assert len(list(call)) == 2 + rec = one(inst, "grpc.ExecutePlan") + assert rec["attrs"]["kind"] == "stream_stream" + assert rec["attrs"]["request_count"] == 2 + assert rec["attrs"]["response_count"] == 2 + + +def test_wrap_channel_disabled_is_passthrough(): + fake = FakeChannel() + assert wrap_channel(fake, None) is fake + assert wrap_channel(fake, NOOP_RECORDER, "inner") is fake + with pytest.raises(ValueError): + GrpcInterceptor(Recorder(lambda r: None), "middle") + + +def test_spans_nest_under_the_current_span(inst): + fake = FakeChannel() + channel = wrap_channel(fake, inst) + with inst.recorder.span("client.to_pandas"): + channel.unary_unary(CONFIG)(proto.ConfigRequest()) + assert ( + one(inst, "grpc.Config")["parent_id"] + == one(inst, "client.to_pandas")["span_id"] + ) + + +@pytest.fixture +def proxied(inst): + """A ProxiedChannel over a FakeChannel, with instrumentation.""" + fake = FakeChannel() + with ( + mock.patch.object(core.proxy, "ManagedSparkSessionProxy") as proxy_cls, + mock.patch.object(core, "DefaultChannelBuilder") as builder_cls, + ): + proxy_cls.return_value.port = 1 + builder_cls.return_value.toChannel.return_value = fake + yield fake, proxy_cls + + +def test_two_layers_with_guard(inst, proxied): + fake, proxy_cls = proxied + fake.handlers[EXECUTE] = lambda req: FakeStreamCall(responses(2)) + guard_calls = [] + + def guard(): + with inst.recorder.span( + "dataproc.get_session", reason="rpc_guard", new_client=True + ): + guard_calls.append(1) + return True + + builder = core.ManagedSparkChannelBuilder( + "sc://example.com:443/;use_ssl=true", + is_active_callback=guard, + instrumentation=inst, + ) + channel = builder.toChannel() + proxy_cls.assert_called_once_with(0, "example.com", recorder=inst.recorder) + list(channel.unary_stream(EXECUTE)(execute_request())) + + outer = [ + s + for s in spans(inst, "grpc.ExecutePlan") + if s["attrs"]["layer"] == "outer" + ] + inner = [ + s + for s in spans(inst, "grpc.ExecutePlan") + if s["attrs"]["layer"] == "inner" + ] + assert len(outer) == 1 and len(inner) == 1 + assert inner[0]["parent_id"] == outer[0]["span_id"] + assert outer[0]["parent_id"] is None + assert guard_calls == [1] + assert outer[0]["attrs"]["guard_ns"] > 0 + assert "guard_ns" not in inner[0]["attrs"] + guard_span = one(inst, "dataproc.get_session") + assert guard_span["parent_id"] == outer[0]["span_id"] + assert guard_span["attrs"]["reason"] == "rpc_guard" + assert outer[0]["attrs"]["response_count"] == 2 + assert inner[0]["attrs"]["response_count"] == 2 + assert outer[0]["duration_ns"] >= inner[0]["duration_ns"] + + +def test_guard_failure_is_an_error_on_the_outer_span(proxied): + fake, _ = proxied + snapshots = [] + inst = Instrumentation( + [ + InMemorySink(), + CallbackSink(lambda r: snapshots.append(json.loads(json.dumps(r)))), + ] + ) + channel = core.ManagedSparkChannelBuilder( + "sc://example.com:443", + is_active_callback=lambda: False, + instrumentation=inst, + ).toChannel() + with pytest.raises(RuntimeError, match="Session not active"): + channel.unary_unary(CONFIG)(proto.ConfigRequest()) + rec = one(inst, "grpc.Config") + assert rec["attrs"]["layer"] == "outer" + assert rec["status"] == "error" + assert rec["error"]["type"] == "RuntimeError" + assert rec["attrs"]["guard_ns"] >= 0 + assert spans(inst, "grpc.Config") == [rec] # inner never ran + (emitted,) = [r for r in snapshots if r["type"] == "span"] + assert emitted["attrs"]["guard_ns"] >= 0 + + +def test_ended_span_ignores_late_set_and_record_exception(inst): + span = inst.recorder.start_span("done", initial=1) + span.end() + span.set(initial=2, extra=3) + span.record_exception(ValueError("too late")) + rec = one(inst, "done") + assert rec["status"] == "ok" + assert rec["error"] is None + assert rec["attrs"] == {"run_id": "r1", "initial": 1} + + +def test_disabled_channel_builder_installs_nothing(proxied): + fake, proxy_cls = proxied + cb = mock.Mock(return_value=True) + channel = core.ManagedSparkChannelBuilder( + "sc://example.com:443", is_active_callback=cb + ).toChannel() + assert isinstance(channel, core.ProxiedChannel) + assert channel._is_active_callback is cb # not wrapped by timed_guard + assert channel._wrapped is fake # no inner interceptor + proxy_cls.assert_called_once_with(0, "example.com", recorder=NOOP_RECORDER) + + +def test_timed_guard(): + assert timed_guard(None) is None + guard = timed_guard(lambda: True) + assert guard() is True + assert guard.__name__ == "" + + +# -------------------------------------------------------------------------- +# Client wrappers +# -------------------------------------------------------------------------- + + +class FakeClient: + + def __init__(self): + self.table = mock.Mock(num_rows=7, num_columns=2) + + def _execute_and_fetch(self, req, observations, self_destruct=False): + return self.table, None, [], [], None + + def to_pandas(self, plan, observations): + table, *_ = self._execute_and_fetch(plan, observations) + return "pdf", "ei" + + +def test_client_to_pandas_wrappers(inst): + client = FakeClient() + install_client_wrappers(client, inst.recorder) + install_client_wrappers(client, inst.recorder) # idempotent + assert client.to_pandas("plan", {}) == ("pdf", "ei") + to_pandas = one(inst, "client.to_pandas") + conv = one(inst, "client.arrow_to_pandas") + assert conv["parent_id"] == to_pandas["span_id"] + assert conv["attrs"]["rows"] == 7 and conv["attrs"]["columns"] == 2 + assert conv["start_mono_ns"] >= to_pandas["start_mono_ns"] + assert ( + conv["start_mono_ns"] + conv["duration_ns"] + <= to_pandas["start_mono_ns"] + to_pandas["duration_ns"] + ) + # A fetch outside to_pandas does not start a conversion span. + client._execute_and_fetch("plan", {}) + assert len(spans(inst, "client.arrow_to_pandas")) == 1 + + uninstall_client_wrappers(client) + uninstall_client_wrappers(client) + assert "to_pandas" not in client.__dict__ + assert "_execute_and_fetch" not in client.__dict__ + client.to_pandas("plan", {}) + assert len(spans(inst, "client.to_pandas")) == 1 + + +def test_client_to_pandas_conversion_error_marks_arrow_to_pandas_error(inst): + class FailingConversionClient(FakeClient): + + def to_pandas(self, plan, observations): + self._execute_and_fetch(plan, observations) + raise ValueError("bad arrow conversion") + + client = FailingConversionClient() + install_client_wrappers(client, inst.recorder) + try: + with pytest.raises(ValueError, match="bad arrow conversion"): + client.to_pandas("plan", {}) + finally: + uninstall_client_wrappers(client) + + to_pandas = one(inst, "client.to_pandas") + conv = one(inst, "client.arrow_to_pandas") + assert to_pandas["status"] == "error" + assert to_pandas["error"]["type"] == "ValueError" + assert conv["status"] == "error" + assert conv["error"]["type"] == "ValueError" + assert conv["parent_id"] == to_pandas["span_id"] + + +def test_client_class_hooks(inst): + import pyarrow as pa + from pyspark.sql.connect.plan import Range + from pyspark.sql.conversion import ArrowTableToRowsConversion + from pyspark.sql.types import LongType, StructField, StructType + + client = FakeClient() + other = FakeClient() + install_client_wrappers(client, inst.recorder) + try: + plan = Range(start=0, end=10, step=1, num_partitions=1).to_proto(client) + assert plan.root.HasField("range") + rec = one(inst, "client.plan_to_proto") + assert rec["attrs"]["plan_bytes"] == plan.ByteSize() + + # Another client (no recorder) goes through the hook untraced. + Range(start=0, end=1, step=1, num_partitions=1).to_proto(other) + assert len(spans(inst, "client.plan_to_proto")) == 1 + + table = pa.table({"a": [1, 2, 3]}) + schema = StructType([StructField("a", LongType())]) + rows = ArrowTableToRowsConversion.convert(table, schema) + assert [r.a for r in rows] == [1, 2, 3] + assert one(inst, "client.arrow_to_rows")["attrs"]["rows"] == 3 + finally: + uninstall_client_wrappers(client) + + # With no instrumented client registered the hooks record nothing. + ArrowTableToRowsConversion.convert(table, schema) + assert len(spans(inst, "client.arrow_to_rows")) == 1 + assert _client._active_recorder() is NOOP_RECORDER + + +def test_install_is_a_noop_when_disabled(): + client = FakeClient() + install_client_wrappers(client, NOOP_RECORDER) + assert "to_pandas" not in client.__dict__ + assert not hasattr(client, _client._RECORDER_ATTR) + + +# -------------------------------------------------------------------------- +# Session wiring (mocked Dataproc API, no network) +# -------------------------------------------------------------------------- + + +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.session_template = "projects/p/locations/r/sessionTemplates/t" + response.runtime_info.endpoints = { + "Spark Connect Server": "sc://spark-connect-server.example.com:443" + } + response.create_time = "2024-01-01T00:00:00Z" + response.state_time = "2024-01-01T00:01:00Z" + response.state_history = [ + { + "state": Session.State.CREATING, + "state_start_time": "2024-01-01T00:00:00Z", + } + ] + 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, + "response": response, + } + + +def _by_name(inst): + out = {} + for rec in spans(inst): + out.setdefault(rec["name"], []).append(rec) + return out + + +def test_session_builder_opt_in(dataproc_mocks): + from google.cloud.managed_spark_connect import ManagedSparkSession + + inst = Instrumentation([InMemorySink()], {"run_id": "r1"}) + session = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .instrumentation(inst) + .getOrCreate() + ) + try: + assert session.instrumentation is inst + assert getattr(session.client, _client._RECORDER_ATTR) is inst.recorder + assert inst.recorder.trace_id == dataproc_mocks["uuid"] + + by_name = _by_name(inst) + root = by_name["session.get_or_create"][0] + assert root["parent_id"] is None + assert root["attrs"] == { + "run_id": "r1", + "mode": "create", + "session_id": dataproc_mocks["session_id"], + "uuid": dataproc_mocks["uuid"], + "runtime_version": "3.0", + "template": "projects/p/locations/r/sessionTemplates/t", + } + for child in ( + "session.create_rpc", + "session.wait_active", + "session.connect", + ): + (rec,) = by_name[child] + assert rec["parent_id"] == root["span_id"], child + wait = by_name["session.wait_active"][0]["attrs"] + assert wait["poll_interval_s"] == 5.0 + assert wait["create_time"].startswith("2024-01-01T00:00:00") + assert wait["state"] == "ACTIVE" + assert json.loads(wait["state_history"])[0]["state"] == "CREATING" + (start,) = by_name["proxy.start"] + assert start["parent_id"] == by_name["session.connect"][0]["span_id"] + assert start["attrs"]["port"] > 0 + assert "session.reuse" not in by_name + # Everything after connect carries the session uuid as trace id. + assert start["trace_id"] == dataproc_mocks["uuid"] + assert root["trace_id"] == dataproc_mocks["uuid"] + finally: + session.stop() + + by_name = _by_name(inst) + (stop,) = by_name["session.stop"] + assert stop["attrs"]["terminate"] is True + assert stop["attrs"]["session_id"] == dataproc_mocks["session_id"] + assert stop["trace_id"] == dataproc_mocks["uuid"] + assert by_name["session.terminate_rpc"][0]["parent_id"] == stop["span_id"] + (wait,) = by_name["session.wait_terminating"] + assert wait["parent_id"] == stop["span_id"] + assert wait["attrs"]["state"] == "TERMINATING" + (poll,) = by_name["dataproc.get_session"] + assert poll["parent_id"] == wait["span_id"] + assert poll["attrs"] == { + "run_id": "r1", + "reason": "stop", + "new_client": False, + } + dataproc_mocks["client"].get_session.assert_called_once() + # stop() closes the instrumentation and detaches the client wrappers. + assert inst.recorder.trace_id == "" + assert not hasattr(session.client, _client._RECORDER_ATTR) + + +def test_session_env_opt_in(dataproc_mocks, monkeypatch, tmp_path): + from google.cloud.managed_spark_connect import ManagedSparkSession + + path = tmp_path / "spans.jsonl" + monkeypatch.setenv(inst_lib.ENV_FILE, str(path)) + session = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .getOrCreate() + ) + try: + assert session.instrumentation is not None + assert [type(s) for s in session.instrumentation.sinks] == [ + JsonlFileSink + ] + finally: + session.stop() + records = [json.loads(l) for l in path.read_text().splitlines()] + assert records[0]["type"] == "meta" + names = {r["name"] for r in records if r["type"] == "span"} + assert {"session.get_or_create", "session.connect", "session.stop"} <= names + + +def test_session_disabled_by_default(dataproc_mocks): + from google.cloud.managed_spark_connect import ManagedSparkSession + + session = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .getOrCreate() + ) + try: + assert session.instrumentation is None + assert not hasattr(session.client, _client._RECORDER_ATTR) + assert isinstance(session.client._channel, core.ProxiedChannel) + assert "to_pandas" not in session.client.__dict__ + finally: + session.stop() + + +def test_session_reuse_spans(dataproc_mocks): + from google.cloud.managed_spark_connect import ManagedSparkSession + + inst = Instrumentation([InMemorySink()]) + builder = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .instrumentation(inst) + ) + first = builder.getOrCreate() + try: + with mock.patch( + "google.cloud.managed_spark_connect.session.get_active_s8s_session_response", + return_value=dataproc_mocks["response"], + ): + second = builder.getOrCreate() + assert second is first + by_name = _by_name(inst) + roots = by_name["session.get_or_create"] + assert [r["attrs"]["mode"] for r in roots] == ["create", "reuse"] + (reuse,) = by_name["session.reuse"] + assert reuse["parent_id"] == roots[1]["span_id"] + assert reuse["attrs"]["found"] is True + (lookup,) = by_name["dataproc.get_session"] + assert lookup["parent_id"] == reuse["span_id"] + assert lookup["attrs"] == {"reason": "reuse", "new_client": True} + finally: + first.stop() + + +def test_session_reuse_fresh_builder_reuses_instrumentation_and_sets_trace_id( + dataproc_mocks, monkeypatch, tmp_path +): + from google.cloud.managed_spark_connect import ManagedSparkSession + + # 1. Env-configured instrumentation: fresh builder reuses active session's + # Instrumentation without emitting a duplicate meta record or orphaning + # the reuse span. + path = tmp_path / "reuse.jsonl" + monkeypatch.setenv(inst_lib.ENV_FILE, str(path)) + first = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .getOrCreate() + ) + try: + with mock.patch( + "google.cloud.managed_spark_connect.session.get_active_s8s_session_response", + return_value=dataproc_mocks["response"], + ): + second = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .getOrCreate() + ) + assert second is first + finally: + first.stop() + + records = [json.loads(line) for line in path.read_text().splitlines()] + assert [r["type"] for r in records].count("meta") == 1 + get_or_create = [ + r + for r in records + if r["type"] == "span" and r["name"] == "session.get_or_create" + ] + assert [r["attrs"]["mode"] for r in get_or_create] == ["create", "reuse"] + assert all(r["trace_id"] == dataproc_mocks["uuid"] for r in get_or_create) + monkeypatch.delenv(inst_lib.ENV_FILE) + + # 2. Programmatic instrumentation: fresh builder without .instrumentation() + # records reuse on the active session's Instrumentation; explicit new + # Instrumentation on reuse still gets trace_id populated. + inst1 = Instrumentation([InMemorySink()]) + inst2 = Instrumentation([InMemorySink()]) + session = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .instrumentation(inst1) + .getOrCreate() + ) + try: + with mock.patch( + "google.cloud.managed_spark_connect.session.get_active_s8s_session_response", + return_value=dataproc_mocks["response"], + ): + reused_default = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .getOrCreate() + ) + reused_explicit = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .instrumentation(inst2) + .getOrCreate() + ) + assert reused_default is session + assert reused_explicit is session + assert [ + r["attrs"]["mode"] for r in _by_name(inst1)["session.get_or_create"] + ] == ["create", "reuse"] + assert inst2.recorder.trace_id == dataproc_mocks["uuid"] + for name in ( + "session.get_or_create", + "session.reuse", + "dataproc.get_session", + ): + (rec,) = _by_name(inst2)[name] + assert rec["trace_id"] == dataproc_mocks["uuid"], name + finally: + session.stop() + + +def test_session_attach_when_local_session_missing_connects_before_versions( + dataproc_mocks, +): + from google.cloud.managed_spark_connect import ManagedSparkSession + + seen_sessions = [] + orig_versions = ManagedSparkSession.Builder._session_versions + + def spy_versions(sess): + seen_sessions.append(sess) + return orig_versions(sess) + + ManagedSparkSession._active_s8s_session_id = dataproc_mocks["session_id"] + session = None + try: + with ( + mock.patch( + "google.cloud.managed_spark_connect.session.get_active_s8s_session_response", + return_value=dataproc_mocks["response"], + ), + mock.patch.object( + ManagedSparkSession.Builder, + "_wait_for_session_available", + return_value=dataproc_mocks["response"], + ), + mock.patch.object( + ManagedSparkSession.Builder, + "_session_versions", + staticmethod(spy_versions), + ), + ): + session = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .getOrCreate() + ) + assert seen_sessions == [session] + assert session is not None + finally: + if session is not None: + session.stop() + ManagedSparkSession._active_s8s_session_id = None + + +# -------------------------------------------------------------------------- +# Real gRPC round trip (in-process server on 127.0.0.1, no network) +# -------------------------------------------------------------------------- + + +@pytest.fixture +def grpc_server(): + from concurrent import futures + + from pyspark.sql.connect.proto import base_pb2_grpc + + class Service(base_pb2_grpc.SparkConnectServiceServicer): + + def ExecutePlan(self, request, context): + yield proto.ExecutePlanResponse(session_id=request.session_id) + if request.operation_id == "fail": + context.abort(grpc.StatusCode.INTERNAL, "boom") + for _ in range(2): + yield proto.ExecutePlanResponse( + session_id=request.session_id, + operation_id=request.operation_id, + ) + + def Config(self, request, context): + return proto.ConfigResponse(session_id=request.session_id) + + server = grpc.server(futures.ThreadPoolExecutor(max_workers=4)) + base_pb2_grpc.add_SparkConnectServiceServicer_to_server(Service(), server) + port = server.add_insecure_port("127.0.0.1:0") + server.start() + try: + yield port + finally: + server.stop(None) + + +def _local_channel(port): + return grpc.insecure_channel( + f"127.0.0.1:{port}", options=[("grpc.enable_http_proxy", 0)] + ) + + +def test_real_grpc_round_trip_through_proxied_channel(inst, grpc_server): + from pyspark.sql.connect.proto import base_pb2_grpc + + guard = mock.Mock(return_value=True) + with ( + mock.patch.object(core.proxy, "ManagedSparkSessionProxy") as proxy_cls, + mock.patch.object(core, "DefaultChannelBuilder") as builder_cls, + ): + proxy_cls.return_value.port = 1 + builder_cls.return_value.toChannel.return_value = _local_channel( + grpc_server + ) + channel = core.ManagedSparkChannelBuilder( + "sc://example.com:443", + is_active_callback=guard, + instrumentation=inst, + ).toChannel() + stub = base_pb2_grpc.SparkConnectServiceStub(channel) + try: + assert ( + stub.Config(proto.ConfigRequest(session_id="s1")).session_id == "s1" + ) + call = stub.ExecutePlan( + proto.ExecutePlanRequest(session_id="s1", operation_id="op1") + ) + assert len(list(call)) == 3 + assert call.code() == grpc.StatusCode.OK + with pytest.raises(grpc.RpcError) as info: + list( + stub.ExecutePlan( + proto.ExecutePlanRequest( + session_id="s1", operation_id="fail" + ) + ) + ) + assert info.value.code() == grpc.StatusCode.INTERNAL + finally: + channel.close() + + assert guard.call_count == 3 + by_name = _by_name(inst) + config = by_name["grpc.Config"] + assert [s["attrs"]["layer"] for s in config] == ["inner", "outer"] + assert config[0]["parent_id"] == config[1]["span_id"] + assert config[1]["attrs"]["grpc_code"] == "OK" + assert config[1]["attrs"]["guard_ns"] > 0 + + execute = by_name["grpc.ExecutePlan"] + ok = [s for s in execute if s["attrs"]["operation_id"] == "op1"] + failed = [s for s in execute if s["attrs"]["operation_id"] == "fail"] + assert len(ok) == 2 and len(failed) == 2 + for rec in ok: + assert rec["attrs"]["response_count"] == 3 + assert rec["attrs"]["grpc_code"] == "OK" + assert [e["name"] for e in rec["events"]] == [ + "first_response", + "status", + ] + for rec in failed: + assert rec["status"] == "error" + assert rec["error"]["grpc_code"] == "INTERNAL" + assert rec["error"]["message"] == "boom" + assert rec["attrs"]["response_count"] == 1 + + +# -------------------------------------------------------------------------- +# Abandoned streams, guard idempotence, attr snapshots, status end time +# -------------------------------------------------------------------------- + + +def test_abandoned_stream_span_is_emitted_when_iterator_is_dropped(inst): + """pyspark stops reading an ExecutePlan stream at result_complete and + drops the iterator without draining it or calling cancel().""" + fake = FakeChannel() + stream = FakeStreamCall(responses(3)) + fake.handlers[EXECUTE] = lambda req: stream + channel = wrap_channel(wrap_channel(fake, inst, "inner"), inst, "outer") + call = channel.unary_stream(EXECUTE)(execute_request()) + next(call) + next(call) + stream._finish() # the server is done; status arrives on the grpc thread + assert spans(inst, "grpc.ExecutePlan") == [] + del call # consumer walks away + recs = spans(inst, "grpc.ExecutePlan") + assert [r["attrs"]["layer"] for r in recs] == ["inner", "outer"] + for rec in recs: + assert rec["attrs"]["response_count"] == 2 + assert rec["attrs"]["grpc_code"] == "OK" + assert [e["name"] for e in rec["events"]] == [ + "first_response", + "status", + ] + assert recs[0]["parent_id"] == recs[1]["span_id"] + + +def test_abandoned_active_stream_does_not_block(inst): + fake = FakeChannel() + fake.handlers[EXECUTE] = lambda req: FakeStreamCall(responses(3)) + channel = wrap_channel(fake, inst) + call = channel.unary_stream(EXECUTE)(execute_request()) + next(call) + del call # RPC still active: code() must not be consulted (it blocks) + (rec,) = spans(inst, "grpc.ExecutePlan") + assert rec["attrs"]["response_count"] == 1 + assert "grpc_code" not in rec["attrs"] + assert [e["name"] for e in rec["events"]] == ["first_response"] + + +def test_late_status_callback_does_not_touch_an_ended_span(inst): + fake = FakeChannel() + stream = FakeStreamCall(responses(1)) + fake.handlers[EXECUTE] = lambda req: stream + call = wrap_channel(fake, inst).unary_stream(EXECUTE)(execute_request()) + next(call) + del call + (rec,) = spans(inst, "grpc.ExecutePlan") + before = json.dumps(rec, sort_keys=True) + stream._finish() # callback fires after the span ended + assert json.dumps(rec, sort_keys=True) == before + + +def test_timed_guard_is_idempotent(): + inner = timed_guard(lambda: True) + assert timed_guard(inner) is inner + assert timed_guard(timed_guard(inner)) is inner + + +def test_guard_wrapped_twice_counts_once(inst, proxied, monkeypatch): + fake, _ = proxied + clock = [1_000_000_000] + + def slow_guard(): + clock[0] += 20_000_000 + return True + + channel = core.ManagedSparkChannelBuilder( + "sc://example.com:443", + is_active_callback=timed_guard(slow_guard), # as scbench does + instrumentation=inst, + ).toChannel() + monkeypatch.setattr("time.perf_counter_ns", lambda: clock[0]) + channel.unary_unary(CONFIG)(proto.ConfigRequest()) + rec = [ + s for s in spans(inst, "grpc.Config") if s["attrs"]["layer"] == "outer" + ][0] + assert rec["attrs"]["guard_ns"] == 20_000_000 + + +def test_global_attributes_are_snapshotted_when_a_span_starts(inst): + early = inst.recorder.start_span("early") + inst.set_attribute("trial_id", "t1") + during = inst.recorder.start_span("during") + inst.remove_attribute("trial_id") + inst.set_attribute("trial_id", "t2") + early.end() + during.end() + assert "trial_id" not in one(inst, "early")["attrs"] + assert one(inst, "during")["attrs"]["trial_id"] == "t1" + # Events are stamped when emitted. + inst.recorder.event("e") + assert [r for r in inst.records if r["type"] == "event"][0]["attrs"][ + "trial_id" + ] == "t2" + + +def test_stream_span_ends_at_status_not_at_release(inst): + """pyspark keeps the finished iterator alive until its ReleaseExecute + has been queued; the span must still end when the status arrived.""" + import time + + fake = FakeChannel() + stream = FakeStreamCall(responses(3)) + fake.handlers[EXECUTE] = lambda req: stream + call = wrap_channel(fake, inst).unary_stream(EXECUTE)(execute_request()) + next(call) + next(call) + stream._finish() # status from the grpc thread: the RPC is over + time.sleep(0.3) # consumer still holds the iterator + released_mono_ns = time.perf_counter_ns() + del call + (rec,) = spans(inst, "grpc.ExecutePlan") + status = [e for e in rec["events"] if e["name"] == "status"][0] + assert ( + rec["duration_ns"] + <= released_mono_ns - rec["start_mono_ns"] - 200_000_000 + ) + assert abs(rec["duration_ns"] - status["offset_ns"]) < 50_000_000 + assert rec["attrs"]["response_count"] == 2 + + +def test_span_end_at_explicit_time(inst): + import time + + span = inst.recorder.start_span("x") + then = time.perf_counter_ns() + time.sleep(0.05) + span.end(end_mono_ns=then) + assert one(inst, "x")["duration_ns"] == then - span.start_mono_ns + early = inst.recorder.start_span("y") + early.end(end_mono_ns=early.start_mono_ns - 10) + assert one(inst, "y")["duration_ns"] == 0 diff --git a/tests/unit/test_proxy.py b/tests/unit/test_proxy.py index 23339391..8aa8172b 100644 --- a/tests/unit/test_proxy.py +++ b/tests/unit/test_proxy.py @@ -149,3 +149,191 @@ def test_proxy_with_timeouts(client_wait, proxy_server_conn, test_message): retry_on_timeouts(proxy_server_conn.recv, 1024).decode() ) assert "\n".join(sent) == "\n".join(received) + + +# -------------------------------------------------------------------------- +# Instrumentation of the proxy (no network) +# -------------------------------------------------------------------------- + + +def _spans(inst, name): + return [ + r for r in inst.records if r["type"] == "span" and r["name"] == name + ] + + +def test_connect_sockets_counts_traffic(): + from google.cloud.managed_spark_connect.instrumentation import ( + InMemorySink, + Instrumentation, + ) + + inst = Instrumentation([InMemorySink()]) + client, proxy_side = socket.socketpair() + bridge_side, backend = socket.socketpair() + span = inst.recorder.start_span("proxy.connection", conn=1) + t = threading.Thread( + target=connect_sockets, + args=[1, proxy_side, bridge_side, span], + daemon=True, + ) + t.start() + client.sendall(b"hello") + assert backend.recv(1024) == b"hello" + backend.sendall(b"world!!") + assert client.recv(1024) == b"world!!" + client.close() + backend.close() + t.join(timeout=5) + assert not t.is_alive() + span.end() + (rec,) = _spans(inst, "proxy.connection") + assert rec["attrs"] == { + "conn": 1, + "bytes_out": 5, + "frames_out": 1, + "bytes_in": 7, + "frames_in": 1, + } + + +def test_connect_tcp_bridge_spans(): + from unittest import mock + + from google.cloud.managed_spark_connect.client import proxy + from google.cloud.managed_spark_connect.instrumentation import ( + InMemorySink, + Instrumentation, + ) + + inst = Instrumentation([InMemorySink()]) + creds = mock.Mock(token="tok") + with ( + mock.patch.object( + proxy.googleauth, "default", return_value=(creds, None) + ), + mock.patch.object(proxy.websocketclient, "connect") as connect, + ): + ws = proxy.connect_tcp_bridge("backend.example.com", inst.recorder, 3) + assert ws is connect.return_value + creds.refresh.assert_called_once() + (parent,) = _spans(inst, "proxy.ws_connect") + assert parent["attrs"] == {"conn": 3, "host": "backend.example.com"} + for name in ("auth.default", "auth.refresh", "ws.handshake"): + (child,) = _spans(inst, name) + assert child["parent_id"] == parent["span_id"] + + +def test_forward_connection_ends_span_on_failure(): + import time + from unittest import mock + + from google.cloud.managed_spark_connect.client import proxy + from google.cloud.managed_spark_connect.instrumentation import ( + InMemorySink, + Instrumentation, + ) + + inst = Instrumentation([InMemorySink()]) + conn = mock.MagicMock() + accepted = (time.perf_counter_ns() - 5_000_000, time.time_ns()) + with mock.patch.object( + proxy, "connect_tcp_bridge", side_effect=OSError("no route") + ): + with pytest.raises(OSError): + proxy.forward_connection( + 1, conn, ("127.0.0.1", 1), "h", inst.recorder, accepted + ) + (rec,) = _spans(inst, "proxy.connection") + assert rec["status"] == "error" + assert rec["error"]["type"] == "OSError" + assert rec["attrs"]["conn"] == 1 + assert rec["start_mono_ns"] == accepted[0] + assert rec["duration_ns"] >= 5_000_000 + + +def test_forward_connection_nests_ws_connect_and_counts(): + from unittest import mock + + from google.cloud.managed_spark_connect.client import proxy + from google.cloud.managed_spark_connect.instrumentation import ( + InMemorySink, + Instrumentation, + ) + + inst = Instrumentation([InMemorySink()]) + client, proxy_side = socket.socketpair() + bridge_side, backend = socket.socketpair() + creds = mock.Mock(token="tok") + + class FakeWs: + + def __enter__(self): + return self + + def __exit__(self, *args): + bridge_side.close() + return False + + def recv(self, timeout=None): + data = bridge_side.recv(1024) + if not data: + raise ConnectionError("closed") + return data.hex() + + def send(self, msg): + bridge_side.sendall(bytes.fromhex(msg)) + + def close(self): + bridge_side.close() + + with ( + mock.patch.object( + proxy.googleauth, "default", return_value=(creds, None) + ), + mock.patch.object( + proxy.websocketclient, "connect", return_value=FakeWs() + ), + ): + t = threading.Thread( + target=proxy.forward_connection, + args=[2, proxy_side, ("127.0.0.1", 1), "h", inst.recorder], + daemon=True, + ) + t.start() + client.sendall(b"ping") + assert backend.recv(1024) == b"ping" + backend.sendall(b"pong!") + assert client.recv(1024) == b"pong!" + client.close() + backend.close() + t.join(timeout=5) + assert not t.is_alive() + + (conn_span,) = _spans(inst, "proxy.connection") + (ws,) = _spans(inst, "proxy.ws_connect") + assert ws["parent_id"] == conn_span["span_id"] + assert _spans(inst, "auth.default")[0]["parent_id"] == ws["span_id"] + assert conn_span["thread"] == ws["thread"] == t.name + assert conn_span["attrs"] == { + "conn": 2, + "bytes_out": 4, + "frames_out": 1, + "bytes_in": 5, + "frames_in": 1, + } + + +def test_proxy_start_span(): + from google.cloud.managed_spark_connect.client import proxy + from google.cloud.managed_spark_connect.instrumentation import ( + InMemorySink, + Instrumentation, + ) + + inst = Instrumentation([InMemorySink()]) + p = proxy.ManagedSparkSessionProxy(0, "h", recorder=inst.recorder) + p.start() + p.stop() + (rec,) = _spans(inst, "proxy.start") + assert rec["attrs"]["port"] == p.port > 0 From ea3810739cfd6781553047abffac47248a8263f1 Mon Sep 17 00:00:00 2001 From: Dave Borowitz Date: Fri, 16 Oct 2026 08:44:00 -0700 Subject: [PATCH 5/5] feat(instrumentation): track in-flight spans; add wait_idle() Recorder.inflight(prefix) lists spans that have started but not ended, and Recorder.wait_idle(timeout_s, prefix='grpc.') (also on Instrumentation) blocks until none match. Benchmarks use it to let the background ReleaseExecute RPCs pyspark issues after a query drain before the next measurement starts, instead of sleeping a fixed amount. --- docs/instrumentation.md | 2 + .../instrumentation/_core.py | 49 +++++++++++++++++++ tests/unit/test_instrumentation.py | 36 ++++++++++++++ 3 files changed, 87 insertions(+) diff --git a/docs/instrumentation.md b/docs/instrumentation.md index fa91a66b..405a4fe3 100644 --- a/docs/instrumentation.md +++ b/docs/instrumentation.md @@ -50,12 +50,14 @@ class Instrumentation: set_attribute(key, value); remove_attribute(key) # global attrs on every subsequent record records -> list[dict] # concatenated InMemorySink records wrap_channel(channel, layer="outer") -> grpc.Channel + wait_idle(timeout_s, prefix="grpc.") -> bool # block until no such span is in flight (e.g. background ReleaseExecute) flush(); close() # idempotent; file sinks reopen on the next write class Recorder: # thread-safe span(name, **attrs) -> ContextManager[Span] # nests under the current span (contextvars) start_span(name, parent=None, **attrs) -> Span # for threads that own their spans; call end() event(name, **attrs); counter(name, delta=1, **attrs) + inflight(prefix="") -> list[Span]; wait_idle(timeout_s, prefix="grpc.") -> bool trace_id: str; enabled: bool class Span: diff --git a/google/cloud/managed_spark_connect/instrumentation/_core.py b/google/cloud/managed_spark_connect/instrumentation/_core.py index 06c1a1f8..666fcac9 100644 --- a/google/cloud/managed_spark_connect/instrumentation/_core.py +++ b/google/cloud/managed_spark_connect/instrumentation/_core.py @@ -163,6 +163,7 @@ def __init__( self._ended = False self._lock = threading.Lock() self._token = None + recorder._track(self) def set(self, **attrs: Any) -> None: """Add or overwrite attributes.""" @@ -209,6 +210,7 @@ def end(self, end_mono_ns: Optional[int] = None) -> None: if self._ended: return self._ended = True + self._recorder._untrack(self) self._recorder._emit( { "type": "span", @@ -307,11 +309,47 @@ def __init__( self._emit = emit self._attributes = attributes self.trace_id = "" + self._inflight: Dict[str, "Span"] = {} + self._idle = threading.Condition() def current_span(self) -> Optional[Span]: """The span that ``span()`` would nest under, if any.""" return _CURRENT_SPAN.get() + def inflight(self, prefix: str = "") -> List[Span]: + """Spans started but not yet ended whose name starts with ``prefix``.""" + with self._idle: + return [ + s for s in self._inflight.values() if s.name.startswith(prefix) + ] + + def wait_idle(self, timeout_s: float, prefix: str = "grpc.") -> bool: + """Block until no span named ``prefix*`` is in flight. + + Returns ``False`` if ``timeout_s`` elapsed first. The default prefix + waits for RPCs, including the background ``ReleaseExecute`` calls + pyspark issues on its thread pool after a query finishes. + """ + deadline = time.monotonic() + timeout_s + with self._idle: + while any( + s.name.startswith(prefix) for s in self._inflight.values() + ): + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + self._idle.wait(remaining) + return True + + def _track(self, span: "Span") -> None: + with self._idle: + self._inflight[span.span_id] = span + + def _untrack(self, span: "Span") -> None: + with self._idle: + if self._inflight.pop(span.span_id, None) is not None: + self._idle.notify_all() + def span(self, name: str, **attrs: Any) -> ContextManager[Span]: parent = _CURRENT_SPAN.get() return Span(self, name, parent.span_id if parent else None, attrs) @@ -380,6 +418,12 @@ def trace_id(self, value: str) -> None: def current_span(self) -> None: return None + def inflight(self, prefix: str = "") -> List[Span]: + return [] + + def wait_idle(self, timeout_s: float, prefix: str = "grpc.") -> bool: + return True + def span(self, name: str, **attrs: Any) -> ContextManager[_NoopSpan]: return _NOOP_SPAN @@ -517,6 +561,11 @@ def wrap_channel(self, channel, layer: str = "outer"): return wrap_channel(channel, self.recorder, layer) + def wait_idle(self, timeout_s: float, prefix: str = "grpc.") -> bool: + """Wait for in-flight spans (RPCs by default); see + :meth:`Recorder.wait_idle`.""" + return self.recorder.wait_idle(timeout_s, prefix) + def flush(self) -> None: for sink in self.sinks: self._call_sink(sink, "flush") diff --git a/tests/unit/test_instrumentation.py b/tests/unit/test_instrumentation.py index 05729a3d..a86ec52a 100644 --- a/tests/unit/test_instrumentation.py +++ b/tests/unit/test_instrumentation.py @@ -1379,6 +1379,8 @@ def test_abandoned_stream_span_is_emitted_when_iterator_is_dropped(inst): "status", ] assert recs[0]["parent_id"] == recs[1]["span_id"] + assert inst.recorder.inflight() == [] + assert inst.wait_idle(0) is True def test_abandoned_active_stream_does_not_block(inst): @@ -1487,3 +1489,37 @@ def test_span_end_at_explicit_time(inst): early = inst.recorder.start_span("y") early.end(end_mono_ns=early.start_mono_ns - 10) assert one(inst, "y")["duration_ns"] == 0 + + +# -------------------------------------------------------------------------- +# In-flight tracking +# -------------------------------------------------------------------------- + + +def test_inflight_and_wait_idle(inst): + import time + + assert inst.recorder.inflight() == [] + assert inst.wait_idle(0) is True + span = inst.recorder.start_span("grpc.ReleaseExecute") + other = inst.recorder.start_span("proxy.connection") + assert [s.name for s in inst.recorder.inflight("grpc.")] == [ + "grpc.ReleaseExecute" + ] + assert len(inst.recorder.inflight()) == 2 + assert inst.wait_idle(0.05) is False + assert inst.wait_idle(0.05, prefix="client.") is True + + def finish_later(): + time.sleep(0.05) + span.end() + + t = threading.Thread(target=finish_later) + t.start() + assert inst.wait_idle(5) is True # woken by span.end(), not the timeout + t.join() + assert inst.recorder.inflight("grpc.") == [] + other.end() + assert inst.recorder.inflight() == [] + assert NOOP_RECORDER.inflight() == [] + assert NOOP_RECORDER.wait_idle(0) is True