ImageApi/tests/otel_propagation.rs

//! Regression test for the tracing migration (#134 / map #135).
//!
//! Origin: spike #133 (go/no-go). It pins the one non-obvious property of
//! `tracing` + `tracing-opentelemetry 0.32` that the whole migration rests on:
//! a span made inside a `tokio::spawn`-ed task nests under the pre-spawn
//! parent (same trace_id, parent_span_id link) **only** when the parent span
//! is explicitly carried across the spawn. Bare `#[instrument]` on its own
//! does NOT do that on a multi-thread runtime — the child orphans onto a fresh
//! trace. That is why every `tokio::spawn` in the codebase must carry the
//! parent span (`.instrument()` or an `enter()` guard); this test fails if
//! that guarantee regresses.
//!
//! Does `tracing` + `tracing-opentelemetry 0.32` propagate OTel context across
//! `.await` AND into a `tokio::spawn`-ed task via `#[instrument]`, so a span
//! made inside the spawned task nests under the pre-spawn parent?
//!
//! Key subtlety (verified by `cargo expand` on tracing-attributes): for an
//! `async fn`, `#[instrument]` generates the span-creation INSIDE the future
//! body — the span is created at FIRST POLL, not at call time. For a
//! `tokio::spawn`ed task the first poll always happens on a runtime worker
//! thread, so the parent is looked up on a thread where it was never
//! entered. This holds for BOTH spawn shapes:
//!
//!   - `tokio::spawn(child())` and
//!   - `tokio::spawn(async move { child().await })`  (the production shape,
//!     see `insight_chat::chat_turn_stream_inner`)
//!
//! Each scenario is driven with `rt.block_on` so the test body always runs on
//! the calling thread while spawned tasks run on (multi-thread) workers,
//! making cross-thread behavior deterministic.
//!
//! VERDICT (Spike #133): ADJUST. Plain `#[instrument]` does NOT propagate
//! across `tokio::spawn` on a multi-thread runtime (S2/S4 export orphaned
//! spans: fresh trace_id, parent_span_id = 0000000000000000) because the
//! span of an `#[instrument]`ed `async fn` is created at FIRST POLL — always
//! on a worker thread for a spawned task — and tracing span context is
//! thread-local. The fix is explicit context carrying into the task:
//! `tokio::spawn(block.instrument(parent_span))` or a `span.enter()` guard at
//! the top of the block (S6/S7 both propagate).
//!
//! GOTCHA discovered while building this: `SubscriberInitExt::set_default` is
//! THREAD-LOCAL — spans from any other thread (plain std or tokio workers)
//! are silently dropped. The subscriber must be installed process-globally
//! (`.try_init()` / `.init()`, i.e. `dispatcher::set_global_default`) for
//! worker threads to see it at all. (The app's existing `init_tracing` uses
//! `.try_init()`, so this is already satisfied in production.)

use opentelemetry::trace::TracerProvider;
use opentelemetry_sdk::trace::{InMemorySpanExporter, SdkTracerProvider, SpanData};
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::registry::Registry;
use tracing_subscriber::util::SubscriberInitExt;
use tracing::Instrument;

const CHILD: &str = "child_in_spawned_task";

#[tracing::instrument]
async fn child_in_spawned_task() {
    // Yield so the remainder of the body really runs as a spawned task.
    tokio::task::yield_now().await;
}

struct Outcome {
    propagated: bool,
    child_trace: String,
    parent_trace: String,
    child_parent_id: String,
    parent_span_id: String,
}

/// Within a scenario's slice there is exactly one parent and one child;
/// identify the parent as the span that is not the child.
fn collect(spans: &[SpanData]) -> Option<Outcome> {
    let child = spans.iter().find(|s| s.name == CHILD)?;
    let parent = spans.iter().find(|s| s.name != CHILD)?;
    Some(Outcome {
        propagated: child.span_context.trace_id() == parent.span_context.trace_id()
            && child.parent_span_id == parent.span_context.span_id(),
        child_trace: child.span_context.trace_id().to_string(),
        parent_trace: parent.span_context.trace_id().to_string(),
        child_parent_id: child.parent_span_id.to_string(),
        parent_span_id: parent.span_context.span_id().to_string(),
    })
}

fn run_scenario<F>(
    exporter: &InMemorySpanExporter,
    make_rt: fn() -> tokio::runtime::Runtime,
    body: F,
) -> Option<Outcome>
where
    F: FnOnce() -> std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send>>,
{
    let before = exporter.get_finished_spans().map(|v| v.len()).unwrap_or(0);
    let rt = make_rt();
    rt.block_on(async {
        let parent = tracing::info_span!("scenario_parent");
        let _entered = parent.enter();
        body().await;
    });
    let all = exporter.get_finished_spans().expect("in-memory exporter");
    collect(&all[before..])
}

fn current_thread_rt() -> tokio::runtime::Runtime {
    tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .expect("current-thread runtime")
}

fn multi_thread_rt() -> tokio::runtime::Runtime {
    tokio::runtime::Builder::new_multi_thread()
        .worker_threads(2)
        .enable_all()
        .build()
        .expect("multi-thread runtime")
}

