diff --git a/CHANGELOG.md b/CHANGELOG.md index 9ea77612..6bbe97d2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,4 @@ +* Add the `ydb.query.session.closed` counter for query session pool closures, labeled by pool name and a standardized closure reason; metrics-enabled clients now advertise `ydb-sdk-metrics/0.2.0` in `x-ydb-sdk-build-info` * Add `SchemeEntry.interrupt_permission_inheritance`, reported by `describe_path` and `list_directory`, telling whether the object inherits permissions from its parents ## 3.31.2 ## diff --git a/docs/observability.rst b/docs/observability.rst index b91fcd4a..2ab8861f 100644 --- a/docs/observability.rst +++ b/docs/observability.rst @@ -323,6 +323,9 @@ adapter maps them to instruments on the ``"ydb.sdk"`` meter): * - ``ydb.query.session.count`` - ObservableUpDownCounter - Current number of open query sessions by pool and ``ydb.query.session.state`` (``idle`` / ``used``). + * - ``ydb.query.session.closed`` + - Counter + - Closed query sessions by pool and closure ``reason``. * - ``ydb.query.session.max`` - ObservableUpDownCounter - Maximum configured number of sessions for a query session pool. @@ -364,6 +367,27 @@ the SDK falls back to the driver connection string ```` (e.g ``QuerySessionPool(..., name="main-pool")`` (sync or async) when several pools share a connection string. Retry metrics are recorded without attributes. +``ydb.query.session.closed`` uses these ``reason`` values: + +* ``pool_idle_timeout`` — the idle-session cleaner removes a session. Python's + ``QuerySessionPool`` has no idle-session timeout lifecycle, so it does not currently + emit this reason. +* ``pool_graceful_shutdown`` — pool shutdown removes a session. +* ``client_timeout`` — a query stream exceeds its client-side transport timeout. +* ``client_cancelled`` — the client closes an unfinished query stream. +* ``attach_closed`` — the server closes the active attach stream. +* ``transport_error`` — an ``UNAVAILABLE`` status or a connection, query-stream, or + attach-stream transport failure retires the session. +* ``node_shutdown`` — the server sends a node shutdown hint. +* ``session_shutdown`` — the server sends a session shutdown hint. +* ``bad_session`` — the server returns ``BAD_SESSION`` or ``SESSION_EXPIRED``. +* ``session_busy`` — the server returns ``SESSION_BUSY``. + +Only an active session managed by a client pool publishes this metric. Failure of the +initial attach handshake does not count as closing an active session, and standalone +``QuerySession`` instances do not publish pool metrics. A session publishes at most one +closure event; the first terminal reason wins. + Writing a Custom Metrics Backend ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/tests/aio/query/test_query_session.py b/tests/aio/query/test_query_session.py index 15d5c1cf..d77b74ab 100644 --- a/tests/aio/query/test_query_session.py +++ b/tests/aio/query/test_query_session.py @@ -73,6 +73,7 @@ async def test_transaction_after_delete_raises(self, session: QuerySession): await session.delete() + assert not session._invalidated with pytest.raises(RuntimeError): session.transaction() diff --git a/tests/observability/test_metrics.py b/tests/observability/test_metrics.py index 66cdefda..b47e0c89 100644 --- a/tests/observability/test_metrics.py +++ b/tests/observability/test_metrics.py @@ -1,3 +1,4 @@ +import asyncio import inspect from contextlib import contextmanager from unittest.mock import AsyncMock, MagicMock @@ -7,6 +8,17 @@ from opentelemetry.sdk.metrics import MeterProvider from opentelemetry.sdk.metrics.export import InMemoryMetricReader +from ydb import issues +from ydb._grpc.common.protos import ydb_query_pb2 +from ydb._grpc.grpcwrapper.common_utils import ServerStatus +from ydb.aio.query.base import AsyncResponseContextIterator +from ydb.aio.query.pool import QuerySessionPool as AsyncQuerySessionPool +from ydb.aio.query.session import QuerySession as AsyncQuerySession +from ydb.observability.metrics import QUERY_SESSION_CLOSED +from ydb.query.base import SyncResponseContextIterator, bad_session_handler +from ydb.query.pool import QuerySessionPool as SyncQuerySessionPool +from ydb.query.session import QuerySession as SyncQuerySession + def _metrics_by_name(reader): data = reader.get_metrics_data() @@ -47,15 +59,30 @@ def _sum_value(reader, name): return _single_point(reader, name).value +def _assert_closed_metric(reader, pool_name, reason): + point = _single_point(reader, QUERY_SESSION_CLOSED) + assert (point.value, point.attributes) == ( + 1, + {"ydb.query.session.pool.name": pool_name, "reason": reason}, + ) + + def _histogram_boundaries_advisory_supported(): return "explicit_bucket_boundaries_advisory" in inspect.signature(Meter.create_histogram).parameters +def test_metrics_build_info_version(metrics_setup): + from ydb.observability.metrics import _metrics_build_info_tokens + + assert _metrics_build_info_tokens() == ["ydb-sdk-metrics/0.2.0"] + + def test_metrics_registry_records_all_instruments(metrics_setup, monkeypatch): from ydb import issues from ydb.observability.metrics import ( CLIENT_OPERATION_DURATION, CLIENT_OPERATION_FAILED, + QUERY_SESSION_CLOSED, QUERY_SESSION_COUNT, QUERY_SESSION_CREATE_TIME, QUERY_SESSION_MAX, @@ -67,6 +94,7 @@ def test_metrics_registry_records_all_instruments(metrics_setup, monkeypatch): ATTEMPT_BUCKETS, DURATION_BUCKETS_SECONDS, RETRY_DURATION_BUCKETS_SECONDS, + SessionMetrics, create_metrics_operation, record_query_session_count, record_query_session_create_time, @@ -83,6 +111,10 @@ def test_metrics_registry_records_all_instruments(metrics_setup, monkeypatch): raise issues.Unavailable("transient") record_query_session_count(2, "main", "used") + session_metrics = SessionMetrics() + session_metrics.pool_name = "main" + session_metrics.count_open() + session_metrics.count_closed("pool_graceful_shutdown") record_query_session_create_time(0.5, "main") record_query_session_max(100, "main") record_query_session_pending_requests(1, "main") @@ -94,6 +126,7 @@ def test_metrics_registry_records_all_instruments(metrics_setup, monkeypatch): assert set(metrics) == { CLIENT_OPERATION_DURATION, CLIENT_OPERATION_FAILED, + QUERY_SESSION_CLOSED, QUERY_SESSION_COUNT, QUERY_SESSION_CREATE_TIME, QUERY_SESSION_MAX, @@ -105,6 +138,7 @@ def test_metrics_registry_records_all_instruments(metrics_setup, monkeypatch): } assert metrics[CLIENT_OPERATION_DURATION].unit == "s" assert metrics[CLIENT_OPERATION_FAILED].unit == "{command}" + assert metrics[QUERY_SESSION_CLOSED].unit == "{session}" assert metrics[QUERY_SESSION_COUNT].unit == "{connection}" assert metrics[QUERY_SESSION_CREATE_TIME].unit == "s" assert metrics[QUERY_SESSION_MAX].unit == "{connection}" @@ -816,6 +850,20 @@ def test_tracing_plugin_attach_context_without_end_on_exit(self, otel_setup): class TestQuerySessionPoolMetricsInstrumentation: + @staticmethod + def _counted_session(pool_name="session-pool", async_=False): + QuerySession = AsyncQuerySession if async_ else SyncQuerySession + session = QuerySession(MagicMock()) + session._session_id = "session-1" + session._session_metrics.pool_name = pool_name + session._session_metrics.count_open() + return session + + @staticmethod + def _failing_status_stream(): + raise RuntimeError("transport error") + yield + def test_query_session_init_metrics_defaults(self): from ydb.query.session import QuerySession @@ -912,7 +960,7 @@ def fake_span_ctx(**kwargs): "ydb.query.session.create_ydb_span", lambda *args, **kwargs: MagicMock(attach_context=fake_span_ctx) ) monkeypatch.setattr(qs, "_create_call", MagicMock()) - monkeypatch.setattr(qs, "_attach", MagicMock()) + monkeypatch.setattr(qs, "_attach", MagicMock(side_effect=qs._session_metrics.count_open)) qs.create() assert qs._session_metrics._counted @@ -934,6 +982,190 @@ def test_sync_session_close_decrements_session_count(self, metrics_setup): assert not qs._session_metrics._counted assert _sum_value(metrics_setup, QUERY_SESSION_COUNT) == -1 + @pytest.mark.asyncio + @pytest.mark.parametrize("async_", [False, True], ids=["sync", "async"]) + async def test_pool_graceful_shutdown_records_closed_session(self, metrics_setup, async_): + QuerySessionPool = AsyncQuerySessionPool if async_ else SyncQuerySessionPool + pool_name = "async-pool" if async_ else "sync-pool" + pool = QuerySessionPool(driver=MagicMock(), size=2, name=pool_name) + session = self._counted_session(pool_name, async_=async_) + session._delete_call = AsyncMock() if async_ else MagicMock() + pool._queue.put_nowait(session) + + if async_: + await pool.stop() + else: + pool.stop() + + _assert_closed_metric(metrics_setup, pool_name, "pool_graceful_shutdown") + + @pytest.mark.parametrize( + ("hint", "reason"), + [ + ("node_shutdown", "node_shutdown"), + ("session_shutdown", "session_shutdown"), + ], + ) + def test_server_shutdown_hint_records_closed_session(self, metrics_setup, hint, reason): + session = self._counted_session("sync-pool") + hint_type = { + "node_shutdown": ydb_query_pb2.NodeShutdownHint, + "session_shutdown": ydb_query_pb2.SessionShutdownHint, + }[hint] + response = ydb_query_pb2.SessionState(status=0, **{hint: hint_type()}) + + session._handle_attach_session_state(response) + session._handle_attach_session_state(response) + + _assert_closed_metric(metrics_setup, "sync-pool", reason) + + @pytest.mark.parametrize( + ("async_", "error", "reason"), + [ + (False, None, "attach_closed"), + (False, RuntimeError("transport error"), "transport_error"), + (True, None, "attach_closed"), + (True, RuntimeError("transport error"), "transport_error"), + ], + ) + @pytest.mark.asyncio + async def test_attach_stream_end_records_closed_session(self, metrics_setup, async_, error, reason): + pool_name = "async-pool" if async_ else "sync-pool" + session = self._counted_session(pool_name, async_=async_) + + if async_: + session._status_stream = MagicMock() + if error is None: + session._status_stream.__aiter__.return_value = [] + else: + session._status_stream.__aiter__.side_effect = error + await session._check_session_status_loop() + else: + status_stream = iter(()) if error is None else self._failing_status_stream() + session._check_session_status_loop(status_stream) + + _assert_closed_metric(metrics_setup, pool_name, reason) + + @pytest.mark.parametrize( + ("status", "reason"), + [ + (issues.StatusCode.BAD_SESSION, "bad_session"), + (issues.StatusCode.SESSION_EXPIRED, "bad_session"), + (issues.StatusCode.SESSION_BUSY, "session_busy"), + (issues.StatusCode.UNAVAILABLE, "transport_error"), + ], + ) + def test_attach_stream_status_records_closed_session(self, metrics_setup, status, reason): + session = self._counted_session("sync-pool") + + session._check_session_status_loop(iter([ServerStatus(status, [])])) + + _assert_closed_metric(metrics_setup, "sync-pool", reason) + + @pytest.mark.parametrize("case", ["attach_status_error", "non_terminal_query_error", "no_client_pool"]) + def test_session_close_not_reported_without_matching_pool_lifecycle(self, metrics_setup, case): + session = self._counted_session(None if case == "no_client_pool" else "sync-pool") + if case == "attach_status_error": + session._check_session_status_loop(iter([ServerStatus(issues.StatusCode.BAD_REQUEST, [])])) + elif case == "non_terminal_query_error": + session._on_execute_stream_error(issues.BadRequest("invalid query")) + else: + session._check_session_status_loop(iter(())) + + assert session.is_active == (case == "non_terminal_query_error") + assert _points(metrics_setup, QUERY_SESSION_CLOSED) == [] + + @pytest.mark.parametrize( + ("error", "reason"), + [ + (issues.DeadlineExceed("deadline"), "client_timeout"), + (issues.Cancelled("cancelled"), "client_cancelled"), + (asyncio.CancelledError("cancelled"), "client_cancelled"), + (issues.SessionBusy("busy"), "session_busy"), + (issues.BadSession("bad session"), "bad_session"), + (issues.SessionExpired("expired"), "bad_session"), + (issues.Unavailable("unavailable"), "transport_error"), + (issues.ConnectionError("connection error"), "transport_error"), + (RuntimeError("transport error"), "transport_error"), + ], + ) + def test_execute_stream_terminal_error_records_closed_session_once(self, metrics_setup, error, reason): + session = self._counted_session("sync-pool") + + session._on_execute_stream_error(error) + session._check_session_status_loop(iter(())) + session._handle_attach_session_state( + ydb_query_pb2.SessionState(status=0, session_shutdown=ydb_query_pb2.SessionShutdownHint()) + ) + session._check_session_status_loop(self._failing_status_stream()) + + assert session.is_closed + _assert_closed_metric(metrics_setup, "sync-pool", reason) + + @pytest.mark.parametrize("iterator_type", ["sync", "async"]) + @pytest.mark.parametrize("finished", [False, True]) + def test_explicit_query_stream_cancel_records_only_unfinished_session(self, metrics_setup, iterator_type, finished): + ResponseContextIterator = ( + SyncResponseContextIterator if iterator_type == "sync" else AsyncResponseContextIterator + ) + session = self._counted_session("sync-pool") + iterator = ResponseContextIterator( + MagicMock(), + MagicMock(), + on_error=session._on_execute_stream_error, + on_finish=MagicMock(), + ) + iterator.it = MagicMock() + + if finished: + iterator._call_on_finish() + iterator.cancel() + iterator.cancel() + + if finished: + assert session.is_active + assert _points(metrics_setup, QUERY_SESSION_CLOSED) == [] + return + + _assert_closed_metric(metrics_setup, "sync-pool", "client_cancelled") + + @pytest.mark.parametrize("error_type", [issues.BadSession, issues.SessionExpired]) + def test_bad_session_response_records_bad_session_once(self, metrics_setup, error_type): + session = self._counted_session("sync-pool") + + @bad_session_handler + def fail(rpc_state, response_pb, query_session): + raise error_type("bad session") + + with pytest.raises(error_type) as error: + fail(None, None, session) + session._on_execute_stream_error(error.value) + + _assert_closed_metric(metrics_setup, "sync-pool", "bad_session") + + @pytest.mark.asyncio + @pytest.mark.parametrize("async_", [False, True], ids=["sync", "async"]) + async def test_initial_attach_error_does_not_record_closed_session(self, metrics_setup, monkeypatch, async_): + QuerySession = AsyncQuerySession if async_ else SyncQuerySession + pool_name = "async-pool" if async_ else "sync-pool" + session = QuerySession(MagicMock()) + session._session_id = "session-1" + session._session_metrics.pool_name = pool_name + mock_type = AsyncMock if async_ else MagicMock + session._attach_call = mock_type(return_value=MagicMock()) + first_response = mock_type(side_effect=RuntimeError("initial attach error")) + module = "ydb.aio.query.session" if async_ else "ydb.query.session" + monkeypatch.setattr(module + "._utilities.get_first_message_with_timeout", first_response) + + with pytest.raises(RuntimeError, match="initial attach error"): + if async_: + await session._attach() + else: + session._attach() + + assert session.is_closed + assert _points(metrics_setup, QUERY_SESSION_CLOSED) == [] + @pytest.mark.asyncio async def test_async_pool_acquire_from_queue_updates_session_count(self, metrics_setup): from ydb.aio.query.pool import QuerySessionPool @@ -1023,7 +1255,7 @@ def fake_span_ctx(**kwargs): "ydb.aio.query.session.create_ydb_span", lambda *args, **kwargs: MagicMock(attach_context=fake_span_ctx) ) monkeypatch.setattr(qs, "_create_call", AsyncMock()) - monkeypatch.setattr(qs, "_attach", AsyncMock()) + monkeypatch.setattr(qs, "_attach", AsyncMock(side_effect=qs._session_metrics.count_open)) await qs.create() assert qs._session_metrics._counted diff --git a/tests/query/test_query_session.py b/tests/query/test_query_session.py index fa85d724..1dce78a6 100644 --- a/tests/query/test_query_session.py +++ b/tests/query/test_query_session.py @@ -71,6 +71,7 @@ def test_transaction_after_delete_raises(self, session: QuerySession): session.delete() + assert not session._invalidated with pytest.raises(RuntimeError): session.transaction() diff --git a/ydb/aio/query/base.py b/ydb/aio/query/base.py index 6c13dfd4..7aa2f229 100644 --- a/ydb/aio/query/base.py +++ b/ydb/aio/query/base.py @@ -1,4 +1,5 @@ from .. import _utilities +from ... import issues class AsyncResponseContextIterator(_utilities.AsyncResponseIterator): @@ -12,6 +13,13 @@ def __init__(self, it, wrapper, on_error=None, on_finish=None): async def __aenter__(self) -> "AsyncResponseContextIterator": return self + def cancel(self): + error = issues.Cancelled("Query stream was cancelled by client") + if self._on_error: + self._on_error(error) + self._call_on_finish(error) + return super().cancel() + async def _next(self): try: return await super()._next() @@ -36,6 +44,7 @@ def _call_on_finish(self, exception=None): if self._on_finish is not None: self._on_finish(exception) self._on_finish = None + self._on_error = None def __del__(self): self._call_on_finish() diff --git a/ydb/aio/query/pool.py b/ydb/aio/query/pool.py index 25f92703..a0e45142 100644 --- a/ydb/aio/query/pool.py +++ b/ydb/aio/query/pool.py @@ -273,6 +273,7 @@ async def stop(self): while True: try: session = self._queue.get_nowait() + session._session_metrics.count_closed("pool_graceful_shutdown") tasks.append(session.delete()) except asyncio.QueueEmpty: break diff --git a/ydb/aio/query/pool_test.py b/ydb/aio/query/pool_test.py index 71d083da..5941e156 100644 --- a/ydb/aio/query/pool_test.py +++ b/ydb/aio/query/pool_test.py @@ -121,6 +121,24 @@ async def test_retry_reacquires_invalidated_session_before_first_use(self): live_session.explain.assert_awaited_once_with("SELECT 1") +class TestQuerySessionDelete(unittest.IsolatedAsyncioTestCase): + async def test_closes_before_delete_call(self): + session = QuerySession(MagicMock()) + session._session_id = "session-id" + session._session_metrics = MagicMock() + + async def delete_call(settings=None): + self.assertTrue(session.is_closed) + session._close_session(invalidate=True, reason="attach_closed") + + session._delete_call = AsyncMock(side_effect=delete_call) + + await session.delete() + + self.assertFalse(session._invalidated) + session._session_metrics.count_closed.assert_called_once_with(None) + + async def _async_empty_iter(): """Async-iterable that yields nothing; usable as a stub for session.execute return value.""" if False: diff --git a/ydb/aio/query/session.py b/ydb/aio/query/session.py index 7f765458..e87178a7 100644 --- a/ydb/aio/query/session.py +++ b/ydb/aio/query/session.py @@ -61,6 +61,8 @@ async def _attach(self) -> None: DEFAULT_INITIAL_RESPONSE_TIMEOUT, ) issues._process_response(first_response) + if not self._closed: + self._session_metrics.count_open() except Exception as e: self._close_session(invalidate=True) raise e @@ -72,11 +74,17 @@ async def _check_session_status_loop(self) -> None: return try: async for status in self._status_stream: - issues._process_response(status) + try: + issues._process_response(status) + except Exception as e: + logger.debug("Attach stream status error: %s, session_id: %s", e, self._session_id) + self._on_attach_stream_status_error(e) + return logger.debug("Attach stream closed, session_id: %s", self._session_id) + self._close_session(invalidate=True, reason="attach_closed") except Exception as e: - logger.debug("Attach stream error: %s, session_id: %s", e, self._session_id) - self._close_session(invalidate=True) + logger.debug("Attach stream transport error: %s, session_id: %s", e, self._session_id) + self._close_session(invalidate=True, reason="transport_error") async def delete(self, settings: Optional[BaseRequestSettings] = None) -> None: """Deletes a Session of Query Service on server side and releases resources. @@ -86,14 +94,14 @@ async def delete(self, settings: Optional[BaseRequestSettings] = None) -> None: if self._closed: return + self._close_session() + if self._session_id: try: await self._delete_call(settings=settings) except Exception: pass - self._close_session() - async def create(self, settings: Optional[BaseRequestSettings] = None) -> "QuerySession": """Creates a Session of Query Service on server side and attaches it. @@ -109,7 +117,6 @@ async def create(self, settings: Optional[BaseRequestSettings] = None) -> "Query await self._create_call(settings=settings) set_peer_attributes(span, self._peer) await self._attach() - self._session_metrics.count_open() return self diff --git a/ydb/observability/metrics.py b/ydb/observability/metrics.py index 95a4d321..245b52ea 100644 --- a/ydb/observability/metrics.py +++ b/ydb/observability/metrics.py @@ -27,6 +27,7 @@ CLIENT_OPERATION_DURATION = "db.client.operation.duration" CLIENT_OPERATION_FAILED = "ydb.client.operation.failed" QUERY_SESSION_COUNT = "ydb.query.session.count" +QUERY_SESSION_CLOSED = "ydb.query.session.closed" QUERY_SESSION_CREATE_TIME = "ydb.query.session.create_time" QUERY_SESSION_PENDING_REQUESTS = "ydb.query.session.pending_requests" QUERY_SESSION_TIMEOUTS = "ydb.query.session.timeouts" @@ -35,7 +36,7 @@ RETRY_ATTEMPTS = "ydb.client.retry.attempts" RETRY_DURATION = "ydb.client.retry.duration" -METRICS_SDK_BUILD_INFO = "ydb-sdk-metrics/0.1.0" +METRICS_SDK_BUILD_INFO = "ydb-sdk-metrics/0.2.0" DURATION_BUCKETS_SECONDS = ( 0.001, @@ -220,7 +221,7 @@ def query_session_pool_name( def _metrics_build_info_tokens() -> List[str]: """Metrics' contribution to the ``x-ydb-sdk-build-info`` header. - Returns ``["ydb-sdk-metrics/0.1.0"]`` once a metrics backend is installed, + Returns ``["ydb-sdk-metrics/0.2.0"]`` once a metrics backend is installed, otherwise an empty list. Aggregated with other features by :func:`ydb.observability.sdk_build_info_tokens`. """ @@ -447,12 +448,13 @@ class SessionMetrics: :attr:`state` and :attr:`pool_name` as the session moves between idle and used. """ - __slots__ = ("pool_name", "state", "_counted") + __slots__ = ("pool_name", "state", "_counted", "_lock") def __init__(self) -> None: self.pool_name: Optional[str] = None self.state: str = "used" self._counted = False + self._lock = threading.Lock() def count_open(self) -> None: if self._counted: @@ -460,11 +462,21 @@ def count_open(self) -> None: self._counted = True record_query_session_count(1, self.pool_name, self.state) - def count_closed(self) -> None: - if not self._counted: - return - self._counted = False + def count_closed(self, reason: Optional[str] = None) -> None: + with self._lock: + if not self._counted: + return + self._counted = False record_query_session_count(-1, self.pool_name, self.state) + if reason is not None and self.pool_name is not None: + _provider.add( + QUERY_SESSION_CLOSED, + 1, + { + "ydb.query.session.pool.name": self.pool_name, + "reason": reason, + }, + ) class _NoopSessionMetrics(SessionMetrics): @@ -473,7 +485,7 @@ class _NoopSessionMetrics(SessionMetrics): def count_open(self) -> None: pass - def count_closed(self) -> None: + def count_closed(self, reason: Optional[str] = None) -> None: pass diff --git a/ydb/opentelemetry/metrics_plugin.py b/ydb/opentelemetry/metrics_plugin.py index e1df71d9..4a25ca31 100644 --- a/ydb/opentelemetry/metrics_plugin.py +++ b/ydb/opentelemetry/metrics_plugin.py @@ -21,6 +21,7 @@ from ydb.observability.metrics import ( CLIENT_OPERATION_DURATION, CLIENT_OPERATION_FAILED, + QUERY_SESSION_CLOSED, QUERY_SESSION_COUNT, QUERY_SESSION_CREATE_TIME, QUERY_SESSION_MAX, @@ -96,6 +97,11 @@ def __init__(self, meter: Meter) -> None: unit="{command}", description="Number of failed YDB client operations.", ), + QUERY_SESSION_CLOSED: meter.create_counter( + QUERY_SESSION_CLOSED, + unit="{session}", + description="Number of closed YDB query sessions.", + ), QUERY_SESSION_TIMEOUTS: meter.create_counter( QUERY_SESSION_TIMEOUTS, unit="{connection}", diff --git a/ydb/query/base.py b/ydb/query/base.py index 12752fe2..a78825dc 100644 --- a/ydb/query/base.py +++ b/ydb/query/base.py @@ -82,6 +82,13 @@ def __init__(self, it, wrapper, on_error=None, on_finish=None): def __enter__(self) -> "SyncResponseContextIterator": return self + def cancel(self): + error = issues.Cancelled("Query stream was cancelled by client") + if self._on_error: + self._on_error(error) + self._call_on_finish(error) + return super().cancel() + def _next(self): try: return super()._next() @@ -105,6 +112,7 @@ def _call_on_finish(self, exception=None): if self._on_finish is not None: self._on_finish(exception) self._on_finish = None + self._on_error = None def __del__(self): self._call_on_finish() @@ -219,8 +227,8 @@ def bad_session_handler(func): def decorator(rpc_state, response_pb, session: "BaseQuerySession", *args, **kwargs): try: return func(rpc_state, response_pb, session, *args, **kwargs) - except issues.BadSession: - session._close_session(invalidate=True) + except (issues.BadSession, issues.SessionExpired): + session._close_session(invalidate=True, reason="bad_session") raise return decorator diff --git a/ydb/query/pool.py b/ydb/query/pool.py index 79f5051a..ccc1a6b0 100644 --- a/ydb/query/pool.py +++ b/ydb/query/pool.py @@ -326,6 +326,7 @@ def stop(self, timeout=None): while True: try: session = self._queue.get_nowait() + session._session_metrics.count_closed("pool_graceful_shutdown") session.delete() except queue.Empty: break diff --git a/ydb/query/pool_test.py b/ydb/query/pool_test.py index 33041ccf..fb00eb1a 100644 --- a/ydb/query/pool_test.py +++ b/ydb/query/pool_test.py @@ -156,6 +156,24 @@ def test_retry_reacquires_invalidated_session_before_first_use(self): live_session.explain.assert_called_once_with("SELECT 1") +class TestQuerySessionDelete(unittest.TestCase): + def test_closes_before_delete_call(self): + session = QuerySession(MagicMock()) + session._session_id = "session-id" + session._session_metrics = MagicMock() + + def delete_call(settings=None): + self.assertTrue(session.is_closed) + session._close_session(invalidate=True, reason="attach_closed") + + session._delete_call = MagicMock(side_effect=delete_call) + + session.delete() + + self.assertFalse(session._invalidated) + session._session_metrics.count_closed.assert_called_once_with(None) + + class TestQuerySessionExecutePoolId(unittest.TestCase): """Test that pool_id flows from session.execute() → _execute_call() → driver.""" diff --git a/ydb/query/session.py b/ydb/query/session.py index b8a7d82e..8e026f5c 100644 --- a/ydb/query/session.py +++ b/ydb/query/session.py @@ -1,4 +1,5 @@ import abc +import asyncio import json import logging import threading @@ -169,18 +170,18 @@ def _handle_attach_session_state(self, response_pb) -> None: if response_pb is None: return - hint = response_pb.WhichOneof("session_hint") - if hint == "node_shutdown": - if self._node_id is not None: - self._driver._pessimize_node(self._node_id) - self._close_session(invalidate=True) - elif hint == "session_shutdown": - self._close_session(invalidate=True) + match response_pb.WhichOneof("session_hint"): + case "node_shutdown": + if self._node_id is not None: + self._driver._pessimize_node(self._node_id) + self._close_session(invalidate=True, reason="node_shutdown") + case "session_shutdown": + self._close_session(invalidate=True, reason="session_shutdown") - def _close_session(self, invalidate: bool = False) -> None: + def _close_session(self, invalidate: bool = False, reason: Optional[str] = None) -> None: if self._closed: return - self._session_metrics.count_closed() + self._session_metrics.count_closed(reason) if invalidate: self._invalidated = True self._closed = True @@ -191,6 +192,19 @@ def _close_session(self, invalidate: bool = False) -> None: except Exception: pass + def _on_attach_stream_status_error(self, e: BaseException) -> None: + reason: Optional[str] + match e: + case issues.BadSession() | issues.SessionExpired(): + reason = "bad_session" + case issues.SessionBusy(): + reason = "session_busy" + case issues.Unavailable() | issues.ConnectionError(): + reason = "transport_error" + case _: + reason = None + self._close_session(invalidate=True, reason=reason) + def _on_execute_stream_error(self, e: BaseException) -> None: # The execute stream is a single gRPC call that carries all of a # query's response parts. If any of these errors surface while reading @@ -205,20 +219,22 @@ def _on_execute_stream_error(self, e: BaseException) -> None: # Accepts BaseException so that asyncio.CancelledError (not an # issues.Error subclass) — the case documented in the bug report — # also invalidates here. - if isinstance(e, issues.Error): - if isinstance( - e, - ( - issues.DeadlineExceed, - issues.SessionBusy, - issues.BadSession, - issues.ConnectionError, - issues.Cancelled, - ), - ): - self._close_session(invalidate=True) - else: - self._close_session(invalidate=True) + match e: + case issues.DeadlineExceed(): + reason = "client_timeout" + case issues.Cancelled() | asyncio.CancelledError(): + reason = "client_cancelled" + case issues.SessionBusy(): + reason = "session_busy" + case issues.BadSession() | issues.SessionExpired(): + reason = "bad_session" + case issues.Unavailable() | issues.ConnectionError(): + reason = "transport_error" + case issues.Error(): + return + case _: + reason = "transport_error" + self._close_session(invalidate=True, reason=reason) # Overloads for _create_call @overload @@ -397,6 +413,8 @@ def _attach(self, first_resp_timeout: int = DEFAULT_INITIAL_RESPONSE_TIMEOUT) -> first_resp_timeout, ) issues._process_response(first_response) + if not self._closed: + self._session_metrics.count_open() except Exception as e: self._close_session(invalidate=True) raise e @@ -411,11 +429,17 @@ def _attach(self, first_resp_timeout: int = DEFAULT_INITIAL_RESPONSE_TIMEOUT) -> def _check_session_status_loop(self, status_stream: _utilities.SyncResponseIterator) -> None: try: for status in status_stream: - issues._process_response(status) + try: + issues._process_response(status) + except Exception as e: + logger.debug("Attach stream status error: %s, session_id: %s", e, self._session_id) + self._on_attach_stream_status_error(e) + return logger.debug("Attach stream closed, session_id: %s", self._session_id) + self._close_session(invalidate=True, reason="attach_closed") except Exception as e: - logger.debug("Attach stream error: %s, session_id: %s", e, self._session_id) - self._close_session(invalidate=True) + logger.debug("Attach stream transport error: %s, session_id: %s", e, self._session_id) + self._close_session(invalidate=True, reason="transport_error") def delete(self, settings: Optional[BaseRequestSettings] = None) -> None: """Deletes a Session of Query Service on server side and releases resources. @@ -425,14 +449,14 @@ def delete(self, settings: Optional[BaseRequestSettings] = None) -> None: if self._closed: return + self._close_session() + if self._session_id: try: self._delete_call(settings=settings) except Exception: pass - self._close_session() - def create(self, settings: Optional[BaseRequestSettings] = None) -> "QuerySession": """Creates a Session of Query Service on server side and attaches it. @@ -448,7 +472,6 @@ def create(self, settings: Optional[BaseRequestSettings] = None) -> "QuerySessio self._create_call(settings=settings) set_peer_attributes(span, self._peer) self._attach() - self._session_metrics.count_open() return self