diff --git a/.changeset/tracing-custom-spans.md b/.changeset/tracing-custom-spans.md new file mode 100644 index 000000000..71c32c805 --- /dev/null +++ b/.changeset/tracing-custom-spans.md @@ -0,0 +1,12 @@ +--- +"workers-rs": minor +--- + +Add `worker::observability` — bindings for Cloudflare Workers [custom spans](https://developers.cloudflare.com/changelog/post/2026-06-16-custom-spans/) (`cloudflare:workers` `enterSpan` / `startActiveSpan`). + +- `enter_span(name, |span| ...)` and `enter_span_async(name, |span| async { ... })` open callback-scoped custom trace spans that nest under the automatic platform spans in the Workers Observability waterfall. +- `start_active_span(name)` opens a span that outlives the callback and is closed by `Span::end` — for streams, and for bridging `tracing`'s separate span create/close. +- `Span::set_attribute` / `Span::is_traced` attach metadata and check sampling. +- `with_active_span` exposes the innermost open span so a `tracing_subscriber::Layer` can forward `tracing` events/fields onto it. + +The new `custom-spans` example ships a `WorkersLayer` that bridges `tracing` span *lifetimes* onto the platform, so `span!` / `#[instrument]` get platform-measured durations with no Workers-specific code at the call site. Addresses #899. diff --git a/Cargo.lock b/Cargo.lock index cb8b17987..4d2f10fdc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -489,6 +489,16 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "console_error_panic_hook" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06aeb73f470f66dcdbf7223caeebb85984942f22f1adb2a088cf9668146bbbc" +dependencies = [ + "cfg-if", + "wasm-bindgen", +] + [[package]] name = "const-oid" version = "0.10.2" @@ -669,6 +679,16 @@ dependencies = [ "worker", ] +[[package]] +name = "custom-spans" +version = "0.1.0" +dependencies = [ + "console_error_panic_hook", + "tracing", + "tracing-subscriber", + "worker", +] + [[package]] name = "data-encoding" version = "2.11.0" diff --git a/examples/custom-spans/Cargo.toml b/examples/custom-spans/Cargo.toml new file mode 100644 index 000000000..e921dd71b --- /dev/null +++ b/examples/custom-spans/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "custom-spans" +version = "0.1.0" +edition = "2021" + +[package.metadata.release] +release = false + +[lib] +crate-type = ["cdylib"] + +[dependencies] +# Binding comes from `worker`; the WorkersLayer (src/layer.rs) needs tracing-subscriber. +worker.workspace = true +tracing = "0.1" +tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] } +console_error_panic_hook = "0.1" diff --git a/examples/custom-spans/README.md b/examples/custom-spans/README.md new file mode 100644 index 000000000..f197e1d8b --- /dev/null +++ b/examples/custom-spans/README.md @@ -0,0 +1,32 @@ +# custom-spans + +Custom trace spans for Workers Observability, in Rust — using +[`worker::observability`](../../worker/src/observability.rs). + +It demonstrates: + +- `enter_span_async("handle_request", |span| async move { … })` — an async + root span around the request handler. +- `enter_span("load_rows", |span| …)` — a nested sync span that auto-parents + under the root via the JS async context. +- `span.set_attribute(...)` and `span.is_traced()`. +- `WorkersLayer` — a `tracing_subscriber::Layer` bridging `tracing` onto the + platform: every `span!` / `#[instrument]` becomes a platform span for its + full lifetime (via `start_active_span` + `Span::end`), and `tracing::info!` + events land as attributes on the span they were emitted in. `summarize()` in + `src/lib.rs` is instrumented that way, with no Workers-specific code. + +Nesting note: a bridged span parents under the nearest enclosing `enter_span`, +not under its `tracing` parent, because the platform derives hierarchy from the +JS async context. Wrap a subtree in `enter_span` where the shape matters. + +Custom spans are recorded only when tracing is enabled in your Worker's +observability config — see `wrangler.toml` (`[observability.traces]`). + +```sh +npx wrangler deploy +``` + +Then open the Worker's **Observability → Traces** view and trigger a request; +`handle_request`, its nested `load_rows` span, and the `tracing`-instrumented +`summarize` span appear in the waterfall next to the automatic `fetch` span. diff --git a/examples/custom-spans/package.json b/examples/custom-spans/package.json new file mode 100644 index 000000000..dfae78872 --- /dev/null +++ b/examples/custom-spans/package.json @@ -0,0 +1,12 @@ +{ + "name": "custom-spans", + "version": "0.0.0", + "private": true, + "scripts": { + "deploy": "cargo install worker-build ; wrangler deploy", + "dev": "cargo install worker-build ; wrangler dev --local" + }, + "devDependencies": { + "wrangler": "^4" + } +} diff --git a/examples/custom-spans/src/layer.rs b/examples/custom-spans/src/layer.rs new file mode 100644 index 000000000..00ec32aac --- /dev/null +++ b/examples/custom-spans/src/layer.rs @@ -0,0 +1,160 @@ +//! `WorkersLayer` — a `tracing_subscriber::Layer` that turns every +//! `tracing::span!` into a real Workers platform span and forwards `tracing` +//! events onto it. +//! +//! ## Why it's shaped this way +//! +//! `tracing` models a span lifetime as two separate operations: the span is +//! created (`on_new_span`) and closed (`on_close`) at unrelated points in time. +//! The original custom-span API was callback-scoped only (`enterSpan(name, +//! cb)`), which a `Layer` can't drive: at `on_new_span` it would have to call +//! `enterSpan` and not return from its callback until the later `on_close`, +//! which a single-threaded Worker can't suspend and resume. +//! +//! [`startActiveSpan()` + `span.end()`][changelog] (2026-07-28) are exactly the +//! imperative pair that lifetime needs, so this layer now bridges span +//! *lifetimes*, not just events: `#[tracing::instrument]` and `span!` produce +//! spans in the trace waterfall with platform-measured durations, with no +//! Workers-specific code at the call site. +//! +//! ## The one limitation +//! +//! Parent-child nesting follows the platform's async context, which is only +//! entered for the instant `startActiveSpan` runs its callback. A bridged span +//! therefore parents under the nearest enclosing span opened by +//! [`worker::observability::enter_span`] / `enter_span_async`, and two nested +//! `tracing` spans come out as siblings under that same parent rather than one +//! inside the other. Wrap a subtree in `enter_span` where the shape matters; +//! closing the gap entirely needs a runtime primitive for attaching to an open +//! span's context. +//! +//! This lives in the example rather than the `worker` crate so `worker` stays +//! free of a `tracing-subscriber` dependency. Copy it into your project, or +//! lift it into `worker` behind a feature if your project wants it there. +//! +//! [changelog]: https://developers.cloudflare.com/changelog/post/2026-07-28-start-active-span/ + +use std::cell::RefCell; +use std::collections::HashMap; + +use tracing::field::{Field, Visit}; +use tracing::span::{Attributes, Id, Record}; +use tracing::{Event, Subscriber}; +use tracing_subscriber::layer::Context; +use tracing_subscriber::registry::LookupSpan; +use tracing_subscriber::Layer; +use worker::observability::{start_active_span, with_active_span, Span}; + +thread_local! { + /// Platform span per live `tracing` span id. A thread-local map rather than + /// the registry's own span extensions because a JS [`Span`] is `!Send` and + /// extensions must be `Send + Sync`; a Worker isolate is single-threaded, + /// so a thread-local is equivalent here. Entries are removed in `on_close`, + /// which `tracing` calls exactly once per span. + static OPEN: RefCell> = RefCell::new(HashMap::new()); +} + +/// Bridges `tracing` onto Workers Observability: each `tracing` span becomes a +/// platform span for its full lifetime, and events land as attributes on the +/// span they were emitted in. Install it on a `tracing_subscriber` registry. +#[derive(Debug, Default, Clone, Copy)] +pub struct WorkersLayer; + +impl Layer for WorkersLayer +where + S: Subscriber + for<'a> LookupSpan<'a>, +{ + fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, _ctx: Context<'_, S>) { + let span = start_active_span(attrs.metadata().name()); + attrs.record(&mut AttrVisitor { + span: &span, + prefix: None, + }); + OPEN.with_borrow_mut(|open| open.insert(id.into_u64(), span)); + } + + fn on_record(&self, id: &Id, values: &Record<'_>, _ctx: Context<'_, S>) { + with_span(id, |span| { + values.record(&mut AttrVisitor { span, prefix: None }) + }); + } + + fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) { + // Events belong to the span they were emitted in; fall back to the + // innermost `enter_span` when there is no enclosing `tracing` span. + let level = event.metadata().level().as_str(); + let recorded = ctx.event_span(event).is_some_and(|s| { + with_span(&s.id(), |span| { + event.record(&mut AttrVisitor { + span, + prefix: Some(level), + }) + }) + .is_some() + }); + + if !recorded { + with_active_span(|span| { + event.record(&mut AttrVisitor { + span, + prefix: Some(level), + }) + }); + } + } + + fn on_close(&self, id: Id, _ctx: Context<'_, S>) { + if let Some(span) = OPEN.with_borrow_mut(|open| open.remove(&id.into_u64())) { + span.end(); + } + } +} + +/// Run `f` against the platform span backing `id`, if it is still open. +fn with_span(id: &Id, f: impl FnOnce(&Span) -> R) -> Option { + OPEN.with_borrow(|open| open.get(&id.into_u64()).map(f)) +} + +/// Writes each visited `tracing` field as a typed `setAttribute` on the +/// platform span. Span fields keep their own names; event fields are prefixed +/// with the level (`"INFO.message"`) so they don't collide with them. +struct AttrVisitor<'a> { + span: &'a Span, + prefix: Option<&'a str>, +} + +impl AttrVisitor<'_> { + fn key(&self, field: &Field) -> String { + match self.prefix { + Some(prefix) => format!("{}.{}", prefix, field.name()), + None => field.name().to_owned(), + } + } +} + +impl Visit for AttrVisitor<'_> { + fn record_bool(&mut self, field: &Field, value: bool) { + self.span.set_attribute(&self.key(field), value); + } + + fn record_i64(&mut self, field: &Field, value: i64) { + self.span.set_attribute(&self.key(field), value); + } + + fn record_u64(&mut self, field: &Field, value: u64) { + self.span.set_attribute(&self.key(field), value); + } + + fn record_f64(&mut self, field: &Field, value: f64) { + self.span.set_attribute(&self.key(field), value); + } + + fn record_str(&mut self, field: &Field, value: &str) { + self.span.set_attribute(&self.key(field), value); + } + + fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) { + self.span + .set_attribute(&self.key(field), format!("{value:?}").as_str()); + } +} diff --git a/examples/custom-spans/src/lib.rs b/examples/custom-spans/src/lib.rs new file mode 100644 index 000000000..96039b8fc --- /dev/null +++ b/examples/custom-spans/src/lib.rs @@ -0,0 +1,65 @@ +//! Custom trace spans for Workers Observability, in Rust. +//! +//! Wraps the request in an async platform span, nests a sync span under it, and +//! lets the `WorkersLayer` turn ordinary `tracing` spans and events into +//! platform spans and attributes. Deploy with `observability.traces` enabled +//! (see `wrangler.toml`) and they appear in the trace waterfall alongside the +//! automatic `fetch` span. + +mod layer; + +use layer::WorkersLayer; +use tracing::{info, instrument}; +use tracing_subscriber::prelude::*; +use worker::observability::{enter_span, enter_span_async}; +use worker::{event, Context, Env, Request, Response, Result}; + +#[event(start)] +fn start() { + console_error_panic_hook::set_once(); + // `try_init` so a hot-reloaded isolate doesn't panic on a second install. + let _ = tracing_subscriber::registry().with(WorkersLayer).try_init(); +} + +#[event(fetch)] +async fn fetch(req: Request, _env: Env, _ctx: Context) -> Result { + let path = req.path(); + + enter_span_async("handle_request", move |span| async move { + span.set_attribute("http.path", path.as_str()); + span.set_attribute("sampled", span.is_traced()); + + // A plain `tracing` event — the layer forwards `user_id` as an + // attribute on `handle_request`. No platform-specific code here. + info!(user_id = 42, "request received"); + + // A nested sync span; auto-parents under `handle_request`. + let rows = enter_span("load_rows", |child| { + let rows = expensive_query(); + child.set_attribute("db.rows", rows); + rows + }); + + // ...and the same work instrumented the ordinary `tracing` way. The + // layer opens a platform span on entry and ends it on close, so this + // gets a platform-measured duration with no Workers-specific code. + let total = summarize(rows); + + info!(rows, "query complete"); + Response::ok(format!("loaded {rows} rows for {path}, total {total}")) + }) + .await +} + +/// Stand-in for real work — a Worker would hit D1 / KV / a binding here. +fn expensive_query() -> u32 { + 1234 +} + +/// Instrumented with plain `tracing`: `WorkersLayer` bridges the span lifetime +/// onto `startActiveSpan` / `span.end()`, and `rows` becomes an attribute. +#[instrument] +fn summarize(rows: u32) -> u32 { + info!("summarizing"); + rows * 2 +} diff --git a/examples/custom-spans/wrangler.toml b/examples/custom-spans/wrangler.toml new file mode 100644 index 000000000..78fcd6ded --- /dev/null +++ b/examples/custom-spans/wrangler.toml @@ -0,0 +1,14 @@ +name = "custom-spans" +main = "build/worker/shim.mjs" +compatibility_date = "2026-06-16" + +[build] +command = "cargo install \"worker-build@^0.8\" && worker-build --release" + +# Custom spans are only recorded when tracing is enabled in observability. +[observability] +enabled = true + +[observability.traces] +enabled = true +head_sampling_rate = 1 diff --git a/worker/src/bindings/mod.rs b/worker/src/bindings/mod.rs index aa5f45d47..9411253d2 100644 --- a/worker/src/bindings/mod.rs +++ b/worker/src/bindings/mod.rs @@ -1 +1,2 @@ pub mod email; +pub mod tracing; diff --git a/worker/src/bindings/tracing.rs b/worker/src/bindings/tracing.rs new file mode 100644 index 000000000..e168bb1f1 --- /dev/null +++ b/worker/src/bindings/tracing.rs @@ -0,0 +1,93 @@ +//! Raw `wasm-bindgen` import of the `cloudflare:workers` `tracing` API. +//! +//! Hand-written (not `ts-gen`-generated) because the safe wrapper in +//! [`crate::observability`] needs a couple of shapes `ts-gen` doesn't express +//! well: the two `enterSpan` callback forms (sync vs. promise-returning) and +//! the `setAttribute` value overloads. +//! +//! Platform surface (`@cloudflare/workers-types`): +//! +//! ```ts +//! interface Tracing { +//! enterSpan( +//! name: string, +//! callback: (span: Span, ...args: A) => T, +//! ...args: A +//! ): T; +//! startActiveSpan( +//! name: string, +//! callback: (span: Span, ...args: A) => T, +//! ...args: A +//! ): T; +//! } +//! declare abstract class Span { +//! get isTraced(): boolean; +//! setAttribute(key: string, value?: boolean | number | string): void; +//! end(): void; +//! } +//! ``` + +use js_sys::Promise; +use wasm_bindgen::closure::ScopedClosure; +use wasm_bindgen::prelude::*; + +#[wasm_bindgen(module = "cloudflare:workers")] +extern "C" { + /// The `tracing` namespace object exported by `cloudflare:workers`. + /// Thread-local because a Worker isolate is single-threaded and the + /// binding is not `Sync`. + #[wasm_bindgen(thread_local_v2, js_name = tracing)] + pub(crate) static TRACING: Tracing; + + #[wasm_bindgen(js_name = Tracing)] + pub(crate) type Tracing; + + /// `enterSpan` with a synchronous callback: the span closes when the + /// callback returns. + #[wasm_bindgen(method, js_name = enterSpan)] + pub(crate) fn enter_span_sync(this: &Tracing, name: &str, cb: &ScopedClosure); + + /// `enterSpan` with an async callback: the callback returns a `Promise` + /// and workerd keeps the span open until it settles. `enterSpan` returns + /// that same promise. + #[wasm_bindgen(method, js_name = enterSpan)] + pub(crate) fn enter_span_async( + this: &Tracing, + name: &str, + cb: &Closure Promise>, + ) -> Promise; + + /// `startActiveSpan`: like `enterSpan`, but the span stays open after the + /// callback returns and is closed by an explicit `end()`. The callback is + /// still invoked once, synchronously, and is the scope in which the new + /// span is the *active* one for auto-nesting. + #[wasm_bindgen(method, js_name = startActiveSpan)] + pub(crate) fn start_active_span( + this: &Tracing, + name: &str, + cb: &ScopedClosure, + ); + + /// A live span handle. Refcounted JS object — cloning is cheap and a clone + /// stays valid while the span is open. + #[wasm_bindgen(js_name = Span)] + #[derive(Debug, Clone)] + pub(crate) type Span; + + // `setAttribute(key, value)` is `boolean | number | string` in JS; bind one + // overload per kind so the safe wrapper stays typed. + #[wasm_bindgen(method, js_name = setAttribute)] + pub(crate) fn set_attribute_bool(this: &Span, key: &str, value: bool); + #[wasm_bindgen(method, js_name = setAttribute)] + pub(crate) fn set_attribute_num(this: &Span, key: &str, value: f64); + #[wasm_bindgen(method, js_name = setAttribute)] + pub(crate) fn set_attribute_str(this: &Span, key: &str, value: &str); + + #[wasm_bindgen(method, getter, js_name = isTraced)] + pub(crate) fn is_traced(this: &Span) -> bool; + + /// Close a span opened with `startActiveSpan`. Idempotent: calls after the + /// first have no effect. Spans opened with `enterSpan` close themselves. + #[wasm_bindgen(method, js_name = end)] + pub(crate) fn end(this: &Span); +} diff --git a/worker/src/lib.rs b/worker/src/lib.rs index 604cf0c09..f48323b44 100644 --- a/worker/src/lib.rs +++ b/worker/src/lib.rs @@ -241,6 +241,7 @@ mod headers; mod http; mod hyperdrive; pub mod kv; +pub mod observability; #[cfg(feature = "queue")] mod queue; mod r2; diff --git a/worker/src/observability.rs b/worker/src/observability.rs new file mode 100644 index 000000000..fff002f31 --- /dev/null +++ b/worker/src/observability.rs @@ -0,0 +1,254 @@ +//! Custom trace spans for Workers Observability. +//! +//! Wraps the `cloudflare:workers` [custom-span API][api]: [`enter_span`] (sync) +//! and [`enter_span_async`] (for `async` handlers) open a span that appears in +//! the trace waterfall alongside the automatic `fetch` / KV / D1 spans, with +//! correct parent-child nesting. Attach metadata with [`Span::set_attribute`]. +//! +//! ```ignore +//! use worker::observability::{enter_span, enter_span_async}; +//! +//! # async fn handler() { +//! enter_span_async("handle_request", |span| async move { +//! span.set_attribute("http.path", "/items"); +//! enter_span("load_rows", |child| { +//! child.set_attribute("db.rows", 42); +//! }); +//! }) +//! .await; +//! # } +//! ``` +//! +//! [`enter_span`] / [`enter_span_async`] are **callback-scoped**: the span opens +//! on the call and closes when the callback returns (sync) or its future +//! resolves (async). For work that outlives a callback — a stream, or a span +//! whose start and end are driven by separate events — [`start_active_span`] +//! opens a span you close yourself with [`Span::end`]. +//! +//! Either way the platform measures the duration, which is why durations are +//! accurate even though guest timer resolution is clamped. +//! +//! ## Bridging the `tracing` crate +//! +//! [`start_active_span`] + [`Span::end`] map onto `tracing`'s two-phase span +//! lifetime (`on_new_span` / `on_close`), so a `tracing_subscriber::Layer` can +//! turn every `tracing::span!` into a real platform span with a real duration; +//! [`with_active_span`] forwards events onto the innermost span opened by +//! `enter_span`. A ready-made layer isn't included here to keep this crate free +//! of a `tracing-subscriber` dependency — see the `custom-spans` example for a +//! `WorkersLayer` you can copy. +//! +//! [api]: https://developers.cloudflare.com/changelog/post/2026-06-16-custom-spans/ + +use std::cell::RefCell; +use std::future::Future; + +use wasm_bindgen::closure::ScopedClosure; +use wasm_bindgen::prelude::*; +use wasm_bindgen_futures::{future_to_promise, JsFuture}; + +use crate::bindings::tracing as raw; + +/// A handle to an open span. Cheap to clone (a refcounted JS object). Valid +/// while the span that produced it is open: the body of the [`enter_span`] / +/// [`enter_span_async`] call that yielded it, or, for [`start_active_span`], +/// until [`Span::end`] is called. +#[derive(Debug, Clone)] +pub struct Span(raw::Span); + +impl Span { + /// Attach an attribute. Accepts any [`AttributeValue`] (`bool`, integer, + /// float, or string). + pub fn set_attribute(&self, key: &str, value: impl AttributeValue) { + value.set_on(self, key); + } + + /// Whether this request is being sampled. Use it to skip building + /// expensive attributes when the span won't be recorded. + pub fn is_traced(&self) -> bool { + self.0.is_traced() + } + + /// Close a span opened by [`start_active_span`]. Idempotent — calls after + /// the first do nothing — so a span whose end is driven by racing events + /// (a stream that either completes or is cancelled) can end from both. + /// + /// A no-op on a span from [`enter_span`] / [`enter_span_async`], which the + /// platform closes when the callback returns. + pub fn end(&self) { + self.0.end(); + } +} + +mod sealed { + pub trait Sealed {} +} + +/// A value that can be attached to a [`Span`] — the JS `boolean | number | +/// string` union. Integers are widened to `f64` (JS `number`); values past +/// 2^53 lose precision. Sealed: implemented for the primitive types only. +pub trait AttributeValue: sealed::Sealed { + #[doc(hidden)] + fn set_on(self, span: &Span, key: &str); +} + +impl sealed::Sealed for bool {} +impl AttributeValue for bool { + fn set_on(self, span: &Span, key: &str) { + span.0.set_attribute_bool(key, self); + } +} + +impl sealed::Sealed for &str {} +impl AttributeValue for &str { + fn set_on(self, span: &Span, key: &str) { + span.0.set_attribute_str(key, self); + } +} + +macro_rules! attribute_value_num { + ($($t:ty),*) => {$( + impl sealed::Sealed for $t {} + impl AttributeValue for $t { + fn set_on(self, span: &Span, key: &str) { + span.0.set_attribute_num(key, self as f64); + } + } + )*}; +} +attribute_value_num!(i8, i16, i32, i64, isize, u8, u16, u32, u64, usize, f32, f64); + +thread_local! { + /// Stack of currently-open platform spans, innermost last. Pushed by + /// `enter_span` / `enter_span_async` for the duration of their body so a + /// `tracing` layer (see [`with_active_span`]) can target the active span. A + /// thread-local rather than the `tracing` registry because a JS [`Span`] is + /// `!Send` and registry extensions must be `Send + Sync`. + static ACTIVE: RefCell> = const { RefCell::new(Vec::new()) }; +} + +fn push_active(span: &Span) { + ACTIVE.with_borrow_mut(|s| s.push(span.clone())); +} + +fn pop_active() { + ACTIVE.with_borrow_mut(|s| { + s.pop(); + }); +} + +/// Run `f` with the innermost open span, if any; returns `None` when no span is +/// open. The hook a `tracing_subscriber::Layer` uses to forward events onto the +/// active platform span (see the `custom-spans` example). `f` must not re-enter +/// this function — a borrow is held for its duration. +pub fn with_active_span(f: impl FnOnce(&Span) -> R) -> Option { + ACTIVE.with_borrow(|stack| stack.last().map(f)) +} + +/// Open a synchronous custom span named `name`, run `f` inside it, and close +/// it when `f` returns. Nests under any span already open on the JS async +/// context — including an enclosing [`enter_span`] — so the platform tree +/// mirrors the call tree. +pub fn enter_span(name: &str, f: impl FnOnce(&Span) -> T) -> T { + let mut f = Some(f); + let mut out: Option = None; + + // `enterSpan` invokes this exactly once, synchronously, before returning, + // so the callback may borrow non-`'static` state. `borrow_mut` encodes + // that the JS function must not outlive this stack frame. + let mut cb = |span: raw::Span| { + let span = Span(span); + push_active(&span); + let f = f.take().expect("enterSpan invoked its callback twice"); + out = Some(f(&span)); + pop_active(); + }; + + { + let scoped = ScopedClosure::borrow_mut(&mut cb); + raw::TRACING.with(|t| t.enter_span_sync(name, &scoped)); + } + + out.expect("enterSpan must invoke its callback synchronously") +} + +/// Open a custom span named `name` that stays open until [`Span::end`] is +/// called, and return its handle. For work whose end isn't a callback return: +/// a stream that outlives the handler, or a `tracing_subscriber::Layer` +/// bridging `tracing`'s separate `on_new_span` / `on_close` (see the +/// `custom-spans` example). +/// +/// Prefer [`enter_span`] / [`enter_span_async`] when the work *is* callback +/// shaped: they can't leak an unclosed span, and anything they open nests under +/// them automatically. +/// +/// # Nesting +/// +/// The new span parents under whichever span is active *at this call*, and is +/// itself the active span only for the instant `startActiveSpan` runs its +/// callback. Work performed after this function returns is therefore NOT inside +/// it as far as the platform's async context is concerned, so spans opened +/// later become siblings rather than children. Reach for [`enter_span`] when +/// you need a subtree. +/// +/// Forgetting to call [`Span::end`] leaves the span open for the rest of the +/// invocation; the platform closes it, but its duration is meaningless. +pub fn start_active_span(name: &str) -> Span { + let mut out: Option = None; + + // Same synchronous-callback contract as `enter_span`, so the closure may + // borrow local state; only the `Span` it hands back outlives the call. + let mut cb = |span: raw::Span| { + out = Some(Span(span)); + }; + + { + let scoped = ScopedClosure::borrow_mut(&mut cb); + raw::TRACING.with(|t| t.start_active_span(name, &scoped)); + } + + out.expect("startActiveSpan must invoke its callback synchronously") +} + +/// Open an asynchronous custom span named `name`, drive `f`'s future inside it, +/// and close it when the future resolves. This is the form `async` request +/// handlers need. +/// +/// Unlike [`enter_span`], the callback returns a `Promise` that outlives the +/// `enterSpan` call (workerd awaits it), so `f` and its future must be +/// `'static`. +pub async fn enter_span_async(name: &str, f: F) -> T +where + F: FnOnce(Span) -> Fut + 'static, + Fut: Future + 'static, + T: 'static, +{ + let result: std::rc::Rc>> = std::rc::Rc::new(RefCell::new(None)); + let sink = result.clone(); + + let mut f = Some(f); + let cb = Closure::wrap(Box::new(move |span: raw::Span| -> js_sys::Promise { + let span = Span(span); + let f = f.take().expect("enterSpan invoked its callback twice"); + let fut = f(span.clone()); + let sink = sink.clone(); + future_to_promise(async move { + push_active(&span); + let value = fut.await; + pop_active(); + *sink.borrow_mut() = Some(value); + Ok(JsValue::UNDEFINED) + }) + }) as Box js_sys::Promise>); + + let promise = raw::TRACING.with(|t| t.enter_span_async(name, &cb)); + // Keep `cb` alive until the span's promise settles, then drop it. + let _ = JsFuture::from(promise).await; + drop(cb); + + std::rc::Rc::try_unwrap(result) + .ok() + .expect("span future outlived its handle") + .into_inner() + .expect("enterSpan async callback must resolve the result") +}