fix(insights): keep parallel insight series in one trace

.semgrep/rules/devex/query-runners-propagate-thread-context.yaml

@@ -0,0 +1,52 @@
+rules:
+    - id: query-runners-propagate-thread-context
+      message: |
+          Bare `threading.Thread` in a query runner. A new thread starts with empty ContextVars, so
+          the worker loses everything the query path reads from context: the OpenTelemetry span
+          context, the ClickHouse query tags, and the warehouse-warning accumulator. Losing the span
+          context means each worker opens its spans with no parent, so one parallel insight reports
+          a separate trace per series instead of a single request trace, and the request trace shows
+          an unexplained gap where the query work should be. Losing the query tags means the
+          worker's ClickHouse queries land in `system.query_log` with no product or feature
+          attribution. A raw thread per work item is also unbounded, so one request can take a web
+          worker's threads and a large share of the shared Postgres connection pool.
+
+          Use `run_in_parallel_threads` from `posthog.hogql_queries.utils.parallel`, which runs the
+          work in a bounded pool and copies the caller's context into every worker. Pass thunks
+          (`functools.partial`) instead of `target=`/`args=`:
+
+              run_in_parallel_threads([partial(run, index, query) for index, query in enumerate(queries)])
+
+          Use `map_in_caller_context` from `posthog.hogql_queries.utils.caller_context` instead when
+          the work returns values rather than writing into a shared result list. It runs serially
+          under `TEST`, so a test that must reach the pool cannot use it.
+
+          A hand-rolled `ThreadPoolExecutor` loses the context the same way, and this rule does not
+          detect that form. Submit a fresh `contextvars.copy_context().run` per task, as
+          `posthog/hogql_queries/utils/timestamp_utils.py` does. One `Context` cannot be entered
+          twice, so hoisting a single `copy_context()` above the loop raises `RuntimeError` on the
+          second task.
+
+          For a thread that genuinely must not inherit the caller's context, add
+          `# nosemgrep: query-runners-propagate-thread-context` with a one-line why.
+
+          Backlog: `products/experiments/backend/hogql_queries/experiment_trends_query_runner.py`
+          still spawns bare threads. Converting it changes which ClickHouse cluster and user its
+          subqueries use, and starts metering their scanned bytes, because that runner propagates no
+          tags today. That needs its own change, so the rule stays at WARNING until then. The
+          new-warnings CI pass still blocks any bare thread a PR adds.
+      severity: WARNING
+      languages: [python]
+      # Anchored at the repo root, which is where CI scans from. The leading slash also opts out of
+      # semgrep's Semgrepignore v2 deprecation warning for unanchored patterns.
+      paths:
+          include:
+              - '/posthog/hogql_queries/**'
+              - '/products/*/backend/hogql_queries/**'
+      pattern-either:
+          - pattern: threading.Thread(...)
+          - pattern: Thread(...)
+      metadata:
+          category: best-practice
+          subcategory:
+              - audit

posthog/hogql_queries/utils/parallel.py