#[test]
fn tracing_opentelemetry_spawn_propagation_matrix() {
    // One process-global subscriber for the whole matrix (set_global_default
    // is one-shot per process; the scenarios run sequentially).
    let exporter = InMemorySpanExporter::default();
    let provider = SdkTracerProvider::builder()
        .with_simple_exporter(exporter.clone())
        .build();
    let tracer = provider.tracer("spike-133");
    let subscriber = Registry::default().with(tracing_opentelemetry::OpenTelemetryLayer::new(tracer));
    subscriber
        .try_init()
        .expect("global subscriber not yet set");

    // S1: single-thread runtime, instrumented call inside the spawned block.
    // Control: spawn and parent share one thread, so the thread-local active
    // span survives into the child's first poll.
    let s1 = run_scenario(&exporter, current_thread_rt, || {
        Box::pin(async {
            let h = tokio::spawn(async move {
                child_in_spawned_task().await;
            });
            h.await.unwrap();
        })
    });

    // S2: multi-thread runtime, instrumented call inside the spawned block —
    // THE PRODUCTION SHAPE (`chat_turn_stream_inner`'s
    // `tokio::spawn(async move { svc.run_streaming_turn(...).await })`).
    // The first poll of the instrumented call happens on a worker thread
    // that never had the parent entered.
    let s2 = run_scenario(&exporter, multi_thread_rt, || {
        Box::pin(async {
            let h = tokio::spawn(async move {
                child_in_spawned_task().await;
            });
            h.await.unwrap();
        })
    });

    // S4: multi-thread runtime, instrumented call at the spawn call site —
    // `tokio::spawn(child_in_spawned_task())`. The span is created on the
    // calling thread where the parent is active.
    let s4 = run_scenario(&exporter, multi_thread_rt, || {
        Box::pin(async {
            let h = tokio::spawn(child_in_spawned_task());
            h.await.unwrap();
        })
    });

    // S5: no spawn at all — plain `.await` on the same task. Sanity control
    // that the layer + exporter wiring works.
    let s5 = run_scenario(&exporter, multi_thread_rt, || {
        Box::pin(async {
            child_in_spawned_task().await;
        })
    });

    // S6: THE FIX — capture the active parent span before the spawn and carry
    // it into the task with `.instrument()`. `.instrument` re-enters the span
    // on every poll (including the first poll on the worker thread), so the
    // `#[instrument]` child created there nests under the captured parent.
    let s6 = run_scenario(&exporter, multi_thread_rt, || {
        Box::pin(async {
            let span = tracing::Span::current();
            let h = tokio::spawn(
                async move {
                    child_in_spawned_task().await;
                }
                .instrument(span),
            );
            h.await.unwrap();
        })
    });

    // S7: alternate fix — re-enter the captured span manually at the top of
    // the spawned block (guard lives for the whole block).
    let s7 = run_scenario(&exporter, multi_thread_rt, || {
        Box::pin(async {
            let span = tracing::Span::current();
            let h = tokio::spawn(async move {
                let _guard = span.enter();
                child_in_spawned_task().await;
            });
            h.await.unwrap();
        })
    });

    let outcomes = vec![
        ("S1 current-thread, in-block call", s1),
        ("S2 multi-thread, in-block call (PRODUCTION)", s2),
        ("S4 multi-thread, direct-call spawn", s4),
        ("S5 multi-thread, same-task await (control)", s5),
        ("S6 multi-thread, spawn + .instrument() (FIX)", s6),
        ("S7 multi-thread, spawn + span.enter() (FIX)", s7),
    ];

    let mut report = String::new();
    for (label, outcome) in &outcomes {
        match outcome {
            Some(o) => {
                let line = format!(
                    "{:<42} propagated={:<5} child_trace={} parent_trace={} child.parent_span_id={} parent.span_id={}",
                    label,
                    o.propagated,
                    o.child_trace,
                    o.parent_trace,
                    o.child_parent_id,
                    o.parent_span_id,
                );
                println!("{}", line);
                report.push_str(&line);
                report.push('\n');
            }
            None => {
                let line = format!(
                    "{:<42} NO SPANS EXPORTED\n",
                    label
                );
                println!("{}", line);
                report.push_str(&line);
                report.push('\n');
            }
        }
    }
    eprintln!("SPIKE-133-REPORT:\n{}", report);

    let get = |prefix: &str| -> &Outcome {
        outcomes
            .iter()
            .find(|(l, _)| l.starts_with(prefix))
            .and_then(|(_, o)| o.as_ref())
            .unwrap_or_else(|| panic!("{prefix} produced no spans:\n{}", report))
    };

    // Wiring sanity: same-task await MUST propagate, else the layer/exporter
    // is broken and nothing below means anything.
    assert!(
        get("S5").propagated,
        "control S5 (same-task await) failed — wiring is broken, not propagation:\n{}",
        report
    );

    // THE FINDING (decisive, production shape): a bare `#[instrument]` fn
    // called from inside a `tokio::spawn`ed block on a multi-thread runtime
    // does NOT nest under the pre-spawn parent — the child gets a fresh
    // trace_id and parent_span_id 0000000000000000 (orphan). We assert the
    // orphan explicitly so the regression is pinned, not hidden.
    assert!(
        !get("S2").propagated,
        "S2 unexpectedly propagated — re-check the finding:\n{}",
        report
    );
    assert!(
        !get("S4").propagated,
        "S4 unexpectedly propagated — re-check the finding:\n{}",
        report
    );

    // Same-thread control: with the parent active on the polling thread the
    // child DOES nest (proves the layer wiring is sound when context is
    // present).
    assert!(
        get("S1").propagated,
        "S1 (current-thread) failed to propagate — wiring is broken:\n{}",
        report
    );

    // THE FIX: carrying the parent span across the spawn (`.instrument()` or
    // an explicit `span.enter()` guard) restores the nesting.
    assert!(
        get("S6").propagated,
        "S6 (spawn + .instrument()) did NOT propagate — the fix is broken:\n{}",
        report
    );
    assert!(
        get("S7").propagated,
        "S7 (spawn + span.enter()) did NOT propagate — the fix is broken:\n{}",
        report
    );
}