diff --git a/docsgpt/api/answer/routes/base.py b/docsgpt/api/answer/routes/base.py index 943efedc..782e11b5 100644 --- a/docsgpt/api/answer/routes/base.py +++ b/docsgpt/api/answer/routes/base.py @@ -110,7 +110,9 @@ def _traced_stream( try: yield from method(*bound.args, **bound.kwargs) finally: - tracing.flush(trace) + # Written on a writer thread so the stream's connection closes + # without waiting on the OTel replay and the INSERT. + tracing.flush(trace, background=True) return wrapper diff --git a/docsgpt/tracing/sink.py b/docsgpt/tracing/sink.py index c5f9ad6c..49a516a8 100644 --- a/docsgpt/tracing/sink.py +++ b/docsgpt/tracing/sink.py @@ -3,29 +3,36 @@ from __future__ import annotations import logging +import threading +from concurrent.futures import Future, ThreadPoolExecutor from typing import Optional from docsgpt.tracing.core import Trace logger = logging.getLogger(__name__) +# Writes handed off by ``flush(background=True)``. Two workers keep up with +# chat traffic (each write is one INSERT); the pool's threads are joined at +# interpreter exit, so queued traces are still written on a clean shutdown. +_writer: Optional[ThreadPoolExecutor] = None +_writer_lock = threading.Lock() -def flush(trace: Optional[Trace], status: Optional[str] = None) -> None: - """Finish ``trace`` and persist it once; later calls are no-ops. + +def _executor() -> ThreadPoolExecutor: + global _writer + with _writer_lock: + if _writer is None: + _writer = ThreadPoolExecutor(max_workers=2, thread_name_prefix="trace-writer") + return _writer + + +def _write(trace: Trace) -> None: + """Export ``trace`` to OTel, then store it; each sink swallows its own failure. OTel replay runs first so the exported trace id can be stored with the - row. Both sinks swallow their own failures: a trace is diagnostic data - and must never fail the request that produced it. - - Args: - trace: The trace to write; ``None`` is accepted and ignored. - status: Final status; defaults to ``error`` when a top-level span - failed, else ``ok``. + row. A trace is diagnostic data and must never fail the request that + produced it. """ - if trace is None or trace.flushed: - return - trace.flushed = True - trace.finish(status) try: from docsgpt.tracing.otel import export_trace @@ -45,6 +52,39 @@ def flush(trace: Optional[Trace], status: Optional[str] = None) -> None: logger.warning("Failed to store trace %s", trace.id, exc_info=True) +def flush( + trace: Optional[Trace], status: Optional[str] = None, *, background: bool = False +) -> Optional[Future]: + """Finish ``trace`` and persist it once; later calls are no-ops. + + The trace is frozen immediately, so nothing recorded afterwards is + included, whichever way it is written. + + Args: + trace: The trace to write; ``None`` is accepted and ignored. + status: Final status; defaults to ``error`` when a top-level span + failed, else ``ok``. + background: Write on a writer thread instead of the caller's. A chat + stream flushes this way so its connection closes without waiting + on the OTel replay and the INSERT. + + Returns: + The pending write when ``background`` is set, else ``None``. + """ + if trace is None or trace.flushed: + return None + trace.flushed = True + trace.finish(status) + if background: + try: + return _executor().submit(_write, trace) + except RuntimeError: + # The pool is shut down (interpreter exit): write inline instead. + pass + _write(trace) + return None + + def discard(trace: Optional[Trace]) -> None: """Drop ``trace`` without writing it (the request was rejected or superseded).""" if trace is None or trace.flushed: diff --git a/tests/api/answer/test_trace_wiring.py b/tests/api/answer/test_trace_wiring.py index 7cf75f7a..d03dec01 100644 --- a/tests/api/answer/test_trace_wiring.py +++ b/tests/api/answer/test_trace_wiring.py @@ -22,7 +22,7 @@ def _captured_flushes(): """Record every flushed trace instead of writing it.""" flushed = [] - def _fake_flush(trace, status=None): + def _fake_flush(trace, status=None, **_kwargs): if trace is None or trace.flushed: return trace.flushed = True diff --git a/tests/e2e/specs/tier-a/execution-traces.spec.ts b/tests/e2e/specs/tier-a/execution-traces.spec.ts index 709585bb..d0474961 100644 --- a/tests/e2e/specs/tier-a/execution-traces.spec.ts +++ b/tests/e2e/specs/tier-a/execution-traces.spec.ts @@ -75,9 +75,10 @@ test.describe('tier-a · execution traces', () => { expect(res.ok()).toBeTruthy(); expect(await res.text()).toContain('"type": "end"'); - // Stored trace: one row for the turn, linked to its message. + // Stored trace: one row for the turn, linked to its message. It is + // written just after the stream closes, so wait for it. + await expect.poll(async () => (await tracesFor(sub)).length).toBe(1); const traces = await tracesFor(sub); - expect(traces).toHaveLength(1); const [trace] = traces; expect(trace.source).toBe('stream'); expect(trace.request_id).toBeTruthy(); @@ -154,6 +155,7 @@ test.describe('tier-a · execution traces', () => { }); expect(await res.text()).toContain('"type": "end"'); + await expect.poll(async () => (await tracesFor(sub)).length).toBe(1); const [trace] = await tracesFor(sub); const byKind = (kind: string) => trace.spans.filter((s) => s.kind === kind); const retrieval = byKind('retrieval'); diff --git a/tests/tracing/test_entry_points.py b/tests/tracing/test_entry_points.py index 95ba3809..5a613892 100644 --- a/tests/tracing/test_entry_points.py +++ b/tests/tracing/test_entry_points.py @@ -22,7 +22,7 @@ def flushed(): """Capture flushed traces instead of writing them.""" captured = [] - def _fake_flush(trace, status=None): + def _fake_flush(trace, status=None, **_kwargs): if trace is None or trace.flushed: return trace.flushed = True diff --git a/tests/tracing/test_sink.py b/tests/tracing/test_sink.py index 7fc7b188..4581434e 100644 --- a/tests/tracing/test_sink.py +++ b/tests/tracing/test_sink.py @@ -98,3 +98,31 @@ class TestDiscard: tracing.discard(trace) tracing.flush(trace) repo.insert.assert_not_called() + + +class TestBackgroundFlush: + def test_background_flush_writes_off_the_calling_thread(self): + import threading + + trace = _trace_with_span() + seen = {} + + def _insert(record): + seen["thread"] = threading.current_thread().name + return True + + with _patched_store() as repo: + repo.insert.side_effect = _insert + future = tracing.flush(trace, background=True) + assert trace.finished # frozen at once, written later + future.result(timeout=5) + assert seen["thread"].startswith("trace-writer") + assert repo.insert.call_count == 1 + + def test_background_flush_is_still_once(self): + trace = _trace_with_span() + with _patched_store() as repo: + future = tracing.flush(trace, background=True) + assert tracing.flush(trace, background=True) is None + future.result(timeout=5) + assert repo.insert.call_count == 1 diff --git a/tests/worker/test_extract_graph.py b/tests/worker/test_extract_graph.py index d6dd5701..10283ad5 100644 --- a/tests/worker/test_extract_graph.py +++ b/tests/worker/test_extract_graph.py @@ -188,7 +188,7 @@ class TestExtractGraphTrace: monkeypatch.setattr("docsgpt.graphrag.extraction.extract_graph_for_source", extract) flushed = [] - def _flush(trace, status=None): + def _flush(trace, status=None, **_kwargs): trace.flushed = True trace.finish(status) flushed.append(trace)