@@ -0,0 +1,80 @@
+"""Run query-runner work concurrently without dropping the caller's ambient context.
+
+A bare `threading.Thread` starts with empty ContextVars, so everything the query path reads from
+context is absent in the worker: the OpenTelemetry span context, the ClickHouse query tags, and the
+warehouse-warning accumulator. Without the span context a worker opens its spans with no parent, so
+one parallel insight reports several disconnected traces instead of a single request trace.
+
+`ThreadPoolExecutor` drops the context the same way, so submitting a bare callable to a pool has
+the same defect as the bare thread it replaces.
+"""
+
+import contextvars
+from collections.abc import Callable, Sequence
+from concurrent.futures import ThreadPoolExecutor
+
+import structlog
+
+from posthog.hogql.constants import INSIGHT_QUERY_FANOUT_CONCURRENCY
+
+from posthog.clickhouse.query_tagging import get_query_tags, query_tags
+
+logger = structlog.get_logger(__name__)
+
+
+def run_in_parallel_threads(
+    work: Sequence[Callable[[], None]],
+    *,
+    max_workers: int = INSIGHT_QUERY_FANOUT_CONCURRENCY,
+    thread_name_prefix: str = "parallel_work",
+) -> None:
+    """Run the callables in a bounded thread pool, then wait for all of them.
+
+    The pool caps how many run at once. One insight request expands into one query per series, and
+    a client controls that expansion, so an unbounded fan-out lets one request take a web worker's
+    threads and a large share of the shared Postgres connection pool.
+
+    Every worker gets its own copy of the context mapping, so rebinding a ContextVar in one worker
+    cannot reach a sibling or the caller. The copy is shallow, so a value reached through the
+    context is still shared, and mutating one in place does reach the caller.
+
+    Warehouse warnings and the query scan accumulator depend on that shallow sharing, because a
+    worker contributes to the caller's object. Query tags do not: this helper gives each worker a
+    private tags copy up front, so even a worker that mutates its tags in place (skipping
+    `tag_queries`, which always replaces the snapshot instead) only ever touches its own copy,
+    never a sibling's or the caller's.
+
+    An exception raised by a work item does not escape to the pool's future and vanish unread. It
+    is collected and re-raised in the caller once every item has finished, so a raising item cannot
+    leave the caller reading a partial result with no signal anything went wrong. If more than one
+    item raises, the caller sees whichever exception is appended first, meaning the first item to
+    finish raising rather than the first item in `work`, and every other collected exception is
+    logged rather than dropped.
+    """
+    if not work:
+        return
+
+    errors: list[Exception] = []
+
+    def run_item(item: Callable[[], None]) -> None:
+        try:
+            # Rebind before running the item, not after: a worker's very first line of code could
+            # mutate the shared QueryTags object in place, so the copy has to exist before that.
+            query_tags.set(get_query_tags().model_copy())
+            item()
+        except Exception as e:
+            errors.append(e)
+
+    with ThreadPoolExecutor(max_workers=min(max_workers, len(work)), thread_name_prefix=thread_name_prefix) as executor:
+        for item in work:
+            # A fresh Context per item, because one Context cannot be entered twice and a pool
+            # thread runs several items in turn.
+            executor.submit(contextvars.copy_context().run, run_item, item)
+
+    if errors:
+        # Surface every secondary error with its traceback before re-raising the first, so a
+        # concurrent failure doesn't disappear into the void. Matches funnels_query_runner's
+        # `_run_in_parallel`.
+        for dropped in errors[1:]:
+            logger.exception("parallel_work_secondary_error", exc_info=dropped)
+        raise errors[0]

posthog/hogql_queries/utils/test/test_parallel.py

@@ -0,0 +1,107 @@
+import threading
+import contextvars
+from functools import partial
+
+import pytest
+
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
+from posthog.clickhouse.query_tagging import get_query_tags, reset_query_tags, tag_queries
+from posthog.hogql_queries.utils.parallel import run_in_parallel_threads
+
+_probe: contextvars.ContextVar[str] = contextvars.ContextVar("parallel_test_probe", default="unset")
+
+
+def test_worker_spans_stay_inside_the_callers_trace() -> None:
+    exporter = InMemorySpanExporter()
+    provider = TracerProvider()
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    tracer = provider.get_tracer("test")
+
+    def work(index: int) -> None:
+        with tracer.start_as_current_span(f"series_{index}"):
+            pass
+
+    with tracer.start_as_current_span("caller") as caller:
+        caller_context = caller.get_span_context()
+        run_in_parallel_threads([partial(work, index) for index in range(3)])
+
+    series = [span for span in exporter.get_finished_spans() if span.name.startswith("series_")]
+    assert len(series) == 3
+    for span in series:
+        # Checked before the span_id read below, so a worker that lost the context reports the
+        # missing parent rather than an AttributeError on None. This also narrows the type, which
+        # is what the union-attr check needs.
+        assert span.parent is not None, f"{span.name} is its own trace root, so the context was lost"
+        assert span.parent.span_id == caller_context.span_id
+        assert span.context.trace_id == caller_context.trace_id
+
+
+def test_workers_read_the_callers_context_and_cannot_rebind_it() -> None:
+    token = _probe.set("from_caller")
+    seen: list[str] = []
+
+    def work(index: int) -> None:
+        seen.append(_probe.get())
+        _probe.set(f"from_worker_{index}")
+
+    try:
+        run_in_parallel_threads([partial(work, index) for index in range(2)])
+
+        assert seen == ["from_caller", "from_caller"]
+        assert _probe.get() == "from_caller"
+    finally:
+        _probe.reset(token)
+
+
+def test_exception_from_a_work_item_is_reraised_to_the_caller() -> None:
+    def boom() -> None:
+        raise ValueError("boom")
+
+    def noop() -> None:
+        pass
+
+    with pytest.raises(ValueError, match="boom"):
+        run_in_parallel_threads([noop, boom])
+
+
+@pytest.mark.parametrize("item_count,max_workers", [(0, 10), (5, 1), (12, 3)])
+def test_every_item_runs_on_no_more_threads_than_the_limit(item_count: int, max_workers: int) -> None:
+    lock = threading.Lock()
+    thread_ids: list[int] = []
+
+    def work() -> None:
+        with lock:
+            thread_ids.append(threading.get_ident())
+
+    run_in_parallel_threads([work] * item_count, max_workers=max_workers)
+
+    assert len(thread_ids) == item_count
+    assert len(set(thread_ids)) <= max_workers
+
+
+def test_worker_mutating_query_tags_in_place_cannot_reach_a_sibling_or_the_caller() -> None:
+    reset_query_tags()
+    tag_queries(team_id=1)
+    # Forces the read to happen strictly after the mutation, so the assertion is deterministic
+    # regardless of thread scheduling: without the per-worker rebind, the mutation below would
+    # reach every reader because all workers would still share the caller's QueryTags object.
+    mutated = threading.Event()
+    seen: list[int | None] = []
+
+    def mutate_without_tag_queries() -> None:
+        # Mutates the object in place instead of going through `tag_queries`, which is exactly
+        # the pattern the per-worker rebind in run_in_parallel_threads has to guard against.
+        get_query_tags().team_id = 999
+        mutated.set()
+
+    def read_team_id_after_mutation() -> None:
+        mutated.wait(timeout=5)
+        seen.append(get_query_tags().team_id)
+
+    run_in_parallel_threads([mutate_without_tag_queries, read_team_id_after_mutation])
+
+    assert seen == [1]
+    assert get_query_tags().team_id == 1

