diff --git a/datasette/database.py b/datasette/database.py index cc35c038..d7cbc99b 100644 --- a/datasette/database.py +++ b/datasette/database.py @@ -1,16 +1,19 @@ import asyncio import atexit +import contextvars import inspect import os import queue import sys import tempfile import threading +import time import uuid from collections import namedtuple from pathlib import Path import sqlite_utils +from opentelemetry import context as otel_context_api from opentelemetry.trace import Status, StatusCode from .inspect import inspect_hash @@ -351,9 +354,15 @@ class Database: return _run() if not write: # Immutable database - no writes can ever occur, so there is no - # write queue to block; run against a fresh read-only connection + # write queue to block; run against a fresh read-only connection. + # A fresh copy_context() is required per submit (not one shared + # copy reused across calls): concurrent execution of the same + # Context raises "RuntimeError: cannot enter context ... already + # entered". This propagates the caller's otel context (e.g. the + # enclosing db.query span) onto the worker thread. + ctx = contextvars.copy_context() return await asyncio.get_running_loop().run_in_executor( - self.ds.executor, _run + self.ds.executor, ctx.run, _run ) # Threaded mode - send to write thread return await self._send_to_write_thread(fn, isolated_connection=True) @@ -458,8 +467,21 @@ class Database: task_id = uuid.uuid5(uuid.NAMESPACE_DNS, "datasette.io") loop = asyncio.get_running_loop() reply_future = loop.create_future() + # Captured here, on the event loop, at enqueue time: the otel + # Context (carrying the enclosing db.query span, if any) and the + # timestamp used to build the db.write.queue_wait span once this + # task is dequeued on the write thread. self._write_queue.put( - WriteTask(fn, task_id, loop, reply_future, isolated_connection, transaction) + WriteTask( + fn, + task_id, + loop, + reply_future, + isolated_connection, + transaction, + otel_context_api.get_current(), + time.time_ns(), + ) ) if block: return await reply_future @@ -473,6 +495,9 @@ class Database: conn = None try: conn = self.connect(write=True) + # This warm-up runs before any write has ever been queued, so + # there is no caller otel context yet to attach - any spans + # created by plugin hooks here are orphans (roots). self.ds._prepare_connection(conn, self.name) except Exception as e: # noqa: BLE001 # Stored and re-raised to whoever queues the next write @@ -487,40 +512,79 @@ class Database: # Best-effort close as the write thread exits pass return - exception = None - result = None - if conn_exception is not None: - exception = conn_exception - elif task.isolated_connection: - try: - isolated_connection = self.connect(write=True) + # Restore the caller's otel context (captured on the event loop + # at enqueue time) so spans created while processing this task + # parent correctly to the request that queued it. Must be + # detached below in `finally` - a leaked token silently poisons + # this thread's ambient context for every write processed after + # it, and a *wrong*-token detach only logs a warning rather than + # raising, so this pairing is load-bearing and easy to get wrong + # silently. + token = otel_context_api.attach(task.otel_context) + try: + exception = None + result = None + # Explicit start_time/end_time rather than a `with` block: + # this span's duration is the time the task actually spent + # waiting in the queue (enqueue -> dequeue), not the near- + # zero time spent constructing/ending the span object here. + tracer.start_span( + "db.write.queue_wait", start_time=task.enqueued_at_ns + ).end(end_time=time.time_ns()) + if conn_exception is not None: + # fn never runs in this branch, so there is nothing to + # wrap in a db.write.execute span. + exception = conn_exception + elif task.isolated_connection: try: - result = task.fn(isolated_connection) - finally: - isolated_connection.close() - try: - self._all_file_connections.remove(isolated_connection) - except ValueError: - # Was probably a memory connection - pass - except Exception as e: # noqa: BLE001 - # Write thread must survive any task failure or the database wedges - sys.stderr.write(f"{e}\n") - sys.stderr.flush() - exception = e - else: - try: - if task.transaction: - with conn: - conn.execute("BEGIN IMMEDIATE") - result = task.fn(conn) - else: - result = task.fn(conn) - except Exception as e: # noqa: BLE001 - sys.stderr.write(f"{e}\n") - sys.stderr.flush() - exception = e - _deliver_write_result(task, result, exception) + with tracer.start_as_current_span("db.write.execute") as span: + span.set_attribute( + "datasette.isolated_connection", + task.isolated_connection, + ) + span.set_attribute( + "datasette.transaction", task.transaction + ) + isolated_connection = self.connect(write=True) + try: + result = task.fn(isolated_connection) + finally: + isolated_connection.close() + try: + self._all_file_connections.remove( + isolated_connection + ) + except ValueError: + # Was probably a memory connection + pass + except Exception as e: # noqa: BLE001 + # Write thread must survive any task failure or the database wedges + sys.stderr.write(f"{e}\n") + sys.stderr.flush() + exception = e + else: + try: + with tracer.start_as_current_span("db.write.execute") as span: + span.set_attribute( + "datasette.isolated_connection", + task.isolated_connection, + ) + span.set_attribute( + "datasette.transaction", task.transaction + ) + if task.transaction: + with conn: + conn.execute("BEGIN IMMEDIATE") + result = task.fn(conn) + else: + result = task.fn(conn) + except Exception as e: # noqa: BLE001 + sys.stderr.write(f"{e}\n") + sys.stderr.flush() + exception = e + _deliver_write_result(task, result, exception) + finally: + otel_context_api.detach(token) async def execute_fn(self, fn): self._check_not_closed() @@ -542,7 +606,13 @@ class Database: with self._pending_execute_futures_lock: self._check_not_closed() - future = self.ds.executor.submit(in_thread) + # A fresh copy_context() is required per submit (not one shared + # copy reused across calls): concurrent execution of the same + # Context raises "RuntimeError: cannot enter context ... + # already entered". This propagates the caller's otel context + # (e.g. the enclosing db.query span) onto the worker thread. + ctx = contextvars.copy_context() + future = self.ds.executor.submit(ctx.run, in_thread) self._pending_execute_futures.add(future) future.add_done_callback(self._remove_pending_execute_future) return await asyncio.wrap_future(future) @@ -564,35 +634,51 @@ class Database: time_limit_ms = custom_time_limit def sql_operation_in_thread(conn): - with sqlite_timelimit(conn, time_limit_ms): - try: - cursor = conn.cursor() - cursor.execute(sql, params if params is not None else {}) - max_returned_rows = self.ds.max_returned_rows - if max_returned_rows == page_size: - max_returned_rows += 1 - if max_returned_rows and truncate: - rows = cursor.fetchmany(max_returned_rows + 1) - truncated = len(rows) > max_returned_rows - rows = rows[:max_returned_rows] - else: - rows = cursor.fetchall() - truncated = False - except (sqlite3.OperationalError, sqlite3.DatabaseError) as e: - if e.args == ("interrupted",): - raise QueryInterrupted(e, sql, params) - if log_sql_errors: - sys.stderr.write( - f"ERROR: conn={conn}, sql = {sql!r}, params = {params}: {e}\n" - ) - sys.stderr.flush() - raise + # This span is created inside the worker thread. Its parent is + # resolved from the ambient otel context, which was propagated + # onto this thread via copy_context() at the executor.submit() + # boundary in execute_fn() (or run_in_executor() for immutable + # databases) - so it parents correctly to the enclosing + # db.query span despite running on a different thread. + # + # Callers passing log_sql_errors=False are probing and treat a + # failure as an expected answer - see the matching handling on the + # db.query span in execute(). Without this, facet suggestion marks + # two spans per text column as failed on every table page. + with tracer.start_as_current_span( + "db.query.execute", + record_exception=log_sql_errors, + set_status_on_exception=log_sql_errors, + ): + with sqlite_timelimit(conn, time_limit_ms): + try: + cursor = conn.cursor() + cursor.execute(sql, params if params is not None else {}) + max_returned_rows = self.ds.max_returned_rows + if max_returned_rows == page_size: + max_returned_rows += 1 + if max_returned_rows and truncate: + rows = cursor.fetchmany(max_returned_rows + 1) + truncated = len(rows) > max_returned_rows + rows = rows[:max_returned_rows] + else: + rows = cursor.fetchall() + truncated = False + except (sqlite3.OperationalError, sqlite3.DatabaseError) as e: + if e.args == ("interrupted",): + raise QueryInterrupted(e, sql, params) + if log_sql_errors: + sys.stderr.write( + f"ERROR: conn={conn}, sql = {sql!r}, params = {params}: {e}\n" + ) + sys.stderr.flush() + raise - if truncate: - return Results(rows, truncated, cursor.description) + if truncate: + return Results(rows, truncated, cursor.description) - else: - return Results(rows, False, cursor.description) + else: + return Results(rows, False, cursor.description) # SIM117 wants these two context managers merged. They are kept nested # deliberately: the hand-rolled tracer's wrapper is on its way out, and @@ -924,16 +1010,26 @@ def _apply_write_wrapper(fn, wrapper_factory, track_event): class WriteTask: __slots__ = ( + "enqueued_at_ns", "fn", "isolated_connection", "loop", + "otel_context", "reply_future", "task_id", "transaction", ) def __init__( - self, fn, task_id, loop, reply_future, isolated_connection, transaction + self, + fn, + task_id, + loop, + reply_future, + isolated_connection, + transaction, + otel_context, + enqueued_at_ns, ): self.fn = fn self.task_id = task_id @@ -941,6 +1037,8 @@ class WriteTask: self.reply_future = reply_future self.isolated_connection = isolated_connection self.transaction = transaction + self.otel_context = otel_context + self.enqueued_at_ns = enqueued_at_ns def _deliver_write_result(task, result, exception): diff --git a/tests/test_internals_database.py b/tests/test_internals_database.py index b1093b1c..ce0d7a53 100644 --- a/tests/test_internals_database.py +++ b/tests/test_internals_database.py @@ -3,11 +3,13 @@ Tests for the datasette.database.Database class """ import asyncio +import threading import uuid from types import SimpleNamespace import pytest import sqlite_utils +from opentelemetry import context as otel_context_api from datasette.app import Datasette from datasette.database import ( @@ -1223,3 +1225,110 @@ async def test_database_close_is_idempotent(tmpdir): # Second call should be a no-op, not raise db.close() ds._internal_database.close() + + +_CONTEXT_LEAK_MARKER_KEY = "otel-context-leak-marker" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("num_sql_threads", (0, 1)) +async def test_write_thread_context_is_detached_between_tasks( + tmp_path, monkeypatch, num_sql_threads +): + """ + The write thread attaches each task's otel Context and must detach it + again before picking up the next task. The thread is persistent and + shared, so a leaked token would grow that thread's context stack for the + rest of the process - and a *wrong*-token detach only logs a warning + rather than raising, so "does it throw" cannot catch either mistake. + + Two things are asserted, because neither alone is sufficient: + + 1. Each task observes the context value that was current on the event + loop when it was queued. This is what fails if the Context is not + carried on WriteTask, or is never attached. It does *not* catch a + missing detach: attach() replaces the current Context wholesale, so a + leftover one from a previous task is simply overwritten. + 2. The write thread's attach depth is identical at the same point in + every task. This is what fails if detach is missing - the stack grows + by one per task - and it holds across a task that raises, because the + detach lives in a `finally`. + + An otel context value is used rather than a plain contextvars.ContextVar: + a plain var set on the event loop never crosses into the write thread, so + the probe would read None every time and the test could not fail. + """ + name = f"context_leak_test_{num_sql_threads}" + db_path = tmp_path / f"{name}.db" + sqlite3.connect(db_path).close() + ds = Datasette([str(db_path)], settings={"num_sql_threads": num_sql_threads}) + db = ds.get_database(name) + await db.execute_write("create table t (id integer primary key)") + + write_thread_name = f"_execute_writes for database {name}" + depth = {"value": 0} + real_attach = otel_context_api.attach + real_detach = otel_context_api.detach + + def counting_attach(context): + token = real_attach(context) + if threading.current_thread().name == write_thread_name: + depth["value"] += 1 + return token + + def counting_detach(token): + real_detach(token) + if threading.current_thread().name == write_thread_name: + depth["value"] -= 1 + + # Patched on the opentelemetry.context module itself, which is what both + # database.py and opentelemetry.trace.use_span() look the functions up on. + monkeypatch.setattr(otel_context_api, "attach", counting_attach) + monkeypatch.setattr(otel_context_api, "detach", counting_detach) + + seen_markers = [] + seen_depths = [] + + def probe(conn): + seen_markers.append(otel_context_api.get_value(_CONTEXT_LEAK_MARKER_KEY)) + seen_depths.append(depth["value"]) + + def failing_probe(conn): + probe(conn) + # Exercises the write thread's exception path: the detach still has + # to happen, which is why it lives in a `finally`. + raise ValueError("deliberate failure inside a write task") + + try: + for i in range(5): + ctx = otel_context_api.set_value(_CONTEXT_LEAK_MARKER_KEY, f"marker-{i}") + token = real_attach(ctx) + try: + if i == 2: + with pytest.raises(ValueError): + await db.execute_write_fn(failing_probe) + else: + await db.execute_write_fn(probe) + finally: + real_detach(token) + + # Sanity check: no marker is active in *this* (event loop) context + # right now, so the final probe is a fair test of the write thread's + # own state rather than something this test forgot to clean up. + assert otel_context_api.get_value(_CONTEXT_LEAK_MARKER_KEY) is None + await db.execute_write_fn(probe) + finally: + db.close() + + assert seen_markers == [ + "marker-0", + "marker-1", + "marker-2", + "marker-3", + "marker-4", + None, + ] + assert len(set(seen_depths)) == 1, ( + f"write thread context stack grew across tasks: {seen_depths} - " + "a token was attached without being detached" + ) diff --git a/tests/test_telemetry.py b/tests/test_telemetry.py index d8e89446..f232a39d 100644 --- a/tests/test_telemetry.py +++ b/tests/test_telemetry.py @@ -2,12 +2,15 @@ import json import sqlite3 import subprocess import sys +import time import pytest +import sqlite_utils from opentelemetry.trace import StatusCode from datasette.app import Datasette -from datasette.telemetry import MAX_SQL_LENGTH, sql_attribute +from datasette.database import Database +from datasette.telemetry import MAX_SQL_LENGTH, sql_attribute, tracer SECRET_PARAM_VALUE = "SUPER_SECRET_PARAM_VALUE_XYZ_123" @@ -33,6 +36,26 @@ def _spans_for_namespace(otel_spans, namespace): ] +def _children_named(otel_spans, name, parent_span_context): + """ + Finished spans called `name` whose parent really is `parent_span_context`. + + Parentage is matched on span id, not on "a span with this name exists" - + a span can exist and still be an unparented root if a thread boundary + dropped the otel context, which is the exact failure these tests exist + to catch. + """ + return [ + span + for span in otel_spans.get_finished_spans() + if span.name == name + and span.parent is not None + and span.parent.span_id == parent_span_context.span_id + and span.parent.trace_id == parent_span_context.trace_id + and span.context.trace_id == parent_span_context.trace_id + ] + + def _all_attribute_values(otel_spans): "Every attribute value across every finished span, for the 'no leaked param values' test." values = [] @@ -267,3 +290,170 @@ async def test_execute_write_many_records_param_sets_not_rows_returned(otel_span # calling this a row count would be a lie. Asserted explicitly because the # attribute really was named datasette.rows_returned at one point. assert "datasette.rows_returned" not in span.attributes + + +# --- Context propagation across thread boundaries -------------------------- +# +# Every assertion below checks parentage (child.parent.span_id == +# expected_parent.span_id, in the same trace), not merely that spans exist. +# Spans can exist and still be wrongly parented - or be unparented roots - if +# a thread boundary drops the otel context, which is exactly the failure mode +# these tests exist to prevent. + + +@pytest.mark.asyncio +async def test_db_query_execute_parents_to_db_query(ds_client, otel_spans): + # execute_fn()'s executor.submit() is thread boundary #1. The + # db.query.execute span is created inside the worker thread; without the + # copy_context() propagation it comes back as an unparented root span + # rather than a child of db.query. + response = await ds_client.get("/fixtures/-/query.json?sql=select+1") + assert response.status_code == 200 + + query_spans = [ + span + for span in _spans_for_namespace(otel_spans, "fixtures") + if span.attributes["db.query.text"] == "select 1" + ] + assert query_spans, "expected a db.query span for 'select 1'" + query_span = query_spans[-1] + + assert [ + span + for span in otel_spans.get_finished_spans() + if span.name == "db.query.execute" + ], "expected at least one db.query.execute span" + children = _children_named(otel_spans, "db.query.execute", query_span.context) + assert len(children) == 1, "expected exactly one db.query.execute child of db.query" + # The execute span is strictly contained by the round-trip span, and the + # gap between the two is the thread-pool wait. + assert query_span.start_time <= children[0].start_time + assert children[0].end_time <= query_span.end_time + + +@pytest.mark.asyncio +async def test_immutable_database_propagates_context(tmp_path, otel_spans): + # Thread boundary #3, the easy one to miss: immutable databases route + # execute_isolated_fn() through loop.run_in_executor() directly rather + # than through the write thread. A span created inside that worker must + # still parent to whatever was current when execute_isolated_fn() was + # awaited, or every immutable-database operation emits orphan roots. + db_path = tmp_path / "t04_immutable.db" + sqlite_utils.Database(str(db_path))["t"].insert({"id": 1}, pk="id") + + ds = Datasette() + db = Database(ds, path=str(db_path), is_mutable=False) + ds.add_database(db, name="t04_immutable") + + def fn(conn): + with tracer.start_as_current_span("t04-child-in-isolated-worker"): + pass + + try: + with tracer.start_as_current_span("t04-parent-on-event-loop") as parent: + parent_context = parent.get_span_context() + await db.execute_isolated_fn(fn) + finally: + ds.remove_database("t04_immutable") + + assert [ + span + for span in otel_spans.get_finished_spans() + if span.name == "t04-child-in-isolated-worker" + ], "expected a span created inside execute_isolated_fn's worker thread" + children = _children_named( + otel_spans, "t04-child-in-isolated-worker", parent_context + ) + assert len(children) == 1 + + +@pytest.mark.asyncio +async def test_write_spans_parent_to_db_query(otel_spans): + # Thread boundary #2: WriteTask -> queue.Queue -> the write thread. + # db.write.queue_wait and db.write.execute are both direct children of + # the db.query span that was current on the event loop at enqueue time, + # so they are siblings rather than nested inside one another. + db = Datasette(memory=True).add_memory_database("t04_write_spans") + await db.execute_write("create table docs (id integer primary key)") + + query_spans = _spans_for_namespace(otel_spans, "t04_write_spans") + assert query_spans, "expected a db.query span from execute_write()" + query_span = query_spans[-1] + + queue_wait_children = _children_named( + otel_spans, "db.write.queue_wait", query_span.context + ) + execute_children = _children_named( + otel_spans, "db.write.execute", query_span.context + ) + assert len(queue_wait_children) == 1 + assert len(execute_children) == 1 + + execute_span = execute_children[0] + assert execute_span.attributes["datasette.isolated_connection"] is False + assert execute_span.attributes["datasette.transaction"] is True + # Siblings, not parent/child: the queue wait is over by the time the + # write begins. + assert queue_wait_children[0].end_time <= execute_span.start_time + + +@pytest.mark.asyncio +async def test_write_queue_wait_duration_reflects_real_wait(otel_spans): + # db.write.queue_wait is built from explicit start/end timestamps - + # task.enqueued_at_ns, captured on the event loop, through to the moment + # the write thread dequeued it. If it were a plain `with` block on the + # write thread it would instead measure the microseconds spent building + # the span object, and this assertion would fail. + ds = Datasette(memory=True) + db = ds.add_memory_database("t04_queue_wait") + await db.execute_write("create table docs (id integer primary key)") + + def slow_write(conn): + time.sleep(0.1) + + # Queue a deliberately slow write without waiting for it, then queue a + # second write immediately behind it: the second task sits in the queue + # for roughly the duration of the first. + _, slow_future = await db._send_to_write_thread(slow_write, block=False) + await db.execute_write("insert into docs (id) values (1)") + await slow_future + + query_spans = [ + span + for span in _spans_for_namespace(otel_spans, "t04_queue_wait") + if span.attributes["db.query.text"] == "insert into docs (id) values (1)" + ] + assert query_spans, "expected a db.query span for the queued-behind insert" + queue_wait_children = _children_named( + otel_spans, "db.write.queue_wait", query_spans[-1].context + ) + assert len(queue_wait_children) == 1 + duration_ns = queue_wait_children[0].end_time - queue_wait_children[0].start_time + # The slow write sleeps 100ms; anything above 10ms is far beyond the + # microseconds a mis-timestamped span would report. + assert duration_ns > 10_000_000, f"queue wait was only {duration_ns}ns" + + +@pytest.mark.asyncio +async def test_suppressed_error_does_not_mark_execute_span(ds_client, otel_spans): + """ + The inner db.query.execute span must honour log_sql_errors too. + + It is created inside the worker thread, so without record_exception / + set_status_on_exception being passed through it would mark every facet + suggestion probe as failed even though the outer db.query span correctly + reports the failure as suppressed. + """ + db = ds_client.ds.get_database("fixtures") + with pytest.raises(sqlite3.OperationalError): + await db.execute(INVALID_SQL, log_sql_errors=False) + + execute_spans = [ + span + for span in otel_spans.get_finished_spans() + if span.name == "db.query.execute" + ] + assert execute_spans + span = execute_spans[-1] + assert span.status.status_code == StatusCode.UNSET + assert not [event for event in span.events if event.name == "exception"]