products/product_analytics/backend/hogql_queries/funnels/funnels_query_runner.py

@@ -1,7 +1,7 @@
 import json
-import threading
 from collections.abc import Callable, Sequence
 from datetime import datetime, timedelta
+from functools import partial
 from math import ceil
 from typing import Any, Optional, TypeVar, cast
 
@@ -25,15 +25,12 @@
     ResolvedDateRangeResponse,
 )
 
-from posthog.hogql import ast, query_stats
+from posthog.hogql import ast
 from posthog.hogql.constants import MAX_BYTES_BEFORE_EXTERNAL_GROUP_BY, HogQLGlobalSettings, LimitContext
 from posthog.hogql.query import execute_hogql_query
-from posthog.hogql.query_stats import QueryStats
 from posthog.hogql.timings import HogQLTimings
 
 from posthog.caching.insights_api import BASE_MINIMUM_INSIGHT_REFRESH_INTERVAL, REDUCED_MINIMUM_INSIGHT_REFRESH_INTERVAL
-from posthog.clickhouse import query_tagging
-from posthog.clickhouse.query_tagging import QueryTags
 from posthog.hogql_queries.query_runner import AnalyticsQueryRunner, ExecutionMode
 from posthog.hogql_queries.utils.breakdowns import (
     ALL_USERS_COHORT_ID,
@@ -43,6 +40,7 @@
     has_breakdown_filter,
     humanize_breakdown_label,
 )
+from posthog.hogql_queries.utils.parallel import run_in_parallel_threads
 from posthog.hogql_queries.utils.query_compare_to_date_range import QueryCompareToDateRange
 from posthog.hogql_queries.utils.query_date_range import QueryDateRange
 from posthog.hogql_queries.utils.query_previous_period_date_range import QueryPreviousPeriodDateRange
@@ -331,20 +329,9 @@ def _run_in_parallel(self, tasks: list[Callable[[HogQLTimings], "_T"]], index_of
         results: list[Optional[_T]] = [None] * len(tasks)
         errors: list[Exception] = []
 
-        def run(
-            index: int,
-            timings: HogQLTimings,
-            query_tags: Optional[QueryTags] = None,
-            stats: Optional[QueryStats] = None,
-        ) -> None:
+        def run(index: int, timings: HogQLTimings) -> None:
             try:
-                # Worker threads start with an empty QueryTags ContextVar — restore the parent's
-                # snapshot so execute_hogql_query has the required feature/product tags.
-                if query_tags is not None:
-                    query_tagging.update_tags(query_tags)
-                # The query scan accumulator is another ContextVar the thread does not inherit.
-                with query_stats.use(stats):
-                    results[index] = tasks[index](timings)
+                results[index] = tasks[index](timings)
             except Exception as exc:  # noqa: BLE001
                 errors.append(exc)
             finally:
@@ -358,19 +345,12 @@ def run(
             for index in range(len(tasks)):
                 run(index, self.timings.clone_for_subquery(index_offset + index))
         else:
-            parent_tags = query_tagging.get_query_tags().model_copy(deep=True)
-            parent_stats = query_stats.get_active()
-            jobs = [
-                threading.Thread(
-                    target=run,
-                    args=(index, self.timings.clone_for_subquery(index_offset + index), parent_tags, parent_stats),
-                )
-                for index in range(len(tasks))
-            ]
-            for job in jobs:
-                job.start()
-            for job in jobs:
-                job.join()
+            run_in_parallel_threads(
+                [
+                    partial(run, index, self.timings.clone_for_subquery(index_offset + index))
+                    for index in range(len(tasks))
+                ]
+            )
 
         if errors:
             # Surface every secondary error with its traceback before re-raising the first,