From d847cb9cb1a5c8bdd836e9021d1300a7b5a14a1a Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Fri, 11 Sep 2026 22:00:07 +0400 Subject: [PATCH 01/21] fix(pegboard-gateway): forward request ray ID to actors --- engine/packages/guard-core/src/request_context.rs | 10 ++++++++++ engine/packages/pegboard-gateway2/src/lib.rs | 4 +++- .../pegboard-gateway3/src/http_stream/handler.rs | 5 ++++- engine/packages/pegboard-gateway3/src/lib.rs | 1 + 4 files changed, 18 insertions(+), 2 deletions(-) diff --git a/engine/packages/guard-core/src/request_context.rs b/engine/packages/guard-core/src/request_context.rs index 3dfe3a75d4..caa4c1d857 100644 --- a/engine/packages/guard-core/src/request_context.rs +++ b/engine/packages/guard-core/src/request_context.rs @@ -2,6 +2,8 @@ use anyhow::{Context, Result}; use hyper::{Method, header::HeaderMap}; use rivet_runner_protocol as protocol; use rivet_util::Id; +use rivet_api_builder::X_RIVET_RAY_ID; +use std::collections::HashMap; use std::{ net::{IpAddr, SocketAddr}, sync::Arc, @@ -90,6 +92,14 @@ impl RequestContext { self.ray_id } + /// Adds this request's ray ID to the headers forwarded to an actor when the + /// caller did not send one, so the actor and gateway use the same ray ID. + pub fn forward_ray(&self, headers: &mut HashMap) { + headers + .entry(X_RIVET_RAY_ID.as_str().to_owned()) + .or_insert_with(|| self.ray_id.to_string()); + } + pub fn req_id(&self) -> Id { self.req_id } diff --git a/engine/packages/pegboard-gateway2/src/lib.rs b/engine/packages/pegboard-gateway2/src/lib.rs index 290ceec819..67ee28c7a9 100644 --- a/engine/packages/pegboard-gateway2/src/lib.rs +++ b/engine/packages/pegboard-gateway2/src/lib.rs @@ -127,7 +127,7 @@ impl PegboardGateway2 { let request_id = req_ctx.in_flight_request_id()?; // Extract request parts - let headers = req + let mut headers = req .headers() .iter() .filter_map(|(name, value)| { @@ -137,6 +137,7 @@ impl PegboardGateway2 { .map(|value_str| (name.to_string(), value_str.to_string())) }) .collect::>(); + req_ctx.forward_ray(&mut headers); // NOTE: Size constraints have already been applied by guard let body_bytes = req @@ -360,6 +361,7 @@ impl PegboardGateway2 { request_headers.insert(name.to_string(), value_str.to_string()); } } + req_ctx.forward_ray(&mut request_headers); let (mut stopped_sub, _) = tokio::try_join!( ctx.subscribe::(("actor_id", self.actor_id)), diff --git a/engine/packages/pegboard-gateway3/src/http_stream/handler.rs b/engine/packages/pegboard-gateway3/src/http_stream/handler.rs index 6bcaa7c014..d9387675c9 100644 --- a/engine/packages/pegboard-gateway3/src/http_stream/handler.rs +++ b/engine/packages/pegboard-gateway3/src/http_stream/handler.rs @@ -6,6 +6,8 @@ use std::{ time::Duration, }; +use std::collections::HashMap; + use anyhow::{Result, anyhow}; use bytes::Bytes; use gas::prelude::*; @@ -71,7 +73,7 @@ impl PegboardGateway3 { req_ctx.request_body_is_end_stream(), ); let request_id = req_ctx.in_flight_request_id()?; - let headers = req_ctx + let mut headers: HashMap = req_ctx .headers() .iter() .filter_map(|(name, value)| { @@ -81,6 +83,7 @@ impl PegboardGateway3 { .map(|value| (name.to_string(), value.to_owned())) }) .collect(); + req_ctx.forward_ray(&mut headers); let (mut stopped_sub, _) = tokio::try_join!( ctx.subscribe::(("actor_id", self.actor_id)), pegboard::utils::ensure_ns_metrics_exporter_for_namespace(ctx, self.namespace_id), diff --git a/engine/packages/pegboard-gateway3/src/lib.rs b/engine/packages/pegboard-gateway3/src/lib.rs index 5ff19be290..2bfa284f1c 100644 --- a/engine/packages/pegboard-gateway3/src/lib.rs +++ b/engine/packages/pegboard-gateway3/src/lib.rs @@ -134,6 +134,7 @@ impl PegboardGateway3 { request_headers.insert(name.to_string(), value_str.to_string()); } } + req_ctx.forward_ray(&mut request_headers); let (mut stopped_sub, _) = tokio::try_join!( ctx.subscribe::(("actor_id", self.actor_id)), From acd152dd8fba11ed5271e03051338fb0addc5f50 Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Tue, 1 Sep 2026 20:29:44 +0400 Subject: [PATCH 02/21] feat(rivetkit): trace actor invocations --- Cargo.lock | 8 + .../packages/rivetkit-core/Cargo.toml | 12 ++ .../packages/rivetkit-core/src/actor/task.rs | 10 +- .../packages/rivetkit-core/src/lib.rs | 1 + .../rivetkit-core/src/registry/dispatch.rs | 2 + .../rivetkit-core/src/registry/http.rs | 14 ++ .../rivetkit-core/src/registry/inspector.rs | 1 + .../rivetkit-core/src/registry/websocket.rs | 1 + .../packages/rivetkit-core/src/telemetry.rs | 128 ++++++++++++++++ .../rivetkit-core/src/telemetry/export.rs | 145 ++++++++++++++++++ .../packages/rivetkit-core/tests/task.rs | 5 + .../packages/rivetkit-napi/src/lib.rs | 20 ++- .../packages/rivetkit-napi/src/registry.rs | 1 + .../packages/rivetkit-napi/src/telemetry.rs | 89 +++++++++++ .../rivetkit/tests/fixtures/otlp-collector.ts | 33 ++++ 15 files changed, 465 insertions(+), 5 deletions(-) create mode 100644 rivetkit-rust/packages/rivetkit-core/src/telemetry.rs create mode 100644 rivetkit-rust/packages/rivetkit-core/src/telemetry/export.rs create mode 100644 rivetkit-typescript/packages/rivetkit-napi/src/telemetry.rs create mode 100644 rivetkit-typescript/packages/rivetkit/tests/fixtures/otlp-collector.ts diff --git a/Cargo.lock b/Cargo.lock index 5ed497b518..306d597014 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3764,6 +3764,7 @@ dependencies = [ "opentelemetry_sdk", "prost 0.13.5", "reqwest 0.12.22", + "serde_json", "thiserror 2.0.12", "tokio", "tonic", @@ -3776,9 +3777,12 @@ version = "0.28.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "56f8870d3024727e99212eb3bb1762ec16e255e3e6f58eeb3dc8db1aa226746d" dependencies = [ + "base64 0.22.1", + "hex", "opentelemetry", "opentelemetry_sdk", "prost 0.13.5", + "serde", "tonic", ] @@ -6313,6 +6317,9 @@ dependencies = [ "include_dir", "js-sys", "nix 0.30.1", + "opentelemetry", + "opentelemetry-otlp", + "opentelemetry_sdk", "parking_lot", "portpicker", "rand 0.8.5", @@ -6342,6 +6349,7 @@ dependencies = [ "tokio-util", "tower-http", "tracing", + "tracing-opentelemetry", "tracing-subscriber", "url", "uuid", diff --git a/rivetkit-rust/packages/rivetkit-core/Cargo.toml b/rivetkit-rust/packages/rivetkit-core/Cargo.toml index b1c5d139f1..09bbec31e2 100644 --- a/rivetkit-rust/packages/rivetkit-core/Cargo.toml +++ b/rivetkit-rust/packages/rivetkit-core/Cargo.toml @@ -15,6 +15,9 @@ default = ["native-runtime"] native-runtime = [ "dep:nix", "dep:reqwest", + "dep:opentelemetry-otlp", + "dep:opentelemetry_sdk", + "dep:tracing-subscriber", "dep:rivetkit-engine-process", "dep:axum", "dep:bytes", @@ -49,6 +52,13 @@ http-body-util = { workspace = true, optional = true } include_dir = { workspace = true } nix = { workspace = true, optional = true, features = ["process"] } parking_lot.workspace = true +opentelemetry = { version = "0.28", default-features = false, features = [ + "trace", + # Lets the SDK report dropped spans and export failures through tracing. + "internal-logs", +] } +opentelemetry-otlp = { version = "0.28", default-features = false, optional = true, features = ["trace", "http-json", "http-proto", "grpc-tonic", "reqwest-blocking-client"] } +opentelemetry_sdk = { version = "0.28", default-features = false, optional = true, features = ["trace", "internal-logs"] } rand.workspace = true reqwest = { workspace = true, optional = true } rusqlite = { workspace = true, optional = true } @@ -74,6 +84,8 @@ tokio-stream = { workspace = true, optional = true } tokio-util.workspace = true tower-http = { workspace = true, optional = true, features = ["fs"] } tracing.workspace = true +tracing-opentelemetry = { version = "0.29", default-features = false } +tracing-subscriber = { workspace = true, optional = true } url.workspace = true vbare.workspace = true diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs index c3cdbe3d22..c6074a80c3 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs @@ -201,6 +201,7 @@ pub enum DispatchCommand { Action { name: String, args: Vec, + incoming: crate::telemetry::IncomingInvocationContext, conn: ConnHandle, reply: oneshot::Sender>>, }, @@ -902,9 +903,12 @@ impl ActorTask { DispatchCommand::Action { name, args, + incoming, conn, reply, } => { + let invocation = + crate::telemetry::ActionInvocationSpan::start(&self.ctx, &name, incoming); tracing::info!( actor_id = %self.ctx.actor_id(), action_name = %name, @@ -938,6 +942,7 @@ impl ActorTask { Ok(result) => { let result = result.map_err(|error| ctx.attach_actor_to_error(error)); + invocation.finish(result.as_ref().err()); tracing::info!( actor_id = %actor_id, action_name = %action_name_for_log, @@ -955,6 +960,7 @@ impl ActorTask { let error = ctx.attach_actor_to_error( ActorLifecycleError::DroppedReply.build(), ); + invocation.finish(Some(&error)); let _ = reply.send(Err(error)); } } @@ -967,7 +973,9 @@ impl ActorTask { ?error, "actor task: failed to enqueue ActorEvent::Action" ); - let _ = reply.send(Err(self.attach_actor_to_error(error))); + let error = self.attach_actor_to_error(error); + invocation.finish(Some(&error)); + let _ = reply.send(Err(error)); self.log_dispatch_command_handled(command_kind, "enqueue_failed"); } } diff --git a/rivetkit-rust/packages/rivetkit-core/src/lib.rs b/rivetkit-rust/packages/rivetkit-core/src/lib.rs index ae6fa5a94e..7b8c84b163 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/lib.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/lib.rs @@ -16,6 +16,7 @@ pub mod registry; pub mod runtime; pub(crate) mod serde_metrics; pub mod serverless; +pub mod telemetry; #[cfg(feature = "native-runtime")] pub mod serverless_http; #[cfg(feature = "native-runtime")] diff --git a/rivetkit-rust/packages/rivetkit-core/src/registry/dispatch.rs b/rivetkit-rust/packages/rivetkit-core/src/registry/dispatch.rs index 8a38fcb0d1..7640a23f59 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/registry/dispatch.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/registry/dispatch.rs @@ -7,6 +7,7 @@ pub(super) async fn dispatch_action_through_task( conn: ConnHandle, name: String, args: Vec, + incoming: crate::telemetry::IncomingInvocationContext, ) -> std::result::Result, ActionDispatchError> { let (reply_tx, reply_rx) = oneshot::channel(); try_send_dispatch_command( @@ -14,6 +15,7 @@ pub(super) async fn dispatch_action_through_task( DispatchCommand::Action { name, args, + incoming, conn, reply: reply_tx, }, diff --git a/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs b/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs index 2d0d566bc1..9158b1bcdb 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs @@ -255,6 +255,20 @@ impl RegistryDispatcher { conn.clone(), action_name.clone(), args, + crate::telemetry::IncomingInvocationContext::from_headers( + request + .headers() + .get("x-rivetkit-ray-id") + .and_then(|value| value.to_str().ok().map(str::to_owned)), + request + .headers() + .get("traceparent") + .and_then(|value| value.to_str().ok()), + request + .headers() + .get("tracestate") + .and_then(|value| value.to_str().ok()), + ), ), ) .await; diff --git a/rivetkit-rust/packages/rivetkit-core/src/registry/inspector.rs b/rivetkit-rust/packages/rivetkit-core/src/registry/inspector.rs index a4dd9f5f96..aae94722ca 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/registry/inspector.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/registry/inspector.rs @@ -344,6 +344,7 @@ impl RegistryDispatcher { conn.clone(), action_name.to_owned(), args, + crate::telemetry::IncomingInvocationContext::default(), ) .await; match &output { diff --git a/rivetkit-rust/packages/rivetkit-core/src/registry/websocket.rs b/rivetkit-rust/packages/rivetkit-core/src/registry/websocket.rs index fd3c5202bf..3fd6f3505b 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/registry/websocket.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/registry/websocket.rs @@ -372,6 +372,7 @@ impl RegistryDispatcher { conn.clone(), request.name.clone(), request.args.into_vec(), + crate::telemetry::IncomingInvocationContext::default(), ) .await { diff --git a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs new file mode 100644 index 0000000000..78e04b6eb5 --- /dev/null +++ b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs @@ -0,0 +1,128 @@ +//! Internal OpenTelemetry spans owned by the actor runtime. + +#[cfg(feature = "native-runtime")] +pub mod export; + +use std::str::FromStr as _; + +use opentelemetry::trace::{ + SpanContext, SpanId, TraceContextExt as _, TraceFlags, TraceId, TraceState, +}; +use tracing_opentelemetry::OpenTelemetrySpanExt as _; + +use crate::{ActorContext, format_actor_key}; + +/// Correlation fields accepted at an invocation boundary. +#[derive(Debug, Default)] +pub struct IncomingInvocationContext { + pub(crate) ray_id: Option, + remote_parent: Option, +} + +impl IncomingInvocationContext { + pub(crate) fn from_headers( + ray_id: Option, + traceparent: Option<&str>, + tracestate: Option<&str>, + ) -> Self { + Self { + ray_id, + remote_parent: parse_remote_parent(traceparent, tracestate), + } + } +} + +/// The single root span for one client action invocation. +#[derive(Debug)] +pub(crate) struct ActionInvocationSpan { + span: Option, +} + +impl ActionInvocationSpan { + pub(crate) fn start( + ctx: &ActorContext, + action_name: &str, + incoming: IncomingInvocationContext, + ) -> Self { + if !tracing::enabled!(target: "rivetkit::telemetry", tracing::Level::INFO) { + return Self { span: None }; + } + + let span = tracing::info_span!( + target: "rivetkit::telemetry", + parent: None, + "rivet.actor.invoke", + otel.kind = "server", + rivet.invocation.type = "action", + rivet.actor.id = %ctx.actor_id(), + rivet.actor.name = %ctx.name(), + rivet.actor.key = %format_actor_key(ctx.key()), + rivet.action.name = %action_name, + rivet.ray.id = tracing::field::Empty, + otel.status_code = tracing::field::Empty, + error.type = tracing::field::Empty, + ); + if let Some(ray_id) = incoming.ray_id.as_deref() { + span.record("rivet.ray.id", ray_id); + } + if let Some(parent) = incoming.remote_parent { + span.set_parent(opentelemetry::Context::new().with_remote_span_context(parent)); + } + + Self { span: Some(span) } + } + + pub(crate) fn finish(mut self, error: Option<&anyhow::Error>) { + let Some(span) = self.span.take() else { + return; + }; + span.record( + "otel.status_code", + if error.is_none() { "OK" } else { "ERROR" }, + ); + if let Some(error) = error { + let error = rivet_error::RivetError::extract(error); + span.record("error.type", format!("{}.{}", error.group(), error.code())); + } + } +} + +impl Drop for ActionInvocationSpan { + fn drop(&mut self) { + let Some(span) = self.span.take() else { + return; + }; + span.record("otel.status_code", "ERROR"); + span.record("error.type", "actor.dropped_reply"); + } +} + +fn parse_remote_parent(traceparent: Option<&str>, tracestate: Option<&str>) -> Option { + let mut fields = traceparent?.split('-'); + let version = fields.next()?; + let trace_id = fields.next()?; + let span_id = fields.next()?; + let flags = fields.next()?; + if fields.next().is_some() + || version.len() != 2 + || version.eq_ignore_ascii_case("ff") + || trace_id.len() != 32 + || span_id.len() != 16 + || flags.len() != 2 + { + return None; + } + + let trace_id = TraceId::from_hex(trace_id).ok()?; + let span_id = SpanId::from_hex(span_id).ok()?; + let flags = u8::from_str_radix(flags, 16).ok()?; + let trace_state = tracestate + .and_then(|value| TraceState::from_str(value).ok()) + .unwrap_or_default(); + let context = SpanContext::new(trace_id, span_id, TraceFlags::new(flags), true, trace_state); + if context.is_valid() { + Some(context) + } else { + None + } +} diff --git a/rivetkit-rust/packages/rivetkit-core/src/telemetry/export.rs b/rivetkit-rust/packages/rivetkit-core/src/telemetry/export.rs new file mode 100644 index 0000000000..35c0051d5a --- /dev/null +++ b/rivetkit-rust/packages/rivetkit-core/src/telemetry/export.rs @@ -0,0 +1,145 @@ +//! Native OTLP export of the runtime's spans. +//! +//! Configuration comes entirely from the standard OpenTelemetry environment +//! variables, so every host that embeds core gets the same behavior by adding +//! [`layer`] to its subscriber and calling [`flush_best_effort`] on shutdown. +//! Core never installs a subscriber itself; which log layers surround the span +//! layer is the host's decision. + +use std::sync::OnceLock; +use std::time::Duration; + +use anyhow::{Context, Result}; +use opentelemetry::KeyValue; +use opentelemetry::trace::TracerProvider as _; +use opentelemetry_otlp::{Protocol, SpanExporter, WithExportConfig as _}; +use opentelemetry_sdk::Resource; +use opentelemetry_sdk::trace::{SdkTracer, SdkTracerProvider}; +use tracing_subscriber::registry::LookupSpan; +use tracing_subscriber::{EnvFilter, Layer}; + +/// Upper bound on the shutdown flush. Long enough for one export round trip +/// to a slow collector, short enough that a stuck collector cannot hold the +/// process open. +const FLUSH_TIMEOUT: Duration = Duration::from_secs(6); + +static PROVIDER: OnceLock = OnceLock::new(); + +/// Builds the span layer when standard OTel environment variables opt in, or +/// nothing when they do not. The layer only sees the runtime's own spans, so a +/// host's log filters do not decide what gets exported. +pub fn layer() -> Result>> +where + S: tracing::Subscriber + for<'a> LookupSpan<'a>, +{ + let Some(tracer) = initialize_if_configured()? else { + return Ok(None); + }; + Ok(Some( + tracing_opentelemetry::layer() + .with_tracer(tracer) + .with_location(false) + .with_threads(false) + .with_tracked_inactivity(false) + .with_filter(EnvFilter::new("rivetkit::telemetry=info")), + )) +} + +/// Builds the OTLP exporter once. A second call reuses the provider so that a +/// host initializing tracing more than once does not open a second pipeline. +fn initialize_if_configured() -> Result> { + if !export_is_configured() { + return Ok(None); + } + if let Some(provider) = PROVIDER.get() { + return Ok(Some(provider.tracer("rivetkit"))); + } + + let exporter = match configured_protocol()? { + Protocol::Grpc => SpanExporter::builder() + .with_tonic() + .build() + .context("build otlp span exporter")?, + protocol @ (Protocol::HttpBinary | Protocol::HttpJson) => SpanExporter::builder() + .with_http() + .with_protocol(protocol) + .build() + .context("build otlp span exporter")?, + }; + let resource = Resource::builder() + // Leave service.version to the application; record the runtime version separately. + .with_attribute(KeyValue::new("rivetkit.version", env!("CARGO_PKG_VERSION"))) + .build(); + let provider = SdkTracerProvider::builder() + .with_resource(resource) + .with_batch_exporter(exporter) + .build(); + let tracer = provider.tracer("rivetkit"); + PROVIDER + .set(provider) + .ok() + .context("tracer provider already initialized")?; + Ok(Some(tracer)) +} + +/// Reads the standard OTLP protocol variables. +/// +/// The exporter's own default comes from a compile-time constant chosen by the +/// enabled cargo features, and neither of its builders reads these variables, +/// so selecting the protocol has to happen here. Enabling `http-json` would +/// otherwise make JSON that compile-time default, which the OTLP specification +/// does not list among the usual defaults. +fn configured_protocol() -> Result { + let configured = std::env::var("OTEL_EXPORTER_OTLP_TRACES_PROTOCOL") + .or_else(|_| std::env::var("OTEL_EXPORTER_OTLP_PROTOCOL")) + .unwrap_or_else(|_| "http/protobuf".to_owned()); + match configured.as_str() { + "grpc" => Ok(Protocol::Grpc), + "http/protobuf" => Ok(Protocol::HttpBinary), + "http/json" => Ok(Protocol::HttpJson), + other => anyhow::bail!( + "native trace export supports grpc, http/protobuf and http/json, got {other:?}" + ), + } +} + +fn export_is_configured() -> bool { + if std::env::var("OTEL_SDK_DISABLED").is_ok_and(|value| value.eq_ignore_ascii_case("true")) + || std::env::var("OTEL_TRACES_EXPORTER") + .is_ok_and(|value| value.eq_ignore_ascii_case("none")) + { + return false; + } + + [ + "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", + "OTEL_EXPORTER_OTLP_ENDPOINT", + ] + .into_iter() + .any(|name| std::env::var_os(name).is_some_and(|value| !value.is_empty())) +} + +/// Exports whatever the batch processor still holds, giving up after +/// [`FLUSH_TIMEOUT`]. Export failures are logged and never returned, because a +/// telemetry problem must not turn a clean shutdown into a failed one. +pub async fn flush_best_effort() { + let Some(provider) = PROVIDER.get().cloned() else { + return; + }; + let flush = tokio::task::spawn_blocking(move || provider.force_flush()); + match tokio::time::timeout(FLUSH_TIMEOUT, flush).await { + Ok(Ok(Ok(()))) => {} + Ok(Ok(Err(error))) => tracing::warn!( + ?error, + "OpenTelemetry trace flush failed; queued spans may be lost" + ), + Ok(Err(error)) => tracing::warn!( + ?error, + "OpenTelemetry trace flush task failed; queued spans may be lost" + ), + Err(_) => tracing::warn!( + timeout = ?FLUSH_TIMEOUT, + "OpenTelemetry trace flush timed out; the collector may be slow or unreachable" + ), + } +} diff --git a/rivetkit-rust/packages/rivetkit-core/tests/task.rs b/rivetkit-rust/packages/rivetkit-core/tests/task.rs index 9e8a56b954..99b62080f1 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/task.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/task.rs @@ -1927,6 +1927,7 @@ pub(crate) mod moved_tests { task.handle_dispatch(DispatchCommand::Action { name: "client-action".to_owned(), args: Vec::new(), + incoming: crate::telemetry::IncomingInvocationContext::default(), conn: client_conn, reply: reply_tx, }) @@ -2030,6 +2031,7 @@ pub(crate) mod moved_tests { task.handle_dispatch(DispatchCommand::Action { name: "slow-action".to_owned(), args: Vec::new(), + incoming: crate::telemetry::IncomingInvocationContext::default(), conn: client_conn, reply: reply_tx, }) @@ -3732,6 +3734,7 @@ pub(crate) mod moved_tests { .send(DispatchCommand::Action { name: "ping".to_owned(), args: Vec::new(), + incoming: crate::telemetry::IncomingInvocationContext::default(), conn: ConnHandle::new("conn-grace", Vec::new(), Vec::new(), false), reply: action_tx, }) @@ -3776,6 +3779,7 @@ pub(crate) mod moved_tests { task.handle_dispatch(DispatchCommand::Action { name: "ping".to_owned(), args: Vec::new(), + incoming: crate::telemetry::IncomingInvocationContext::default(), conn: ConnHandle::new("conn-finalize", Vec::new(), Vec::new(), false), reply: reply_tx, }) @@ -4533,6 +4537,7 @@ pub(crate) mod moved_tests { .send(DispatchCommand::Action { name: "ping".to_owned(), args: Vec::new(), + incoming: crate::telemetry::IncomingInvocationContext::default(), conn: ConnHandle::new("conn-log-flow", Vec::new(), Vec::new(), false), reply: action_tx, }) diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/lib.rs b/rivetkit-typescript/packages/rivetkit-napi/src/lib.rs index 1c1c1b0a86..30a8760639 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/lib.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/lib.rs @@ -9,6 +9,7 @@ pub mod napi_actor_events; pub mod queue; pub mod registry; pub mod schedule; +mod telemetry; pub mod types; pub mod websocket; @@ -16,7 +17,7 @@ use std::sync::Once; use rivet_error::RivetError as RivetTransportError; use rivetkit_core::error::public_error_status_code; -use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; +use tracing_subscriber::{Layer as _, layer::SubscriberExt, util::SubscriberInitExt}; static INIT_TRACING: Once = Once::new(); pub(crate) const BRIDGE_RIVET_ERROR_PREFIX: &str = "__RIVET_ERROR_JSON__:"; @@ -115,10 +116,15 @@ pub(crate) fn init_tracing(log_level: Option<&str>) { .or_else(|| std::env::var("RUST_LOG").ok()) .unwrap_or_else(|| "warn".to_string()); + let log_filter = format!("{filter},rivetkit::telemetry=off"); let log_format = LogFormat::from_env(); + let (otel_layer, otel_error) = match rivetkit_core::telemetry::export::layer() { + Ok(layer) => (layer, None), + Err(error) => (None, Some(error)), + }; tracing_subscriber::registry() - .with(tracing_subscriber::EnvFilter::new(&filter)) + .with(otel_layer) .with(match log_format { LogFormat::Logfmt => Some( tracing_logfmt::builder() @@ -128,7 +134,8 @@ pub(crate) fn init_tracing(log_level: Option<&str>) { .with_location(env_flag("RUST_LOG_LOCATION")) .with_module_path(env_flag("RUST_LOG_MODULE_PATH")) .with_ansi_color(env_flag("RUST_LOG_ANSI_COLOR")) - .layer(), + .layer() + .with_filter(tracing_subscriber::EnvFilter::new(&log_filter)), ), LogFormat::Gcp => None, }) @@ -136,10 +143,15 @@ pub(crate) fn init_tracing(log_level: Option<&str>) { LogFormat::Logfmt => None, LogFormat::Gcp => Some( tracing_stackdriver::layer() - .with_source_location(env_flag("RUST_LOG_LOCATION")), + .with_source_location(env_flag("RUST_LOG_LOCATION")) + .with_filter(tracing_subscriber::EnvFilter::new(&log_filter)), ), }) .init(); + + if let Some(error) = otel_error { + tracing::warn!(?error, "OpenTelemetry trace export could not be initialized"); + } }); } diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/registry.rs b/rivetkit-typescript/packages/rivetkit-napi/src/registry.rs index c410453cc7..97d9cd0195 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/registry.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/registry.rs @@ -322,6 +322,7 @@ impl CoreRegistry { // `wait_ready()` may have armed its waiter while `serve()` was still // registering. Wake it after the state transition so it observes shutdown. self.serving_envoy_ready.notify_waiters(); + rivetkit_core::telemetry::export::flush_best_effort().await; Ok(()) } diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/telemetry.rs b/rivetkit-typescript/packages/rivetkit-napi/src/telemetry.rs new file mode 100644 index 0000000000..376b56b78b --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit-napi/src/telemetry.rs @@ -0,0 +1,89 @@ +//! Node-side telemetry glue. Export itself lives in `rivetkit_core::telemetry::export`. + +/// Forwards the OpenTelemetry SDK's own diagnostics to the JavaScript logger. +/// +/// The SDK reports dropped spans and export failures through Rust `tracing`, +/// which prints to stdout in a different format from the actor's Pino logs. +/// This layer hands those events to a JS callback instead, so an operator sees +/// them alongside everything else the actor logs. +pub(crate) mod sdk_log_bridge { + use std::sync::OnceLock; + + use napi::bindgen_prelude::*; + use napi::threadsafe_function::{ErrorStrategy, ThreadSafeCallContext, ThreadsafeFunction}; + use tracing::field::{Field, Visit}; + use tracing_subscriber::Layer; + use tracing_subscriber::layer::Context; + + /// One SDK diagnostic, flattened for the JavaScript side. + pub(crate) struct SdkLogEvent { + pub(crate) name: String, + pub(crate) message: String, + } + + static SINK: OnceLock> = OnceLock::new(); + + /// Installs the JavaScript sink. Only the first call takes effect, matching + /// the one-shot initialization of the tracing subscriber itself. + /// + /// The threadsafe function is unreferenced. A referenced one counts as live + /// work on the Node event loop, so a process that had registered the sink + /// would never exit on its own. Warnings still cross while the application + /// is running; the sink just stops being a reason to keep running. + pub(crate) fn install(env: Env, callback: JsFunction) -> Result<()> { + let mut tsfn = + callback.create_threadsafe_function(0, |ctx: ThreadSafeCallContext| { + let mut object = ctx.env.create_object()?; + object.set("name", ctx.value.name)?; + object.set("message", ctx.value.message)?; + Ok(vec![object.into_unknown()]) + })?; + tsfn.unref(&env)?; + let _ = SINK.set(tsfn); + Ok(()) + } + + #[derive(Default)] + struct FieldCollector { + name: String, + message: String, + } + + impl Visit for FieldCollector { + fn record_str(&mut self, field: &Field, value: &str) { + match field.name() { + "name" => self.name = value.to_owned(), + "message" => self.message = value.to_owned(), + _ => {} + } + } + + fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) { + let rendered = format!("{value:?}"); + match field.name() { + "name" => self.name = rendered, + "message" => self.message = rendered, + _ => {} + } + } + } + + pub(crate) struct SdkLogLayer; + + impl Layer for SdkLogLayer { + fn on_event(&self, event: &tracing::Event<'_>, _ctx: Context<'_, S>) { + let Some(sink) = SINK.get() else { + return; + }; + let mut fields = FieldCollector::default(); + event.record(&mut fields); + sink.call( + SdkLogEvent { + name: fields.name, + message: fields.message, + }, + napi::threadsafe_function::ThreadsafeFunctionCallMode::NonBlocking, + ); + } + } +} diff --git a/rivetkit-typescript/packages/rivetkit/tests/fixtures/otlp-collector.ts b/rivetkit-typescript/packages/rivetkit/tests/fixtures/otlp-collector.ts new file mode 100644 index 0000000000..944b7214c1 --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/tests/fixtures/otlp-collector.ts @@ -0,0 +1,33 @@ +import { createServer } from "node:http"; + +export interface OtlpCollector { + readonly endpoint: string; + spans(): Buffer[]; + close(): Promise; +} + +export async function startOtlpCollector(port: number): Promise { + const exports: Buffer[] = []; + const server = createServer((request, response) => { + const chunks: Buffer[] = []; + request.on("data", (chunk: Buffer) => chunks.push(chunk)); + request.on("end", () => { + exports.push(Buffer.concat(chunks)); + response.writeHead(200, { "content-type": "application/json" }); + response.end(); + }); + }); + + await new Promise((resolve) => + server.listen(port, "127.0.0.1", resolve), + ); + + return { + endpoint: `http://127.0.0.1:${port}/v1/traces`, + spans: () => exports, + close: () => + new Promise((resolve, reject) => { + server.close((error) => (error ? reject(error) : resolve())); + }), + }; +} From 58c21dec7694ce0d18bd5733aa8f0a161e9d7efd Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Tue, 1 Sep 2026 20:42:55 +0400 Subject: [PATCH 03/21] feat(rivetkit): trace SQLite operations --- .../rivetkit-core/src/actor/context.rs | 52 +++- .../rivetkit-core/src/actor/messages.rs | 4 + .../rivetkit-core/src/actor/sqlite/mod.rs | 84 ++++++- .../rivetkit-core/src/actor/sqlite/tx.rs | 30 ++- .../packages/rivetkit-core/src/actor/task.rs | 2 + .../packages/rivetkit-core/src/lib.rs | 3 + .../packages/rivetkit-core/src/telemetry.rs | 225 +++++++++++++++--- .../packages/rivetkit-core/tests/context.rs | 1 + rivetkit-rust/packages/rivetkit/src/event.rs | 4 + rivetkit-rust/packages/rivetkit/src/start.rs | 1 + .../tests/integration_canned_events.rs | 1 + .../packages/rivetkit-napi/index.d.ts | 1 + .../rivetkit-napi/src/actor_context.rs | 7 +- .../rivetkit-napi/src/actor_factory.rs | 6 +- .../rivetkit-napi/src/napi_actor_events.rs | 4 + .../rivetkit-napi/tests/napi_actor_events.rs | 3 + .../rivetkit/src/registry/napi-runtime.ts | 48 +++- .../packages/rivetkit/src/registry/native.ts | 36 +-- .../packages/rivetkit/src/registry/runtime.ts | 1 + .../rivetkit/src/registry/wasm-runtime.ts | 8 + .../tests/fixtures/napi-runtime-server.ts | 3 + .../rivetkit/tests/runtime-parity.test.ts | 6 +- .../rivetkit/tests/wasm-runtime.test.ts | 10 +- 23 files changed, 463 insertions(+), 77 deletions(-) diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs index cfc312a148..ed8d9faf23 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs @@ -62,7 +62,12 @@ use crate::types::{ActorKey, ConnId, ListOpts, format_actor_key}; /// and on the returned runtime objects like `SqliteDb`, schedule APIs, /// queue APIs, `ConnHandle`, and `WebSocket`. #[derive(Clone)] -pub struct ActorContext(pub(crate) Arc); +pub struct ActorContext( + pub(crate) Arc, + // Telemetry of the invocation this handle serves. `None` on the actor-owned + // handle and on any handle created outside an invocation. + pub(crate) Option, +); #[derive(Clone)] pub struct ActorKv { @@ -172,6 +177,7 @@ pub(crate) struct ActorContextInner { hibernated_connection_liveness_override: RwLock, Vec)>>>, pub(super) metrics: ActorMetrics, diagnostics: ActorDiagnostics, + telemetry_identity: Arc, actor_id: String, name: String, key: ActorKey, @@ -242,6 +248,31 @@ impl ActorKv { } impl ActorContext { + /// Returns a handle bound to `telemetry`, so schedules and SQLite work done + /// through it are attributed to that invocation. + pub fn with_invocation_telemetry( + mut self, + telemetry: Option, + ) -> Self { + self.1 = telemetry; + self + } + + /// Returns the SQLite handle bound to this handle's invocation. + pub fn invocation_sql(&self) -> crate::actor::sqlite::SqliteDb { + self.0.sql.clone().with_invocation_telemetry(self.1.clone()) + } + + pub(crate) fn invocation_telemetry(&self) -> Option<&crate::ActorInvocationTelemetry> { + self.1.as_ref() + } + + /// Returns whether two handles belong to the same running actor generation. + #[doc(hidden)] + pub fn is_same_instance(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) + } + #[cfg(test)] pub(crate) fn new( actor_id: impl Into, @@ -300,7 +331,7 @@ impl ActorContext { let shutdown_deadline = CancellationToken::new(); let sleep = SleepState::new(config.clone()); let user_kv = ActorKv { sql: sql.clone() }; - let ctx = Self(Arc::new(ActorContextInner { + let inner = Arc::new(ActorContextInner { legacy_kv, user_kv, sql, @@ -387,11 +418,17 @@ impl ActorContext { hibernated_connection_liveness_override: RwLock::new(None), metrics, diagnostics, + telemetry_identity: Arc::new(crate::telemetry::ActorTelemetryIdentity { + actor_id: actor_id.clone(), + actor_name: name.clone(), + actor_key: crate::types::format_actor_key(&key), + }), actor_id, name, key, region, - })); + }); + let ctx = Self(inner, None); ctx.configure_sleep_hooks(); ctx } @@ -911,6 +948,12 @@ impl ActorContext { &self.0.metrics } + /// Identity fields shared by every invocation on this actor. Built once so a + /// span does not re-allocate them per action. + pub(crate) fn telemetry_identity(&self) -> Arc { + self.0.telemetry_identity.clone() + } + pub(crate) fn record_user_task_started(&self, kind: UserTaskKind) { self.0.metrics.begin_user_task(kind); } @@ -1337,7 +1380,7 @@ impl ActorContext { } pub(crate) fn from_weak(weak: &Weak) -> Option { - weak.upgrade().map(Self) + weak.upgrade().map(|inner| Self(inner, None)) } #[doc(hidden)] @@ -1759,6 +1802,7 @@ impl ActorContext { args, conn: None, scheduled_fire: Some(scheduled_fire), + invocation_telemetry: None, reply: Reply::from(reply_tx), }, "scheduled_action", diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs index 702fa8f6ec..b8382520ab 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs @@ -392,6 +392,10 @@ pub enum ActorEvent { args: Vec, conn: Option, scheduled_fire: Option, + /// Telemetry of the invocation this action runs as, for the host + /// runtime to bind onto the context it hands the action. Absent when + /// the dispatch opened no invocation. + invocation_telemetry: Option, reply: Reply>, }, HttpRequest { diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/mod.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/mod.rs index c5cc659d00..9f16a05161 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/mod.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/mod.rs @@ -1,4 +1,5 @@ use std::collections::HashSet; +use std::future::Future; use std::io::Cursor; use std::sync::{ Arc, @@ -22,6 +23,9 @@ use serde_json::{Map as JsonMap, Value as JsonValue}; use tokio::sync::Mutex as AsyncMutex; #[cfg(feature = "sqlite-local")] use tokio::task::JoinHandle; +use tracing::Instrument as _; + +use crate::telemetry::SqliteOperation; #[cfg(feature = "sqlite-local")] mod envoy_sqlite_transport; @@ -234,6 +238,7 @@ pub struct SqliteDb { /// always sets up sqlite storage under the hood, so handle/actor_id are /// not a reliable signal for whether the user opted in; this flag is. enabled: bool, + invocation_telemetry: Option, #[cfg(feature = "sqlite-local")] // Forced-sync: native SQLite handles are used inside spawn_blocking and // synchronous diagnostic accessors. @@ -263,6 +268,7 @@ impl Default for SqliteDb { SqliteBackend::RemoteEnvoy }, enabled: false, + invocation_telemetry: None, #[cfg(feature = "sqlite-local")] db: Default::default(), #[cfg(feature = "sqlite-local")] @@ -299,6 +305,7 @@ impl SqliteDb { generation, backend: select_sqlite_backend(remote_sqlite)?, enabled, + invocation_telemetry: None, #[cfg(feature = "sqlite-local")] db: Default::default(), #[cfg(feature = "sqlite-local")] @@ -351,6 +358,34 @@ impl SqliteDb { self.backend } + #[doc(hidden)] + pub fn with_invocation_telemetry( + mut self, + telemetry: Option, + ) -> Self { + self.invocation_telemetry = telemetry; + self + } + + /// Runs one SQLite operation inside a `rivet.sqlite.*` span when the + /// current invocation is traced. + pub(super) async fn traced( + &self, + operation: SqliteOperation, + future: impl Future>, + ) -> Result { + let Some(mut span) = self + .invocation_telemetry + .as_ref() + .and_then(|telemetry| telemetry.start_sqlite(operation)) + else { + return future.await; + }; + let result = future.instrument(span.span()).await; + span.finish(result.as_ref().err()); + result + } + pub async fn get_pages( &self, request: protocol::SqliteGetPagesRequest, @@ -513,6 +548,11 @@ impl SqliteDb { pub async fn exec(&self, sql: impl Into) -> Result { let sql = sql.into(); + self.traced(SqliteOperation::Exec, self.exec_untraced(sql)) + .await + } + + async fn exec_untraced(&self, sql: String) -> Result { let sql_for_log = sql.clone(); #[cfg(feature = "sqlite-local")] let started_at = self @@ -566,6 +606,15 @@ impl SqliteDb { params: Option>, ) -> Result { let sql = sql.into(); + self.traced(SqliteOperation::Query, self.query_untraced(sql, params)) + .await + } + + async fn query_untraced( + &self, + sql: String, + params: Option>, + ) -> Result { let sql_for_log = sql.clone(); let binding_count = bind_param_count(¶ms); #[cfg(feature = "sqlite-local")] @@ -626,6 +675,15 @@ impl SqliteDb { params: Option>, ) -> Result { let sql = sql.into(); + self.traced(SqliteOperation::Run, self.run_untraced(sql, params)) + .await + } + + async fn run_untraced( + &self, + sql: String, + params: Option>, + ) -> Result { let sql_for_log = sql.clone(); let binding_count = bind_param_count(¶ms); #[cfg(feature = "sqlite-local")] @@ -685,6 +743,15 @@ impl SqliteDb { params: Option>, ) -> Result { let sql = sql.into(); + self.traced(SqliteOperation::Execute, self.execute_untraced(sql, params)) + .await + } + + async fn execute_untraced( + &self, + sql: String, + params: Option>, + ) -> Result { let sql_for_log = sql.clone(); let binding_count = bind_param_count(¶ms); #[cfg(feature = "sqlite-local")] @@ -736,6 +803,17 @@ impl SqliteDb { pub async fn execute_batch( &self, statements: Vec, + ) -> Result> { + self.traced( + SqliteOperation::ExecuteBatch, + self.execute_batch_untraced(statements), + ) + .await + } + + async fn execute_batch_untraced( + &self, + statements: Vec, ) -> Result> { let statement_count = statements.len(); let binding_count = statements @@ -749,7 +827,11 @@ impl SqliteDb { } } else { async { - let transaction = self.begin_transaction(None).await?; + let transaction = self + .clone() + .with_invocation_telemetry(None) + .begin_transaction(None) + .await?; let mut results = Vec::with_capacity(statements.len()); for statement in statements { match transaction.execute(statement.sql, statement.params).await { diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/tx.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/tx.rs index 1f011fcff7..1102395883 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/tx.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/tx.rs @@ -22,6 +22,7 @@ use tokio_util::sync::CancellationToken; #[cfg(not(target_arch = "wasm32"))] use crate::runtime::RuntimeSpawner; +use crate::telemetry::SqliteOperation; #[cfg(feature = "sqlite-local")] use super::profiling::{FINGERPRINT_FORMAT_VERSION, TransactionProfile}; @@ -277,7 +278,7 @@ impl SqliteDb { } let db = self.clone(); - run_detached_transaction_task( + let task = run_detached_transaction_task( async move { db.begin_transaction_profiled_inner( key, @@ -289,8 +290,8 @@ impl SqliteDb { .await }, "sqlite transaction begin task failed", - ) - .await + ); + self.traced(SqliteOperation::TransactionBegin, task).await } #[cfg(test)] @@ -495,11 +496,11 @@ impl SqliteDb { async fn transaction_exec(&self, key: &str, sql: String) -> Result { let db = self.clone(); let key = key.to_owned(); - run_detached_transaction_task( + let task = run_detached_transaction_task( async move { db.transaction_exec_inner(&key, sql).await }, "sqlite transaction exec task failed", - ) - .await + ); + self.traced(SqliteOperation::TransactionExec, task).await } async fn transaction_exec_inner(&self, key: &str, sql: String) -> Result { @@ -547,11 +548,11 @@ impl SqliteDb { ) -> Result { let db = self.clone(); let key = key.to_owned(); - run_detached_transaction_task( + let task = run_detached_transaction_task( async move { db.transaction_execute_inner(&key, sql, params).await }, "sqlite transaction execute task failed", - ) - .await + ); + self.traced(SqliteOperation::TransactionExecute, task).await } async fn transaction_execute_inner( @@ -616,11 +617,16 @@ impl SqliteDb { async fn finish_transaction(&self, key: &str, commit: bool) -> Result<()> { let db = self.clone(); let key = key.to_owned(); - run_detached_transaction_task( + let operation = if commit { + SqliteOperation::TransactionCommit + } else { + SqliteOperation::TransactionRollback + }; + let task = run_detached_transaction_task( async move { db.finish_transaction_inner(&key, commit).await }, "sqlite transaction finish task failed", - ) - .await + ); + self.traced(operation, task).await } async fn finish_transaction_inner(&self, key: &str, commit: bool) -> Result<()> { diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs index c6074a80c3..315afc9dc5 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs @@ -909,6 +909,7 @@ impl ActorTask { } => { let invocation = crate::telemetry::ActionInvocationSpan::start(&self.ctx, &name, incoming); + let invocation_telemetry = invocation.telemetry(); tracing::info!( actor_id = %self.ctx.actor_id(), action_name = %name, @@ -925,6 +926,7 @@ impl ActorTask { args, conn: Some(conn), scheduled_fire: None, + invocation_telemetry: Some(invocation_telemetry), reply: Reply::from(tracked_reply_tx), }, ) { diff --git a/rivetkit-rust/packages/rivetkit-core/src/lib.rs b/rivetkit-rust/packages/rivetkit-core/src/lib.rs index 7b8c84b163..7191bab8c8 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/lib.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/lib.rs @@ -17,6 +17,9 @@ pub mod runtime; pub(crate) mod serde_metrics; pub mod serverless; pub mod telemetry; +// Internal bridge types consumed by the NAPI and Wasm runtime adapters. +#[doc(hidden)] +pub use telemetry::ActorInvocationTelemetry; #[cfg(feature = "native-runtime")] pub mod serverless_http; #[cfg(feature = "native-runtime")] diff --git a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs index 78e04b6eb5..4499d1fdba 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs @@ -4,13 +4,15 @@ pub mod export; use std::str::FromStr as _; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use opentelemetry::trace::{ SpanContext, SpanId, TraceContextExt as _, TraceFlags, TraceId, TraceState, }; use tracing_opentelemetry::OpenTelemetrySpanExt as _; -use crate::{ActorContext, format_actor_key}; +use crate::ActorContext; /// Correlation fields accepted at an invocation boundary. #[derive(Debug, Default)] @@ -35,6 +37,88 @@ impl IncomingInvocationContext { /// The single root span for one client action invocation. #[derive(Debug)] pub(crate) struct ActionInvocationSpan { + telemetry: ActorInvocationTelemetry, +} + +/// Opaque invocation context carried across foreign-runtime adapters. +#[doc(hidden)] +#[derive(Clone, Debug)] +pub struct ActorInvocationTelemetry(Arc); + +/// Identity fields that do not change while an actor is alive. Built once per +/// actor and shared by every invocation, so starting one does not re-allocate +/// them. +#[derive(Debug)] +pub(crate) struct ActorTelemetryIdentity { + pub(crate) actor_id: String, + pub(crate) actor_name: String, + pub(crate) actor_key: String, +} + +/// Shared invocation state. The span never changes once the invocation starts, +/// so only the terminal record needs guarding: `finished` lets exactly one of +/// the explicit finish path and the drop path record a status. +#[derive(Debug)] +struct InvocationInner { + span: Option, + finished: AtomicBool, + identity: Arc, +} + +const OPERATION_ABANDONED_ERROR_TYPE: &str = "actor.operation_abandoned"; + +/// The closed set of SQLite operations that get a span. +/// +/// Both names are `&'static str`, so starting one of these spans allocates +/// nothing. Adding an operation is a compile error here rather than a silently +/// wrong span name. +#[derive(Clone, Copy, Debug)] +pub(crate) enum SqliteOperation { + Exec, + Execute, + ExecuteBatch, + Query, + Run, + TransactionBegin, + TransactionExec, + TransactionExecute, + TransactionCommit, + TransactionRollback, +} + +impl SqliteOperation { + fn as_str(self) -> &'static str { + match self { + Self::Exec => "exec", + Self::Execute => "execute", + Self::ExecuteBatch => "execute_batch", + Self::Query => "query", + Self::Run => "run", + Self::TransactionBegin => "transaction.begin", + Self::TransactionExec => "transaction.exec", + Self::TransactionExecute => "transaction.execute", + Self::TransactionCommit => "transaction.commit", + Self::TransactionRollback => "transaction.rollback", + } + } + + fn span_name(self) -> &'static str { + match self { + Self::Exec => "rivet.sqlite.exec", + Self::Execute => "rivet.sqlite.execute", + Self::ExecuteBatch => "rivet.sqlite.execute_batch", + Self::Query => "rivet.sqlite.query", + Self::Run => "rivet.sqlite.run", + Self::TransactionBegin => "rivet.sqlite.transaction.begin", + Self::TransactionExec => "rivet.sqlite.transaction.exec", + Self::TransactionExecute => "rivet.sqlite.transaction.execute", + Self::TransactionCommit => "rivet.sqlite.transaction.commit", + Self::TransactionRollback => "rivet.sqlite.transaction.rollback", + } + } +} + +pub(crate) struct SqliteOperationSpan { span: Option, } @@ -44,56 +128,139 @@ impl ActionInvocationSpan { action_name: &str, incoming: IncomingInvocationContext, ) -> Self { - if !tracing::enabled!(target: "rivetkit::telemetry", tracing::Level::INFO) { - return Self { span: None }; + let identity = ctx.telemetry_identity(); + let span = tracing::enabled!(target: "rivetkit::telemetry", tracing::Level::INFO).then(|| { + let span = tracing::info_span!( + target: "rivetkit::telemetry", + parent: None, + "rivet.actor.invoke", + otel.kind = "server", + rivet.invocation.type = "action", + rivet.actor.id = %identity.actor_id, + rivet.actor.name = %identity.actor_name, + rivet.actor.key = %identity.actor_key, + rivet.action.name = %action_name, + rivet.ray.id = tracing::field::Empty, + otel.status_code = tracing::field::Empty, + error.type = tracing::field::Empty, + ); + if let Some(ray_id) = incoming.ray_id.as_deref() { + span.record("rivet.ray.id", ray_id); + } + if let Some(parent) = incoming.remote_parent { + span.set_parent(opentelemetry::Context::new().with_remote_span_context(parent)); + } + span + }); + + Self { + telemetry: ActorInvocationTelemetry::new(span, identity), } + } + + pub(crate) fn telemetry(&self) -> ActorInvocationTelemetry { + self.telemetry.clone() + } + + pub(crate) fn finish(self, error: Option<&anyhow::Error>) { + self.telemetry.finish(error); + } +} + +impl Drop for ActionInvocationSpan { + fn drop(&mut self) { + self.telemetry.finish_dropped(); + } +} +impl ActorInvocationTelemetry { + fn new( + span: Option, + identity: Arc, + ) -> Self { + Self(Arc::new(InvocationInner { + span, + finished: AtomicBool::new(false), + identity, + })) + } + + pub(crate) fn start_sqlite(&self, operation: SqliteOperation) -> Option { + let parent = self.0.span.as_ref()?; let span = tracing::info_span!( target: "rivetkit::telemetry", - parent: None, - "rivet.actor.invoke", - otel.kind = "server", - rivet.invocation.type = "action", - rivet.actor.id = %ctx.actor_id(), - rivet.actor.name = %ctx.name(), - rivet.actor.key = %format_actor_key(ctx.key()), - rivet.action.name = %action_name, - rivet.ray.id = tracing::field::Empty, + parent: parent, + "rivet.sqlite.operation", + otel.name = operation.span_name(), + otel.kind = "internal", + rivet.operation.system = "sqlite", + rivet.operation.name = operation.as_str(), + rivet.actor.id = %self.0.identity.actor_id, + rivet.actor.name = %self.0.identity.actor_name, + rivet.actor.key = %self.0.identity.actor_key, otel.status_code = tracing::field::Empty, error.type = tracing::field::Empty, ); - if let Some(ray_id) = incoming.ray_id.as_deref() { - span.record("rivet.ray.id", ray_id); - } - if let Some(parent) = incoming.remote_parent { - span.set_parent(opentelemetry::Context::new().with_remote_span_context(parent)); + Some(SqliteOperationSpan { span: Some(span) }) + } + + fn finish(&self, error: Option<&anyhow::Error>) { + let Some(span) = self.take_span() else { + return; + }; + record_outcome(span, error); + } + + fn finish_dropped(&self) { + let Some(span) = self.take_span() else { + return; + }; + span.record("otel.status_code", "ERROR"); + span.record("error.type", "actor.dropped_reply"); + } + + /// Claims the terminal record, so the finish and drop paths cannot both + /// record a status for the same invocation. + fn take_span(&self) -> Option<&tracing::Span> { + if self.0.finished.swap(true, Ordering::AcqRel) { + return None; } + self.0.span.as_ref() + } +} - Self { span: Some(span) } +impl SqliteOperationSpan { + pub(crate) fn span(&self) -> tracing::Span { + self.span.as_ref().expect("sqlite span is present").clone() } - pub(crate) fn finish(mut self, error: Option<&anyhow::Error>) { + pub(crate) fn finish(&mut self, error: Option<&anyhow::Error>) { let Some(span) = self.span.take() else { return; }; - span.record( - "otel.status_code", - if error.is_none() { "OK" } else { "ERROR" }, - ); - if let Some(error) = error { - let error = rivet_error::RivetError::extract(error); - span.record("error.type", format!("{}.{}", error.group(), error.code())); - } + record_outcome(&span, error); } } -impl Drop for ActionInvocationSpan { +impl Drop for SqliteOperationSpan { fn drop(&mut self) { let Some(span) = self.span.take() else { return; }; span.record("otel.status_code", "ERROR"); - span.record("error.type", "actor.dropped_reply"); + span.record("error.type", OPERATION_ABANDONED_ERROR_TYPE); + } +} + +/// Records the terminal status and error identity of a finished span. +fn record_outcome(span: &tracing::Span, error: Option<&anyhow::Error>) { + span.record( + "otel.status_code", + if error.is_none() { "OK" } else { "ERROR" }, + ); + if let Some(error) = error { + let error = rivet_error::RivetError::extract(error); + span.record("error.type", format!("{}.{}", error.group(), error.code())); } } diff --git a/rivetkit-rust/packages/rivetkit-core/tests/context.rs b/rivetkit-rust/packages/rivetkit-core/tests/context.rs index 00d4b28582..6e8e070d12 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/context.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/context.rs @@ -1137,6 +1137,7 @@ mod moved_tests { args, conn, scheduled_fire, + invocation_telemetry: _, reply, } => { assert_eq!(name, "tick"); diff --git a/rivetkit-rust/packages/rivetkit/src/event.rs b/rivetkit-rust/packages/rivetkit/src/event.rs index 6a0b259e7e..e7f6ff8ead 100644 --- a/rivetkit-rust/packages/rivetkit/src/event.rs +++ b/rivetkit-rust/packages/rivetkit/src/event.rs @@ -95,6 +95,9 @@ impl RuntimeEvent { args, conn, scheduled_fire, + // The typed runtime hands actions the actor-wide `Ctx` built + // at start, so it has no per-invocation handle to bind this to. + invocation_telemetry: _, reply, } => Self::Action(ActionCall { name, @@ -1500,6 +1503,7 @@ mod tests { args: Vec::new(), conn: None, scheduled_fire: None, + invocation_telemetry: None, reply: reply_tx.into(), }) .expect("queue action event"); diff --git a/rivetkit-rust/packages/rivetkit/src/start.rs b/rivetkit-rust/packages/rivetkit/src/start.rs index fb229ab7c5..8a5eb88f77 100644 --- a/rivetkit-rust/packages/rivetkit/src/start.rs +++ b/rivetkit-rust/packages/rivetkit/src/start.rs @@ -2486,6 +2486,7 @@ mod tests { args: args.to_vec(), conn, scheduled_fire: None, + invocation_telemetry: None, reply: reply_tx.into(), }) .expect("send action event"); diff --git a/rivetkit-rust/packages/rivetkit/tests/integration_canned_events.rs b/rivetkit-rust/packages/rivetkit/tests/integration_canned_events.rs index e76aa2e5bf..5061648bf9 100644 --- a/rivetkit-rust/packages/rivetkit/tests/integration_canned_events.rs +++ b/rivetkit-rust/packages/rivetkit/tests/integration_canned_events.rs @@ -119,6 +119,7 @@ async fn send_action(event_tx: &mpsc::UnboundedSender, name: &str) - args: Vec::new(), conn: None, scheduled_fire: None, + invocation_telemetry: None, reply: reply_tx.into(), }) .expect("send action event"); diff --git a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts index 6e980d29c6..ac09cbe5b5 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts +++ b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts @@ -306,6 +306,7 @@ export declare class ActorContext { endOnStateChange(): void kv(): Kv sql(): JsNativeDatabase + sameActorInstance(other: ActorContext): boolean provisionActorRuntimeSocket(): Promise schedule(): Schedule queue(): Queue diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs b/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs index 49478854b7..b626bbab75 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs @@ -273,11 +273,16 @@ impl ActorContext { #[napi] pub fn sql(&self) -> JsNativeDatabase { JsNativeDatabase::new( - self.inner.sql().clone(), + self.inner.invocation_sql(), Some(self.inner.actor_id().to_owned()), ) } + #[napi] + pub fn same_actor_instance(&self, other: &ActorContext) -> bool { + self.inner.is_same_instance(&other.inner) + } + #[napi] pub async fn provision_actor_runtime_socket( &self, diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/actor_factory.rs b/rivetkit-typescript/packages/rivetkit-napi/src/actor_factory.rs index d01c1c0187..93a828007f 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/actor_factory.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/actor_factory.rs @@ -196,6 +196,7 @@ pub(crate) struct ConnectionPayload { #[derive(Clone)] pub(crate) struct ActionPayload { pub(crate) ctx: CoreActorContext, + pub(crate) telemetry: Option, pub(crate) conn: Option, pub(crate) name: String, pub(crate) args: Vec, @@ -871,7 +872,10 @@ fn build_connection_payload( fn build_action_payload(env: &Env, payload: ActionPayload) -> napi::Result> { let mut object = env.create_object()?; - object.set("ctx", ActorContext::new(payload.ctx))?; + object.set( + "ctx", + ActorContext::new(payload.ctx.with_invocation_telemetry(payload.telemetry)), + )?; match payload.conn { Some(conn) => object.set("conn", ConnHandle::new(conn))?, None => object.set("conn", env.get_null()?)?, diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs b/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs index 619f8d49b2..e6aa7d0e34 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs @@ -396,6 +396,7 @@ pub(crate) async fn dispatch_event( args, conn, scheduled_fire, + invocation_telemetry, reply, } => { tracing::info!( @@ -430,6 +431,7 @@ pub(crate) async fn dispatch_event( call_action( &callback, &ctx, + invocation_telemetry, conn, name.clone(), args.clone(), @@ -1177,6 +1179,7 @@ async fn call_run( async fn call_action( callback: &crate::actor_factory::CallbackTsfn, ctx: &ActorContext, + telemetry: Option, conn: Option, name: String, args: Vec, @@ -1189,6 +1192,7 @@ async fn call_action( callback, ActionPayload { ctx: ctx.inner().clone(), + telemetry, conn, name, args, diff --git a/rivetkit-typescript/packages/rivetkit-napi/tests/napi_actor_events.rs b/rivetkit-typescript/packages/rivetkit-napi/tests/napi_actor_events.rs index 171d0dc2fb..293810edde 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/tests/napi_actor_events.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/tests/napi_actor_events.rs @@ -168,6 +168,7 @@ mod moved_tests { args: vec![1, 2, 3], conn: None, scheduled_fire: None, + invocation_telemetry: None, reply: tx.into(), }, &bindings, @@ -490,6 +491,7 @@ mod moved_tests { args: Vec::new(), conn: None, scheduled_fire: None, + invocation_telemetry: None, reply: first_tx.into(), }) .expect("first action event should send"); @@ -499,6 +501,7 @@ mod moved_tests { args: Vec::new(), conn: None, scheduled_fire: None, + invocation_telemetry: None, reply: second_tx.into(), }) .expect("second action event should send"); diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts index 3d589d5242..950c242de4 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts @@ -1,3 +1,4 @@ +import type { AsyncLocalStorage } from "node:async_hooks"; import type { ActorContext as NativeActorContext, NapiActorFactory as NativeActorFactory, @@ -255,17 +256,37 @@ export class NapiCoreRuntime implements CoreRuntime { #bindings: NativeBindings; #sql = new WeakMap(); + #invocationContext: AsyncLocalStorage; - constructor(bindings: NativeBindings) { + constructor( + bindings: NativeBindings, + invocationContext: AsyncLocalStorage, + ) { this.#bindings = bindings; + this.#invocationContext = invocationContext; + } + + #actorContextForOperation(owner: ActorContextHandle): NativeActorContext { + const ownerCtx = asNativeActorContext(owner); + const active = this.#invocationContext.getStore(); + if (active?.sameActorInstance(ownerCtx)) { + return active; + } + return ownerCtx; } + // Cache only the actor-owned handle, which is closed on sleep. + // Invocation handles carry per-call context and must not be reused. #actorSql(ctx: ActorContextHandle): NapiSqlDatabase { - const nativeCtx = asNativeActorContext(ctx); - let database = this.#sql.get(nativeCtx); + const ownerCtx = asNativeActorContext(ctx); + const activeCtx = this.#actorContextForOperation(ctx); + if (activeCtx !== ownerCtx) { + return activeCtx.sql(); + } + let database = this.#sql.get(ownerCtx); if (!database) { - database = nativeCtx.sql(); - this.#sql.set(nativeCtx, database); + database = ownerCtx.sql(); + this.#sql.set(ownerCtx, database); } return database; } @@ -552,6 +573,10 @@ export class NapiCoreRuntime implements CoreRuntime { return asNativeActorContext(ctx).actorId(); } + runWithActorInvocationContext(ctx: ActorContextHandle, run: () => T): T { + return this.#invocationContext.run(asNativeActorContext(ctx), run); + } + actorName(ctx: ActorContextHandle): string { return asNativeActorContext(ctx).name(); } @@ -871,11 +896,7 @@ export class NapiCoreRuntime implements CoreRuntime { async actorSqlClose(ctx: ActorContextHandle): Promise { const nativeCtx = asNativeActorContext(ctx); - const database = this.#sql.get(nativeCtx); - if (!database) { - return; - } - + const database = this.#sql.get(nativeCtx) ?? nativeCtx.sql(); this.#sql.delete(nativeCtx); await database.close(); } @@ -1173,9 +1194,12 @@ export async function loadNapiRuntime(): Promise<{ // would snapshot the native `.node` addon into the deploy and 413. The // computed specifier keeps it opaque to static analysis so it is never // bundled. Enforced by scripts/ci/check-edge-native-closure.mjs. - const bindings = await import(["@rivetkit", "rivetkit-napi"].join("/")); + const [{ AsyncLocalStorage }, bindings] = await Promise.all([ + import("node:async_hooks"), + import(["@rivetkit", "rivetkit-napi"].join("/")), + ]); return { bindings, - runtime: new NapiCoreRuntime(bindings), + runtime: new NapiCoreRuntime(bindings, new AsyncLocalStorage()), }; } diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts index 4d3a55b45e..a2e01aab35 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts @@ -5312,21 +5312,29 @@ export function buildNativeFactory( conn != null ? makeConnCtx(ctx, conn, undefined, cancelToken) : makeActorCtx(ctx, undefined, cancelToken); - try { - return encodeValue( - await handler( - actorCtx, - ...validateActionArgs( - schemaConfig.actionInputSchemas, - name, - decodeArgs(args), + const runAction = async () => { + try { + return encodeValue( + await handler( + actorCtx, + ...validateActionArgs( + schemaConfig.actionInputSchemas, + name, + decodeArgs(args), + ), + ...(scheduledFire + ? [scheduledFire] + : []), ), - ...(scheduledFire ? [scheduledFire] : []), - ), - ); - } finally { - await actorCtx.dispose(); - } + ); + } finally { + await actorCtx.dispose(); + } + }; + return await runtime.runWithActorInvocationContext( + ctx, + runAction, + ); }, ), ]), diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts index 383cec1d1b..76029c7cea 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts @@ -540,6 +540,7 @@ export interface CoreRuntime { writes: RuntimeWorkflowKvWrite[], ): Promise; actorId(ctx: ActorContextHandle): string; + runWithActorInvocationContext(ctx: ActorContextHandle, run: () => T): T; actorName(ctx: ActorContextHandle): string; actorKey(ctx: ActorContextHandle): RuntimeActorKeySegment[]; actorRegion(ctx: ActorContextHandle): string; diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts index b0fb460ccf..c374433642 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts @@ -535,6 +535,14 @@ export class WasmCoreRuntime implements CoreRuntime { return callHandle(asWasmActorContext(ctx), "actorId"); } + runWithActorInvocationContext( + _ctx: ActorContextHandle, + run: () => T, + ): T { + // Wasm does not yet carry invocation telemetry across its runtime boundary. + return run(); + } + actorName(ctx: ActorContextHandle): string { return callHandle(asWasmActorContext(ctx), "name"); } diff --git a/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts b/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts index 9355ac780a..c601dbd3d5 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts @@ -113,6 +113,9 @@ const integrationActor = actor({ count: c.state.count, }; }, + sqliteFailure: async (c) => { + await c.db.execute("SELECT value FROM missing_trace_test_table"); + }, stateSnapshot: async (c) => { const kvValue = await c.kv.get("count"); return { diff --git a/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts b/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts index afe2e1db3a..a6f2149b70 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts @@ -1,3 +1,4 @@ +import { AsyncLocalStorage } from "node:async_hooks"; import { describe, expect, test } from "vitest"; import { BRIDGE_RIVET_ERROR_PREFIX, @@ -335,7 +336,10 @@ function createRuntimeCase(kind: CoreRuntime["kind"]): RuntimeCase { scenario, runtime: kind === "napi" - ? new NapiCoreRuntime(fakeNapiBindings(scenario) as never) + ? new NapiCoreRuntime( + fakeNapiBindings(scenario) as never, + new AsyncLocalStorage(), + ) : new WasmCoreRuntime(fakeWasmBindings(scenario)), }; } diff --git a/rivetkit-typescript/packages/rivetkit/tests/wasm-runtime.test.ts b/rivetkit-typescript/packages/rivetkit/tests/wasm-runtime.test.ts index a9bf9d42d2..c8209a363d 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/wasm-runtime.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/wasm-runtime.test.ts @@ -1,3 +1,4 @@ +import { AsyncLocalStorage } from "node:async_hooks"; import { describe, expect, test, vi } from "vitest"; import { BRIDGE_RIVET_ERROR_PREFIX, RivetError } from "@/actor/errors"; import { actor } from "@/actor/mod"; @@ -240,7 +241,9 @@ describe("WasmCoreRuntime", () => { const acceptRuntime = (_runtime: CoreRuntime) => {}; acceptRuntime(new WasmCoreRuntime(fakeWasmBindings())); - acceptRuntime(new NapiCoreRuntime({} as never)); + acceptRuntime( + new NapiCoreRuntime({} as never, new AsyncLocalStorage()), + ); }); test("maps raw wasm registry, factory, and cancellation handles", () => { @@ -370,7 +373,10 @@ describe("WasmCoreRuntime", () => { } as unknown as ActorContextHandle; expect( - new NapiCoreRuntime({} as never).actorQueueMaxSize(context), + new NapiCoreRuntime( + {} as never, + new AsyncLocalStorage(), + ).actorQueueMaxSize(context), ).toBe(maxSize); expect( new WasmCoreRuntime(fakeWasmBindings()).actorQueueMaxSize(context), From 884389e0eeaa83de63317bd66eb023acf1c5291c Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Tue, 1 Sep 2026 20:58:41 +0400 Subject: [PATCH 04/21] feat(rivetkit): propagate actor trace context --- pnpm-lock.yaml | 3 + .../rivetkit-core/src/actor/context.rs | 22 +- .../packages/rivetkit-core/src/lib.rs | 4 +- .../rivetkit-core/src/registry/http.rs | 15 +- .../packages/rivetkit-core/src/telemetry.rs | 236 +++++++++++++++--- .../packages/rivetkit-napi/index.d.ts | 14 ++ .../rivetkit-napi/src/actor_context.rs | 48 +++- .../packages/rivetkit/package.json | 1 + .../rivetkit/src/client/actor-handle.ts | 36 ++- .../packages/rivetkit/src/client/client.ts | 22 +- .../src/common/actor-router-consts.ts | 3 + .../src/common/actor-telemetry-context.ts | 21 ++ .../rivetkit/src/common/otel-context.ts | 40 +++ .../src/engine-client/actor-http-client.ts | 16 ++ .../rivetkit/src/registry/napi-runtime.ts | 20 +- .../packages/rivetkit/src/registry/native.ts | 16 +- .../packages/rivetkit/src/registry/runtime.ts | 19 ++ .../rivetkit/src/registry/wasm-runtime.ts | 7 + 18 files changed, 463 insertions(+), 80 deletions(-) create mode 100644 rivetkit-typescript/packages/rivetkit/src/common/actor-telemetry-context.ts create mode 100644 rivetkit-typescript/packages/rivetkit/src/common/otel-context.ts diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index e8d9fcae90..50dc6369c5 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -3538,6 +3538,9 @@ importers: '@hono/zod-openapi': specifier: ^1.1.5 version: 1.1.5(hono@4.11.9)(zod@4.1.13) + '@opentelemetry/api': + specifier: ^1.1.0 + version: 1.9.0 '@rivet-dev/agent-os-core': specifier: ^0.1.1 version: 0.1.1(pyodide@0.28.3) diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs index ed8d9faf23..8f00eaf5e4 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs @@ -263,6 +263,12 @@ impl ActorContext { self.0.sql.clone().with_invocation_telemetry(self.1.clone()) } + /// Returns correlation for the invocation this handle serves, absent when + /// the handle is not bound to one or tracing is disabled. + pub fn invocation_trace_context(&self) -> Option { + self.1.as_ref()?.trace_context() + } + pub(crate) fn invocation_telemetry(&self) -> Option<&crate::ActorInvocationTelemetry> { self.1.as_ref() } @@ -751,9 +757,17 @@ impl ActorContext { false } + /// Runs `future` to completion after the current reply, without blocking + /// it. Work started from an invocation keeps that invocation's span open + /// until it settles, so its SQLite operations and logs stay attributed to + /// the request that started them. #[cfg(not(feature = "wasm-runtime"))] pub fn wait_until(&self, future: impl Future + Send + 'static) { - self.spawn_work(ActorWorkKind::WaitUntil, future); + let invocation = self.1.as_ref().map(crate::ActorInvocationTelemetry::hold_open); + self.spawn_work(ActorWorkKind::WaitUntil, async move { + future.await; + drop(invocation); + }); } #[cfg(not(feature = "wasm-runtime"))] @@ -763,7 +777,11 @@ impl ActorContext { #[cfg(feature = "wasm-runtime")] pub fn wait_until(&self, future: impl Future + 'static) { - self.spawn_work(ActorWorkKind::WaitUntil, future); + let invocation = self.1.as_ref().map(crate::ActorInvocationTelemetry::hold_open); + self.spawn_work(ActorWorkKind::WaitUntil, async move { + future.await; + drop(invocation); + }); } #[cfg(feature = "wasm-runtime")] diff --git a/rivetkit-rust/packages/rivetkit-core/src/lib.rs b/rivetkit-rust/packages/rivetkit-core/src/lib.rs index 7191bab8c8..d789dd9849 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/lib.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/lib.rs @@ -19,7 +19,9 @@ pub mod serverless; pub mod telemetry; // Internal bridge types consumed by the NAPI and Wasm runtime adapters. #[doc(hidden)] -pub use telemetry::ActorInvocationTelemetry; +pub use telemetry::{ + ActorInvocationSpanContext, ActorInvocationTelemetry, ActorInvocationTraceContext, +}; #[cfg(feature = "native-runtime")] pub mod serverless_http; #[cfg(feature = "native-runtime")] diff --git a/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs b/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs index 9158b1bcdb..f85838037d 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs @@ -255,20 +255,7 @@ impl RegistryDispatcher { conn.clone(), action_name.clone(), args, - crate::telemetry::IncomingInvocationContext::from_headers( - request - .headers() - .get("x-rivetkit-ray-id") - .and_then(|value| value.to_str().ok().map(str::to_owned)), - request - .headers() - .get("traceparent") - .and_then(|value| value.to_str().ok()), - request - .headers() - .get("tracestate") - .and_then(|value| value.to_str().ok()), - ), + crate::telemetry::IncomingInvocationContext::from_http_headers(request.headers()), ), ) .await; diff --git a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs index 4499d1fdba..013db97002 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs @@ -5,11 +5,12 @@ pub mod export; use std::str::FromStr as _; use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use opentelemetry::trace::{ SpanContext, SpanId, TraceContextExt as _, TraceFlags, TraceId, TraceState, }; +use parking_lot::Mutex; use tracing_opentelemetry::OpenTelemetrySpanExt as _; use crate::ActorContext; @@ -21,6 +22,9 @@ pub struct IncomingInvocationContext { remote_parent: Option, } +/// Header carrying the caller's ray into an actor. +pub(crate) const HEADER_RIVETKIT_RAY_ID: &str = "x-rivetkit-ray-id"; + impl IncomingInvocationContext { pub(crate) fn from_headers( ray_id: Option, @@ -32,6 +36,35 @@ impl IncomingInvocationContext { remote_parent: parse_remote_parent(traceparent, tracestate), } } + + /// Reads the ray ID and W3C trace context an HTTP request carries. Every + /// HTTP entry point into an actor reads them through here, so they all + /// apply the same bounds. + pub(crate) fn from_http_headers(headers: &http::HeaderMap) -> Self { + Self::from_headers( + invocation_ray_id(headers), + headers.get("traceparent").and_then(|value| value.to_str().ok()), + headers.get("tracestate").and_then(|value| value.to_str().ok()), + ) + } +} + +/// Reads the caller's ray id. The header is untrusted, so it is bounded to +/// 128 characters of `[A-Za-z0-9_-]`; anything else counts as absent and the +/// invocation mints a fresh ray instead. +fn invocation_ray_id(headers: &http::HeaderMap) -> Option { + headers + .get(HEADER_RIVETKIT_RAY_ID)? + .to_str() + .ok() + .filter(|value| { + !value.is_empty() + && value.len() <= 128 + && value + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_')) + }) + .map(str::to_owned) } /// The single root span for one client action invocation. @@ -55,18 +88,51 @@ pub(crate) struct ActorTelemetryIdentity { pub(crate) actor_key: String, } -/// Shared invocation state. The span never changes once the invocation starts, -/// so only the terminal record needs guarding: `finished` lets exactly one of -/// the explicit finish path and the drop path record a status. +/// Shared invocation state. `finished` lets exactly one of the finish and +/// drop paths record the terminal status, and marks the invocation closed even +/// when tracing is off and there is no span. The span slot is emptied, which +/// is what exports the span, once the status is recorded and no work handed +/// to `wait_until` from this invocation is still running. `pending_work` +/// counts that work, so the invocation stays active for it after the reply. #[derive(Debug)] struct InvocationInner { - span: Option, + ray_id: Option, + // Forced-sync: the slot is emptied from `Drop` paths and sync accessors, + // and the guard is never held across an await. + span: Mutex>, finished: AtomicBool, + pending_work: AtomicUsize, identity: Arc, } const OPERATION_ABANDONED_ERROR_TYPE: &str = "actor.operation_abandoned"; +/// Keeps an invocation open while one piece of `wait_until` work runs. The +/// span is released when the last guard drops after the terminal status has +/// been recorded, so work that settles before the reply changes nothing. +pub(crate) struct InvocationWorkGuard(ActorInvocationTelemetry); + +/// Active actor invocation fields exposed to foreign-runtime adapters. +#[doc(hidden)] +#[derive(Clone, Debug)] +pub struct ActorInvocationTraceContext { + pub ray_id: Option, + /// Present only while the invocation runs inside a valid span. + pub span: Option, +} + +/// W3C span context of the current invocation span. A span context is either +/// complete or absent, so these fields are never optional individually. +#[doc(hidden)] +#[derive(Clone, Debug)] +pub struct ActorInvocationSpanContext { + pub trace_id: String, + pub span_id: String, + pub trace_flags: u8, + pub traceparent: String, + pub tracestate: Option, +} + /// The closed set of SQLite operations that get a span. /// /// Both names are `&'static str`, so starting one of these spans allocates @@ -129,32 +195,33 @@ impl ActionInvocationSpan { incoming: IncomingInvocationContext, ) -> Self { let identity = ctx.telemetry_identity(); - let span = tracing::enabled!(target: "rivetkit::telemetry", tracing::Level::INFO).then(|| { - let span = tracing::info_span!( - target: "rivetkit::telemetry", - parent: None, - "rivet.actor.invoke", - otel.kind = "server", - rivet.invocation.type = "action", - rivet.actor.id = %identity.actor_id, - rivet.actor.name = %identity.actor_name, - rivet.actor.key = %identity.actor_key, - rivet.action.name = %action_name, - rivet.ray.id = tracing::field::Empty, - otel.status_code = tracing::field::Empty, - error.type = tracing::field::Empty, - ); - if let Some(ray_id) = incoming.ray_id.as_deref() { - span.record("rivet.ray.id", ray_id); - } - if let Some(parent) = incoming.remote_parent { - span.set_parent(opentelemetry::Context::new().with_remote_span_context(parent)); - } - span - }); + let ray_id = incoming.ray_id; + let span = if tracing::enabled!(target: "rivetkit::telemetry", tracing::Level::INFO) { + let span = tracing::info_span!( + target: "rivetkit::telemetry", + parent: None, + "rivet.actor.invoke", + otel.kind = "server", + rivet.invocation.type = "action", + rivet.actor.id = %identity.actor_id, + rivet.actor.name = %identity.actor_name, + rivet.actor.key = %identity.actor_key, + rivet.action.name = %action_name, + rivet.ray.id = tracing::field::Empty, + otel.status_code = tracing::field::Empty, + error.type = tracing::field::Empty, + ); + span.record("rivet.ray.id", ray_id.as_deref()); + if let Some(parent) = incoming.remote_parent { + span.set_parent(opentelemetry::Context::new().with_remote_span_context(parent)); + } + Some(span) + } else { + None + }; Self { - telemetry: ActorInvocationTelemetry::new(span, identity), + telemetry: ActorInvocationTelemetry::new(ray_id, span, identity), } } @@ -175,26 +242,73 @@ impl Drop for ActionInvocationSpan { impl ActorInvocationTelemetry { fn new( + ray_id: Option, span: Option, identity: Arc, ) -> Self { Self(Arc::new(InvocationInner { - span, + ray_id, + span: Mutex::new(span), finished: AtomicBool::new(false), + pending_work: AtomicUsize::new(0), identity, })) } + /// Registers work that outlives the reply, so the invocation span stays + /// open and keeps parenting operations until the returned guard drops. + pub(crate) fn hold_open(&self) -> InvocationWorkGuard { + self.0.pending_work.fetch_add(1, Ordering::SeqCst); + InvocationWorkGuard(self.clone()) + } + + /// Returns correlation fields only while this actor invocation is active. + #[doc(hidden)] + pub fn trace_context(&self) -> Option { + let active = self.active()?; + let span = active.span.lock().clone().and_then(|span| { + let context = span.context(); + let context_span = context.span(); + let span_context = context_span.span_context(); + if !span_context.is_valid() { + return None; + } + let tracestate = span_context.trace_state().header(); + Some(ActorInvocationSpanContext { + trace_id: span_context.trace_id().to_string(), + span_id: span_context.span_id().to_string(), + trace_flags: span_context.trace_flags().to_u8(), + traceparent: format!( + "00-{}-{}-{:02x}", + span_context.trace_id(), + span_context.span_id(), + span_context.trace_flags().to_u8(), + ), + tracestate: if tracestate.is_empty() { + None + } else { + Some(tracestate) + }, + }) + }); + + Some(ActorInvocationTraceContext { + ray_id: active.ray_id.clone(), + span, + }) + } + pub(crate) fn start_sqlite(&self, operation: SqliteOperation) -> Option { - let parent = self.0.span.as_ref()?; + let parent = self.active()?.span.lock().clone()?; let span = tracing::info_span!( target: "rivetkit::telemetry", - parent: parent, + parent: &parent, "rivet.sqlite.operation", otel.name = operation.span_name(), otel.kind = "internal", rivet.operation.system = "sqlite", rivet.operation.name = operation.as_str(), + rivet.ray.id = self.0.ray_id.as_deref(), rivet.actor.id = %self.0.identity.actor_id, rivet.actor.name = %self.0.identity.actor_name, rivet.actor.key = %self.0.identity.actor_key, @@ -205,27 +319,71 @@ impl ActorInvocationTelemetry { } fn finish(&self, error: Option<&anyhow::Error>) { - let Some(span) = self.take_span() else { + let Some(span) = self.claim_terminal() else { return; }; - record_outcome(span, error); + record_outcome(&span, error); + self.mark_reply_sent(&span); + self.release_span_if_settled(); } fn finish_dropped(&self) { - let Some(span) = self.take_span() else { + let Some(span) = self.claim_terminal() else { return; }; span.record("otel.status_code", "ERROR"); span.record("error.type", "actor.dropped_reply"); + self.mark_reply_sent(&span); + self.release_span_if_settled(); + } + + /// Borrows the invocation while it is still open: before its status is + /// recorded, or after it while `wait_until` work from it still runs. A + /// settled invocation yields nothing, so late SQLite work and retained + /// handles cannot attach to a span that has already ended. + fn active(&self) -> Option<&InvocationInner> { + let open = !self.0.finished.load(Ordering::SeqCst) + || self.0.pending_work.load(Ordering::SeqCst) > 0; + if open { Some(&*self.0) } else { None } } /// Claims the terminal record, so the finish and drop paths cannot both - /// record a status for the same invocation. - fn take_span(&self) -> Option<&tracing::Span> { - if self.0.finished.swap(true, Ordering::AcqRel) { + /// record a status for the same invocation. The span stays in its slot + /// until `release_span_if_settled` empties it. + fn claim_terminal(&self) -> Option { + if self.0.finished.swap(true, Ordering::SeqCst) { return None; } - self.0.span.as_ref() + self.0.span.lock().clone() + } + + /// Marks the moment the caller got its answer when the span will outlive + /// it, so the reply point stays visible inside a span that is still + /// running `wait_until` work. + fn mark_reply_sent(&self, span: &tracing::Span) { + if self.0.pending_work.load(Ordering::SeqCst) > 0 { + tracing::info!(target: "rivetkit::telemetry", parent: span, "reply sent"); + } + } + + /// Ends the span, which exports it, unless `wait_until` work is still + /// holding the invocation open. The last guard to drop ends it instead. + /// Both sides set their flag before reading the other's, so the two + /// cannot each see the other as still pending and leave the span behind. + fn release_span_if_settled(&self) { + if self.0.pending_work.load(Ordering::SeqCst) == 0 { + self.0.span.lock().take(); + } + } +} + +impl Drop for InvocationWorkGuard { + fn drop(&mut self) { + let inner = &self.0.0; + let was_last = inner.pending_work.fetch_sub(1, Ordering::SeqCst) == 1; + if was_last && inner.finished.load(Ordering::SeqCst) { + inner.span.lock().take(); + } } } diff --git a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts index ac09cbe5b5..07102087e1 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts +++ b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts @@ -8,6 +8,19 @@ export interface JsActorKeySegment { stringValue?: string numberValue?: number } +/** Active actor invocation correlation exposed to the TypeScript runtime adapter. */ +export interface JsActorInvocationTraceContext { + rayId?: string + span?: JsActorInvocationSpanContext +} +/** W3C span context of the current invocation span, present only when tracing is active. */ +export interface JsActorInvocationSpanContext { + traceId: string + spanId: string + traceFlags: number + traceparent: string + tracestate?: string +} export interface JsHttpRequest { method: string uri: string @@ -307,6 +320,7 @@ export declare class ActorContext { kv(): Kv sql(): JsNativeDatabase sameActorInstance(other: ActorContext): boolean + invocationTraceContext(): JsActorInvocationTraceContext | null provisionActorRuntimeSocket(): Promise schedule(): Schedule queue(): Queue diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs b/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs index b626bbab75..917e159367 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs @@ -18,8 +18,9 @@ use napi_derive::napi; use parking_lot::Mutex; use rivetkit_core::types::ActorKeySegment; use rivetkit_core::{ - ActorContext as CoreActorContext, ActorWorkKind, ConnHandle as CoreConnHandle, KeepAwakeRegion, - Request as CoreRequest, RequestSaveOpts, StateDelta, WebSocketCallbackRegion, WorkflowKvWrite, + ActorContext as CoreActorContext, ActorInvocationSpanContext, ActorInvocationTraceContext, + ActorWorkKind, ConnHandle as CoreConnHandle, KeepAwakeRegion, Request as CoreRequest, + RequestSaveOpts, StateDelta, WebSocketCallbackRegion, WorkflowKvWrite, }; use scc::HashMap as SccHashMap; use tokio::sync::mpsc::UnboundedSender; @@ -79,6 +80,44 @@ pub struct JsActorKeySegment { pub number_value: Option, } +/// Active actor invocation correlation exposed to the TypeScript runtime adapter. +#[napi(object)] +pub struct JsActorInvocationTraceContext { + pub ray_id: Option, + pub span: Option, +} + +/// W3C span context of the current invocation span, present only when tracing is active. +#[napi(object)] +pub struct JsActorInvocationSpanContext { + pub trace_id: String, + pub span_id: String, + pub trace_flags: u8, + pub traceparent: String, + pub tracestate: Option, +} + +impl From for JsActorInvocationTraceContext { + fn from(value: ActorInvocationTraceContext) -> Self { + Self { + ray_id: value.ray_id, + span: value.span.map(JsActorInvocationSpanContext::from), + } + } +} + +impl From for JsActorInvocationSpanContext { + fn from(value: ActorInvocationSpanContext) -> Self { + Self { + trace_id: value.trace_id, + span_id: value.span_id, + trace_flags: value.trace_flags, + traceparent: value.traceparent, + tracestate: value.tracestate, + } + } +} + #[napi(object)] pub struct JsHttpRequest { pub method: String, @@ -283,6 +322,11 @@ impl ActorContext { self.inner.is_same_instance(&other.inner) } + #[napi] + pub fn invocation_trace_context(&self) -> Option { + self.inner.invocation_trace_context().map(Into::into) + } + #[napi] pub async fn provision_actor_runtime_socket( &self, diff --git a/rivetkit-typescript/packages/rivetkit/package.json b/rivetkit-typescript/packages/rivetkit/package.json index ac27eda52e..38549ec69b 100644 --- a/rivetkit-typescript/packages/rivetkit/package.json +++ b/rivetkit-typescript/packages/rivetkit/package.json @@ -209,6 +209,7 @@ }, "dependencies": { "@hono/zod-openapi": "^1.1.5", + "@opentelemetry/api": "^1.1.0", "@rivet-dev/agent-os-core": "^0.1.1", "@rivet-dev/services": "^0.1.5", "@rivetkit/bare-ts": "^0.6.2", diff --git a/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts b/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts index e41aa6366d..d3569adced 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts @@ -3,6 +3,9 @@ import type { ActorSpecifier } from "@/actor/errors"; import { HEADER_CONN_PARAMS, HEADER_ENCODING, + HEADER_RIVETKIT_RAY_ID, + HEADER_TRACEPARENT, + HEADER_TRACESTATE, } from "@/common/actor-router-consts"; import { isRequestLike } from "@/common/fetch-like"; import type * as protocol from "@/common/client-protocol"; @@ -24,6 +27,7 @@ import { AsyncMutex } from "@/common/database/shared"; import type { Encoding, JsonCompatValue } from "@/common/encoding"; import { deconstructError } from "@/common/utils"; import type { EngineControlClient } from "@/engine-client/driver"; +import type { CurrentActorInvocation } from "@/registry/runtime"; import { decodeCborCompat, deserializeWithEncoding, @@ -82,6 +86,7 @@ export class ActorHandleRaw { #resolvedActorId?: string; #resolvingActorId?: Promise; #queueSendMutex = new AsyncMutex(); + #currentActorInvocation?: CurrentActorInvocation; /** * Do not call this directly. @@ -99,6 +104,7 @@ export class ActorHandleRaw { actorResolutionState: ActorResolutionState, gatewayOptions: ActorGatewayOptions = {}, signal?: AbortSignal, + currentActorInvocation?: CurrentActorInvocation, ) { this.#client = client; this.#driver = driver; @@ -108,6 +114,7 @@ export class ActorHandleRaw { this.#params = params; this.#getParams = getParams; this.#signal = signal; + this.#currentActorInvocation = currentActorInvocation; } async #resolveConnectionParams(): Promise { @@ -320,6 +327,24 @@ export class ActorHandleRaw { name: opts.name, encoding: this.#encoding, }); + const invocation = this.#currentActorInvocation?.(); + const headers: Record = { + [HEADER_ENCODING]: this.#encoding, + }; + if (this.#params !== undefined) { + headers[HEADER_CONN_PARAMS] = JSON.stringify(this.#params); + } + if (invocation) { + headers[HEADER_RIVETKIT_RAY_ID] = invocation.rayId; + if (invocation.span) { + headers[HEADER_TRACEPARENT] = + invocation.span.traceparent; + if (invocation.span.tracestate) { + headers[HEADER_TRACESTATE] = + invocation.span.tracestate; + } + } + } const output = await sendHttpRequest< protocol.HttpActionRequest, protocol.HttpActionResponse, @@ -330,16 +355,7 @@ export class ActorHandleRaw { >({ url: `http://actor/action/${encodeURIComponent(opts.name)}`, method: "POST", - headers: { - [HEADER_ENCODING]: this.#encoding, - ...(this.#params !== undefined - ? { - [HEADER_CONN_PARAMS]: JSON.stringify( - this.#params, - ), - } - : {}), - }, + headers, body: opts.args, encoding: this.#encoding, customFetch: async (request) => diff --git a/rivetkit-typescript/packages/rivetkit/src/client/client.ts b/rivetkit-typescript/packages/rivetkit/src/client/client.ts index 439403871e..f717bc4751 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/client.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/client.ts @@ -3,6 +3,7 @@ import type { ActorQuery } from "@/client/query"; import type { Encoding } from "@/common/encoding"; import type { EngineControlClient } from "@/engine-client/driver"; import type { Registry } from "@/registry"; +import type { CurrentActorInvocation } from "@/registry/runtime"; import type { ActorActionFunction, ActorGatewayOptions } from "./actor-common"; import { type ActorConn, @@ -181,6 +182,13 @@ export const CREATE_ACTOR_CONN_PROXY = Symbol("createActorConnProxy"); * @template A The actors map type that defines the available actors. * @see {@link https://rivet.dev/docs/manage|Create & Manage Actors} */ +export interface ClientRawOptions { + encoding?: Encoding; + gateway?: ActorGatewayOptions; + /** Supplies the calling actor's invocation so actor-to-actor clients propagate its trace and ray ID. */ + currentActorInvocation?: CurrentActorInvocation; +} + export class ClientRaw { #disposed = false; @@ -189,19 +197,20 @@ export class ClientRaw { #driver: EngineControlClient; #encodingKind: Encoding; #gatewayOptions: ActorGatewayOptions; + #currentActorInvocation?: CurrentActorInvocation; /** * Creates an instance of Client. */ public constructor( driver: EngineControlClient, - encoding: Encoding | undefined, - gatewayOptions: ActorGatewayOptions = {}, + options: ClientRawOptions = {}, ) { this.#driver = driver; - this.#encodingKind = encoding ?? "bare"; - this.#gatewayOptions = gatewayOptions; + this.#encodingKind = options.encoding ?? "bare"; + this.#gatewayOptions = options.gateway ?? {}; + this.#currentActorInvocation = options.currentActorInvocation; } /** @@ -403,6 +412,7 @@ export class ClientRaw { actorQuery, this.#gatewayOptions, signal, + this.#currentActorInvocation, ); } @@ -459,9 +469,9 @@ export type AnyClient = Client>; export function createClientWithDriver>( driver: EngineControlClient, - config: { encoding?: Encoding; gateway?: ActorGatewayOptions } = {}, + options: ClientRawOptions = {}, ): Client { - const client = new ClientRaw(driver, config.encoding, config.gateway); + const client = new ClientRaw(driver, options); // Create proxy for accessing actors by name return new Proxy(client, { diff --git a/rivetkit-typescript/packages/rivetkit/src/common/actor-router-consts.ts b/rivetkit-typescript/packages/rivetkit/src/common/actor-router-consts.ts index edefc6cc8f..22c1efbcd6 100644 --- a/rivetkit-typescript/packages/rivetkit/src/common/actor-router-consts.ts +++ b/rivetkit-typescript/packages/rivetkit/src/common/actor-router-consts.ts @@ -20,6 +20,9 @@ export const HEADER_ACTOR_GENERATION = "x-rivet-actor-generation"; export const HEADER_ACTOR_KEY = "x-rivet-actor-key"; export const HEADER_RIVET_TOKEN = "x-rivet-token"; +export const HEADER_RIVET_RAY_ID = "x-rivet-ray-id"; +export const HEADER_TRACEPARENT = "traceparent"; +export const HEADER_TRACESTATE = "tracestate"; // MARK: Manager Gateway Headers export const HEADER_RIVET_TARGET = "x-rivet-target"; diff --git a/rivetkit-typescript/packages/rivetkit/src/common/actor-telemetry-context.ts b/rivetkit-typescript/packages/rivetkit/src/common/actor-telemetry-context.ts new file mode 100644 index 0000000000..bdc61ab950 --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/src/common/actor-telemetry-context.ts @@ -0,0 +1,21 @@ +/** Correlation owned by the currently executing Core actor invocation. */ +export interface ActorInvocationTraceContext { + /** Ray ID the invocation received, absent when nothing upstream supplied one. */ + readonly rayId?: string; + /** Core-owned invocation span context, absent when tracing is disabled. */ + readonly span?: ActorInvocationSpanContext; +} + +/** W3C span context of the Core invocation span. */ +export interface ActorInvocationSpanContext { + /** W3C trace identifier. */ + readonly traceId: string; + /** W3C span identifier for the Core invocation. */ + readonly spanId: string; + /** OpenTelemetry trace flags encoded as an integer. */ + readonly traceFlags: number; + /** Serialized W3C Trace Context for the Core invocation. */ + readonly traceparent: string; + /** Optional vendor trace state inherited by the Core invocation. */ + readonly tracestate?: string; +} diff --git a/rivetkit-typescript/packages/rivetkit/src/common/otel-context.ts b/rivetkit-typescript/packages/rivetkit/src/common/otel-context.ts new file mode 100644 index 0000000000..e7e89b1de2 --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/src/common/otel-context.ts @@ -0,0 +1,40 @@ +import { + type Context, + context, + createTraceState, + isSpanContextValid, + trace, +} from "@opentelemetry/api"; +import type { ActorInvocationSpanContext } from "./actor-telemetry-context"; + +/** + * Runs `run` with the Core invocation span as the active OpenTelemetry span, + * so application spans started inside an actor callback nest under it. With + * no span, or an invalid one, `run` executes unchanged. + */ +export function runWithActorInvocationSpan( + invocation: ActorInvocationSpanContext | undefined, + run: () => T, +): T { + if (!invocation) return run(); + + let parent: Context; + try { + const spanContext = { + traceId: invocation.traceId, + spanId: invocation.spanId, + traceFlags: invocation.traceFlags, + traceState: invocation.tracestate + ? createTraceState(invocation.tracestate) + : undefined, + isRemote: false, + }; + if (!isSpanContextValid(spanContext)) return run(); + parent = trace.setSpanContext(context.active(), spanContext); + } catch { + // Invalid telemetry must not prevent the action from running. + return run(); + } + + return context.with(parent, run); +} diff --git a/rivetkit-typescript/packages/rivetkit/src/engine-client/actor-http-client.ts b/rivetkit-typescript/packages/rivetkit/src/engine-client/actor-http-client.ts index 2fc62b41a8..563fd90f8c 100644 --- a/rivetkit-typescript/packages/rivetkit/src/engine-client/actor-http-client.ts +++ b/rivetkit-typescript/packages/rivetkit/src/engine-client/actor-http-client.ts @@ -4,6 +4,9 @@ import { HEADER_RIVET_SKIP_READY_WAIT, HEADER_RIVET_TARGET, HEADER_RIVET_TOKEN, + HEADER_RIVET_RAY_ID, + HEADER_TRACEPARENT, + HEADER_TRACESTATE, } from "@/common/actor-router-consts"; import { type GatewayRequestOptions, shouldSkipReadyWait } from "./driver"; @@ -55,6 +58,19 @@ function buildGuardHeaders( for (const [key, value] of Object.entries(runConfig.headers)) { headers.set(key, value as string); } + // Invocation headers are per action call. Apply the active request last so + // static client configuration cannot retain or override an earlier action. + for (const name of [ + HEADER_RIVET_RAY_ID, + HEADER_TRACEPARENT, + HEADER_TRACESTATE, + ]) { + headers.delete(name); + const value = actorRequest.headers.get(name); + if (value !== null) { + headers.set(name, value); + } + } // Add guard-specific headers if (runConfig.token) { headers.set(HEADER_RIVET_TOKEN, runConfig.token); diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts index 950c242de4..753d89389c 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts @@ -8,6 +8,8 @@ import type { HttpResponseBodyStream as NativeHttpResponseBodyStream, WebSocket as NativeWebSocket, } from "@rivetkit/rivetkit-napi"; +import type { ActorInvocationTraceContext } from "@/common/actor-telemetry-context"; +import { runWithActorInvocationSpan } from "@/common/otel-context"; import type { ActorContextHandle, ActorFactoryHandle, @@ -574,7 +576,23 @@ export class NapiCoreRuntime implements CoreRuntime { } runWithActorInvocationContext(ctx: ActorContextHandle, run: () => T): T { - return this.#invocationContext.run(asNativeActorContext(ctx), run); + const nativeCtx = asNativeActorContext(ctx); + const traceContext = this.#actorInvocationTraceContext(nativeCtx); + return this.#invocationContext.run(nativeCtx, () => + runWithActorInvocationSpan(traceContext?.span, run), + ); + } + + actorInvocationTraceContext( + ctx: ActorContextHandle, + ): ActorInvocationTraceContext | undefined { + return this.#actorInvocationTraceContext(asNativeActorContext(ctx)); + } + + #actorInvocationTraceContext( + ctx: NativeActorContext, + ): ActorInvocationTraceContext | undefined { + return ctx.invocationTraceContext() ?? undefined; } actorName(ctx: ActorContextHandle): string { diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts index a2e01aab35..f0d0fbb66e 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts @@ -3995,12 +3995,18 @@ export function buildNativeFactory( events: config.events, queues: config.queues, }; - const createClient = () => + const createClient = (ctx: ActorContextHandle) => createClientWithDriver( new RemoteEngineControlClient( convertRegistryConfigToClientConfig(registryConfig), ), - { encoding: "bare" }, + { + encoding: "bare", + currentActorInvocation: () => + callNativeSync(() => + runtime.actorInvocationTraceContext(ctx), + ), + }, ); const run = getRunFunction(config.run); const runHandlerCoordinator = @@ -4044,7 +4050,7 @@ export function buildNativeFactory( new ActorContextHandleAdapter( runtime, ctx, - createClient, + () => createClient(ctx), schemaConfig, databaseProvider, request, @@ -4063,7 +4069,7 @@ export function buildNativeFactory( runtime, ctx, conn, - createClient, + () => createClient(ctx), schemaConfig, databaseProvider, request, @@ -5373,7 +5379,7 @@ export function buildNativeFactory( runtime, ctx, conn, - createClient, + () => createClient(ctx), schemaConfig, databaseProvider, jsRequest, diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts index 76029c7cea..a66abb199a 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts @@ -1,3 +1,4 @@ +import type { ActorInvocationTraceContext } from "@/common/actor-telemetry-context"; import type { SqliteNativeMetrics, SqliteProfilingOptions, @@ -29,6 +30,11 @@ export interface RuntimeActorKeySegment { numberValue?: number; } +/** Resolves correlation at operation time so retained clients cannot freeze stale context. */ +export type CurrentActorInvocation = () => + | ActorInvocationTraceContext + | undefined; + export interface RuntimeHttpRequest { method: string; uri: string; @@ -540,7 +546,20 @@ export interface CoreRuntime { writes: RuntimeWorkflowKvWrite[], ): Promise; actorId(ctx: ActorContextHandle): string; + /** + * Runs one actor callback with `ctx` as the current invocation: operations + * on retained handles for the same actor resolve to it, and its Core span + * is the active OpenTelemetry span for the duration of `run`. + */ runWithActorInvocationContext(ctx: ActorContextHandle, run: () => T): T; + /** + * Correlation of the invocation currently executing for this actor, or + * `undefined` outside an invocation or after it finished. A sampled-out + * invocation can still expose valid span context for propagation. + */ + actorInvocationTraceContext( + ctx: ActorContextHandle, + ): ActorInvocationTraceContext | undefined; actorName(ctx: ActorContextHandle): string; actorKey(ctx: ActorContextHandle): RuntimeActorKeySegment[]; actorRegion(ctx: ActorContextHandle): string; diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts index c374433642..a89a9d3fd4 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts @@ -1,4 +1,5 @@ import { decodeBridgeRivetError, RivetError } from "@/actor/errors"; +import type { ActorInvocationTraceContext } from "@/common/actor-telemetry-context"; import type { WasmRuntimeBindings, WasmRuntimeConfig, @@ -543,6 +544,12 @@ export class WasmCoreRuntime implements CoreRuntime { return run(); } + actorInvocationTraceContext( + _ctx: ActorContextHandle, + ): ActorInvocationTraceContext | undefined { + return undefined; + } + actorName(ctx: ActorContextHandle): string { return callHandle(asWasmActorContext(ctx), "name"); } From 6fb50d50db6efc7bb13e758e052dc0b8a192c241 Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Tue, 1 Sep 2026 21:18:42 +0400 Subject: [PATCH 05/21] feat(rivetkit): record invocation metrics --- .../rivetkit-core/src/actor/context.rs | 22 +++- .../rivetkit-core/src/actor/metrics.rs | 119 +++++++++++++++++- .../packages/rivetkit-core/src/actor/task.rs | 25 ++-- .../packages/rivetkit-core/tests/task.rs | 1 + 4 files changed, 155 insertions(+), 12 deletions(-) diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs index 8f00eaf5e4..4a9e4cbc5e 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs @@ -41,7 +41,7 @@ use crate::actor::internal_storage; use crate::actor::kv::LegacyActorKv; use crate::actor::lifecycle_hooks::Reply; use crate::actor::messages::{ActorEvent, Request, StateDelta, WorkflowKvWrite}; -use crate::actor::metrics::ActorMetrics; +use crate::actor::metrics::{ActorMetrics, InvocationStatus, InvocationType}; use crate::actor::queue::{QueueInspectorUpdateCallback, QueueMetadata, QueueWaitActivityCallback}; use crate::actor::schedule::{InternalKeepAwakeCallback, LocalAlarmCallback}; use crate::actor::sleep::{CanSleep, SleepState}; @@ -323,8 +323,11 @@ impl ActorContext { let mut sql = sql; #[cfg(feature = "sqlite-local")] sql.set_profiling_config(config.sqlite_profiling.clone()); - let metrics = - ActorMetrics::new_with_sqlite_profiling(name.clone(), config.sqlite_profiling.clone()); + let metrics = ActorMetrics::new_for_actor( + name.clone(), + config.actions.iter().map(|action| action.name.clone()), + config.sqlite_profiling.clone(), + ); #[cfg(feature = "sqlite-local")] sql.set_vfs_metrics(Arc::new(metrics.clone())); let diagnostics = ActorDiagnostics::new(actor_id.clone()); @@ -1814,6 +1817,8 @@ impl ActorContext { let (reply_tx, reply_rx) = oneshot::channel(); let mut dispatch_error = None; + let mut invocation_status = InvocationStatus::Ok; + let mut action_ran = true; match ctx.try_send_actor_event( ActorEvent::Action { name: action.clone(), @@ -1828,6 +1833,7 @@ impl ActorContext { Ok(()) => match reply_rx.await { Ok(Ok(_)) => {} Ok(Err(error)) => { + invocation_status = InvocationStatus::from_error(&error); dispatch_error = Some(error); tracing::error!( error = ?dispatch_error.as_ref().expect("just assigned"), @@ -1837,6 +1843,7 @@ impl ActorContext { ); } Err(error) => { + invocation_status = InvocationStatus::Dropped; dispatch_error = Some(error.into()); tracing::error!( error = ?dispatch_error.as_ref().expect("just assigned"), @@ -1847,6 +1854,7 @@ impl ActorContext { } }, Err(error) => { + action_ran = false; dispatch_error = Some(error); tracing::error!( error = ?dispatch_error.as_ref().expect("just assigned"), @@ -1856,6 +1864,14 @@ impl ActorContext { ); } } + if action_ran { + ctx.metrics().record_invocation( + &action_name, + InvocationType::Scheduled, + invocation_status, + started_at.elapsed(), + ); + } ctx.finish_schedule_dispatch(&event_id, history_id, dispatch_error.as_ref()) .await; diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/metrics.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/metrics.rs index 97892f8734..eb4846d206 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/metrics.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/metrics.rs @@ -1,4 +1,4 @@ -use std::collections::BTreeMap; +use std::collections::{BTreeMap, BTreeSet}; use std::fmt; use std::sync::atomic::{AtomicBool, AtomicU8, Ordering}; #[cfg(feature = "sqlite-local")] @@ -23,6 +23,7 @@ use crate::time::Instant; const ACTOR_LABELS: &[&str] = &["actor_name"]; const INBOX_LABELS: &[&str] = &["actor_name", "inbox"]; const USER_TASK_LABELS: &[&str] = &["actor_name", "kind"]; +const INVOCATION_LABELS: &[&str] = &["actor_name", "action_name", "invocation_type", "result"]; const WORK_LABELS: &[&str] = &["actor_name", "kind"]; const SHUTDOWN_LABELS: &[&str] = &["actor_name", "reason"]; const STATE_MUTATION_LABELS: &[&str] = &["actor_name", "reason"]; @@ -129,9 +130,52 @@ pub(crate) struct StartupTimer { finished: bool, } +#[derive(Clone, Copy, Debug)] +pub(crate) enum InvocationType { + Action, + Scheduled, +} + +impl InvocationType { + fn as_label(self) -> &'static str { + match self { + Self::Action => "action", + Self::Scheduled => "scheduled", + } + } +} + +#[derive(Clone, Copy, Debug)] +pub(crate) enum InvocationStatus { + Ok, + Error, + Dropped, +} + +impl InvocationStatus { + /// Classifies a failed invocation, distinguishing dropped replies from user or runtime errors. + pub(crate) fn from_error(error: &anyhow::Error) -> Self { + let structured = rivet_error::RivetError::extract(error); + if structured.group() == "actor" && structured.code() == "dropped_reply" { + Self::Dropped + } else { + Self::Error + } + } + + fn as_label(self) -> &'static str { + match self { + Self::Ok => "ok", + Self::Error => "error", + Self::Dropped => "dropped", + } + } +} + #[derive(Debug)] struct ActorMetricInner { labels: ActorMetricLabels, + action_names: BTreeSet, #[cfg(feature = "sqlite-local")] sqlite_profiling: crate::SqliteProfilingConfig, #[cfg(feature = "sqlite-local")] @@ -182,6 +226,8 @@ struct ActorMetricCollectors { inbox_depth: IntGaugeVec, user_tasks_active: IntGaugeVec, user_task_duration_seconds: HistogramVec, + invocations_total: IntCounterVec, + invocation_duration_seconds: HistogramVec, http_requests_active: IntGaugeVec, keep_awake_active: IntGaugeVec, shutdown_tasks_active: IntGaugeVec, @@ -1091,6 +1137,25 @@ impl ActorMetricCollectors { USER_TASK_LABELS, ) .expect("create actor_user_task_duration_seconds histogram"); + let invocations_total = IntCounterVec::new( + Opts::new( + "rivetkit_actor_invocations_total", + "completed actor invocations", + ), + INVOCATION_LABELS, + ) + .expect("create actor_invocations_total counter"); + let invocation_duration_seconds = HistogramVec::new( + HistogramOpts::new( + "rivetkit_actor_invocation_duration_seconds", + "actor invocation duration in seconds", + ) + // Invocations land in the hundreds of microseconds, which the + // Prometheus default buckets collapse into their first bucket. + .buckets(rivet_metrics::MICRO_BUCKETS.to_vec()), + INVOCATION_LABELS, + ) + .expect("create actor_invocation_duration_seconds histogram"); let http_requests_active = IntGaugeVec::new( Opts::new( "rivetkit_actor_http_requests_active", @@ -1407,6 +1472,11 @@ impl ActorMetricCollectors { register_metric(&rivet_metrics::REGISTRY, inbox_depth.clone()); register_metric(&rivet_metrics::REGISTRY, user_tasks_active.clone()); register_metric(&rivet_metrics::REGISTRY, user_task_duration_seconds.clone()); + register_metric(&rivet_metrics::REGISTRY, invocations_total.clone()); + register_metric( + &rivet_metrics::REGISTRY, + invocation_duration_seconds.clone(), + ); register_metric(&rivet_metrics::REGISTRY, http_requests_active.clone()); register_metric(&rivet_metrics::REGISTRY, keep_awake_active.clone()); register_metric(&rivet_metrics::REGISTRY, shutdown_tasks_active.clone()); @@ -1521,6 +1591,8 @@ impl ActorMetricCollectors { inbox_depth, user_tasks_active, user_task_duration_seconds, + invocations_total, + invocation_duration_seconds, http_requests_active, keep_awake_active, shutdown_tasks_active, @@ -1584,12 +1656,25 @@ impl ActorMetricCollectors { impl ActorMetrics { pub(crate) fn new(actor_name: impl Into) -> Self { - Self::new_with_sqlite_profiling(actor_name, crate::SqliteProfilingConfig::default()) + Self::new_for_actor( + actor_name, + std::iter::empty(), + crate::SqliteProfilingConfig::default(), + ) } + #[cfg(all(test, feature = "sqlite-local"))] pub(crate) fn new_with_sqlite_profiling( actor_name: impl Into, _sqlite_profiling: crate::SqliteProfilingConfig, + ) -> Self { + Self::new_for_actor(actor_name, std::iter::empty(), _sqlite_profiling) + } + + pub(crate) fn new_for_actor( + actor_name: impl Into, + action_names: impl IntoIterator, + _sqlite_profiling: crate::SqliteProfilingConfig, ) -> Self { let labels = ActorMetricLabels { actor_name: actor_name.into(), @@ -1610,6 +1695,7 @@ impl ActorMetrics { Self { inner: Arc::new(ActorMetricInner { labels, + action_names: action_names.into_iter().collect(), #[cfg(feature = "sqlite-local")] sqlite_profiling: _sqlite_profiling, #[cfg(feature = "sqlite-local")] @@ -1864,6 +1950,35 @@ impl ActorMetrics { .observe(duration.as_secs_f64()); } + pub(crate) fn record_invocation( + &self, + action_name: &str, + invocation_type: InvocationType, + result: InvocationStatus, + duration: Duration, + ) { + let actor_labels = self.actor_labels(); + // Action names arrive from callers, so an undeclared one would mint a new + // label series per value. `_OTHER` is the fallback OpenTelemetry defines + // for exactly this, and it cannot collide with a declared action name. + let action_name = if self.inner.action_names.contains(action_name) { + action_name + } else { + "_OTHER" + }; + let labels = [ + actor_labels[0], + action_name, + invocation_type.as_label(), + result.as_label(), + ]; + METRICS.invocations_total.with_label_values(&labels).inc(); + METRICS + .invocation_duration_seconds + .with_label_values(&labels) + .observe(duration.as_secs_f64()); + } + pub(crate) fn set_http_requests_active(&self, count: usize) { let labels = self.actor_labels(); let mut state = self.inner.state.lock(); diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs index 315afc9dc5..0e2626d43f 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs @@ -51,7 +51,7 @@ use crate::actor::messages::{ ActorEvent, ActorHttpResponse, QueueSendResult, Request, SerializeStateReason, StateDelta, WorkflowKvWrite, }; -use crate::actor::metrics::startup_phase::StartupPhase; +use crate::actor::metrics::{InvocationStatus, InvocationType, startup_phase::StartupPhase}; use crate::actor::state::{PersistedActor, RequestSaveOpts}; use crate::actor::task_types::ShutdownKind; use crate::actor::work_registry::ActorWorkKind; @@ -909,6 +909,7 @@ impl ActorTask { } => { let invocation = crate::telemetry::ActionInvocationSpan::start(&self.ctx, &name, incoming); + let invocation_started_at = Instant::now(); let invocation_telemetry = invocation.telemetry(); tracing::info!( actor_id = %self.ctx.actor_id(), @@ -940,18 +941,21 @@ impl ActorTask { let actor_id = self.ctx.actor_id().to_owned(); let ctx = self.ctx.clone(); self.ctx.spawn_work(ActorWorkKind::Action, async move { - match tracked_reply_rx.await { + let (result, status) = match tracked_reply_rx.await { Ok(result) => { let result = result.map_err(|error| ctx.attach_actor_to_error(error)); - invocation.finish(result.as_ref().err()); + let status = match result.as_ref() { + Ok(_) => InvocationStatus::Ok, + Err(error) => InvocationStatus::from_error(error), + }; tracing::info!( actor_id = %actor_id, action_name = %action_name_for_log, ok = result.is_ok(), "actor task: tracked reply received, forwarding" ); - let _ = reply.send(result); + (result, status) } Err(_) => { tracing::warn!( @@ -962,10 +966,17 @@ impl ActorTask { let error = ctx.attach_actor_to_error( ActorLifecycleError::DroppedReply.build(), ); - invocation.finish(Some(&error)); - let _ = reply.send(Err(error)); + (Err(error), InvocationStatus::Dropped) } - } + }; + ctx.metrics().record_invocation( + &action_name_for_log, + InvocationType::Action, + status, + invocation_started_at.elapsed(), + ); + invocation.finish(result.as_ref().err()); + let _ = reply.send(result); }); } Err(error) => { diff --git a/rivetkit-rust/packages/rivetkit-core/tests/task.rs b/rivetkit-rust/packages/rivetkit-core/tests/task.rs index 99b62080f1..f59f835f5e 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/task.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/task.rs @@ -1,4 +1,5 @@ pub(crate) mod moved_tests { + use anyhow::anyhow; use std::collections::{BTreeMap, HashMap}; use std::path::PathBuf; use std::process::Command; From 1ebefff32199956ca30993c157bade0df1283cca Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Wed, 2 Sep 2026 02:58:45 +0400 Subject: [PATCH 06/21] feat(rivetkit): add trace context to logs --- .../packages/rivetkit/src/registry/native.ts | 27 ++++++++++++++++++- .../tests/fixtures/napi-runtime-server.ts | 7 +++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts index f0d0fbb66e..5bce0db6e5 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts @@ -8,6 +8,7 @@ import { type ActorCron, type ActorCronEveryOptions, type ActorCronSetOptions, + type ActorLogger, type ActorSchedule, CONN_STATE_MANAGER_SYMBOL, type CronFire, @@ -46,6 +47,7 @@ import { } from "@/client/client"; import { convertRegistryConfigToClientConfig } from "@/client/config"; import { HEADER_CONN_PARAMS } from "@/common/actor-router-consts"; +import type { ActorInvocationTraceContext } from "@/common/actor-telemetry-context"; import type { AnyDatabaseProvider, SqliteProfilingOptions, @@ -2720,6 +2722,7 @@ export class ActorContextHandleAdapter { #db?: unknown; #dispatchCancelToken?: CancellationTokenHandle; #kv?: NativeKvAdapter; + #log?: ActorLogger; #queue?: NativeQueueAdapter; #request?: Request; #schedule?: NativeScheduleAdapter; @@ -2956,8 +2959,30 @@ export class ActorContextHandleAdapter { return this.#connMap; } + #invocationTraceContext(): ActorInvocationTraceContext | undefined { + return callNativeSync(() => + this.#runtime.actorInvocationTraceContext(this.#ctx), + ); + } + get log() { - return logger(); + if (!this.#log) { + // Actor fields follow the camelCase used by the rest of the + // TypeScript logs. trace_id and span_id stay snake_case because + // that is what OpenTelemetry log correlation tooling looks for. + const invocation = this.#invocationTraceContext(); + this.#log = logger().child({ + actorId: this.actorId, + actorName: this.name, + actorKey: this.key, + ...(invocation?.rayId && { rayId: invocation.rayId }), + ...(invocation?.span && { + trace_id: invocation.span.traceId, + span_id: invocation.span.spanId, + }), + }); + } + return this.#log; } get abortSignal(): AbortSignal { diff --git a/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts b/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts index c601dbd3d5..a4588b60e3 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts @@ -60,6 +60,13 @@ const integrationActor = actor({ getCount: async (c) => { return c.state.count; }, + logContext: async (c, correlationToken: string) => { + c.log.warn( + { correlation_token: correlationToken }, + "native actor log context", + ); + return correlationToken; + }, validatedAction: async (_c, payload: { amount: number }) => { return payload.amount; }, From ad9b9bd81cbef048dd3e659beb9059a5c2b7570b Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Wed, 2 Sep 2026 00:56:54 +0400 Subject: [PATCH 07/21] feat(rivetkit): pass outbound trace context from client calls --- .../rivetkit/src/client/actor-handle.ts | 23 +++----- .../rivetkit/src/client/outbound-telemetry.ts | 57 +++++++++++++++++++ .../packages/rivetkit/src/client/raw-utils.ts | 3 + .../src/common/actor-telemetry-context.ts | 9 +++ .../rivetkit/src/common/otel-context.ts | 49 +++++++++++++++- 5 files changed, 125 insertions(+), 16 deletions(-) create mode 100644 rivetkit-typescript/packages/rivetkit/src/client/outbound-telemetry.ts diff --git a/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts b/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts index d3569adced..8e614ec88e 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts @@ -3,9 +3,6 @@ import type { ActorSpecifier } from "@/actor/errors"; import { HEADER_CONN_PARAMS, HEADER_ENCODING, - HEADER_RIVETKIT_RAY_ID, - HEADER_TRACEPARENT, - HEADER_TRACESTATE, } from "@/common/actor-router-consts"; import { isRequestLike } from "@/common/fetch-like"; import type * as protocol from "@/common/client-protocol"; @@ -57,6 +54,7 @@ import { type ClientRaw, CREATE_ACTOR_CONN_PROXY } from "./client"; import { ActorError, isSchedulingError } from "./errors"; import { retryOnLifecycleBoundary } from "./lifecycle-errors"; import { logger } from "./log"; +import { outboundTelemetryHeaders } from "./outbound-telemetry"; import { createQueueSender, type QueueSendNoWaitOptions, @@ -327,24 +325,15 @@ export class ActorHandleRaw { name: opts.name, encoding: this.#encoding, }); - const invocation = this.#currentActorInvocation?.(); const headers: Record = { [HEADER_ENCODING]: this.#encoding, + ...outboundTelemetryHeaders( + this.#currentActorInvocation?.(), + ), }; if (this.#params !== undefined) { headers[HEADER_CONN_PARAMS] = JSON.stringify(this.#params); } - if (invocation) { - headers[HEADER_RIVETKIT_RAY_ID] = invocation.rayId; - if (invocation.span) { - headers[HEADER_TRACEPARENT] = - invocation.span.traceparent; - if (invocation.span.tracestate) { - headers[HEADER_TRACESTATE] = - invocation.span.tracestate; - } - } - } const output = await sendHttpRequest< protocol.HttpActionRequest, protocol.HttpActionResponse, @@ -696,6 +685,9 @@ export class ActorHandleRaw { skipReadyWait, }, ); + const telemetryHeaders = outboundTelemetryHeaders( + this.#currentActorInvocation?.(), + ); for (let attempt = 0; attempt < maxAttempts; attempt++) { let actorId: string | undefined; @@ -714,6 +706,7 @@ export class ActorHandleRaw { clonesInputBody ? input.clone() : input, requestInit, gatewayOptions, + telemetryHeaders, ); const retry = await this.#shouldRetryRawFetchResponse( response, diff --git a/rivetkit-typescript/packages/rivetkit/src/client/outbound-telemetry.ts b/rivetkit-typescript/packages/rivetkit/src/client/outbound-telemetry.ts new file mode 100644 index 0000000000..abe93b056c --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/src/client/outbound-telemetry.ts @@ -0,0 +1,57 @@ +import { + HEADER_RIVET_RAY_ID, + HEADER_TRACEPARENT, + HEADER_TRACESTATE, +} from "@/common/actor-router-consts"; +import type { ActorInvocationTraceContext } from "@/common/actor-telemetry-context"; +import { readActiveRayId, readActiveTraceHeaders } from "@/common/otel-context"; + +/** + * Headers that carry a caller's ray ID and trace context into an actor. One + * rule for every kind of outbound call, so an action call and a queue send + * made from the same place land in the same trace under the same ray ID. + * + * The ray ID is the calling invocation's when the caller is itself inside an + * actor, else the one placed in OpenTelemetry baggage by the surrounding + * request handler. The trace context is the application span active in this + * JavaScript context, else the calling actor's own Core invocation span. + */ +export function outboundTelemetryHeaders( + invocation: ActorInvocationTraceContext | undefined, +): Record { + const headers: Record = {}; + const rayId = invocation?.rayId ?? readActiveRayId(); + if (rayId) { + headers[HEADER_RIVET_RAY_ID] = rayId; + } + const traceHeaders = readActiveTraceHeaders() ?? invocation?.span; + if (traceHeaders) { + headers[HEADER_TRACEPARENT] = traceHeaders.traceparent; + if (traceHeaders.tracestate) { + headers[HEADER_TRACESTATE] = traceHeaders.tracestate; + } + } + return headers; +} + +/** + * Adds outbound telemetry headers to a request whose caller may have set + * some already. A header the caller set wins. `traceparent` and `tracestate` + * describe one span between them, so when the caller set either, both stay + * as the caller set them and neither is added. + */ +export function addOutboundTelemetryHeaders( + headers: Headers, + telemetry: Record, +): void { + const callerSetTraceContext = + headers.has(HEADER_TRACEPARENT) || headers.has(HEADER_TRACESTATE); + for (const [name, value] of Object.entries(telemetry)) { + const isTraceContext = + name === HEADER_TRACEPARENT || name === HEADER_TRACESTATE; + if (isTraceContext ? callerSetTraceContext : headers.has(name)) { + continue; + } + headers.set(name, value); + } +} diff --git a/rivetkit-typescript/packages/rivetkit/src/client/raw-utils.ts b/rivetkit-typescript/packages/rivetkit/src/client/raw-utils.ts index de89ec1b96..b71d45a05f 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/raw-utils.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/raw-utils.ts @@ -12,6 +12,7 @@ import type { } from "@/engine-client/driver"; import { ActorError } from "./errors"; import { logger } from "./log"; +import { addOutboundTelemetryHeaders } from "./outbound-telemetry"; /** Buffer a one-shot ReadableStream body to bytes so every retry attempt can re-send it. */ export async function prepareRetryableInit( @@ -37,6 +38,7 @@ export async function rawHttpFetch( input: string | URL | Request, init?: RequestInit, options: GatewayRequestOptions = {}, + telemetryHeaders: Record = {}, ): Promise { // Extract path and merge init options let path: string; @@ -111,6 +113,7 @@ export async function rawHttpFetch( if (params) { proxyRequestHeaders.set(HEADER_CONN_PARAMS, JSON.stringify(params)); } + addOutboundTelemetryHeaders(proxyRequestHeaders, telemetryHeaders); // Forward the request to the actor const proxyRequest = new Request(url, { diff --git a/rivetkit-typescript/packages/rivetkit/src/common/actor-telemetry-context.ts b/rivetkit-typescript/packages/rivetkit/src/common/actor-telemetry-context.ts index bdc61ab950..fec4e65275 100644 --- a/rivetkit-typescript/packages/rivetkit/src/common/actor-telemetry-context.ts +++ b/rivetkit-typescript/packages/rivetkit/src/common/actor-telemetry-context.ts @@ -19,3 +19,12 @@ export interface ActorInvocationSpanContext { /** Optional vendor trace state inherited by the Core invocation. */ readonly tracestate?: string; } + +/** Formats a W3C `traceparent` header from its span identifiers. */ +export function formatTraceparent( + traceId: string, + spanId: string, + traceFlags: number, +): string { + return `00-${traceId}-${spanId}-${traceFlags.toString(16).padStart(2, "0")}`; +} diff --git a/rivetkit-typescript/packages/rivetkit/src/common/otel-context.ts b/rivetkit-typescript/packages/rivetkit/src/common/otel-context.ts index e7e89b1de2..b478afcd6f 100644 --- a/rivetkit-typescript/packages/rivetkit/src/common/otel-context.ts +++ b/rivetkit-typescript/packages/rivetkit/src/common/otel-context.ts @@ -3,9 +3,56 @@ import { context, createTraceState, isSpanContextValid, + propagation, trace, } from "@opentelemetry/api"; -import type { ActorInvocationSpanContext } from "./actor-telemetry-context"; +import { + type ActorInvocationSpanContext, + formatTraceparent, +} from "./actor-telemetry-context"; + +/** W3C headers derived from the active JavaScript OTel context. */ +export interface ActiveTraceHeaders { + /** W3C Trace Context identifying the active trace and span. */ + readonly traceparent: string; + /** Optional vendor trace state associated with the active span. */ + readonly tracestate?: string; +} + +/** Returns the active W3C trace context, when an OTel provider has installed one. */ +export function readActiveTraceHeaders(): ActiveTraceHeaders | undefined { + const spanContext = trace.getSpanContext(context.active()); + if (!spanContext || !isSpanContextValid(spanContext)) return undefined; + + const tracestate = spanContext.traceState?.serialize(); + return { + traceparent: formatTraceparent( + spanContext.traceId, + spanContext.spanId, + spanContext.traceFlags, + ), + ...(tracestate ? { tracestate } : {}), + }; +} + +/** W3C Baggage key that carries a ray ID through application code. */ +export const RAY_BAGGAGE_KEY = "rivet.ray.id"; + +const RAY_ID_PATTERN = /^[A-Za-z0-9_-]{1,128}$/; + +/** + * Returns the ray ID carried in the active OpenTelemetry baggage, so a request + * handler that received a ray ID can pass it to the actors it calls. The value is + * bounded by the same rule Core applies at the actor edge: 1 to 128 characters + * of `[A-Za-z0-9_-]`. Anything else counts as absent. + */ +export function readActiveRayId(): string | undefined { + const rayId = propagation + .getBaggage(context.active()) + ?.getEntry(RAY_BAGGAGE_KEY)?.value; + if (rayId === undefined || !RAY_ID_PATTERN.test(rayId)) return undefined; + return rayId; +} /** * Runs `run` with the Core invocation span as the active OpenTelemetry span, From a93aff770324f2cf2ff6a4276b23b560803225a7 Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Wed, 9 Sep 2026 22:48:30 +0400 Subject: [PATCH 08/21] feat(rivetkit): send trace context from the Rust client --- Cargo.lock | 3 + .../packages/client-protocol/src/lib.rs | 1 + .../client-protocol/src/telemetry_headers.rs | 37 ++++++ rivetkit-rust/packages/client/Cargo.toml | 3 + rivetkit-rust/packages/client/src/client.rs | 15 +++ .../packages/client/src/remote_manager.rs | 89 +++++++++++++ rivetkit-rust/packages/client/tests/bare.rs | 123 ++++++++++++++++++ .../packages/rivetkit-core/src/telemetry.rs | 31 ++--- 8 files changed, 283 insertions(+), 19 deletions(-) create mode 100644 rivetkit-rust/packages/client-protocol/src/telemetry_headers.rs diff --git a/Cargo.lock b/Cargo.lock index 306d597014..bf7a273262 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6265,6 +6265,8 @@ dependencies = [ "bytes", "fs_extra", "futures-util", + "opentelemetry", + "opentelemetry_sdk", "parking_lot", "portpicker", "reqwest 0.12.22", @@ -6279,6 +6281,7 @@ dependencies = [ "tokio-test", "tokio-tungstenite", "tracing", + "tracing-opentelemetry", "tracing-subscriber", "tungstenite", "urlencoding", diff --git a/rivetkit-rust/packages/client-protocol/src/lib.rs b/rivetkit-rust/packages/client-protocol/src/lib.rs index 70f35bb01b..faf8453c77 100644 --- a/rivetkit-rust/packages/client-protocol/src/lib.rs +++ b/rivetkit-rust/packages/client-protocol/src/lib.rs @@ -1,4 +1,5 @@ pub mod generated; +pub mod telemetry_headers; pub mod versioned; // Re-export latest. diff --git a/rivetkit-rust/packages/client-protocol/src/telemetry_headers.rs b/rivetkit-rust/packages/client-protocol/src/telemetry_headers.rs new file mode 100644 index 0000000000..dc26a1e1cd --- /dev/null +++ b/rivetkit-rust/packages/client-protocol/src/telemetry_headers.rs @@ -0,0 +1,37 @@ +//! Headers that carry a caller's ray ID and W3C trace context into an actor, +//! shared by every Rust peer of the actor HTTP surface: the runtime that reads +//! them and the clients that send them. + +/// Bounded correlation string that follows work across actors, surviving +/// sampling and trace-root boundaries. +pub const HEADER_RIVET_RAY_ID: &str = "x-rivet-ray-id"; +pub const HEADER_TRACEPARENT: &str = "traceparent"; +pub const HEADER_TRACESTATE: &str = "tracestate"; + +/// W3C Baggage key that carries a ray ID through application code, so a request +/// handler that received a ray ID can hand it to the actors it calls. +pub const RAY_BAGGAGE_KEY: &str = "rivet.ray.id"; + +const RAY_ID_MAX_LEN: usize = 128; + +/// Returns `value` when it is a ray ID the runtime accepts: 1 to 128 characters +/// of `[A-Za-z0-9_-]`. The value arrives from a caller, so anything else +/// counts as absent. +pub fn bounded_ray_id(value: &str) -> Option<&str> { + let valid = !value.is_empty() + && value.len() <= RAY_ID_MAX_LEN + && value + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_')); + if valid { Some(value) } else { None } +} + +/// Formats a version 00 W3C `traceparent` from the parts of a span context, +/// so every Rust peer writes the header the same way. +pub fn format_traceparent( + trace_id: impl std::fmt::Display, + span_id: impl std::fmt::Display, + trace_flags: u8, +) -> String { + format!("00-{trace_id}-{span_id}-{trace_flags:02x}") +} diff --git a/rivetkit-rust/packages/client/Cargo.toml b/rivetkit-rust/packages/client/Cargo.toml index edcadab797..7dc2cf96f8 100644 --- a/rivetkit-rust/packages/client/Cargo.toml +++ b/rivetkit-rust/packages/client/Cargo.toml @@ -13,6 +13,7 @@ anyhow = "1.0" base64 = "0.22.1" bytes = { workspace = true } futures-util = "0.3.31" +opentelemetry = { version = "0.28", default-features = false, features = ["trace"] } parking_lot.workspace = true reqwest = { version = "0.12.12", default-features = false, features = ["json", "charset", "http2", "macos-system-configuration", "rustls-tls-native-roots", "rustls-tls-webpki-roots"] } rivetkit-client-protocol.workspace = true @@ -24,12 +25,14 @@ serde_json = "1.0" tokio = { version = "1", features = ["full"] } tokio-tungstenite = { version = "0.26.1", features = ["rustls-tls-native-roots", "rustls-tls-webpki-roots", "handshake"] } tracing = "0.1.41" +tracing-opentelemetry = { version = "0.29", default-features = false } tungstenite = "0.26.2" urlencoding = "2.1.3" vbare = "0.0.4" [dev-dependencies] axum = { workspace = true, features = ["ws"] } +opentelemetry_sdk = { version = "0.28", default-features = false, features = ["trace"] } tracing-subscriber = { version = "0.3.19", features = ["env-filter", "std", "registry"]} tempfile = "3.10.1" tokio-test = "0.4.3" diff --git a/rivetkit-rust/packages/client/src/client.rs b/rivetkit-rust/packages/client/src/client.rs index 79a8a622b4..7578be94fb 100644 --- a/rivetkit-rust/packages/client/src/client.rs +++ b/rivetkit-rust/packages/client/src/client.rs @@ -46,6 +46,9 @@ pub struct ClientConfig { pub encoding: EncodingKind, pub transport: TransportKind, pub headers: Option>, + /// Ray sent on every actor request made outside an active baggage + /// context, so work this client causes can be found by one string. + pub ray_id: Option, pub max_input_size: Option, pub disable_metadata_lookup: bool, } @@ -60,6 +63,7 @@ impl ClientConfig { encoding: EncodingKind::Bare, transport: TransportKind::WebSocket, headers: None, + ray_id: None, max_input_size: None, disable_metadata_lookup: false, } @@ -107,6 +111,16 @@ impl ClientConfig { self } + /// Sets the ray sent on every actor request. It must be 1 to 128 + /// characters of `[A-Za-z0-9_-]`, the same bound the runtime applies; a + /// value outside it is dropped with a warning when the client is built. + /// A ray carried in the active OpenTelemetry baggage under `rivet.ray.id` + /// takes precedence per call. + pub fn ray_id(mut self, ray_id: impl Into) -> Self { + self.ray_id = Some(ray_id.into()); + self + } + pub fn max_input_size(mut self, max_input_size: usize) -> Self { self.max_input_size = Some(max_input_size); self @@ -153,6 +167,7 @@ impl Client { config.namespace, config.pool_name, config.headers, + config.ray_id, config.max_input_size, config.disable_metadata_lookup, ); diff --git a/rivetkit-rust/packages/client/src/remote_manager.rs b/rivetkit-rust/packages/client/src/remote_manager.rs index adf46e666d..780078fc5e 100644 --- a/rivetkit-rust/packages/client/src/remote_manager.rs +++ b/rivetkit-rust/packages/client/src/remote_manager.rs @@ -1,15 +1,22 @@ use anyhow::{anyhow, Context, Result}; use base64::{engine::general_purpose, engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; use bytes::Bytes; +use opentelemetry::baggage::BaggageExt as _; +use opentelemetry::trace::TraceContextExt as _; use reqwest::{ header::{HeaderMap, HeaderName, HeaderValue, USER_AGENT}, Method, }; +use rivetkit_client_protocol::telemetry_headers::{ + HEADER_RIVET_RAY_ID, HEADER_TRACEPARENT, HEADER_TRACESTATE, RAY_BAGGAGE_KEY, bounded_ray_id, + format_traceparent, +}; use serde::{Deserialize, Serialize}; use serde_cbor; use std::{collections::HashMap, str::FromStr, sync::Arc}; use tokio::sync::OnceCell; use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tracing_opentelemetry::OpenTelemetrySpanExt as _; use crate::{ common::{ @@ -29,6 +36,7 @@ pub struct RemoteManager { namespace: String, pool_name: String, headers: HashMap, + ray_id: Option, max_input_size: usize, disable_metadata_lookup: bool, resolved_config: Arc>, @@ -126,6 +134,7 @@ impl RemoteManager { namespace: default_namespace(), pool_name: default_pool_name(), headers: HashMap::new(), + ray_id: None, max_input_size: default_max_input_size(), disable_metadata_lookup: false, resolved_config: Arc::new(OnceCell::new()), @@ -139,15 +148,27 @@ impl RemoteManager { namespace: Option, pool_name: Option, headers: Option>, + ray_id: Option, max_input_size: Option, disable_metadata_lookup: bool, ) -> Self { + let ray_id = ray_id.and_then(|ray_id| { + let bounded = bounded_ray_id(&ray_id).map(str::to_owned); + if bounded.is_none() { + tracing::warn!( + len = ray_id.len(), + "dropping configured ray id; it must be 1 to 128 characters of [A-Za-z0-9_-]" + ); + } + bounded + }); Self { endpoint, token, namespace: namespace.unwrap_or_else(default_namespace), pool_name: pool_name.unwrap_or_else(default_pool_name), headers: headers.unwrap_or_default(), + ray_id, max_input_size: max_input_size.unwrap_or_else(default_max_input_size), disable_metadata_lookup, resolved_config: Arc::new(OnceCell::new()), @@ -214,6 +235,11 @@ impl RemoteManager { req = req.header(USER_AGENT, USER_AGENT_VALUE); for (key, value) in &self.headers { + // Ray ID and trace context are per call, so a configured value cannot + // pin stale context on every request. Matches the TypeScript client. + if is_telemetry_header(key) { + continue; + } let name = HeaderName::from_str(key) .with_context(|| format!("invalid configured header name `{key}`"))?; let value = HeaderValue::from_str(value) @@ -493,6 +519,19 @@ impl RemoteManager { let mut req = self.apply_common_headers_with(builder, &config)?; + // Per-call context wins over configured headers, and headers the + // caller passed for this request win over both, matching the + // TypeScript client. + let mut headers = headers; + let caller_set_trace_context = + headers.contains_key(HEADER_TRACEPARENT) || headers.contains_key(HEADER_TRACESTATE); + for (name, value) in self.telemetry_headers()? { + let is_trace_context = name == HEADER_TRACEPARENT || name == HEADER_TRACESTATE; + if is_trace_context && caller_set_trace_context { + continue; + } + headers.entry(name).or_insert(value); + } req = req.headers(headers); if let Some(body_data) = body { @@ -503,6 +542,50 @@ impl RemoteManager { Ok(res) } + /// Headers that carry the caller's trace context and ray ID into the actor, + /// read from the `tracing` span current at the call. Without a registered + /// OpenTelemetry layer the span carries no context and only a configured + /// ray ID is sent. + fn telemetry_headers(&self) -> Result> { + let mut headers = Vec::with_capacity(3); + let context = tracing::Span::current().context(); + let span = context.span(); + let span_context = span.span_context(); + if span_context.is_valid() { + let traceparent = format_traceparent( + span_context.trace_id(), + span_context.span_id(), + span_context.trace_flags().to_u8(), + ); + headers.push(( + HeaderName::from_static(HEADER_TRACEPARENT), + HeaderValue::from_str(&traceparent).context("format traceparent header")?, + )); + let tracestate = span_context.trace_state().header(); + if !tracestate.is_empty() { + headers.push(( + HeaderName::from_static(HEADER_TRACESTATE), + HeaderValue::from_str(&tracestate).context("format tracestate header")?, + )); + } + } + let baggage_ray = context + .baggage() + .get(RAY_BAGGAGE_KEY) + .map(|value| value.as_str().into_owned()); + let ray_id = baggage_ray + .as_deref() + .and_then(bounded_ray_id) + .or(self.ray_id.as_deref()); + if let Some(ray_id) = ray_id { + headers.push(( + HeaderName::from_static(HEADER_RIVET_RAY_ID), + HeaderValue::from_str(ray_id).context("format ray id header")?, + )); + } + Ok(headers) + } + pub fn gateway_url(&self, query: &ActorQuery) -> Result { let config = self.base_config(); match query { @@ -823,3 +906,9 @@ fn default_pool_name() -> String { fn default_max_input_size() -> usize { 4 * 1024 } + +fn is_telemetry_header(name: &str) -> bool { + [HEADER_RIVET_RAY_ID, HEADER_TRACEPARENT, HEADER_TRACESTATE] + .iter() + .any(|header| header.eq_ignore_ascii_case(name)) +} diff --git a/rivetkit-rust/packages/client/tests/bare.rs b/rivetkit-rust/packages/client/tests/bare.rs index 6900a4b91d..e2f6e949ba 100644 --- a/rivetkit-rust/packages/client/tests/bare.rs +++ b/rivetkit-rust/packages/client/tests/bare.rs @@ -64,6 +64,11 @@ struct ConfigHeaderTestState { saw_raw_websocket: Arc, } +#[derive(Clone)] +struct TelemetryHeaderState { + seen: Arc>>>, +} + #[derive(Clone)] struct MetadataLookupState { saw_metadata: Arc, @@ -547,6 +552,96 @@ async fn config_headers_are_sent_on_http_and_websocket_paths() { server.abort(); } +#[tokio::test] +async fn active_span_and_ray_reach_the_actor_over_configured_headers() { + use opentelemetry::trace::{TraceContextExt as _, TracerProvider as _}; + use opentelemetry::{Context as OtelContext, KeyValue, baggage::BaggageExt as _}; + use tracing::Instrument as _; + use tracing_opentelemetry::OpenTelemetrySpanExt as _; + use tracing_subscriber::layer::SubscriberExt as _; + + // Register the application tracing layer without an exporter. + let provider = opentelemetry_sdk::trace::SdkTracerProvider::builder().build(); + let subscriber = tracing_subscriber::registry() + .with(tracing_opentelemetry::layer().with_tracer(provider.tracer("test"))); + let _subscriber = tracing::subscriber::set_default(subscriber); + + let state = TelemetryHeaderState { + seen: Arc::new(std::sync::Mutex::new(Vec::new())), + }; + let app = Router::new() + .route( + "/gateway/{actor_id}/action/{action}", + post(action_capturing_telemetry_headers), + ) + .with_state(state.clone()); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + + // A configured `traceparent` must not pin stale context, and the + // configured ray applies only when no baggage carries one. + let client = Client::new( + ClientConfig::new(endpoint(addr)) + .disable_metadata_lookup(true) + .header( + "traceparent", + "00-000000000000000000000000000000aa-00000000000000aa-01", + ) + .header("x-rivet-ray-id", "stale-configured-ray") + .ray_id("ray-from-config"), + ); + let actor = client + .get_or_create( + "counter", + vec!["telemetry-headers".to_owned()], + GetOrCreateOptions::default(), + ) + .unwrap(); + + let span = tracing::info_span!("caller"); + let expected_span = span.context().span().span_context().clone(); + let output = async { actor.action("increment", vec![json!(2)]).await.unwrap() } + .instrument(span) + .await; + assert_eq!(output, json!({ "count": 3 })); + + let baggage = + OtelContext::current_with_baggage(vec![KeyValue::new("rivet.ray.id", "ray-from-baggage")]); + let baggage_span = tracing::info_span!(parent: None, "handler"); + baggage_span.set_parent(baggage); + let output = async { actor.action("increment", vec![json!(2)]).await.unwrap() } + .instrument(baggage_span) + .await; + assert_eq!(output, json!({ "count": 3 })); + + let seen = state.seen.lock().unwrap().clone(); + assert_eq!(seen.len(), 2); + assert_eq!( + seen[0].get("traceparent").map(String::as_str), + Some( + format!( + "00-{}-{}-01", + expected_span.trace_id(), + expected_span.span_id() + ) + .as_str() + ) + ); + assert_eq!( + seen[0].get("x-rivet-ray-id").map(String::as_str), + Some("ray-from-config") + ); + assert_eq!( + seen[1].get("x-rivet-ray-id").map(String::as_str), + Some("ray-from-baggage") + ); + + server.abort(); +} + #[tokio::test] async fn max_input_size_checks_raw_query_input_before_base64url_encoding() { let client = Client::new( @@ -1169,6 +1264,34 @@ async fn action_with_config_header( .await } +async fn action_capturing_telemetry_headers( + State(state): State, + Path((actor_id, action_name)): Path<(String, String)>, + headers: HeaderMap, + body: Bytes, +) -> impl IntoResponse { + state.seen.lock().unwrap().push( + headers + .iter() + .filter_map(|(name, value)| { + Some((name.as_str().to_owned(), value.to_str().ok()?.to_owned())) + }) + .collect(), + ); + action( + State(TestState { + saw_bare_action: Arc::new(AtomicBool::new(false)), + saw_bare_queue: Arc::new(AtomicBool::new(false)), + saw_raw_fetch: Arc::new(AtomicBool::new(false)), + saw_raw_websocket: Arc::new(AtomicBool::new(false)), + }), + Path((actor_id, action_name)), + headers, + body, + ) + .await +} + async fn action_for_disable_metadata( Path((actor_id, action_name)): Path<(String, String)>, headers: HeaderMap, diff --git a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs index 013db97002..32dca43a70 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs @@ -22,8 +22,7 @@ pub struct IncomingInvocationContext { remote_parent: Option, } -/// Header carrying the caller's ray into an actor. -pub(crate) const HEADER_RIVETKIT_RAY_ID: &str = "x-rivetkit-ray-id"; +pub(crate) use rivetkit_client_protocol::telemetry_headers::HEADER_RIVET_RAY_ID; impl IncomingInvocationContext { pub(crate) fn from_headers( @@ -43,28 +42,22 @@ impl IncomingInvocationContext { pub(crate) fn from_http_headers(headers: &http::HeaderMap) -> Self { Self::from_headers( invocation_ray_id(headers), - headers.get("traceparent").and_then(|value| value.to_str().ok()), - headers.get("tracestate").and_then(|value| value.to_str().ok()), + headers + .get("traceparent") + .and_then(|value| value.to_str().ok()), + headers + .get("tracestate") + .and_then(|value| value.to_str().ok()), ) } } -/// Reads the caller's ray id. The header is untrusted, so it is bounded to -/// 128 characters of `[A-Za-z0-9_-]`; anything else counts as absent and the -/// invocation mints a fresh ray instead. +/// Reads the caller's ray ID. The header is untrusted, so it is bounded by +/// the rule shared with the clients that send it; anything else counts as +/// absent. fn invocation_ray_id(headers: &http::HeaderMap) -> Option { - headers - .get(HEADER_RIVETKIT_RAY_ID)? - .to_str() - .ok() - .filter(|value| { - !value.is_empty() - && value.len() <= 128 - && value - .bytes() - .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_')) - }) - .map(str::to_owned) + let value = headers.get(HEADER_RIVET_RAY_ID)?.to_str().ok()?; + rivetkit_client_protocol::telemetry_headers::bounded_ray_id(value).map(str::to_owned) } /// The single root span for one client action invocation. From 563a438105f68f6dbede458a8bd51b3f14581b24 Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Wed, 2 Sep 2026 09:40:49 +0400 Subject: [PATCH 09/21] feat(rivetkit): parent SQLite and schedule spans to active application context --- .../rivetkit-core/src/actor/context.rs | 14 ++++ .../packages/rivetkit-core/src/telemetry.rs | 43 +++++++++-- .../packages/rivetkit-napi/index.d.ts | 5 ++ .../rivetkit-napi/src/actor_context.rs | 16 ++++ .../rivetkit/src/registry/napi-runtime.ts | 73 ++++++++++++------- .../rivetkit/tests/runtime-parity.test.ts | 8 ++ 6 files changed, 126 insertions(+), 33 deletions(-) diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs index 4a9e4cbc5e..e7e10c00b9 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs @@ -258,6 +258,20 @@ impl ActorContext { self } + /// Returns a handle for the same invocation whose spans parent to the + /// application span the host runtime has active, given as W3C + /// `traceparent` and `tracestate`. A handle that serves no invocation is + /// returned unchanged, because it opens no spans. + #[doc(hidden)] + pub fn with_application_span(&self, traceparent: Option<&str>, tracestate: Option<&str>) -> Self { + Self( + self.0.clone(), + self.1 + .as_ref() + .map(|telemetry| telemetry.with_application_span(traceparent, tracestate)), + ) + } + /// Returns the SQLite handle bound to this handle's invocation. pub fn invocation_sql(&self) -> crate::actor::sqlite::SqliteDb { self.0.sql.clone().with_invocation_telemetry(self.1.clone()) diff --git a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs index 32dca43a70..6d201a0172 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs @@ -67,9 +67,14 @@ pub(crate) struct ActionInvocationSpan { } /// Opaque invocation context carried across foreign-runtime adapters. +/// +/// The second field is the application span the host runtime had active when +/// it resolved this handle. Core cannot see the host's span stack, so a span +/// Core opens through this handle parents there when it is set and to the +/// invocation span otherwise. Every clone of a handle shares one invocation. #[doc(hidden)] #[derive(Clone, Debug)] -pub struct ActorInvocationTelemetry(Arc); +pub struct ActorInvocationTelemetry(Arc, Option); /// Identity fields that do not change while an actor is alive. Built once per /// actor and shared by every invocation, so starting one does not re-allocate @@ -239,13 +244,30 @@ impl ActorInvocationTelemetry { span: Option, identity: Arc, ) -> Self { - Self(Arc::new(InvocationInner { - ray_id, - span: Mutex::new(span), - finished: AtomicBool::new(false), - pending_work: AtomicUsize::new(0), - identity, - })) + Self( + Arc::new(InvocationInner { + ray_id, + span: Mutex::new(span), + finished: AtomicBool::new(false), + pending_work: AtomicUsize::new(0), + identity, + }), + None, + ) + } + + /// Returns a handle for the same invocation whose spans parent to the + /// application span identified by `traceparent` and `tracestate`. Invalid + /// or absent context yields a handle that parents to the invocation span. + pub(crate) fn with_application_span( + &self, + traceparent: Option<&str>, + tracestate: Option<&str>, + ) -> Self { + Self( + self.0.clone(), + parse_remote_parent(traceparent, tracestate), + ) } /// Registers work that outlives the reply, so the invocation span stays @@ -308,6 +330,11 @@ impl ActorInvocationTelemetry { otel.status_code = tracing::field::Empty, error.type = tracing::field::Empty, ); + if let Some(application_span) = &self.1 { + span.set_parent( + opentelemetry::Context::new().with_remote_span_context(application_span.clone()), + ); + } Some(SqliteOperationSpan { span: Some(span) }) } diff --git a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts index 07102087e1..9a2de32aef 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts +++ b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts @@ -320,6 +320,11 @@ export declare class ActorContext { kv(): Kv sql(): JsNativeDatabase sameActorInstance(other: ActorContext): boolean + /** + * Returns a handle for the same invocation whose Core spans parent to the + * application span active in JavaScript, given as W3C headers. + */ + withApplicationSpan(traceparent?: string | undefined | null, tracestate?: string | undefined | null): ActorContext invocationTraceContext(): JsActorInvocationTraceContext | null provisionActorRuntimeSocket(): Promise schedule(): Schedule diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs b/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs index 917e159367..48d3c6acaa 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs @@ -322,6 +322,22 @@ impl ActorContext { self.inner.is_same_instance(&other.inner) } + /// Returns a handle for the same invocation whose Core spans parent to the + /// application span active in JavaScript, given as W3C headers. + #[napi] + pub fn with_application_span( + &self, + traceparent: Option, + tracestate: Option, + ) -> ActorContext { + ActorContext { + inner: self + .inner + .with_application_span(traceparent.as_deref(), tracestate.as_deref()), + shared: self.shared.clone(), + } + } + #[napi] pub fn invocation_trace_context(&self) -> Option { self.inner.invocation_trace_context().map(Into::into) diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts index 753d89389c..76f0638a8b 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts @@ -9,7 +9,10 @@ import type { WebSocket as NativeWebSocket, } from "@rivetkit/rivetkit-napi"; import type { ActorInvocationTraceContext } from "@/common/actor-telemetry-context"; -import { runWithActorInvocationSpan } from "@/common/otel-context"; +import { + readActiveTraceHeaders, + runWithActorInvocationSpan, +} from "@/common/otel-context"; import type { ActorContextHandle, ActorFactoryHandle, @@ -259,6 +262,9 @@ export class NapiCoreRuntime implements CoreRuntime { #bindings: NativeBindings; #sql = new WeakMap(); #invocationContext: AsyncLocalStorage; + // `traceparent` of each invocation's own Core span, so an operation can + // tell that span apart from an application span without a native call. + #invocationTraceparent = new WeakMap(); constructor( bindings: NativeBindings, @@ -268,13 +274,26 @@ export class NapiCoreRuntime implements CoreRuntime { this.#invocationContext = invocationContext; } + // Core cannot read the active JS span; bind it to the operation handle here. #actorContextForOperation(owner: ActorContextHandle): NativeActorContext { const ownerCtx = asNativeActorContext(owner); const active = this.#invocationContext.getStore(); - if (active?.sameActorInstance(ownerCtx)) { + if (!active?.sameActorInstance(ownerCtx)) { + return ownerCtx; + } + const applicationSpan = readActiveTraceHeaders(); + // Avoid a native call when Core already has the active parent. + if ( + !applicationSpan || + applicationSpan.traceparent === + this.#invocationTraceparent.get(active) + ) { return active; } - return ownerCtx; + return active.withApplicationSpan( + applicationSpan.traceparent, + applicationSpan.tracestate ?? null, + ); } // Cache only the actor-owned handle, which is closed on sleep. @@ -577,22 +596,22 @@ export class NapiCoreRuntime implements CoreRuntime { runWithActorInvocationContext(ctx: ActorContextHandle, run: () => T): T { const nativeCtx = asNativeActorContext(ctx); - const traceContext = this.#actorInvocationTraceContext(nativeCtx); + const span = nativeCtx.invocationTraceContext()?.span; + if (span) { + this.#invocationTraceparent.set(nativeCtx, span.traceparent); + } return this.#invocationContext.run(nativeCtx, () => - runWithActorInvocationSpan(traceContext?.span, run), + runWithActorInvocationSpan(span, run), ); } actorInvocationTraceContext( ctx: ActorContextHandle, ): ActorInvocationTraceContext | undefined { - return this.#actorInvocationTraceContext(asNativeActorContext(ctx)); - } - - #actorInvocationTraceContext( - ctx: NativeActorContext, - ): ActorInvocationTraceContext | undefined { - return ctx.invocationTraceContext() ?? undefined; + return ( + this.#actorContextForOperation(ctx).invocationTraceContext() ?? + undefined + ); } actorName(ctx: ActorContextHandle): string { @@ -643,7 +662,7 @@ export class NapiCoreRuntime implements CoreRuntime { } actorWaitUntil(ctx: ActorContextHandle, promise: Promise): void { - asNativeActorContext(ctx).waitUntil(promise); + this.#actorContextForOperation(ctx).waitUntil(promise); } async actorWaitForTrackedShutdownWork( @@ -1041,7 +1060,7 @@ export class NapiCoreRuntime implements CoreRuntime { actionName: string, args: RuntimeBytes, ): Promise { - return await asNativeActorContext(ctx) + return await this.#actorContextForOperation(ctx) .schedule() .after(durationMs, actionName, toNapiBuffer(args)); } @@ -1052,13 +1071,13 @@ export class NapiCoreRuntime implements CoreRuntime { actionName: string, args: RuntimeBytes, ): Promise { - return await asNativeActorContext(ctx) + return await this.#actorContextForOperation(ctx) .schedule() .at(timestampMs, actionName, toNapiBuffer(args)); } async actorScheduleCancel(ctx: ActorContextHandle, id: string) { - return await asNativeActorContext(ctx).schedule().cancel(id); + return await this.#actorContextForOperation(ctx).schedule().cancel(id); } async actorScheduleGet( @@ -1066,12 +1085,13 @@ export class NapiCoreRuntime implements CoreRuntime { id: string, ): Promise { return ( - (await asNativeActorContext(ctx).schedule().get(id)) ?? undefined + (await this.#actorContextForOperation(ctx).schedule().get(id)) ?? + undefined ); } async actorScheduleList(ctx: ActorContextHandle) { - return await asNativeActorContext(ctx).schedule().list(); + return await this.#actorContextForOperation(ctx).schedule().list(); } async actorCronSet( @@ -1083,7 +1103,7 @@ export class NapiCoreRuntime implements CoreRuntime { args: RuntimeBytes, maxHistory: number | undefined, ) { - await asNativeActorContext(ctx) + await this.#actorContextForOperation(ctx) .schedule() .cronSet( name, @@ -1103,7 +1123,7 @@ export class NapiCoreRuntime implements CoreRuntime { args: RuntimeBytes, maxHistory: number | undefined, ) { - await asNativeActorContext(ctx) + await this.#actorContextForOperation(ctx) .schedule() .cronEvery( name, @@ -1118,20 +1138,23 @@ export class NapiCoreRuntime implements CoreRuntime { ctx: ActorContextHandle, name: string, ): Promise { - return ((await asNativeActorContext(ctx).schedule().cronGet(name)) ?? - undefined) as RuntimeCronJobInfo | undefined; + return ((await this.#actorContextForOperation(ctx) + .schedule() + .cronGet(name)) ?? undefined) as RuntimeCronJobInfo | undefined; } async actorCronList( ctx: ActorContextHandle, ): Promise { - return (await asNativeActorContext(ctx) + return (await this.#actorContextForOperation(ctx) .schedule() .cronList()) as RuntimeCronJobInfo[]; } async actorCronDelete(ctx: ActorContextHandle, name: string) { - return await asNativeActorContext(ctx).schedule().cronDelete(name); + return await this.#actorContextForOperation(ctx) + .schedule() + .cronDelete(name); } async actorCronHistory( @@ -1139,7 +1162,7 @@ export class NapiCoreRuntime implements CoreRuntime { name: string, limit: number | undefined, ): Promise { - return (await asNativeActorContext(ctx) + return (await this.#actorContextForOperation(ctx) .schedule() .cronHistory(name, limit)) as RuntimeCronFire[]; } diff --git a/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts b/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts index a6f2149b70..c1adaad769 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts @@ -158,6 +158,14 @@ class FakeActorContext { return this.runtimeBag; } + invocationTraceContext(): undefined { + return undefined; + } + + sameActorInstance(other: FakeActorContext): boolean { + return this === other; + } + actorId(): string { return "parity-actor"; } From d4b3b5f6c2c7627f2314e735e552bd61d3716655 Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Wed, 2 Sep 2026 09:50:04 +0400 Subject: [PATCH 10/21] feat(rivetkit-core): trace scheduled invocations --- .../rivetkit-core/src/actor/context.rs | 32 ++--- .../rivetkit-core/src/actor/metrics.rs | 34 +++-- .../packages/rivetkit-core/src/actor/task.rs | 21 +-- .../packages/rivetkit-core/src/telemetry.rs | 125 +++++++++++++----- 4 files changed, 135 insertions(+), 77 deletions(-) diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs index e7e10c00b9..e6791220d0 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs @@ -41,7 +41,7 @@ use crate::actor::internal_storage; use crate::actor::kv::LegacyActorKv; use crate::actor::lifecycle_hooks::Reply; use crate::actor::messages::{ActorEvent, Request, StateDelta, WorkflowKvWrite}; -use crate::actor::metrics::{ActorMetrics, InvocationStatus, InvocationType}; +use crate::actor::metrics::ActorMetrics; use crate::actor::queue::{QueueInspectorUpdateCallback, QueueMetadata, QueueWaitActivityCallback}; use crate::actor::schedule::{InternalKeepAwakeCallback, LocalAlarmCallback}; use crate::actor::sleep::{CanSleep, SleepState}; @@ -1826,20 +1826,20 @@ impl ActorContext { self.track_shutdown_task(async move { let _internal_keep_awake_region = internal_keep_awake_region; ctx.record_user_task_started(UserTaskKind::ScheduledAction); - let started_at = Instant::now(); + let user_task_started_at = Instant::now(); let action_name = action.clone(); + let invocation = crate::telemetry::ActorInvocation::start_scheduled(&ctx, &action_name); + let invocation_telemetry = invocation.telemetry(); let (reply_tx, reply_rx) = oneshot::channel(); let mut dispatch_error = None; - let mut invocation_status = InvocationStatus::Ok; - let mut action_ran = true; match ctx.try_send_actor_event( ActorEvent::Action { name: action.clone(), args, conn: None, scheduled_fire: Some(scheduled_fire), - invocation_telemetry: None, + invocation_telemetry: Some(invocation_telemetry), reply: Reply::from(reply_tx), }, "scheduled_action", @@ -1847,7 +1847,6 @@ impl ActorContext { Ok(()) => match reply_rx.await { Ok(Ok(_)) => {} Ok(Err(error)) => { - invocation_status = InvocationStatus::from_error(&error); dispatch_error = Some(error); tracing::error!( error = ?dispatch_error.as_ref().expect("just assigned"), @@ -1856,9 +1855,9 @@ impl ActorContext { "scheduled event execution failed" ); } - Err(error) => { - invocation_status = InvocationStatus::Dropped; - dispatch_error = Some(error.into()); + Err(_) => { + // Use the same dropped-reply error for tracing, metrics, and schedule history. + dispatch_error = Some(ActorLifecycleError::DroppedReply.build()); tracing::error!( error = ?dispatch_error.as_ref().expect("just assigned"), event_id, @@ -1868,7 +1867,6 @@ impl ActorContext { } }, Err(error) => { - action_ran = false; dispatch_error = Some(error); tracing::error!( error = ?dispatch_error.as_ref().expect("just assigned"), @@ -1878,14 +1876,7 @@ impl ActorContext { ); } } - if action_ran { - ctx.metrics().record_invocation( - &action_name, - InvocationType::Scheduled, - invocation_status, - started_at.elapsed(), - ); - } + invocation.finish(dispatch_error.as_ref()); ctx.finish_schedule_dispatch(&event_id, history_id, dispatch_error.as_ref()) .await; @@ -1903,7 +1894,10 @@ impl ActorContext { } } - ctx.record_user_task_finished(UserTaskKind::ScheduledAction, started_at.elapsed()); + ctx.record_user_task_finished( + UserTaskKind::ScheduledAction, + user_task_started_at.elapsed(), + ); }); } diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/metrics.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/metrics.rs index eb4846d206..46129142eb 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/metrics.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/metrics.rs @@ -137,12 +137,21 @@ pub(crate) enum InvocationType { } impl InvocationType { - fn as_label(self) -> &'static str { + pub(crate) fn as_label(self) -> &'static str { match self { Self::Action => "action", Self::Scheduled => "scheduled", } } + + /// OpenTelemetry span kind for this invocation. An action is entered from + /// outside the actor, while a scheduled fire originates inside it. + pub(crate) fn otel_kind(self) -> &'static str { + match self { + Self::Action => "server", + Self::Scheduled => "internal", + } + } } #[derive(Clone, Copy, Debug)] @@ -1950,6 +1959,20 @@ impl ActorMetrics { .observe(duration.as_secs_f64()); } + /// Folds an undeclared action name down to a bounded placeholder. + /// + /// Action names arrive from callers, so using one verbatim would mint a new + /// series per value wherever the name becomes a dimension. `_OTHER` is the + /// fallback OpenTelemetry defines for exactly this, and it cannot collide + /// with a declared action name. + pub(crate) fn label_action_name<'a>(&'a self, action_name: &'a str) -> &'a str { + if self.inner.action_names.contains(action_name) { + action_name + } else { + "_OTHER" + } + } + pub(crate) fn record_invocation( &self, action_name: &str, @@ -1958,14 +1981,7 @@ impl ActorMetrics { duration: Duration, ) { let actor_labels = self.actor_labels(); - // Action names arrive from callers, so an undeclared one would mint a new - // label series per value. `_OTHER` is the fallback OpenTelemetry defines - // for exactly this, and it cannot collide with a declared action name. - let action_name = if self.inner.action_names.contains(action_name) { - action_name - } else { - "_OTHER" - }; + let action_name = self.label_action_name(action_name); let labels = [ actor_labels[0], action_name, diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs index 0e2626d43f..e6d6b8697b 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs @@ -51,7 +51,7 @@ use crate::actor::messages::{ ActorEvent, ActorHttpResponse, QueueSendResult, Request, SerializeStateReason, StateDelta, WorkflowKvWrite, }; -use crate::actor::metrics::{InvocationStatus, InvocationType, startup_phase::StartupPhase}; +use crate::actor::metrics::startup_phase::StartupPhase; use crate::actor::state::{PersistedActor, RequestSaveOpts}; use crate::actor::task_types::ShutdownKind; use crate::actor::work_registry::ActorWorkKind; @@ -908,8 +908,7 @@ impl ActorTask { reply, } => { let invocation = - crate::telemetry::ActionInvocationSpan::start(&self.ctx, &name, incoming); - let invocation_started_at = Instant::now(); + crate::telemetry::ActorInvocation::start_action(&self.ctx, &name, incoming); let invocation_telemetry = invocation.telemetry(); tracing::info!( actor_id = %self.ctx.actor_id(), @@ -941,21 +940,17 @@ impl ActorTask { let actor_id = self.ctx.actor_id().to_owned(); let ctx = self.ctx.clone(); self.ctx.spawn_work(ActorWorkKind::Action, async move { - let (result, status) = match tracked_reply_rx.await { + let result = match tracked_reply_rx.await { Ok(result) => { let result = result.map_err(|error| ctx.attach_actor_to_error(error)); - let status = match result.as_ref() { - Ok(_) => InvocationStatus::Ok, - Err(error) => InvocationStatus::from_error(error), - }; tracing::info!( actor_id = %actor_id, action_name = %action_name_for_log, ok = result.is_ok(), "actor task: tracked reply received, forwarding" ); - (result, status) + result } Err(_) => { tracing::warn!( @@ -966,15 +961,9 @@ impl ActorTask { let error = ctx.attach_actor_to_error( ActorLifecycleError::DroppedReply.build(), ); - (Err(error), InvocationStatus::Dropped) + Err(error) } }; - ctx.metrics().record_invocation( - &action_name_for_log, - InvocationType::Action, - status, - invocation_started_at.elapsed(), - ); invocation.finish(result.as_ref().err()); let _ = reply.send(result); }); diff --git a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs index 6d201a0172..149e32045e 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs @@ -14,6 +14,8 @@ use parking_lot::Mutex; use tracing_opentelemetry::OpenTelemetrySpanExt as _; use crate::ActorContext; +use crate::actor::metrics::{ActorMetrics, InvocationStatus, InvocationType}; +use crate::time::Instant; /// Correlation fields accepted at an invocation boundary. #[derive(Debug, Default)] @@ -60,10 +62,14 @@ fn invocation_ray_id(headers: &http::HeaderMap) -> Option { rivetkit_client_protocol::telemetry_headers::bounded_ray_id(value).map(str::to_owned) } -/// The single root span for one client action invocation. +/// Owns the complete lifecycle of one actor invocation. #[derive(Debug)] -pub(crate) struct ActionInvocationSpan { +pub(crate) struct ActorInvocation { telemetry: ActorInvocationTelemetry, + metrics: ActorMetrics, + action_name: String, + invocation_type: InvocationType, + started_at: Instant, } /// Opaque invocation context carried across foreign-runtime adapters. @@ -186,21 +192,49 @@ pub(crate) struct SqliteOperationSpan { span: Option, } -impl ActionInvocationSpan { - pub(crate) fn start( +impl ActorInvocation { + pub(crate) fn start_action( ctx: &ActorContext, action_name: &str, incoming: IncomingInvocationContext, + ) -> Self { + Self::start( + ctx, + action_name, + InvocationType::Action, + incoming.ray_id, + incoming.remote_parent, + ) + } + + pub(crate) fn start_scheduled(ctx: &ActorContext, action_name: &str) -> Self { + Self::start( + ctx, + action_name, + InvocationType::Scheduled, + None, + None, + ) + } + + fn start( + ctx: &ActorContext, + action_name: &str, + invocation_type: InvocationType, + ray_id: Option, + parent: Option, ) -> Self { let identity = ctx.telemetry_identity(); - let ray_id = incoming.ray_id; + // Use bounded names for both spans and metrics. + let action_name = ctx.metrics().label_action_name(action_name).to_owned(); let span = if tracing::enabled!(target: "rivetkit::telemetry", tracing::Level::INFO) { let span = tracing::info_span!( target: "rivetkit::telemetry", parent: None, "rivet.actor.invoke", - otel.kind = "server", - rivet.invocation.type = "action", + otel.name = %format!("{}/{}", identity.actor_name, action_name), + otel.kind = invocation_type.otel_kind(), + rivet.invocation.type = invocation_type.as_label(), rivet.actor.id = %identity.actor_id, rivet.actor.name = %identity.actor_name, rivet.actor.key = %identity.actor_key, @@ -210,7 +244,7 @@ impl ActionInvocationSpan { error.type = tracing::field::Empty, ); span.record("rivet.ray.id", ray_id.as_deref()); - if let Some(parent) = incoming.remote_parent { + if let Some(parent) = parent { span.set_parent(opentelemetry::Context::new().with_remote_span_context(parent)); } Some(span) @@ -220,6 +254,10 @@ impl ActionInvocationSpan { Self { telemetry: ActorInvocationTelemetry::new(ray_id, span, identity), + metrics: ctx.metrics().clone(), + action_name, + invocation_type, + started_at: Instant::now(), } } @@ -227,14 +265,54 @@ impl ActionInvocationSpan { self.telemetry.clone() } - pub(crate) fn finish(self, error: Option<&anyhow::Error>) { - self.telemetry.finish(error); + pub(crate) fn finish(mut self, error: Option<&anyhow::Error>) { + self.finish_with_status( + error.map_or(InvocationStatus::Ok, InvocationStatus::from_error), + error, + ); + } + + fn finish_with_status(&mut self, status: InvocationStatus, error: Option<&anyhow::Error>) { + let Some(span) = self.telemetry.claim_terminal() else { + return; + }; + self.record_finished(span, status, error); + } + + /// Records the terminal metric and span status of an invocation whose + /// completion the caller has already claimed through `claim_terminal`. + /// The metric measures what the caller waited for, so it is recorded here + /// even when `wait_until` work keeps the span open past this point. + fn record_finished( + &self, + span: Option, + status: InvocationStatus, + error: Option<&anyhow::Error>, + ) { + self.metrics.record_invocation( + &self.action_name, + self.invocation_type, + status, + self.started_at.elapsed(), + ); + if let Some(span) = span { + record_outcome(&span, error); + self.telemetry.mark_reply_sent(&span); + } + self.telemetry.release_span_if_settled(); } } -impl Drop for ActionInvocationSpan { +impl Drop for ActorInvocation { fn drop(&mut self) { - self.telemetry.finish_dropped(); + // `finish` consumes the invocation, so this runs on the completed path + // too. Claim the terminal record first, so the dropped-reply error is + // only built for an invocation that really was dropped. + let Some(span) = self.telemetry.claim_terminal() else { + return; + }; + let error = crate::error::ActorLifecycle::DroppedReply.build(); + self.record_finished(span, InvocationStatus::Dropped, Some(&error)); } } @@ -338,25 +416,6 @@ impl ActorInvocationTelemetry { Some(SqliteOperationSpan { span: Some(span) }) } - fn finish(&self, error: Option<&anyhow::Error>) { - let Some(span) = self.claim_terminal() else { - return; - }; - record_outcome(&span, error); - self.mark_reply_sent(&span); - self.release_span_if_settled(); - } - - fn finish_dropped(&self) { - let Some(span) = self.claim_terminal() else { - return; - }; - span.record("otel.status_code", "ERROR"); - span.record("error.type", "actor.dropped_reply"); - self.mark_reply_sent(&span); - self.release_span_if_settled(); - } - /// Borrows the invocation while it is still open: before its status is /// recorded, or after it while `wait_until` work from it still runs. A /// settled invocation yields nothing, so late SQLite work and retained @@ -370,11 +429,11 @@ impl ActorInvocationTelemetry { /// Claims the terminal record, so the finish and drop paths cannot both /// record a status for the same invocation. The span stays in its slot /// until `release_span_if_settled` empties it. - fn claim_terminal(&self) -> Option { + fn claim_terminal(&self) -> Option> { if self.0.finished.swap(true, Ordering::SeqCst) { return None; } - self.0.span.lock().clone() + Some(self.0.span.lock().clone()) } /// Marks the moment the caller got its answer when the span will outlive From d72a4132a228e447a5c57aeca68562810367ba9f Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Wed, 2 Sep 2026 09:51:28 +0400 Subject: [PATCH 11/21] feat(rivetkit-core): persist schedule trace origins --- .../packages/actor-persist/src/versioned.rs | 48 ++++ .../rivetkit-core/src/actor/context.rs | 6 +- .../src/actor/internal_storage/mod.rs | 7 + .../src/actor/internal_storage/queries.rs | 13 +- .../src/actor/internal_storage/schema.rs | 2 +- .../rivetkit-core/src/actor/schedule.rs | 253 ++++++++++++++---- .../packages/rivetkit-core/src/telemetry.rs | 102 +++++-- .../rivetkit-core/tests/sql_efficiency.rs | 36 ++- 8 files changed, 381 insertions(+), 86 deletions(-) diff --git a/rivetkit-rust/packages/actor-persist/src/versioned.rs b/rivetkit-rust/packages/actor-persist/src/versioned.rs index d64bd2d9e0..d8dde9dbde 100644 --- a/rivetkit-rust/packages/actor-persist/src/versioned.rs +++ b/rivetkit-rust/packages/actor-persist/src/versioned.rs @@ -1,4 +1,5 @@ use anyhow::{Result, bail}; +use serde::{Deserialize, Serialize}; use vbare::OwnedVersionedData; use crate::generated::{v1, v2, v3, v4}; @@ -500,6 +501,53 @@ pub enum RunWakeAt { V1(Option), } +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +pub struct ScheduleTraceContextData { + pub ray_id: Option, + pub traceparent: Option, + pub tracestate: Option, +} + +pub enum ScheduleTraceContext { + V1(ScheduleTraceContextData), +} + +impl OwnedVersionedData for ScheduleTraceContext { + type Latest = ScheduleTraceContextData; + + fn wrap_latest(latest: Self::Latest) -> Self { + Self::V1(latest) + } + + fn unwrap_latest(self) -> Result { + match self { + Self::V1(data) => Ok(data), + } + } + + fn deserialize_version(payload: &[u8], version: u16) -> Result { + match version { + 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), + _ => bail!("invalid schedule trace context version: {version}"), + } + } + + fn serialize_version(self, version: u16) -> Result> { + match (self, version) { + (Self::V1(data), 1) => serde_bare::to_vec(&data).map_err(Into::into), + (_, version) => bail!("unexpected schedule trace context version: {version}"), + } + } + + fn deserialize_converters() -> Vec Result> { + Vec:: Result>::new() + } + + fn serialize_converters() -> Vec Result> { + Vec:: Result>::new() + } +} + impl OwnedVersionedData for RunWakeAt { type Latest = Option; diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs index e6791220d0..6073122f8c 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs @@ -1828,7 +1828,11 @@ impl ActorContext { ctx.record_user_task_started(UserTaskKind::ScheduledAction); let user_task_started_at = Instant::now(); let action_name = action.clone(); - let invocation = crate::telemetry::ActorInvocation::start_scheduled(&ctx, &action_name); + let invocation = crate::telemetry::ActorInvocation::start_scheduled( + &ctx, + &action_name, + dispatch.origin, + ); let invocation_telemetry = invocation.telemetry(); let (reply_tx, reply_rx) = oneshot::channel(); diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/mod.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/mod.rs index 78020880a2..a8faac3c7a 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/mod.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/mod.rs @@ -149,6 +149,10 @@ pub(crate) async fn import_legacy_actor_snapshot( sql: UPSERT_ACTOR_STATE_SQL.to_owned(), params: Some(vec![BindParam::Blob(actor.state.clone())]), }, + SqliteBatchStatement { + sql: RESET_SCHEDULE_TRACE_CONTEXTS_SQL.to_owned(), + params: None, + }, SqliteBatchStatement { sql: RESET_SCHEDULES_FOR_LEGACY_IMPORT_SQL.to_owned(), params: None, @@ -1163,6 +1167,9 @@ pub(crate) async fn clear_imported_storage(db: &SqliteDb, actor_id: &str) -> Res .await .with_context(|| format!("clear partially imported {table} rows"))?; } + db.execute(RESET_SCHEDULE_TRACE_CONTEXTS_SQL, None) + .await + .context("clear imported schedule trace contexts")?; Ok(()) } diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/queries.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/queries.rs index dad1417c00..babe13ec6b 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/queries.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/queries.rs @@ -14,6 +14,7 @@ pub(crate) const DELETE_CONN_STATE_SQL: &str = "DELETE FROM _rivet_conn_state WH pub(crate) const DELETE_CONN_SQL: &str = "DELETE FROM _rivet_conns WHERE conn_id = ?"; pub(crate) const RESET_SCHEDULES_FOR_LEGACY_IMPORT_SQL: &str = "DELETE FROM _rivet_schedule_events"; +pub(crate) const RESET_SCHEDULE_TRACE_CONTEXTS_SQL: &str = "DELETE FROM _rivet_meta WHERE key >= 'schedule_trace_context:' AND key < 'schedule_trace_context;'"; pub(crate) const INSERT_SCHEDULE_EVENT_SQL: &str = "INSERT INTO _rivet_schedule_events (event_id, trigger_at, action, args, kind, cron_expression, timezone, interval_ms, last_started_at, max_history) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"; pub(crate) const UPSERT_RECURRING_SCHEDULE_SQL: &str = "INSERT INTO _rivet_schedule_events (event_id, trigger_at, action, args, kind, cron_expression, timezone, interval_ms, last_started_at, max_history) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(event_id) DO UPDATE SET trigger_at = excluded.trigger_at, action = excluded.action, args = excluded.args, kind = excluded.kind, cron_expression = excluded.cron_expression, timezone = excluded.timezone, interval_ms = excluded.interval_ms, max_history = excluded.max_history"; pub(crate) const CANCEL_SCHEDULE_SQL: &str = @@ -29,7 +30,10 @@ pub(crate) const LIST_CRONS_SQL: &str = "SELECT event_id, trigger_at, action, ar pub(crate) const CRON_HISTORY_SQL: &str = "SELECT action, scheduled_at, fired_at, finished_at, result, error_group, error_code, error_message, error_metadata FROM _rivet_schedule_history WHERE schedule_id = ? ORDER BY fired_at DESC, id DESC LIMIT ?"; pub(crate) const LOAD_SCHEDULE_SQL: &str = "SELECT event_id, trigger_at, action, args, kind, cron_expression, timezone, interval_ms, last_started_at, max_history FROM _rivet_schedule_events WHERE event_id = ?"; pub(crate) const COUNT_SCHEDULES_SQL: &str = "SELECT COUNT(*) FROM _rivet_schedule_events"; -pub(crate) const TAKE_DUE_SCHEDULES_SQL: &str = "SELECT event_id, trigger_at, action, args, kind, cron_expression, timezone, interval_ms, last_started_at, max_history FROM _rivet_schedule_events WHERE trigger_at <= ? ORDER BY trigger_at, event_id"; +pub(crate) const TAKE_DUE_SCHEDULES_SQL: &str = "SELECT event_id, trigger_at, action, args, kind, cron_expression, timezone, interval_ms, last_started_at, max_history, (SELECT value FROM _rivet_meta WHERE key = 'schedule_trace_context:' || event_id) FROM _rivet_schedule_events WHERE trigger_at <= ? ORDER BY trigger_at, event_id"; +pub(crate) const UPSERT_SCHEDULE_TRACE_CONTEXT_SQL: &str = "INSERT INTO _rivet_meta (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value"; +pub(crate) const DELETE_SCHEDULE_TRACE_CONTEXT_SQL: &str = "DELETE FROM _rivet_meta WHERE key = ?"; +pub(crate) const DELETE_ORPHAN_SCHEDULE_TRACE_CONTEXT_SQL: &str = "DELETE FROM _rivet_meta WHERE key = ? AND NOT EXISTS (SELECT 1 FROM _rivet_schedule_events WHERE event_id = ?)"; pub(crate) const ADVANCE_SKIPPED_SCHEDULE_SQL: &str = "UPDATE _rivet_schedule_events SET trigger_at = ? WHERE event_id = ?"; pub(crate) const ADVANCE_SCHEDULE_SQL: &str = @@ -52,6 +56,13 @@ pub(crate) fn claim_one_shots_sql(event_count: usize) -> String { ) } +pub(crate) fn delete_schedule_trace_contexts_sql(event_count: usize) -> String { + let placeholders = std::iter::repeat_n("?", event_count) + .collect::>() + .join(", "); + format!("DELETE FROM _rivet_meta WHERE key IN ({placeholders})") +} + pub(crate) const LOAD_QUEUE_NEXT_ID_SQL: &str = "SELECT queue_next_id FROM _rivet_runtime WHERE id = 1"; pub(crate) const LOAD_QUEUE_STATS_SQL: &str = "SELECT COUNT(*), MAX(id) FROM _rivet_queue"; diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/schema.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/schema.rs index 41c23e44b2..0e44b07679 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/schema.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/schema.rs @@ -13,7 +13,7 @@ const SCHEMA_VERSION_KEY: &str = "schema_version"; // interrupted imports can be detected and retried. Fixed core-owned logical // metadata may also live here when adding a column would break older runtimes' // ability to open the database. This is not a general-purpose runtime KV store. -// W[bootstrap + core metadata only | point upsert | <100 B | 1-page map] +// W[bootstrap + bounded core metadata | point upsert | schedule metadata capped by max_schedules] pub(crate) const CREATE_META_TABLE: &str = r#" CREATE TABLE IF NOT EXISTS _rivet_meta ( key TEXT PRIMARY KEY, diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/schedule.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/schedule.rs index f54026a1b2..671981070e 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/schedule.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/schedule.rs @@ -11,6 +11,7 @@ use futures::future::BoxFuture; use futures::future::{AbortHandle, Abortable}; use rivet_envoy_client::handle::EnvoyHandle; use rivet_error::RivetError; +use rivetkit_actor_persist::versioned::{ScheduleTraceContext, ScheduleTraceContextData}; use serde::{Deserialize, Serialize}; use tokio::runtime::Handle; use tokio::sync::oneshot; @@ -19,8 +20,12 @@ use uuid::Uuid; use crate::actor::context::ActorContext; use crate::actor::internal_storage::queries::*; +use crate::actor::persist::{ + decode_latest_with_embedded_version, encode_latest_with_embedded_version, +}; use crate::error::{ScheduleRuntimeError, client_error_message, client_error_metadata}; use crate::sqlite::{BindParam, ColumnValue, SqliteBatchStatement}; +use crate::telemetry::{ActorInvocationTelemetry, TraceOrigin}; use crate::time::{SystemTime, UNIX_EPOCH, sleep}; const CRON_ID_PREFIX: &str = "cron:"; @@ -34,6 +39,8 @@ pub const MAX_ACTOR_HISTORY: i64 = 10_000; pub const MIN_INTERVAL_MS: i64 = 5_000; const DEFAULT_HISTORY_LIMIT: i64 = 20; const CLAIM_ONE_SHOT_BATCH_SIZE: usize = 128; +const SCHEDULE_TRACE_CONTEXT_META_PREFIX: &str = "schedule_trace_context:"; +const SCHEDULE_TRACE_CONTEXT_VERSION: u16 = 1; pub(crate) const GLOBAL_HISTORY_PRUNE_INTERVAL: usize = 100; pub(crate) const GLOBAL_HISTORY_RETAINED_ROWS: i64 = MAX_ACTOR_HISTORY - GLOBAL_HISTORY_PRUNE_INTERVAL as i64; @@ -148,6 +155,7 @@ pub(crate) struct DueScheduleDispatch { pub args: Vec, pub fire: ScheduledFireInfo, pub history_id: Option, + pub origin: TraceOrigin, } #[derive(Clone, Debug)] @@ -183,6 +191,14 @@ impl ActorContext { system_now_timestamp_ms() } + /// Trace origin of the invocation defining this schedule, empty when the + /// caller is not inside a traced invocation. + fn schedule_trace_origin(&self) -> TraceOrigin { + self.invocation_telemetry() + .map(ActorInvocationTelemetry::trace_origin) + .unwrap_or_default() + } + pub async fn after( &self, duration: Duration, @@ -195,25 +211,29 @@ impl ActorContext { } pub async fn at(&self, timestamp_ms: i64, action_name: &str, args: &[u8]) -> Result { + let origin = self.schedule_trace_origin(); let _mutation = self.0.schedule_mutation_lock.lock().await; self.ensure_schedule_capacity(false).await?; let event_id = Uuid::new_v4().to_string(); + let schedule_params = vec![ + BindParam::Text(event_id.clone()), + BindParam::Integer(timestamp_ms), + BindParam::Text(action_name.to_owned()), + args_param(args), + BindParam::Integer(ScheduleKind::At.as_i64()), + BindParam::Null, + BindParam::Null, + BindParam::Null, + BindParam::Null, + BindParam::Integer(0), + ]; + let mut statements = vec![SqliteBatchStatement { + sql: INSERT_SCHEDULE_EVENT_SQL.to_owned(), + params: Some(schedule_params), + }]; + append_schedule_trace_context_upsert(&mut statements, &event_id, origin)?; self.sql() - .execute( - INSERT_SCHEDULE_EVENT_SQL, - Some(vec![ - BindParam::Text(event_id.clone()), - BindParam::Integer(timestamp_ms), - BindParam::Text(action_name.to_owned()), - args_param(args), - BindParam::Integer(ScheduleKind::At.as_i64()), - BindParam::Null, - BindParam::Null, - BindParam::Null, - BindParam::Null, - BindParam::Integer(0), - ]), - ) + .execute_batch(statements) .await .context("insert one-shot schedule")?; self.mark_schedule_dirty(); @@ -224,18 +244,21 @@ impl ActorContext { pub async fn cancel_schedule(&self, event_id: &str) -> Result { let _mutation = self.0.schedule_mutation_lock.lock().await; - let result = self + let results = self .sql() - .execute( - CANCEL_SCHEDULE_SQL, - Some(vec![ - BindParam::Text(event_id.to_owned()), - BindParam::Integer(ScheduleKind::At.as_i64()), - ]), - ) + .execute_batch(vec![ + SqliteBatchStatement { + sql: CANCEL_SCHEDULE_SQL.to_owned(), + params: Some(vec![ + BindParam::Text(event_id.to_owned()), + BindParam::Integer(ScheduleKind::At.as_i64()), + ]), + }, + delete_orphan_schedule_trace_context(event_id), + ]) .await .context("cancel one-shot schedule")?; - let removed = result.changes > 0; + let removed = results.first().is_some_and(|result| result.changes > 0); if removed { self.mark_schedule_dirty(); self.record_schedules_updated(); @@ -304,6 +327,7 @@ impl ActorContext { args: &[u8], max_history: Option, ) -> Result<()> { + let origin = self.schedule_trace_origin(); validate_name(name)?; let timezone = timezone.unwrap_or("UTC"); let timezone_parsed = parse_timezone(timezone)?; @@ -334,6 +358,7 @@ impl ActorContext { Some(timezone), None, max_history, + origin, ) .await?; self.prune_schedule_history(&event_id, max_history).await?; @@ -350,6 +375,7 @@ impl ActorContext { args: &[u8], max_history: Option, ) -> Result<()> { + let origin = self.schedule_trace_origin(); validate_name(name)?; if interval_ms < MIN_INTERVAL_MS { return Err(ScheduleRuntimeError::InvalidInterval { @@ -382,6 +408,7 @@ impl ActorContext { None, Some(interval_ms), max_history, + origin, ) .await?; self.prune_schedule_history(&event_id, max_history).await?; @@ -402,23 +429,26 @@ impl ActorContext { timezone: Option<&str>, interval_ms: Option, max_history: i64, + origin: TraceOrigin, ) -> Result<()> { + let mut statements = vec![SqliteBatchStatement { + sql: UPSERT_RECURRING_SCHEDULE_SQL.to_owned(), + params: Some(vec![ + BindParam::Text(event_id.to_owned()), + BindParam::Integer(trigger_at), + BindParam::Text(action_name.to_owned()), + args_param(args), + BindParam::Integer(kind.as_i64()), + optional_text_param(cron_expression), + optional_text_param(timezone), + optional_i64_param(interval_ms), + BindParam::Null, + BindParam::Integer(max_history), + ]), + }]; + append_schedule_trace_context_upsert(&mut statements, event_id, origin)?; self.sql() - .execute( - UPSERT_RECURRING_SCHEDULE_SQL, - Some(vec![ - BindParam::Text(event_id.to_owned()), - BindParam::Integer(trigger_at), - BindParam::Text(action_name.to_owned()), - args_param(args), - BindParam::Integer(kind.as_i64()), - optional_text_param(cron_expression), - optional_text_param(timezone), - optional_i64_param(interval_ms), - BindParam::Null, - BindParam::Integer(max_history), - ]), - ) + .execute_batch(statements) .await .context("upsert recurring schedule")?; Ok(()) @@ -442,10 +472,11 @@ impl ActorContext { SqliteBatchStatement { sql: DELETE_CRON_SQL.to_owned(), params: Some(vec![ - BindParam::Text(event_id), + BindParam::Text(event_id.clone()), BindParam::Integer(ScheduleKind::At.as_i64()), ]), }, + delete_orphan_schedule_trace_context(&event_id), ]) .await .context("delete recurring schedule and history")?; @@ -463,19 +494,23 @@ impl ActorContext { pub(crate) async fn cron_delete_if_action(&self, name: &str, action: &str) -> Result { validate_name(name)?; let _mutation = self.0.schedule_mutation_lock.lock().await; - let result = self + let event_id = cron_event_id(name); + let results = self .sql() - .execute( - DELETE_CRON_IF_ACTION_SQL, - Some(vec![ - BindParam::Text(cron_event_id(name)), - BindParam::Integer(ScheduleKind::At.as_i64()), - BindParam::Text(action.to_owned()), - ]), - ) + .execute_batch(vec![ + SqliteBatchStatement { + sql: DELETE_CRON_IF_ACTION_SQL.to_owned(), + params: Some(vec![ + BindParam::Text(event_id.clone()), + BindParam::Integer(ScheduleKind::At.as_i64()), + BindParam::Text(action.to_owned()), + ]), + }, + delete_orphan_schedule_trace_context(&event_id), + ]) .await .context("delete recurring schedule with matching action")?; - let removed = result.changes > 0; + let removed = results.first().is_some_and(|result| result.changes > 0); if removed { self.mark_schedule_dirty(); self.record_schedules_updated(); @@ -594,24 +629,38 @@ impl ActorContext { let due_schedules = result .rows .iter() - .map(|row| read_stored_schedule(row)) + .map(|row| read_due_schedule(row)) .collect::>>()?; let claim_statements = due_schedules .iter() - .filter(|event| event.kind == ScheduleKind::At) + .filter(|(event, _)| event.kind == ScheduleKind::At) .collect::>() .chunks(CLAIM_ONE_SHOT_BATCH_SIZE) - .map(|events| { + .flat_map(|events| { let mut params = Vec::with_capacity(events.len() * 2 + 1); params.push(BindParam::Integer(ScheduleKind::At.as_i64())); - for event in events { + for (event, _) in events { params.push(BindParam::Text(event.event_id.clone())); params.push(BindParam::Integer(event.trigger_at)); } - SqliteBatchStatement { - sql: claim_one_shots_sql(events.len()), - params: Some(params), - } + let delete_contexts = SqliteBatchStatement { + sql: delete_schedule_trace_contexts_sql(events.len()), + params: Some( + events + .iter() + .map(|(event, _)| { + BindParam::Text(schedule_trace_context_key(&event.event_id)) + }) + .collect(), + ), + }; + [ + delete_contexts, + SqliteBatchStatement { + sql: claim_one_shots_sql(events.len()), + params: Some(params), + }, + ] }) .collect::>(); if !claim_statements.is_empty() { @@ -621,7 +670,7 @@ impl ActorContext { .context("claim due one-shot schedules")?; } let mut dispatches = Vec::new(); - for event in due_schedules { + for (event, origin) in due_schedules { if event.kind == ScheduleKind::At { dispatches.push(DueScheduleDispatch { event_id: event.event_id.clone(), @@ -635,6 +684,7 @@ impl ActorContext { fired_at: now_ms, }, history_id: None, + origin, }); continue; } @@ -708,6 +758,7 @@ impl ActorContext { fired_at: now_ms, }, history_id, + origin, }); } self.mark_schedule_dirty(); @@ -1353,6 +1404,90 @@ fn read_stored_schedule(row: &[ColumnValue]) -> Result { } } +fn read_due_schedule(row: &[ColumnValue]) -> Result<(StoredSchedule, TraceOrigin)> { + let event = read_stored_schedule(row)?; + let origin = read_optional_blob(row, 10, "schedule trace context")? + .and_then(|payload| { + decode_latest_with_embedded_version::( + &payload, + "schedule trace context", + ) + .inspect_err(|error| { + tracing::warn!( + event_id = %event.event_id, + ?error, + "ignoring undecodable schedule trace context" + ); + }) + .ok() + }) + .map_or_else(TraceOrigin::default, TraceOrigin::from); + Ok((event, origin)) +} + +impl From for TraceOrigin { + fn from(context: ScheduleTraceContextData) -> Self { + Self { + ray_id: context.ray_id, + traceparent: context.traceparent, + tracestate: context.tracestate, + } + } +} + +impl From for ScheduleTraceContextData { + fn from(origin: TraceOrigin) -> Self { + Self { + ray_id: origin.ray_id, + traceparent: origin.traceparent, + tracestate: origin.tracestate, + } + } +} + +/// Stores the defining invocation's trace context beside a schedule row, or +/// clears a stale one when the definer carried no context. +fn append_schedule_trace_context_upsert( + statements: &mut Vec, + event_id: &str, + origin: TraceOrigin, +) -> Result<()> { + let key = schedule_trace_context_key(event_id); + if origin.is_empty() { + statements.push(SqliteBatchStatement { + sql: DELETE_SCHEDULE_TRACE_CONTEXT_SQL.to_owned(), + params: Some(vec![BindParam::Text(key)]), + }); + return Ok(()); + } + let payload = encode_latest_with_embedded_version::( + ScheduleTraceContextData::from(origin), + SCHEDULE_TRACE_CONTEXT_VERSION, + "schedule trace context", + )?; + statements.push(SqliteBatchStatement { + sql: UPSERT_SCHEDULE_TRACE_CONTEXT_SQL.to_owned(), + params: Some(vec![BindParam::Text(key), BindParam::Blob(payload)]), + }); + Ok(()) +} + +/// Removes a schedule's trace context once its row is gone. Runs after the +/// row delete in the same batch so a mismatched delete leaves the context alone. +fn delete_orphan_schedule_trace_context(event_id: &str) -> SqliteBatchStatement { + SqliteBatchStatement { + sql: DELETE_ORPHAN_SCHEDULE_TRACE_CONTEXT_SQL.to_owned(), + params: Some(vec![ + BindParam::Text(schedule_trace_context_key(event_id)), + BindParam::Text(event_id.to_owned()), + ]), + } +} + +fn schedule_trace_context_key(event_id: &str) -> String { + format!("{SCHEDULE_TRACE_CONTEXT_META_PREFIX}{event_id}") +} + fn read_cron_fire(row: &[ColumnValue]) -> Result { let error_group = read_optional_text(row, 5, "error_group")?; let error_code = read_optional_text(row, 6, "error_code")?; diff --git a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs index 149e32045e..3ef5659075 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs @@ -11,6 +11,7 @@ use opentelemetry::trace::{ SpanContext, SpanId, TraceContextExt as _, TraceFlags, TraceId, TraceState, }; use parking_lot::Mutex; +use rivetkit_client_protocol::telemetry_headers::format_traceparent; use tracing_opentelemetry::OpenTelemetrySpanExt as _; use crate::ActorContext; @@ -116,6 +117,24 @@ const OPERATION_ABANDONED_ERROR_TYPE: &str = "actor.operation_abandoned"; /// been recorded, so work that settles before the reply changes nothing. pub(crate) struct InvocationWorkGuard(ActorInvocationTelemetry); +/// Where later work came from: the ray ID of the invocation that caused it and +/// the span that was active there. Persisted beside schedules and queue +/// messages so the work they cause can link back to its origin. +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct TraceOrigin { + pub(crate) ray_id: Option, + pub(crate) traceparent: Option, + pub(crate) tracestate: Option, +} + +impl TraceOrigin { + /// True when there is nothing to persist: the work was caused outside any + /// traced invocation. + pub fn is_empty(&self) -> bool { + self.ray_id.is_none() && self.traceparent.is_none() && self.tracestate.is_none() + } +} + /// Active actor invocation fields exposed to foreign-runtime adapters. #[doc(hidden)] #[derive(Clone, Debug)] @@ -204,16 +223,24 @@ impl ActorInvocation { InvocationType::Action, incoming.ray_id, incoming.remote_parent, + None, ) } - pub(crate) fn start_scheduled(ctx: &ActorContext, action_name: &str) -> Self { + pub(crate) fn start_scheduled( + ctx: &ActorContext, + action_name: &str, + origin: TraceOrigin, + ) -> Self { + let origin_parent = + parse_remote_parent(origin.traceparent.as_deref(), origin.tracestate.as_deref()); Self::start( ctx, action_name, InvocationType::Scheduled, + origin.ray_id, None, - None, + origin_parent, ) } @@ -223,6 +250,7 @@ impl ActorInvocation { invocation_type: InvocationType, ray_id: Option, parent: Option, + link: Option, ) -> Self { let identity = ctx.telemetry_identity(); // Use bounded names for both spans and metrics. @@ -247,6 +275,9 @@ impl ActorInvocation { if let Some(parent) = parent { span.set_parent(opentelemetry::Context::new().with_remote_span_context(parent)); } + if let Some(link) = link { + span.add_link(link); + } Some(span) } else { None @@ -362,27 +393,7 @@ impl ActorInvocationTelemetry { let span = active.span.lock().clone().and_then(|span| { let context = span.context(); let context_span = context.span(); - let span_context = context_span.span_context(); - if !span_context.is_valid() { - return None; - } - let tracestate = span_context.trace_state().header(); - Some(ActorInvocationSpanContext { - trace_id: span_context.trace_id().to_string(), - span_id: span_context.span_id().to_string(), - trace_flags: span_context.trace_flags().to_u8(), - traceparent: format!( - "00-{}-{}-{:02x}", - span_context.trace_id(), - span_context.span_id(), - span_context.trace_flags().to_u8(), - ), - tracestate: if tracestate.is_empty() { - None - } else { - Some(tracestate) - }, - }) + w3c_span_context(context_span.span_context()) }); Some(ActorInvocationTraceContext { @@ -391,6 +402,31 @@ impl ActorInvocationTelemetry { }) } + /// Trace origin work caused by this invocation records: the invocation's + /// ray ID, and the application span active in the host runtime at that + /// moment, or the invocation span when there was none. Work that links + /// back to it then points at the code that caused it rather than at the + /// whole invocation around that code. + pub(crate) fn trace_origin(&self) -> TraceOrigin { + let Some(context) = self.trace_context() else { + return TraceOrigin::default(); + }; + let span = self + .1 + .as_ref() + .and_then(w3c_span_context) + .or(context.span); + let (traceparent, tracestate) = match span { + Some(span) => (Some(span.traceparent), span.tracestate), + None => (None, None), + }; + TraceOrigin { + ray_id: context.ray_id, + traceparent, + tracestate, + } + } + pub(crate) fn start_sqlite(&self, operation: SqliteOperation) -> Option { let parent = self.active()?.span.lock().clone()?; let span = tracing::info_span!( @@ -489,6 +525,26 @@ impl Drop for SqliteOperationSpan { } } +/// W3C fields of a span context, or nothing when it is not valid and so +/// carries nothing worth propagating. +fn w3c_span_context(span_context: &SpanContext) -> Option { + if !span_context.is_valid() { + return None; + } + let tracestate = span_context.trace_state().header(); + Some(ActorInvocationSpanContext { + trace_id: span_context.trace_id().to_string(), + span_id: span_context.span_id().to_string(), + trace_flags: span_context.trace_flags().to_u8(), + traceparent: format_traceparent( + span_context.trace_id(), + span_context.span_id(), + span_context.trace_flags().to_u8(), + ), + tracestate: (!tracestate.is_empty()).then_some(tracestate), + }) +} + /// Records the terminal status and error identity of a finished span. fn record_outcome(span: &tracing::Span, error: Option<&anyhow::Error>) { span.record( diff --git a/rivetkit-rust/packages/rivetkit-core/tests/sql_efficiency.rs b/rivetkit-rust/packages/rivetkit-core/tests/sql_efficiency.rs index eedba335de..eb335bd614 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/sql_efficiency.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/sql_efficiency.rs @@ -294,6 +294,12 @@ fn query_catalog() -> Vec { bound: "the legacy actor snapshot containing the source schedule vector is capped at 256 KiB", }]), }, + QueryCase { + id: "migration.reset_schedule_trace_contexts", + sql: internal_storage::RESET_SCHEDULE_TRACE_CONTEXTS_SQL.into(), + params: vec![], + expectation: indexed(None, &["_rivet_meta"]), + }, QueryCase { id: "queue.next_id", sql: internal_storage::LOAD_QUEUE_NEXT_ID_SQL.into(), @@ -488,6 +494,21 @@ fn query_catalog() -> Vec { params: vec![text("at:00000000"), 0_i64.into()], expectation: indexed(None, all_schedules), }, + QueryCase { + id: "schedule.delete_orphan_trace_context", + sql: queries::DELETE_ORPHAN_SCHEDULE_TRACE_CONTEXT_SQL.into(), + params: vec![ + text("schedule_trace_context:at:00000000"), + text("at:00000000"), + ], + expectation: indexed(None, &["_rivet_meta", "_rivet_schedule_events"]), + }, + QueryCase { + id: "schedule.delete_trace_context", + sql: queries::DELETE_SCHEDULE_TRACE_CONTEXT_SQL.into(), + params: vec![text("schedule_trace_context:at:00000000")], + expectation: indexed(None, &["_rivet_meta"]), + }, QueryCase { id: "schedule.get_one_shot", sql: queries::GET_SCHEDULED_EVENT_SQL.into(), @@ -558,7 +579,10 @@ fn query_catalog() -> Vec { id: "schedule.due", sql: queries::TAKE_DUE_SCHEDULES_SQL.into(), params: vec![5_i64.into()], - expectation: indexed(Some("_rivet_schedule_events_trigger_at"), all_schedules), + expectation: indexed( + Some("_rivet_schedule_events_trigger_at"), + &["_rivet_schedule_events", "_rivet_meta"], + ), }, QueryCase { id: "schedule.claim_one_shots", @@ -574,6 +598,16 @@ fn query_catalog() -> Vec { ], expectation: indexed(None, all_schedules), }, + QueryCase { + id: "schedule.delete_claimed_trace_contexts", + sql: queries::delete_schedule_trace_contexts_sql(3), + params: vec![ + text("schedule_trace_context:at:00000000"), + text("schedule_trace_context:at:00000003"), + text("schedule_trace_context:at:00000006"), + ], + expectation: indexed(None, &["_rivet_meta"]), + }, QueryCase { id: "schedule.advance_skipped", sql: queries::ADVANCE_SKIPPED_SCHEDULE_SQL.into(), From f9cb8bfd9e90bbaf0b293d5855e7a9dfd94a83eb Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Wed, 9 Sep 2026 22:24:59 +0400 Subject: [PATCH 12/21] feat(rivetkit-core): trace raw HTTP request invocations --- .../rivetkit-core/src/actor/context.rs | 1 + .../rivetkit-core/src/actor/messages.rs | 11 + .../rivetkit-core/src/actor/metrics.rs | 20 +- .../packages/rivetkit-core/src/actor/task.rs | 98 ++++--- .../rivetkit-core/src/actor/work_registry.rs | 9 + .../packages/rivetkit-core/src/telemetry.rs | 127 +++++++-- .../tests/integration/counter.rs | 6 +- .../integration/sqlite_corruption_fuzz.rs | 6 +- rivetkit-rust/packages/rivetkit/src/event.rs | 6 +- rivetkit-rust/packages/rivetkit/src/start.rs | 10 +- .../packages/rivetkit-napi/src/http.rs | 6 +- .../rivetkit-napi/src/napi_actor_events.rs | 10 +- .../packages/rivetkit/src/registry/native.ts | 256 ++++++++++-------- 13 files changed, 383 insertions(+), 183 deletions(-) diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs index 6073122f8c..badc4490e6 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs @@ -1576,6 +1576,7 @@ impl ActorContext { } let region = match kind { ActorWorkKind::Action => self.internal_keep_awake_region(), + ActorWorkKind::DispatchReply => self.internal_keep_awake_region(), ActorWorkKind::KeepAwake => self.keep_awake_region_state(), ActorWorkKind::InternalKeepAwake => self.internal_keep_awake_region(), ActorWorkKind::WaitUntil => return None, diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs index b8382520ab..8389c5746c 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs @@ -313,6 +313,15 @@ pub enum ActorHttpResponse { Stream(StreamingResponse), } +impl ActorHttpResponse { + pub fn status(&self) -> u16 { + match self { + Self::Buffered(response) => response.status().as_u16(), + Self::Stream(response) => response.status, + } + } +} + impl From for ActorHttpResponse { fn from(value: Response) -> Self { Self::Buffered(value) @@ -400,6 +409,8 @@ pub enum ActorEvent { }, HttpRequest { request: Request, + /// Telemetry of the invocation this request runs as. See `Action`. + invocation_telemetry: Option, reply: Reply, }, QueueSend { diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/metrics.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/metrics.rs index 46129142eb..4a660e72b3 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/metrics.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/metrics.rs @@ -130,10 +130,12 @@ pub(crate) struct StartupTimer { finished: bool, } -#[derive(Clone, Copy, Debug)] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] pub(crate) enum InvocationType { Action, Scheduled, + /// A raw HTTP request served by the actor's `onRequest` handler. + Request, } impl InvocationType { @@ -141,14 +143,16 @@ impl InvocationType { match self { Self::Action => "action", Self::Scheduled => "scheduled", + Self::Request => "request", } } - /// OpenTelemetry span kind for this invocation. An action is entered from - /// outside the actor, while a scheduled fire originates inside it. + /// OpenTelemetry span kind for this invocation. An action or a raw HTTP request + /// is entered from outside the actor, while a scheduled fire originates + /// inside it. pub(crate) fn otel_kind(self) -> &'static str { match self { - Self::Action => "server", + Self::Action | Self::Request => "server", Self::Scheduled => "internal", } } @@ -1973,18 +1977,20 @@ impl ActorMetrics { } } + /// `action_label` is already bounded: the caller folds action names + /// through `label_action_name`, and a request invocation uses a fixed + /// name, which the fold would otherwise turn into `_OTHER`. pub(crate) fn record_invocation( &self, - action_name: &str, + action_label: &str, invocation_type: InvocationType, result: InvocationStatus, duration: Duration, ) { let actor_labels = self.actor_labels(); - let action_name = self.label_action_name(action_name); let labels = [ actor_labels[0], - action_name, + action_label, invocation_type.as_label(), result.as_label(), ]; diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs index e6d6b8697b..8a85a08fdd 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs @@ -57,6 +57,7 @@ use crate::actor::task_types::ShutdownKind; use crate::actor::work_registry::ActorWorkKind; use crate::error::{ActorLifecycle as ActorLifecycleError, ActorRuntime}; use crate::runtime::RuntimeSpawner; +use crate::telemetry::{ActorInvocation, IncomingInvocationContext}; #[cfg(test)] use crate::time::sleep; use crate::time::{Instant, sleep_until, timeout}; @@ -907,8 +908,7 @@ impl ActorTask { conn, reply, } => { - let invocation = - crate::telemetry::ActorInvocation::start_action(&self.ctx, &name, incoming); + let invocation = ActorInvocation::start_action(&self.ctx, &name, incoming); let invocation_telemetry = invocation.telemetry(); tracing::info!( actor_id = %self.ctx.actor_id(), @@ -939,34 +939,23 @@ impl ActorTask { self.log_dispatch_command_handled(command_kind, "enqueued"); let actor_id = self.ctx.actor_id().to_owned(); let ctx = self.ctx.clone(); - self.ctx.spawn_work(ActorWorkKind::Action, async move { - let result = match tracked_reply_rx.await { - Ok(result) => { - let result = - result.map_err(|error| ctx.attach_actor_to_error(error)); - tracing::info!( - actor_id = %actor_id, - action_name = %action_name_for_log, - ok = result.is_ok(), - "actor task: tracked reply received, forwarding" - ); - result - } - Err(_) => { - tracing::warn!( - actor_id = %actor_id, - action_name = %action_name_for_log, - "actor task: tracked reply dropped before completion" - ); - let error = ctx.attach_actor_to_error( - ActorLifecycleError::DroppedReply.build(), - ); - Err(error) - } - }; - invocation.finish(result.as_ref().err()); - let _ = reply.send(result); - }); + self.forward_tracked_reply( + ActorWorkKind::Action, + tracked_reply_rx, + reply, + move |result| { + let result = + result.map_err(|error| ctx.attach_actor_to_error(error)); + tracing::info!( + actor_id = %actor_id, + action_name = %action_name_for_log, + ok = result.is_ok(), + "actor task: tracked reply received, forwarding" + ); + invocation.finish(result.as_ref().err()); + result + }, + ); } Err(error) => { tracing::warn!( @@ -1010,17 +999,33 @@ impl ActorTask { } }, DispatchCommand::Http { request, reply } => { + let incoming = IncomingInvocationContext::from_http_headers(request.headers()); + let invocation = ActorInvocation::start_request(&self.ctx, &request, incoming); + let invocation_telemetry = invocation.telemetry(); + let (tracked_reply_tx, tracked_reply_rx) = oneshot::channel(); match self.send_actor_event( "dispatch_http", ActorEvent::HttpRequest { request, - reply: Reply::from(reply), + invocation_telemetry: Some(invocation_telemetry), + reply: Reply::from(tracked_reply_tx), }, ) { Ok(()) => { self.log_dispatch_command_handled(command_kind, "enqueued"); + self.forward_tracked_reply( + ActorWorkKind::DispatchReply, + tracked_reply_rx, + reply, + move |result| { + invocation.finish_request(result.as_ref()); + result + }, + ); } - Err(_error) => { + Err(error) => { + invocation.finish(Some(&error)); + let _ = reply.send(Err(error)); self.log_dispatch_command_handled(command_kind, "enqueue_failed"); } } @@ -1082,6 +1087,35 @@ impl ActorTask { } } + /// Waits for the runtime adapter's reply to one dispatched invocation and + /// forwards it to the caller. `on_reply` sees the reply first, finishes the + /// invocation, and returns what the caller gets. A reply channel closed + /// without an answer counts as a dropped reply, so every arm reports that + /// case the same way. + fn forward_tracked_reply( + &self, + kind: ActorWorkKind, + tracked_reply_rx: oneshot::Receiver>, + reply: oneshot::Sender>, + on_reply: impl FnOnce(Result) -> Result + Send + 'static, + ) { + let actor_id = self.ctx.actor_id().to_owned(); + self.ctx.spawn_work(kind, async move { + let result = match tracked_reply_rx.await { + Ok(result) => result, + Err(_) => { + tracing::warn!( + actor_id = %actor_id, + "actor task: tracked reply dropped before completion; the runtime adapter dropped its reply handle" + ); + Err(ActorLifecycleError::DroppedReply.build()) + } + }; + let result = on_reply(result); + let _ = reply.send(result); + }); + } + fn log_dispatch_command_handled(&self, command: &'static str, outcome: &'static str) { tracing::debug!( actor_id = %self.ctx.actor_id(), diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/work_registry.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/work_registry.rs index df3bfabf26..028042ef0a 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/work_registry.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/work_registry.rs @@ -17,6 +17,8 @@ use crate::actor::task_types::UserTaskKind; pub enum ActorWorkKind { /// Dispatched action work awaiting its reply from the runtime adapter. Action, + /// Non-action invocation work awaiting its reply from the runtime adapter. + DispatchReply, /// User work that keeps the actor out of idle sleep while it runs. KeepAwake, /// Runtime-owned work that should behave like keep-awake without exposing a user API. @@ -54,6 +56,12 @@ impl ActorWorkKind { aborts_at_shutdown_deadline: true, user_task_kind: Some(UserTaskKind::Action), }, + ActorWorkKind::DispatchReply => ActorWorkPolicy { + blocks_idle_sleep: true, + drains_shutdown_grace: true, + aborts_at_shutdown_deadline: true, + user_task_kind: None, + }, ActorWorkKind::KeepAwake => ActorWorkPolicy { blocks_idle_sleep: true, drains_shutdown_grace: true, @@ -97,6 +105,7 @@ impl ActorWorkKind { pub(crate) fn label(self) -> &'static str { match self { ActorWorkKind::Action => "action", + ActorWorkKind::DispatchReply => "dispatch_reply", ActorWorkKind::KeepAwake => "keep_awake", ActorWorkKind::InternalKeepAwake => "internal_keep_awake", ActorWorkKind::WaitUntil => "wait_until", diff --git a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs index 3ef5659075..96acc89b37 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs @@ -63,6 +63,19 @@ fn invocation_ray_id(headers: &http::HeaderMap) -> Option { rivetkit_client_protocol::telemetry_headers::bounded_ray_id(value).map(str::to_owned) } +/// Name a request invocation is reported under, in place of an action name. +/// It cannot collide with an action, because the metric and span carry the +/// invocation type beside it. +const REQUEST_INVOCATION_NAME: &str = "onRequest"; + +/// What an invocation ran, which decides its name and the attributes that +/// identify it on the span. +#[derive(Clone, Copy, Debug)] +enum InvocationSubject<'a> { + Action(&'a str), + Request { method: &'a str }, +} + /// Owns the complete lifecycle of one actor invocation. #[derive(Debug)] pub(crate) struct ActorInvocation { @@ -219,7 +232,7 @@ impl ActorInvocation { ) -> Self { Self::start( ctx, - action_name, + InvocationSubject::Action(action_name), InvocationType::Action, incoming.ray_id, incoming.remote_parent, @@ -236,7 +249,7 @@ impl ActorInvocation { parse_remote_parent(origin.traceparent.as_deref(), origin.tracestate.as_deref()); Self::start( ctx, - action_name, + InvocationSubject::Action(action_name), InvocationType::Scheduled, origin.ray_id, None, @@ -244,9 +257,29 @@ impl ActorInvocation { ) } + /// Starts the invocation for one raw HTTP request served by `onRequest`. + /// The span is named after the handler rather than the path, because a + /// path is caller-supplied and would make the name a cardinality surface. + pub(crate) fn start_request( + ctx: &ActorContext, + request: &crate::actor::messages::Request, + incoming: IncomingInvocationContext, + ) -> Self { + Self::start( + ctx, + InvocationSubject::Request { + method: request.method().as_str(), + }, + InvocationType::Request, + incoming.ray_id, + incoming.remote_parent, + None, + ) + } + fn start( ctx: &ActorContext, - action_name: &str, + subject: InvocationSubject<'_>, invocation_type: InvocationType, ray_id: Option, parent: Option, @@ -254,7 +287,11 @@ impl ActorInvocation { ) -> Self { let identity = ctx.telemetry_identity(); // Use bounded names for both spans and metrics. - let action_name = ctx.metrics().label_action_name(action_name).to_owned(); + let (action_name, http_method) = match subject { + InvocationSubject::Action(name) => (ctx.metrics().label_action_name(name), None), + InvocationSubject::Request { method } => (REQUEST_INVOCATION_NAME, Some(method)), + }; + let action_name = action_name.to_owned(); let span = if tracing::enabled!(target: "rivetkit::telemetry", tracing::Level::INFO) { let span = tracing::info_span!( target: "rivetkit::telemetry", @@ -266,12 +303,18 @@ impl ActorInvocation { rivet.actor.id = %identity.actor_id, rivet.actor.name = %identity.actor_name, rivet.actor.key = %identity.actor_key, - rivet.action.name = %action_name, + rivet.action.name = tracing::field::Empty, rivet.ray.id = tracing::field::Empty, + http.request.method = tracing::field::Empty, + http.response.status_code = tracing::field::Empty, otel.status_code = tracing::field::Empty, error.type = tracing::field::Empty, ); span.record("rivet.ray.id", ray_id.as_deref()); + match http_method { + Some(method) => span.record("http.request.method", method), + None => span.record("rivet.action.name", &action_name), + }; if let Some(parent) = parent { span.set_parent(opentelemetry::Context::new().with_remote_span_context(parent)); } @@ -303,11 +346,41 @@ impl ActorInvocation { ); } + /// Finishes a request invocation with the HTTP status the handler + /// answered, which is recorded on the span beside the outcome. A 5xx + /// answer counts as a failed invocation with the status as its error + /// identity, following the HTTP server span convention. An error means no + /// response was produced, so only the error identity is recorded. + pub(crate) fn finish_request( + mut self, + response: std::result::Result<&crate::actor::messages::ActorHttpResponse, &anyhow::Error>, + ) { + match response { + Ok(response) => { + let status = response.status(); + self.telemetry.record_http_status(status); + if status >= 500 { + self.finish_with_failure(InvocationFailure::HttpStatus(status)); + } else { + self.finish_with_status(InvocationStatus::Ok, None); + } + } + Err(error) => self.finish(Some(error)), + } + } + fn finish_with_status(&mut self, status: InvocationStatus, error: Option<&anyhow::Error>) { let Some(span) = self.telemetry.claim_terminal() else { return; }; - self.record_finished(span, status, error); + self.record_finished(span, status, error.map(InvocationFailure::Error)); + } + + fn finish_with_failure(&mut self, failure: InvocationFailure<'_>) { + let Some(span) = self.telemetry.claim_terminal() else { + return; + }; + self.record_finished(span, InvocationStatus::Error, Some(failure)); } /// Records the terminal metric and span status of an invocation whose @@ -318,7 +391,7 @@ impl ActorInvocation { &self, span: Option, status: InvocationStatus, - error: Option<&anyhow::Error>, + failure: Option>, ) { self.metrics.record_invocation( &self.action_name, @@ -327,13 +400,27 @@ impl ActorInvocation { self.started_at.elapsed(), ); if let Some(span) = span { - record_outcome(&span, error); + match failure { + Some(InvocationFailure::HttpStatus(status)) => { + span.record("otel.status_code", "ERROR"); + span.record("error.type", status.to_string()); + } + Some(InvocationFailure::Error(error)) => record_outcome(&span, Some(error)), + None => record_outcome(&span, None), + } self.telemetry.mark_reply_sent(&span); } self.telemetry.release_span_if_settled(); } } +/// Why an invocation is recorded as failed: an error crossing the runtime +/// boundary, or a request the handler answered with a server error status. +enum InvocationFailure<'a> { + Error(&'a anyhow::Error), + HttpStatus(u16), +} + impl Drop for ActorInvocation { fn drop(&mut self) { // `finish` consumes the invocation, so this runs on the completed path @@ -343,7 +430,11 @@ impl Drop for ActorInvocation { return; }; let error = crate::error::ActorLifecycle::DroppedReply.build(); - self.record_finished(span, InvocationStatus::Dropped, Some(&error)); + self.record_finished( + span, + InvocationStatus::Dropped, + Some(InvocationFailure::Error(&error)), + ); } } @@ -373,10 +464,14 @@ impl ActorInvocationTelemetry { traceparent: Option<&str>, tracestate: Option<&str>, ) -> Self { - Self( - self.0.clone(), - parse_remote_parent(traceparent, tracestate), - ) + Self(self.0.clone(), parse_remote_parent(traceparent, tracestate)) + } + + /// Records the status a request invocation answered with. + fn record_http_status(&self, status: u16) { + if let Some(span) = self.0.span.lock().as_ref() { + span.record("http.response.status_code", status); + } } /// Registers work that outlives the reply, so the invocation span stays @@ -411,11 +506,7 @@ impl ActorInvocationTelemetry { let Some(context) = self.trace_context() else { return TraceOrigin::default(); }; - let span = self - .1 - .as_ref() - .and_then(w3c_span_context) - .or(context.span); + let span = self.1.as_ref().and_then(w3c_span_context).or(context.span); let (traceparent, tracestate) = match span { Some(span) => (Some(span.traceparent), span.tracestate), None => (None, None), diff --git a/rivetkit-rust/packages/rivetkit-core/tests/integration/counter.rs b/rivetkit-rust/packages/rivetkit-core/tests/integration/counter.rs index a2c5b1247e..8575aea9a7 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/integration/counter.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/integration/counter.rs @@ -86,7 +86,11 @@ fn counter_factory() -> ActorFactory { ActorEvent::RunGracefulCleanup { reason: _, reply } => { reply.send(Ok(())); } - ActorEvent::HttpRequest { request: _, reply } => { + ActorEvent::HttpRequest { + request: _, + invocation_telemetry: _, + reply, + } => { reply.send(Err(anyhow::anyhow!("http requests are not handled"))); } ActorEvent::QueueSend { diff --git a/rivetkit-rust/packages/rivetkit-core/tests/integration/sqlite_corruption_fuzz.rs b/rivetkit-rust/packages/rivetkit-core/tests/integration/sqlite_corruption_fuzz.rs index c8b9c2c5e6..3bf89e3955 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/integration/sqlite_corruption_fuzz.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/integration/sqlite_corruption_fuzz.rs @@ -579,7 +579,11 @@ fn sqlite_fuzz_factory() -> ActorFactory { ActorEvent::RunGracefulCleanup { reason: _, reply } => { reply.send(Ok(())); } - ActorEvent::HttpRequest { request: _, reply } => { + ActorEvent::HttpRequest { + request: _, + invocation_telemetry: _, + reply, + } => { reply.send(Err(anyhow::anyhow!("http requests are not handled"))); } ActorEvent::QueueSend { diff --git a/rivetkit-rust/packages/rivetkit/src/event.rs b/rivetkit-rust/packages/rivetkit/src/event.rs index e7f6ff8ead..82ec55df3d 100644 --- a/rivetkit-rust/packages/rivetkit/src/event.rs +++ b/rivetkit-rust/packages/rivetkit/src/event.rs @@ -106,7 +106,11 @@ impl RuntimeEvent { scheduled_fire, reply: Some(reply), }), - ActorEvent::HttpRequest { request, reply } => Self::Http(HttpCall { + ActorEvent::HttpRequest { + request, + invocation_telemetry: _, + reply, + } => Self::Http(HttpCall { request: Some(request), reply: Some(reply), }), diff --git a/rivetkit-rust/packages/rivetkit/src/start.rs b/rivetkit-rust/packages/rivetkit/src/start.rs index 8a5eb88f77..ed6b2fb75d 100644 --- a/rivetkit-rust/packages/rivetkit/src/start.rs +++ b/rivetkit-rust/packages/rivetkit/src/start.rs @@ -385,7 +385,11 @@ async fn handle_actor_event( } } } - ActorEvent::HttpRequest { request, reply } => { + ActorEvent::HttpRequest { + request, + invocation_telemetry: _, + reply, + } => { if let Some(http_pools) = http_pools { let class = A::classify_http_request(&request); let Some(permit) = http_pools.try_acquire(class) else { @@ -1224,6 +1228,7 @@ mod tests { let (reply_tx, reply_rx) = oneshot::channel(); tx.send(ActorEvent::HttpRequest { request: rivetkit_core::Request::default(), + invocation_telemetry: None, reply: reply_tx.into(), }) .expect("send http event"); @@ -1248,6 +1253,7 @@ mod tests { let (reply_tx, reply_rx) = oneshot::channel(); tx.send(ActorEvent::HttpRequest { request: Request::default(), + invocation_telemetry: None, reply: reply_tx.into(), }) .expect("send http event"); @@ -1286,6 +1292,7 @@ mod tests { tx.send(ActorEvent::HttpRequest { request: Request::from_parts("GET", "/standard", Default::default(), Vec::new()) .expect("standard request"), + invocation_telemetry: None, reply: standard_tx.into(), }) .expect("send standard http event"); @@ -1293,6 +1300,7 @@ mod tests { tx.send(ActorEvent::HttpRequest { request: Request::from_parts("GET", "/live", Default::default(), Vec::new()) .expect("live request"), + invocation_telemetry: None, reply: live_tx.into(), }) .expect("send live http event"); diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/http.rs b/rivetkit-typescript/packages/rivetkit-napi/src/http.rs index 80068e10ef..ae89bb04ad 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/http.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/http.rs @@ -30,6 +30,7 @@ pub struct JsHttpResponse { #[derive(Clone)] pub(crate) struct HttpRequestPayload { pub(crate) ctx: CoreActorContext, + pub(crate) telemetry: Option, pub(crate) request: Request, pub(crate) cancel_token: Option, pub(crate) response_stream: Option, @@ -220,7 +221,10 @@ pub(crate) fn build_http_request_payload( payload: HttpRequestPayload, ) -> napi::Result> { let mut object = env.create_object()?; - object.set("ctx", ActorContext::new(payload.ctx))?; + object.set( + "ctx", + ActorContext::new(payload.ctx.with_invocation_telemetry(payload.telemetry)), + )?; object.set("request", build_request_object(env, payload.request)?)?; match payload.response_stream { Some(response_stream) => object.set("responseBodyStream", response_stream)?, diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs b/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs index e6aa7d0e34..91aebc7602 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs @@ -462,7 +462,11 @@ pub(crate) async fn dispatch_event( } }); } - ActorEvent::HttpRequest { request, reply } => { + ActorEvent::HttpRequest { + request, + invocation_telemetry, + reply, + } => { let Some(callback) = bindings.on_request.clone() else { reply.send(Err(missing_callback("onRequest"))); return; @@ -476,7 +480,7 @@ pub(crate) async fn dispatch_event( "Action timed out", None, timeout, - call_http_request(&callback, &ctx, request), + call_http_request(&callback, &ctx, invocation_telemetry, request), ) .await }); @@ -1226,6 +1230,7 @@ async fn call_on_before_action_response( async fn call_http_request( callback: &crate::actor_factory::CallbackTsfn, ctx: &ActorContext, + telemetry: Option, request: rivetkit_core::Request, ) -> Result { let request_cancel_token = request.cancellation_token(); @@ -1234,6 +1239,7 @@ async fn call_http_request( callback, HttpRequestPayload { ctx: ctx.inner().clone(), + telemetry, request, cancel_token: Some(request_cancel_token), response_stream: None, diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts index 5bce0db6e5..e3d4f5c674 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts @@ -5072,129 +5072,147 @@ export function buildNativeFactory( try { const { ctx, request, cancelToken, responseBodyStream } = unwrapTsfnPayload(error, payload); - const inspectorResponse = - await maybeHandleNativeInspectorRequest(ctx, request); - if (inspectorResponse) { - await cancelNativeHttpRequestBody(request.bodyStream); - return ( - await convertNativeHttpResponse( - inspectorResponse, - responseBodyStream, - ) - ).response; - } - - if (typeof config.onRequest !== "function") { - await cancelNativeHttpRequestBody(request.bodyStream); - return ( - await convertNativeHttpResponse( - new Response(null, { status: 404 }), - responseBodyStream, - ) - ).response; - } + return await runtime.runWithActorInvocationContext( + ctx, + async () => { + const inspectorResponse = + await maybeHandleNativeInspectorRequest( + ctx, + request, + ); + if (inspectorResponse) { + await cancelNativeHttpRequestBody( + request.bodyStream, + ); + return ( + await convertNativeHttpResponse( + inspectorResponse, + responseBodyStream, + ) + ).response; + } - const requestAbortController = new AbortController(); - const handlerRequest = buildNativeHttpRequest({ - ...request, - abortController: requestAbortController, - }); - const rawConnParams = - handlerRequest.headers.get(HEADER_CONN_PARAMS); - let requestCtx: - | ReturnType - | undefined; - let conn: ConnHandle | undefined; - let removeRequestAbortListener: (() => void) | undefined; - let cleanupDeferredToBody = false; - let cleanedUp = false; - const cleanupRequest = async () => { - if (cleanedUp) return; - cleanedUp = true; - removeRequestAbortListener?.(); - try { - await requestCtx?.dispose(); - } finally { - if (conn) { - await runtime.connDisconnect(conn); + if (typeof config.onRequest !== "function") { + await cancelNativeHttpRequestBody( + request.bodyStream, + ); + return ( + await convertNativeHttpResponse( + new Response(null, { status: 404 }), + responseBodyStream, + ) + ).response; } - } - }; - try { - const connParams = validateConnParams( - schemaConfig.connParamsSchema, - rawConnParams - ? JSON.parse(rawConnParams) - : undefined, - ); - conn = await callNative(() => - runtime.actorConnectConn( - ctx, - encodeValue(connParams), - request, - ), - ); - requestCtx = makeConnCtx( - ctx, - conn, - handlerRequest, - cancelToken, - ); - const ctxAbortSignal = requestCtx.abortSignal; - const abortRequest = () => - requestAbortController.abort(ctxAbortSignal.reason); - if (ctxAbortSignal.aborted) { - abortRequest(); - } else { - ctxAbortSignal.addEventListener( - "abort", - abortRequest, - { once: true }, - ); - removeRequestAbortListener = () => - ctxAbortSignal.removeEventListener( - "abort", - abortRequest, + + const requestAbortController = + new AbortController(); + const handlerRequest = buildNativeHttpRequest({ + ...request, + abortController: requestAbortController, + }); + const rawConnParams = + handlerRequest.headers.get(HEADER_CONN_PARAMS); + let requestCtx: + | ReturnType + | undefined; + let conn: ConnHandle | undefined; + let removeRequestAbortListener: + | (() => void) + | undefined; + let cleanupDeferredToBody = false; + let cleanedUp = false; + const cleanupRequest = async () => { + if (cleanedUp) return; + cleanedUp = true; + removeRequestAbortListener?.(); + try { + await requestCtx?.dispose(); + } finally { + if (conn) { + await runtime.connDisconnect(conn); + } + } + }; + try { + const connParams = validateConnParams( + schemaConfig.connParamsSchema, + rawConnParams + ? JSON.parse(rawConnParams) + : undefined, ); - } - const response = await config.onRequest( - requestCtx, - handlerRequest, - ); - if (!isResponseLike(response)) { - throw new Error( - "onRequest handler must return a Response", - ); - } - const conversion = await convertNativeHttpResponse( - response, - responseBodyStream, - ); - if (conversion.bodyCompletion) { - cleanupDeferredToBody = true; - void conversion.bodyCompletion - .then(cleanupRequest) - .catch((cleanupError) => { - logger().error({ - msg: "failed to clean up native streaming http request", - error: cleanupError, - }); - }); - } - return conversion.response; - } finally { - try { - // Handler completion ends upload ownership even when - // the Web Request body is locked or partly consumed. - await cancelNativeHttpRequestBody( - request.bodyStream, - ); - } finally { - if (!cleanupDeferredToBody) { - await cleanupRequest(); + conn = await callNative(() => + runtime.actorConnectConn( + ctx, + encodeValue(connParams), + request, + ), + ); + requestCtx = makeConnCtx( + ctx, + conn, + handlerRequest, + cancelToken, + ); + const ctxAbortSignal = requestCtx.abortSignal; + const abortRequest = () => + requestAbortController.abort( + ctxAbortSignal.reason, + ); + if (ctxAbortSignal.aborted) { + abortRequest(); + } else { + ctxAbortSignal.addEventListener( + "abort", + abortRequest, + { once: true }, + ); + removeRequestAbortListener = () => + ctxAbortSignal.removeEventListener( + "abort", + abortRequest, + ); + } + const response = await config.onRequest( + requestCtx, + handlerRequest, + ); + if (!isResponseLike(response)) { + throw new Error( + "onRequest handler must return a Response", + ); + } + const conversion = + await convertNativeHttpResponse( + response, + responseBodyStream, + ); + if (conversion.bodyCompletion) { + cleanupDeferredToBody = true; + void conversion.bodyCompletion + .then(cleanupRequest) + .catch((cleanupError) => { + logger().error({ + msg: "failed to clean up native streaming http request", + error: cleanupError, + }); + }); + } + return conversion.response; + } finally { + try { + // Handler completion ends upload ownership even when + // the Web Request body is locked or partly consumed. + await cancelNativeHttpRequestBody( + request.bodyStream, + ); + } finally { + if (!cleanupDeferredToBody) { + await cleanupRequest(); + } + } } - } - } + }, + ); } catch (error) { logger().error({ msg: "native onRequest failed", From fc47ec89658cd8f7e5be67b894be1ae75e67f7cc Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Wed, 9 Sep 2026 22:35:18 +0400 Subject: [PATCH 13/21] feat(rivetkit-core): trace queue sends and persist message origins --- .../packages/actor-persist/src/versioned.rs | 50 ++++++ .../rivetkit-core/src/actor/config.rs | 13 ++ .../rivetkit-core/src/actor/context.rs | 17 +- .../src/actor/internal_storage/mod.rs | 76 +++++++-- .../src/actor/internal_storage/queries.rs | 15 +- .../rivetkit-core/src/actor/messages.rs | 2 + .../rivetkit-core/src/actor/metrics.rs | 29 +++- .../packages/rivetkit-core/src/actor/queue.rs | 85 +++++++++- .../packages/rivetkit-core/src/actor/task.rs | 55 ++++-- .../packages/rivetkit-core/src/lib.rs | 4 +- .../rivetkit-core/src/registry/http.rs | 3 + .../packages/rivetkit-core/src/telemetry.rs | 158 ++++++++++++++---- .../tests/integration/counter.rs | 1 + .../integration/sqlite_corruption_fuzz.rs | 1 + .../rivetkit-core/tests/sql_efficiency.rs | 18 ++ rivetkit-rust/packages/rivetkit/src/event.rs | 1 + rivetkit-rust/packages/rivetkit/src/start.rs | 1 + .../packages/rivetkit-napi/index.d.ts | 4 + .../rivetkit-napi/src/actor_factory.rs | 21 ++- .../rivetkit-napi/src/napi_actor_events.rs | 2 + .../packages/rivetkit-wasm/src/lib.rs | 13 ++ .../rivetkit/src/client/actor-conn.ts | 2 + .../rivetkit/src/client/actor-handle.ts | 4 + .../packages/rivetkit/src/client/queue.ts | 3 + .../rivetkit/src/registry/napi-runtime.ts | 13 +- .../packages/rivetkit/src/registry/native.ts | 118 +++++++------ .../packages/rivetkit/src/registry/runtime.ts | 1 + 27 files changed, 564 insertions(+), 146 deletions(-) diff --git a/rivetkit-rust/packages/actor-persist/src/versioned.rs b/rivetkit-rust/packages/actor-persist/src/versioned.rs index d8dde9dbde..a2e2c0c462 100644 --- a/rivetkit-rust/packages/actor-persist/src/versioned.rs +++ b/rivetkit-rust/packages/actor-persist/src/versioned.rs @@ -548,6 +548,56 @@ impl OwnedVersionedData for ScheduleTraceContext { } } +/// Trace origin of a queue message, stored beside its `_rivet_queue` row. It +/// has its own versioned type so its persisted meaning cannot be conflated with +/// `ScheduleTraceContext` even though both currently encode the same fields. +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +pub struct QueueTraceContextData { + pub ray_id: Option, + pub traceparent: Option, + pub tracestate: Option, +} + +pub enum QueueTraceContext { + V1(QueueTraceContextData), +} + +impl OwnedVersionedData for QueueTraceContext { + type Latest = QueueTraceContextData; + + fn wrap_latest(latest: Self::Latest) -> Self { + Self::V1(latest) + } + + fn unwrap_latest(self) -> Result { + match self { + Self::V1(data) => Ok(data), + } + } + + fn deserialize_version(payload: &[u8], version: u16) -> Result { + match version { + 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), + _ => bail!("invalid queue trace context version: {version}"), + } + } + + fn serialize_version(self, version: u16) -> Result> { + match (self, version) { + (Self::V1(data), 1) => serde_bare::to_vec(&data).map_err(Into::into), + (_, version) => bail!("unexpected queue trace context version: {version}"), + } + } + + fn deserialize_converters() -> Vec Result> { + Vec:: Result>::new() + } + + fn serialize_converters() -> Vec Result> { + Vec:: Result>::new() + } +} + impl OwnedVersionedData for RunWakeAt { type Latest = Option; diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/config.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/config.rs index e9519f8824..ae941ce7b7 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/config.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/config.rs @@ -55,6 +55,13 @@ pub struct ActionDefinition { pub name: String, } +/// A queue the actor declares. Names are bounded by this set wherever a queue +/// name becomes a telemetry dimension. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct QueueDefinition { + pub name: String, +} + /// Experimental SQLite profiling configuration. /// /// This entire configuration surface, including every field, is subject to @@ -178,6 +185,7 @@ pub struct ActorConfig { pub max_outgoing_message_size: u32, pub overrides: Option, pub actions: Vec, + pub queues: Vec, /// Author-declared inspector tab entries (custom tabs + built-in /// hides). Validated upstream (Zod / builder). pub inspector_tabs: Vec, @@ -212,6 +220,7 @@ pub struct ActorConfigInput { pub max_incoming_message_size: Option, pub max_outgoing_message_size: Option, pub actions: Option>, + pub queues: Option>, pub inspector_tabs: Option>, } @@ -289,6 +298,9 @@ impl ActorConfig { if let Some(actions) = config.actions { actor_config.actions = actions; } + if let Some(queues) = config.queues { + actor_config.queues = queues; + } if let Some(tabs) = config.inspector_tabs { actor_config.inspector_tabs = tabs; } @@ -364,6 +376,7 @@ impl Default for ActorConfig { max_outgoing_message_size: DEFAULT_MAX_OUTGOING_MESSAGE_SIZE, overrides: None, actions: Vec::new(), + queues: Vec::new(), inspector_tabs: Vec::new(), } } diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs index badc4490e6..0e257d3558 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs @@ -263,7 +263,11 @@ impl ActorContext { /// `traceparent` and `tracestate`. A handle that serves no invocation is /// returned unchanged, because it opens no spans. #[doc(hidden)] - pub fn with_application_span(&self, traceparent: Option<&str>, tracestate: Option<&str>) -> Self { + pub fn with_application_span( + &self, + traceparent: Option<&str>, + tracestate: Option<&str>, + ) -> Self { Self( self.0.clone(), self.1 @@ -340,6 +344,7 @@ impl ActorContext { let metrics = ActorMetrics::new_for_actor( name.clone(), config.actions.iter().map(|action| action.name.clone()), + config.queues.iter().map(|queue| queue.name.clone()), config.sqlite_profiling.clone(), ); #[cfg(feature = "sqlite-local")] @@ -780,7 +785,10 @@ impl ActorContext { /// the request that started them. #[cfg(not(feature = "wasm-runtime"))] pub fn wait_until(&self, future: impl Future + Send + 'static) { - let invocation = self.1.as_ref().map(crate::ActorInvocationTelemetry::hold_open); + let invocation = self + .1 + .as_ref() + .map(crate::ActorInvocationTelemetry::hold_open); self.spawn_work(ActorWorkKind::WaitUntil, async move { future.await; drop(invocation); @@ -794,7 +802,10 @@ impl ActorContext { #[cfg(feature = "wasm-runtime")] pub fn wait_until(&self, future: impl Future + 'static) { - let invocation = self.1.as_ref().map(crate::ActorInvocationTelemetry::hold_open); + let invocation = self + .1 + .as_ref() + .map(crate::ActorInvocationTelemetry::hold_open); self.spawn_work(ActorWorkKind::WaitUntil, async move { future.await; drop(invocation); diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/mod.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/mod.rs index a8faac3c7a..cf7223b02b 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/mod.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/mod.rs @@ -381,15 +381,18 @@ pub(crate) async fn load_queue_metadata(db: &SqliteDb) -> Result }) } +/// Writes a queue message and, when the sender carried one, its encoded trace +/// context, in one batch so a message never exists without its origin. pub(crate) async fn persist_queue_message( db: &SqliteDb, id: u64, next_id: u64, message: &PersistedQueueMessage, + trace_context: Option>, ) -> Result<()> { let id = i64::try_from(id).context("queue message id exceeds sqlite integer range")?; let next_id = i64::try_from(next_id).context("queue next id exceeds sqlite integer range")?; - db.execute_batch(vec![ + let mut statements = vec![ SqliteBatchStatement { sql: INSERT_QUEUE_MESSAGE_SQL.to_owned(), params: Some(vec![ @@ -407,12 +410,26 @@ pub(crate) async fn persist_queue_message( BindParam::Integer(next_id), ]), }, - ]) - .await - .context("persist internal queue message")?; + ]; + if let Some(trace_context) = trace_context { + statements.push(SqliteBatchStatement { + sql: UPSERT_QUEUE_TRACE_CONTEXT_SQL.to_owned(), + params: Some(vec![ + BindParam::Text(queue_trace_context_key(id)), + BindParam::Blob(trace_context), + ]), + }); + } + db.execute_batch(statements) + .await + .context("persist internal queue message")?; Ok(()) } +pub(crate) fn queue_trace_context_key(id: i64) -> String { + format!("queue_trace_context:{id}") +} + /// Persists imported queue rows without rewriting `queue_next_id` for every /// message. The importer writes that singleton once after every row is copied. pub(crate) async fn persist_queue_messages( @@ -646,6 +663,7 @@ fn decode_queue_message_rows(rows: &[Vec]) -> Result Result< return Ok(0); } - let mut statements = Vec::with_capacity(ids.len()); + // Each message row is followed by the delete of its trace context, so the + // two go in one batch and the row count below reads every other result. + let mut statements = Vec::with_capacity(ids.len() * 2); for id in ids { + let id = i64::try_from(*id).context("queue message id exceeds sqlite integer range")?; statements.push(SqliteBatchStatement { sql: DELETE_QUEUE_MESSAGE_SQL.to_owned(), - params: Some(vec![BindParam::Integer( - i64::try_from(*id).context("queue message id exceeds sqlite integer range")?, - )]), + params: Some(vec![BindParam::Integer(id)]), + }); + statements.push(SqliteBatchStatement { + sql: DELETE_QUEUE_TRACE_CONTEXT_SQL.to_owned(), + params: Some(vec![BindParam::Text(queue_trace_context_key(id))]), }); } let results = db .execute_batch(statements) .await .context("delete internal queue messages")?; - results.into_iter().try_fold(0u32, |deleted, result| { - let changes = u32::try_from(result.changes) - .context("deleted queue message count is outside u32 range")?; - deleted - .checked_add(changes) - .context("deleted queue message count exceeds u32 range") - }) + results + .into_iter() + .step_by(2) + .try_fold(0u32, |deleted, result| { + let changes = u32::try_from(result.changes) + .context("deleted queue message count is outside u32 range")?; + deleted + .checked_add(changes) + .context("deleted queue message count exceeds u32 range") + }) } pub(crate) async fn reset_queue(db: &SqliteDb) -> Result<()> { - db.execute(RESET_QUEUE_SQL, None) - .await - .context("reset internal queue")?; + db.execute_batch(vec![ + SqliteBatchStatement { + sql: RESET_QUEUE_SQL.to_owned(), + params: None, + }, + SqliteBatchStatement { + sql: RESET_QUEUE_TRACE_CONTEXTS_SQL.to_owned(), + params: None, + }, + ]) + .await + .context("reset internal queue")?; Ok(()) } @@ -689,6 +724,8 @@ pub(crate) async fn reset_queue(db: &SqliteDb) -> Result<()> { pub(crate) struct QueueMessageRow { pub id: u64, pub message: PersistedQueueMessage, + /// Encoded `QueueTraceContext`, when the sender carried one. + pub trace_context: Option>, } pub(crate) async fn user_kv_batch_get( @@ -1170,6 +1207,9 @@ pub(crate) async fn clear_imported_storage(db: &SqliteDb, actor_id: &str) -> Res db.execute(RESET_SCHEDULE_TRACE_CONTEXTS_SQL, None) .await .context("clear imported schedule trace contexts")?; + db.execute(RESET_QUEUE_TRACE_CONTEXTS_SQL, None) + .await + .context("clear imported queue trace contexts")?; Ok(()) } diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/queries.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/queries.rs index babe13ec6b..a5a94508b0 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/queries.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/queries.rs @@ -66,11 +66,10 @@ pub(crate) fn delete_schedule_trace_contexts_sql(event_count: usize) -> String { pub(crate) const LOAD_QUEUE_NEXT_ID_SQL: &str = "SELECT queue_next_id FROM _rivet_runtime WHERE id = 1"; pub(crate) const LOAD_QUEUE_STATS_SQL: &str = "SELECT COUNT(*), MAX(id) FROM _rivet_queue"; -pub(crate) const LOAD_QUEUE_MESSAGES_SQL: &str = - "SELECT id, name, body, created_at FROM _rivet_queue ORDER BY id"; -pub(crate) const LOAD_QUEUE_MESSAGES_LIMITED_SQL: &str = - "SELECT id, name, body, created_at FROM _rivet_queue ORDER BY id LIMIT ?"; -pub(crate) const LOAD_QUEUE_MESSAGES_FOR_NAME_SQL: &str = "SELECT id, name, body, created_at FROM _rivet_queue INDEXED BY _rivet_queue_name_id WHERE name = ? ORDER BY id LIMIT ?"; +// The trace-origin subquery uses a primary-key lookup per returned message. +pub(crate) const LOAD_QUEUE_MESSAGES_SQL: &str = "SELECT id, name, body, created_at, (SELECT value FROM _rivet_meta WHERE key = 'queue_trace_context:' || id) FROM _rivet_queue ORDER BY id"; +pub(crate) const LOAD_QUEUE_MESSAGES_LIMITED_SQL: &str = "SELECT id, name, body, created_at, (SELECT value FROM _rivet_meta WHERE key = 'queue_trace_context:' || id) FROM _rivet_queue ORDER BY id LIMIT ?"; +pub(crate) const LOAD_QUEUE_MESSAGES_FOR_NAME_SQL: &str = "SELECT id, name, body, created_at, (SELECT value FROM _rivet_meta WHERE key = 'queue_trace_context:' || id) FROM _rivet_queue INDEXED BY _rivet_queue_name_id WHERE name = ? ORDER BY id LIMIT ?"; pub(crate) const HAS_QUEUE_MESSAGES_SQL: &str = "SELECT 1 FROM _rivet_queue LIMIT 1"; pub(crate) const HAS_QUEUE_MESSAGES_FOR_NAME_SQL: &str = "SELECT 1 FROM _rivet_queue INDEXED BY _rivet_queue_name_id WHERE name = ? LIMIT 1"; @@ -82,13 +81,17 @@ pub(crate) fn load_queue_messages_by_ids_sql(id_count: usize) -> String { .collect::>() .join(", "); format!( - "SELECT id, name, body, created_at FROM _rivet_queue WHERE id IN ({placeholders}) ORDER BY id" + "SELECT id, name, body, created_at, (SELECT value FROM _rivet_meta WHERE key = 'queue_trace_context:' || id) FROM _rivet_queue WHERE id IN ({placeholders}) ORDER BY id" ) } pub(crate) const INSERT_QUEUE_MESSAGE_SQL: &str = "INSERT OR REPLACE INTO _rivet_queue (id, name, body, created_at) VALUES (?, ?, ?, ?)"; pub(crate) const DELETE_QUEUE_MESSAGE_SQL: &str = "DELETE FROM _rivet_queue WHERE id = ?"; pub(crate) const RESET_QUEUE_SQL: &str = "DELETE FROM _rivet_queue"; +pub(crate) const UPSERT_QUEUE_TRACE_CONTEXT_SQL: &str = "INSERT INTO _rivet_meta (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value"; +pub(crate) const DELETE_QUEUE_TRACE_CONTEXT_SQL: &str = "DELETE FROM _rivet_meta WHERE key = ?"; +pub(crate) const RESET_QUEUE_TRACE_CONTEXTS_SQL: &str = + "DELETE FROM _rivet_meta WHERE key >= 'queue_trace_context:' AND key < 'queue_trace_context;'"; pub(crate) const DELETE_USER_KV_SQL: &str = "DELETE FROM _rivet_user_kv WHERE key = ?"; pub(crate) const DELETE_USER_KV_RANGE_SQL: &str = diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs index 8389c5746c..a252c2d0cf 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs @@ -420,6 +420,8 @@ pub enum ActorEvent { request: Request, wait: bool, timeout_ms: Option, + /// Telemetry of the invocation this send runs as. See `Action`. + invocation_telemetry: Option, reply: Reply, }, WebSocketOpen { diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/metrics.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/metrics.rs index 4a660e72b3..b01601cea4 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/metrics.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/metrics.rs @@ -136,6 +136,8 @@ pub(crate) enum InvocationType { Scheduled, /// A raw HTTP request served by the actor's `onRequest` handler. Request, + /// A message sent into one of the actor's queues from outside it. + QueueSend, } impl InvocationType { @@ -144,16 +146,18 @@ impl InvocationType { Self::Action => "action", Self::Scheduled => "scheduled", Self::Request => "request", + Self::QueueSend => "queue_send", } } /// OpenTelemetry span kind for this invocation. An action or a raw HTTP request - /// is entered from outside the actor, while a scheduled fire originates - /// inside it. + /// is entered from outside the actor, a scheduled fire originates inside + /// it, and a queue send produces a message the actor consumes later. pub(crate) fn otel_kind(self) -> &'static str { match self { Self::Action | Self::Request => "server", Self::Scheduled => "internal", + Self::QueueSend => "producer", } } } @@ -189,6 +193,7 @@ impl InvocationStatus { struct ActorMetricInner { labels: ActorMetricLabels, action_names: BTreeSet, + queue_names: BTreeSet, #[cfg(feature = "sqlite-local")] sqlite_profiling: crate::SqliteProfilingConfig, #[cfg(feature = "sqlite-local")] @@ -1672,6 +1677,7 @@ impl ActorMetrics { Self::new_for_actor( actor_name, std::iter::empty(), + std::iter::empty(), crate::SqliteProfilingConfig::default(), ) } @@ -1681,12 +1687,18 @@ impl ActorMetrics { actor_name: impl Into, _sqlite_profiling: crate::SqliteProfilingConfig, ) -> Self { - Self::new_for_actor(actor_name, std::iter::empty(), _sqlite_profiling) + Self::new_for_actor( + actor_name, + std::iter::empty(), + std::iter::empty(), + _sqlite_profiling, + ) } pub(crate) fn new_for_actor( actor_name: impl Into, action_names: impl IntoIterator, + queue_names: impl IntoIterator, _sqlite_profiling: crate::SqliteProfilingConfig, ) -> Self { let labels = ActorMetricLabels { @@ -1709,6 +1721,7 @@ impl ActorMetrics { inner: Arc::new(ActorMetricInner { labels, action_names: action_names.into_iter().collect(), + queue_names: queue_names.into_iter().collect(), #[cfg(feature = "sqlite-local")] sqlite_profiling: _sqlite_profiling, #[cfg(feature = "sqlite-local")] @@ -1977,6 +1990,16 @@ impl ActorMetrics { } } + /// Folds an undeclared queue name the same way `label_action_name` folds + /// an action name, for the same reason: it arrives from the caller. + pub(crate) fn label_queue_name<'a>(&'a self, queue_name: &'a str) -> &'a str { + if self.inner.queue_names.contains(queue_name) { + queue_name + } else { + "_OTHER" + } + } + /// `action_label` is already bounded: the caller folds action names /// through `label_action_name`, and a request invocation uses a fixed /// name, which the fold would otherwise turn into `_OTHER`. diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/queue.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/queue.rs index 68ed09adaa..f549f0f485 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/queue.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/queue.rs @@ -10,6 +10,7 @@ use crate::time::{Instant, SystemTime, UNIX_EPOCH, sleep}; use anyhow::{Context, Result}; use rivet_error::RivetError; +use rivetkit_actor_persist::versioned::{QueueTraceContext, QueueTraceContextData}; use rivetkit_actor_persist::{generated::v4 as persist_v4, versioned as persist_versioned}; use serde::{Deserialize, Serialize}; #[cfg(not(target_arch = "wasm32"))] @@ -26,6 +27,9 @@ use crate::actor::persist::{ use crate::actor::task_types::UserTaskKind; #[cfg(target_arch = "wasm32")] use crate::error::ActorRuntime; +use crate::telemetry::{self, TraceOrigin}; + +const QUEUE_TRACE_CONTEXT_VERSION: u16 = 1; #[derive(Clone, Debug, Default)] pub struct QueueNextOpts { @@ -101,6 +105,10 @@ pub struct QueueMessage { pub name: String, pub body: Vec, pub created_at: i64, + /// Ray ID and span of the invocation that sent the message, which the span + /// covering its receipt links back to. Empty when the sender carried no + /// trace context. + pub trace_origin: TraceOrigin, completion: Option, } @@ -110,9 +118,30 @@ pub struct CompletableQueueMessage { pub name: String, pub body: Vec, pub created_at: i64, + pub trace_origin: TraceOrigin, completion: CompletionHandle, } +impl From for TraceOrigin { + fn from(context: QueueTraceContextData) -> Self { + Self { + ray_id: context.ray_id, + traceparent: context.traceparent, + tracestate: context.tracestate, + } + } +} + +impl From for QueueTraceContextData { + fn from(origin: TraceOrigin) -> Self { + Self { + ray_id: origin.ray_id, + traceparent: origin.traceparent, + tracestate: origin.tracestate, + } + } +} + #[derive(Clone)] struct CompletionHandle(Arc); @@ -281,6 +310,24 @@ impl ActorContext { in_flight_at: None, }; let encoded_message = encode_queue_message(&persisted).context("encode queue message")?; + // A send from inside an invocation records that invocation as the + // message's origin. A send from outside one has nothing to record. + let trace_origin = self + .invocation_telemetry() + .map(crate::ActorInvocationTelemetry::trace_origin) + .unwrap_or_default(); + let encoded_trace_context = if trace_origin.is_empty() { + None + } else { + Some( + encode_latest_with_embedded_version::( + QueueTraceContextData::from(trace_origin.clone()), + QUEUE_TRACE_CONTEXT_VERSION, + "queue trace context", + ) + .context("encode queue trace context")?, + ) + }; let config = self.config(); if encoded_message.len() > config.max_queue_message_size as usize { @@ -323,9 +370,14 @@ impl ActorContext { false }; - let persist_result = - internal_storage::persist_queue_message(self.sql(), id, metadata.next_id, &persisted) - .await; + let persist_result = internal_storage::persist_queue_message( + self.sql(), + id, + metadata.next_id, + &persisted, + encoded_trace_context, + ) + .await; if let Err(error) = persist_result { metadata.next_id = id; @@ -350,6 +402,7 @@ impl ActorContext { name: name.to_owned(), body: body.to_vec(), created_at, + trace_origin, completion: None, }) } @@ -653,6 +706,12 @@ impl ActorContext { return Ok(Vec::new()); } + // Keep receipt spans open until the batch is handed back. + let _receive_spans: Vec = selected + .iter() + .map(|message| telemetry::start_queue_receive(self, message)) + .collect(); + if completable { let queue_size = self.0.queue_metadata.lock().await.size; self.0 @@ -927,6 +986,7 @@ impl QueueMessage { name: self.name, body: self.body, created_at: self.created_at, + trace_origin: self.trace_origin, completion, }) } @@ -947,6 +1007,7 @@ impl CompletableQueueMessage { name: self.name, body: self.body, created_at: self.created_at, + trace_origin: self.trace_origin, completion: Some(self.completion), } } @@ -1051,11 +1112,29 @@ fn normalize_names(names: Option>) -> Option> { } fn queue_message_from_row(row: internal_storage::QueueMessageRow) -> QueueMessage { + // A malformed trace context is a telemetry defect, not a queue defect, so + // the message is still delivered and only its origin is lost. + let trace_origin = row + .trace_context + .as_deref() + .and_then(|payload| { + decode_latest_with_embedded_version::(payload, "queue trace context") + .inspect_err(|error| { + tracing::warn!( + message_id = row.id, + ?error, + "ignoring undecodable queue trace context" + ); + }) + .ok() + }) + .map_or_else(TraceOrigin::default, TraceOrigin::from); QueueMessage { id: row.id, name: row.message.name, body: row.message.body, created_at: row.message.created_at, + trace_origin, completion: None, } } diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs index 8a85a08fdd..e5bcdbcf41 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs @@ -209,6 +209,7 @@ pub enum DispatchCommand { QueueSend { name: String, body: Vec, + incoming: crate::telemetry::IncomingInvocationContext, conn: ConnHandle, request: Request, wait: bool, @@ -974,30 +975,48 @@ impl ActorTask { DispatchCommand::QueueSend { name, body, + incoming, conn, request, wait, timeout_ms, reply, - } => match self.send_actor_event( - "dispatch_queue_send", - ActorEvent::QueueSend { - name, - body, - conn, - request, - wait, - timeout_ms, - reply: Reply::from(reply), - }, - ) { - Ok(()) => { - self.log_dispatch_command_handled(command_kind, "enqueued"); - } - Err(_error) => { - self.log_dispatch_command_handled(command_kind, "enqueue_failed"); + } => { + let invocation = ActorInvocation::start_queue_send(&self.ctx, &name, incoming); + let invocation_telemetry = invocation.telemetry(); + let (tracked_reply_tx, tracked_reply_rx) = oneshot::channel(); + match self.send_actor_event( + "dispatch_queue_send", + ActorEvent::QueueSend { + name, + body, + conn, + request, + wait, + timeout_ms, + invocation_telemetry: Some(invocation_telemetry), + reply: Reply::from(tracked_reply_tx), + }, + ) { + Ok(()) => { + self.log_dispatch_command_handled(command_kind, "enqueued"); + self.forward_tracked_reply( + ActorWorkKind::DispatchReply, + tracked_reply_rx, + reply, + move |result| { + invocation.finish(result.as_ref().err()); + result + }, + ); + } + Err(error) => { + invocation.finish(Some(&error)); + let _ = reply.send(Err(error)); + self.log_dispatch_command_handled(command_kind, "enqueue_failed"); + } } - }, + } DispatchCommand::Http { request, reply } => { let incoming = IncomingInvocationContext::from_http_headers(request.headers()); let invocation = ActorInvocation::start_request(&self.ctx, &request, incoming); diff --git a/rivetkit-rust/packages/rivetkit-core/src/lib.rs b/rivetkit-rust/packages/rivetkit-core/src/lib.rs index d789dd9849..2fa505c912 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/lib.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/lib.rs @@ -20,7 +20,7 @@ pub mod telemetry; // Internal bridge types consumed by the NAPI and Wasm runtime adapters. #[doc(hidden)] pub use telemetry::{ - ActorInvocationSpanContext, ActorInvocationTelemetry, ActorInvocationTraceContext, + ActorInvocationSpanContext, ActorInvocationTelemetry, ActorInvocationTraceContext, TraceOrigin, }; #[cfg(feature = "native-runtime")] pub mod serverless_http; @@ -130,7 +130,7 @@ pub use actor::{kv, sqlite}; pub use actor::action::ActionDispatchError; pub use actor::config::{ ActionDefinition, ActorConfig, ActorConfigInput, ActorConfigOverrides, CanHibernateWebSocket, - SqliteProfilingConfig, SqliteProfilingConfigInput, + QueueDefinition, SqliteProfilingConfig, SqliteProfilingConfigInput, }; pub use actor::connection::ConnHandle; pub use actor::context::{ diff --git a/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs b/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs index f85838037d..ae1b12664a 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs @@ -369,12 +369,15 @@ impl RegistryDispatcher { } }; + let incoming = + crate::telemetry::IncomingInvocationContext::from_http_headers(request.headers()); let (reply_tx, reply_rx) = oneshot::channel(); let dispatch_result = try_send_dispatch_command( &instance.dispatch, DispatchCommand::QueueSend { name: queue_name, body: queue_request.body, + incoming, conn: conn.clone(), request, wait: queue_request.wait, diff --git a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs index 96acc89b37..6c934cdce2 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs @@ -16,6 +16,7 @@ use tracing_opentelemetry::OpenTelemetrySpanExt as _; use crate::ActorContext; use crate::actor::metrics::{ActorMetrics, InvocationStatus, InvocationType}; +use crate::actor::queue::QueueMessage; use crate::time::Instant; /// Correlation fields accepted at an invocation boundary. @@ -68,12 +69,17 @@ fn invocation_ray_id(headers: &http::HeaderMap) -> Option { /// invocation type beside it. const REQUEST_INVOCATION_NAME: &str = "onRequest"; +/// Name a queue send invocation is reported under. The queue itself is an +/// attribute, so one series covers every queue of an actor. +const QUEUE_SEND_INVOCATION_NAME: &str = "queue.send"; + /// What an invocation ran, which decides its name and the attributes that /// identify it on the span. #[derive(Clone, Copy, Debug)] enum InvocationSubject<'a> { Action(&'a str), Request { method: &'a str }, + QueueSend { queue: &'a str }, } /// Owns the complete lifecycle of one actor invocation. @@ -135,9 +141,9 @@ pub(crate) struct InvocationWorkGuard(ActorInvocationTelemetry); /// messages so the work they cause can link back to its origin. #[derive(Clone, Debug, Default, Eq, PartialEq)] pub struct TraceOrigin { - pub(crate) ray_id: Option, - pub(crate) traceparent: Option, - pub(crate) tracestate: Option, + pub ray_id: Option, + pub traceparent: Option, + pub tracestate: Option, } impl TraceOrigin { @@ -257,6 +263,24 @@ impl ActorInvocation { ) } + /// Starts the invocation for one message sent into `queue_name` from + /// outside the actor. It ends when the send is acknowledged, or when the + /// sender's wait for a completion ends. + pub(crate) fn start_queue_send( + ctx: &ActorContext, + queue_name: &str, + incoming: IncomingInvocationContext, + ) -> Self { + Self::start( + ctx, + InvocationSubject::QueueSend { queue: queue_name }, + InvocationType::QueueSend, + incoming.ray_id, + incoming.remote_parent, + None, + ) + } + /// Starts the invocation for one raw HTTP request served by `onRequest`. /// The span is named after the handler rather than the path, because a /// path is caller-supplied and would make the name a cardinality surface. @@ -287,41 +311,48 @@ impl ActorInvocation { ) -> Self { let identity = ctx.telemetry_identity(); // Use bounded names for both spans and metrics. - let (action_name, http_method) = match subject { - InvocationSubject::Action(name) => (ctx.metrics().label_action_name(name), None), - InvocationSubject::Request { method } => (REQUEST_INVOCATION_NAME, Some(method)), + let (action_name, http_method, queue_name) = match subject { + InvocationSubject::Action(name) => (ctx.metrics().label_action_name(name), None, None), + InvocationSubject::Request { method } => (REQUEST_INVOCATION_NAME, Some(method), None), + InvocationSubject::QueueSend { queue } => ( + QUEUE_SEND_INVOCATION_NAME, + None, + Some(ctx.metrics().label_queue_name(queue)), + ), }; let action_name = action_name.to_owned(); let span = if tracing::enabled!(target: "rivetkit::telemetry", tracing::Level::INFO) { - let span = tracing::info_span!( - target: "rivetkit::telemetry", - parent: None, - "rivet.actor.invoke", - otel.name = %format!("{}/{}", identity.actor_name, action_name), - otel.kind = invocation_type.otel_kind(), - rivet.invocation.type = invocation_type.as_label(), - rivet.actor.id = %identity.actor_id, - rivet.actor.name = %identity.actor_name, - rivet.actor.key = %identity.actor_key, - rivet.action.name = tracing::field::Empty, - rivet.ray.id = tracing::field::Empty, - http.request.method = tracing::field::Empty, - http.response.status_code = tracing::field::Empty, - otel.status_code = tracing::field::Empty, - error.type = tracing::field::Empty, - ); - span.record("rivet.ray.id", ray_id.as_deref()); - match http_method { - Some(method) => span.record("http.request.method", method), - None => span.record("rivet.action.name", &action_name), - }; - if let Some(parent) = parent { - span.set_parent(opentelemetry::Context::new().with_remote_span_context(parent)); - } - if let Some(link) = link { - span.add_link(link); - } - Some(span) + let span = tracing::info_span!( + target: "rivetkit::telemetry", + parent: None, + "rivet.actor.invoke", + otel.name = %format!("{}/{}", identity.actor_name, action_name), + otel.kind = invocation_type.otel_kind(), + rivet.invocation.type = invocation_type.as_label(), + rivet.actor.id = %identity.actor_id, + rivet.actor.name = %identity.actor_name, + rivet.actor.key = %identity.actor_key, + rivet.action.name = tracing::field::Empty, + rivet.ray.id = tracing::field::Empty, + http.request.method = tracing::field::Empty, + http.response.status_code = tracing::field::Empty, + rivet.queue.name = tracing::field::Empty, + otel.status_code = tracing::field::Empty, + error.type = tracing::field::Empty, + ); + span.record("rivet.ray.id", ray_id.as_deref()); + match (http_method, queue_name) { + (Some(method), _) => span.record("http.request.method", method), + (None, Some(queue)) => span.record("rivet.queue.name", queue), + (None, None) => span.record("rivet.action.name", &action_name), + }; + if let Some(parent) = parent { + span.set_parent(opentelemetry::Context::new().with_remote_span_context(parent)); + } + if let Some(link) = link { + span.add_link(link); + } + Some(span) } else { None }; @@ -636,6 +667,63 @@ fn w3c_span_context(span_context: &SpanContext) -> Option tracing::Span { + if !tracing::enabled!(target: "rivetkit::telemetry", tracing::Level::INFO) { + return tracing::Span::none(); + } + let identity = ctx.telemetry_identity(); + let invocation = ctx + .1 + .as_ref() + .and_then(|telemetry| telemetry.active().map(|inner| (telemetry, inner))); + let span = tracing::info_span!( + target: "rivetkit::telemetry", + parent: None, + "rivet.queue.receive", + otel.name = %format!("{}/queue.receive", identity.actor_name), + otel.kind = "consumer", + rivet.actor.id = %identity.actor_id, + rivet.actor.name = %identity.actor_name, + rivet.actor.key = %identity.actor_key, + rivet.queue.name = ctx.metrics().label_queue_name(&message.name), + rivet.ray.id = tracing::field::Empty, + ); + match invocation { + Some((telemetry, inner)) => { + span.record("rivet.ray.id", inner.ray_id.as_deref()); + let parent = match &telemetry.1 { + Some(application_span) => Some( + opentelemetry::Context::new() + .with_remote_span_context(application_span.clone()), + ), + None => inner.span.lock().as_ref().map(|span| span.context()), + }; + if let Some(parent) = parent { + span.set_parent(parent); + } + } + None => { + if let Some(ray_id) = &message.trace_origin.ray_id { + span.record("rivet.ray.id", ray_id); + } + } + } + if let Some(link) = parse_remote_parent( + message.trace_origin.traceparent.as_deref(), + message.trace_origin.tracestate.as_deref(), + ) { + span.add_link(link); + } + span +} + /// Records the terminal status and error identity of a finished span. fn record_outcome(span: &tracing::Span, error: Option<&anyhow::Error>) { span.record( diff --git a/rivetkit-rust/packages/rivetkit-core/tests/integration/counter.rs b/rivetkit-rust/packages/rivetkit-core/tests/integration/counter.rs index 8575aea9a7..35caea7fd0 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/integration/counter.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/integration/counter.rs @@ -100,6 +100,7 @@ fn counter_factory() -> ActorFactory { request: _, wait: _, timeout_ms: _, + invocation_telemetry: _, reply, } => { reply.send(Err(anyhow::anyhow!("queue sends are not handled"))); diff --git a/rivetkit-rust/packages/rivetkit-core/tests/integration/sqlite_corruption_fuzz.rs b/rivetkit-rust/packages/rivetkit-core/tests/integration/sqlite_corruption_fuzz.rs index 3bf89e3955..7ebc9622e1 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/integration/sqlite_corruption_fuzz.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/integration/sqlite_corruption_fuzz.rs @@ -593,6 +593,7 @@ fn sqlite_fuzz_factory() -> ActorFactory { request: _, wait: _, timeout_ms: _, + invocation_telemetry: _, reply, } => { reply.send(Err(anyhow::anyhow!("queue sends are not handled"))); diff --git a/rivetkit-rust/packages/rivetkit-core/tests/sql_efficiency.rs b/rivetkit-rust/packages/rivetkit-core/tests/sql_efficiency.rs index eb335bd614..4aa43b35df 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/sql_efficiency.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/sql_efficiency.rs @@ -382,6 +382,24 @@ fn query_catalog() -> Vec { params: vec![1_i64.into()], expectation: indexed(None, &["_rivet_queue"]), }, + QueryCase { + id: "queue.upsert_trace_context", + sql: internal_storage::UPSERT_QUEUE_TRACE_CONTEXT_SQL.into(), + params: vec![text("queue_trace_context:1"), Value::Blob(vec![1])], + expectation: indexed(None, &["_rivet_meta"]), + }, + QueryCase { + id: "queue.delete_trace_context", + sql: internal_storage::DELETE_QUEUE_TRACE_CONTEXT_SQL.into(), + params: vec![text("queue_trace_context:1")], + expectation: indexed(None, &["_rivet_meta"]), + }, + QueryCase { + id: "queue.reset_trace_contexts", + sql: internal_storage::RESET_QUEUE_TRACE_CONTEXTS_SQL.into(), + params: vec![], + expectation: indexed(None, &["_rivet_meta"]), + }, QueryCase { id: "queue.reset", sql: internal_storage::RESET_QUEUE_SQL.into(), diff --git a/rivetkit-rust/packages/rivetkit/src/event.rs b/rivetkit-rust/packages/rivetkit/src/event.rs index 82ec55df3d..5cc75f8c64 100644 --- a/rivetkit-rust/packages/rivetkit/src/event.rs +++ b/rivetkit-rust/packages/rivetkit/src/event.rs @@ -121,6 +121,7 @@ impl RuntimeEvent { request, wait, timeout_ms, + invocation_telemetry: _, reply, } => Self::QueueSend(QueueSend { name, diff --git a/rivetkit-rust/packages/rivetkit/src/start.rs b/rivetkit-rust/packages/rivetkit/src/start.rs index ed6b2fb75d..19deea7bee 100644 --- a/rivetkit-rust/packages/rivetkit/src/start.rs +++ b/rivetkit-rust/packages/rivetkit/src/start.rs @@ -2515,6 +2515,7 @@ mod tests { request: rivetkit_core::Request::default(), wait: true, timeout_ms: None, + invocation_telemetry: None, reply: reply_tx.into(), }) .expect("send queue event"); diff --git a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts index 9a2de32aef..eee15955e2 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts +++ b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts @@ -64,6 +64,9 @@ export interface JsQueueSendResult { export interface JsActionDefinition { name: string } +export interface JsQueueDefinition { + name: string +} /** * One entry in the actor's `inspector.tabs[]` declaration. Either a * custom-tab descriptor (id + label + source dir) or a built-in modifier @@ -133,6 +136,7 @@ export interface JsActorConfig { maxIncomingMessageSize?: number maxOutgoingMessageSize?: number actions?: Array + queues?: Array inspectorTabs?: Array } export interface JsBindParam { diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/actor_factory.rs b/rivetkit-typescript/packages/rivetkit-napi/src/actor_factory.rs index 93a828007f..025b63ce5d 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/actor_factory.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/actor_factory.rs @@ -11,7 +11,7 @@ use rivet_error::{ActorSpecifier, RivetError, RivetErrorKind}; use rivetkit_core::inspector::InspectorTabEntry; use rivetkit_core::{ ActionDefinition, ActorConfig, ActorConfigInput, ActorContext as CoreActorContext, - ActorFactory as CoreActorFactory, ConnHandle as CoreConnHandle, Request, + ActorFactory as CoreActorFactory, ConnHandle as CoreConnHandle, QueueDefinition, Request, SqliteProfilingConfigInput, WebSocket as CoreWebSocket, }; @@ -53,6 +53,12 @@ pub struct JsActionDefinition { pub name: String, } +#[napi(object)] +#[derive(Clone, Default)] +pub struct JsQueueDefinition { + pub name: String, +} + /// One entry in the actor's `inspector.tabs[]` declaration. Either a /// custom-tab descriptor (id + label + source dir) or a built-in modifier /// (id + hidden=true). Validation already happened on the TS side; the @@ -124,6 +130,7 @@ pub struct JsActorConfig { pub max_incoming_message_size: Option, pub max_outgoing_message_size: Option, pub actions: Option>, + pub queues: Option>, pub inspector_tabs: Option>, } @@ -155,6 +162,7 @@ pub(crate) struct MigratePayload { #[derive(Clone)] pub(crate) struct QueueSendPayload { pub(crate) ctx: CoreActorContext, + pub(crate) telemetry: Option, pub(crate) conn: CoreConnHandle, pub(crate) request: Request, pub(crate) name: String, @@ -805,7 +813,10 @@ fn build_queue_send_payload( payload: QueueSendPayload, ) -> napi::Result> { let mut object = env.create_object()?; - object.set("ctx", ActorContext::new(payload.ctx))?; + object.set( + "ctx", + ActorContext::new(payload.ctx.with_invocation_telemetry(payload.telemetry)), + )?; object.set("conn", ConnHandle::new(payload.conn))?; object.set("request", build_request_object(env, payload.request)?)?; object.set("name", payload.name)?; @@ -1026,6 +1037,12 @@ impl From for ActorConfigInput { .map(|action| ActionDefinition { name: action.name }) .collect() }), + queues: value.queues.map(|queues| { + queues + .into_iter() + .map(|queue| QueueDefinition { name: queue.name }) + .collect() + }), inspector_tabs: value.inspector_tabs.map(|tabs| { tabs.into_iter() .map(|tab| { diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs b/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs index 91aebc7602..d841261f0d 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs @@ -492,6 +492,7 @@ pub(crate) async fn dispatch_event( request, wait, timeout_ms, + invocation_telemetry, reply, } => { let Some(callback) = bindings.on_queue_send.clone() else { @@ -514,6 +515,7 @@ pub(crate) async fn dispatch_event( &callback, QueueSendPayload { ctx: ctx.inner().clone(), + telemetry: invocation_telemetry, conn, request, name, diff --git a/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs b/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs index f1efda5650..4f68f3630b 100644 --- a/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs +++ b/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs @@ -177,6 +177,12 @@ pub struct WasmActionDefinition { pub name: String, } +#[derive(Clone, Default, serde::Deserialize)] +#[serde(default, rename_all = "camelCase")] +pub struct WasmQueueDefinition { + pub name: String, +} + /// Experimental SQLite profiling configuration. This entire configuration /// surface is subject to change without notice. #[derive(Clone, Default, serde::Deserialize)] @@ -229,6 +235,7 @@ pub struct WasmActorConfig { pub max_incoming_message_size: Option, pub max_outgoing_message_size: Option, pub actions: Option>, + pub queues: Option>, } impl From for ActorConfigInput { @@ -265,6 +272,12 @@ impl From for ActorConfigInput { .map(|action| rivetkit_core::ActionDefinition { name: action.name }) .collect() }), + queues: config.queues.map(|queues| { + queues + .into_iter() + .map(|queue| rivetkit_core::QueueDefinition { name: queue.name }) + .collect() + }), // Custom inspector tabs serve assets from a filesystem `root`, which is a // native/server feature that has no meaning in a browser wasm host. inspector_tabs: None, diff --git a/rivetkit-typescript/packages/rivetkit/src/client/actor-conn.ts b/rivetkit-typescript/packages/rivetkit/src/client/actor-conn.ts index db8edb885f..a106156c19 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/actor-conn.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/actor-conn.ts @@ -50,6 +50,7 @@ import { ACTOR_CONNS_SYMBOL, type ClientRaw } from "./client"; import * as errors from "./errors"; import { isRetryableLifecycleReconnectSignal } from "./lifecycle-errors"; import { logger } from "./log"; +import { outboundTelemetryHeaders } from "./outbound-telemetry"; import { createQueueSender, type QueueSendNoWaitOptions, @@ -229,6 +230,7 @@ export class ActorConnRaw { this.#queueSender = createQueueSender({ encoding: this.#encoding, params: this.#params, + telemetryHeaders: () => outboundTelemetryHeaders(undefined), customFetch: async (request: Request) => { return await this.#driver.sendRequest( getGatewayTarget(this.#actorResolutionState), diff --git a/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts b/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts index 8e614ec88e..d571828f85 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts @@ -171,6 +171,10 @@ export class ActorHandleRaw { return await createQueueSender({ encoding: this.#encoding, params: this.#params, + telemetryHeaders: () => + outboundTelemetryHeaders( + this.#currentActorInvocation?.(), + ), customFetch: async (request: Request) => { return await this.#driver.sendRequest( target, diff --git a/rivetkit-typescript/packages/rivetkit/src/client/queue.ts b/rivetkit-typescript/packages/rivetkit/src/client/queue.ts index a55fc86852..c081a78b0f 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/queue.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/queue.ts @@ -54,6 +54,8 @@ export interface QueueSendResult { interface QueueSenderOptions { encoding: Encoding; params: unknown; + /** Ray ID and trace context headers, read per send. */ + telemetryHeaders: () => Record; customFetch: (request: Request) => Promise; } @@ -90,6 +92,7 @@ export function createQueueSender( method: "POST", headers: { [HEADER_ENCODING]: senderOptions.encoding, + ...senderOptions.telemetryHeaders(), ...(senderOptions.params !== undefined ? { [HEADER_CONN_PARAMS]: JSON.stringify( diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts index 76f0638a8b..4b3acbc7d0 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts @@ -71,6 +71,9 @@ type NapiSqlTransaction = Awaited< type NapiActorStateTransaction = Awaited< ReturnType >; +type NapiQueueMessage = Awaited< + ReturnType["send"]> +>; function asNativeRegistry(handle: RegistryHandle): NativeCoreRegistry { return handle as unknown as NativeCoreRegistry; @@ -239,7 +242,7 @@ function toNapiKvEntry(entry: RuntimeKvEntry): { }; } -function toNapiQueueMessage(message: RuntimeQueueMessage): RuntimeQueueMessage { +function toNapiQueueMessage(message: NapiQueueMessage): RuntimeQueueMessage { return { id: () => message.id(), name: () => message.name(), @@ -948,7 +951,7 @@ export class NapiCoreRuntime implements CoreRuntime { body: RuntimeBytes, ): Promise { return toNapiQueueMessage( - await asNativeActorContext(ctx) + await this.#actorContextForOperation(ctx) .queue() .send(name, toNapiBuffer(body)), ); @@ -959,7 +962,7 @@ export class NapiCoreRuntime implements CoreRuntime { options?: RuntimeQueueNextBatchOptions | undefined | null, signal?: CancellationTokenHandle | undefined | null, ): Promise { - const messages = await asNativeActorContext(ctx) + const messages = await this.#actorContextForOperation(ctx) .queue() .nextBatch( options, @@ -975,7 +978,7 @@ export class NapiCoreRuntime implements CoreRuntime { signal?: CancellationTokenHandle | undefined | null, ): Promise { return toNapiQueueMessage( - await asNativeActorContext(ctx) + await this.#actorContextForOperation(ctx) .queue() .waitForNames( names, @@ -1022,7 +1025,7 @@ export class NapiCoreRuntime implements CoreRuntime { options?: RuntimeQueueEnqueueAndWaitOptions | undefined | null, signal?: CancellationTokenHandle | undefined | null, ): Promise { - return await asNativeActorContext(ctx) + return await this.#actorContextForOperation(ctx) .queue() .enqueueAndWait( name, diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts index e3d4f5c674..a93b0d06e2 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts @@ -3906,6 +3906,9 @@ function buildActorConfig( actions: Object.keys(flattenActionHandlers(config.actions)) .sort() .map((name) => ({ name })), + queues: Object.keys(config.queues ?? {}) + .sort() + .map((name) => ({ name })), inspectorTabs: buildInspectorTabs(config.inspector, runtimeKind), }; } @@ -5431,62 +5434,75 @@ export function buildNativeFactory( cancelToken, run !== undefined, ); - try { - if ( - !schemaConfig.queues || - !hasSchemaConfigKey(schemaConfig.queues, name) - ) { - return { status: "completed" }; - } - - const canPublish = getQueueCanPublish( - schemaConfig.queues, - name, - ); - if (canPublish && !(await canPublish(actorCtx))) { - throw forbiddenError(); - } - - const decodedBody = decodeValue(body); - if (wait) { + return await runtime.runWithActorInvocationContext( + ctx, + async () => { try { - const response = - await actorCtx.queue.enqueueAndWait( - name, - decodedBody, - { - timeout: - timeoutMs === undefined || - timeoutMs === null - ? undefined - : Number(timeoutMs), - }, - ); - return { - status: "completed", - response: - response === undefined - ? undefined - : encodeValue(response), - }; - } catch (error) { if ( - (error as { group?: string; code?: string }) - .group === "queue" && - (error as { group?: string; code?: string }) - .code === "timed_out" + !schemaConfig.queues || + !hasSchemaConfigKey(schemaConfig.queues, name) ) { - return { status: "timedOut" }; + return { status: "completed" }; } - throw error; - } - } - await actorCtx.queue.send(name, decodedBody); - return { status: "completed" }; - } finally { - await actorCtx.dispose(); - } + const canPublish = getQueueCanPublish( + schemaConfig.queues, + name, + ); + if (canPublish && !(await canPublish(actorCtx))) { + throw forbiddenError(); + } + + const decodedBody = decodeValue(body); + if (wait) { + try { + const response = + await actorCtx.queue.enqueueAndWait( + name, + decodedBody, + { + timeout: + timeoutMs === undefined || + timeoutMs === null + ? undefined + : Number(timeoutMs), + }, + ); + return { + status: "completed", + response: + response === undefined + ? undefined + : encodeValue(response), + }; + } catch (error) { + if ( + ( + error as { + group?: string; + code?: string; + } + ).group === "queue" && + ( + error as { + group?: string; + code?: string; + } + ).code === "timed_out" + ) { + return { status: "timedOut" }; + } + throw error; + } + } + + await actorCtx.queue.send(name, decodedBody); + return { status: "completed" }; + } finally { + await actorCtx.dispose(); + } + }, + ); }, ), serializeState: wrapNativeCallback( diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts index a66abb199a..dbe68a6364 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts @@ -324,6 +324,7 @@ export interface RuntimeActorConfig { preloadMaxWorkflowBytes?: number; preloadMaxConnectionsBytes?: number; actions?: Array<{ name: string }>; + queues?: Array<{ name: string }>; inspectorTabs?: Array; } From 2fcaad7cc7f37450af8d89fc7a0c6ea6cd59e7a9 Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Wed, 2 Sep 2026 09:56:53 +0400 Subject: [PATCH 14/21] test(rivetkit): cover schedule trace origins --- .../tests/fixtures/napi-runtime-server.ts | 7 +++++ .../tests/napi-runtime-integration.test.ts | 27 +++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts b/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts index a4588b60e3..64c767fe2e 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts @@ -120,6 +120,13 @@ const integrationActor = actor({ count: c.state.count, }; }, + scheduleTrace: async (c, correlationToken: string) => { + await c.schedule.after(50, "scheduledTrace", correlationToken); + return correlationToken; + }, + scheduledTrace: async (c, correlationToken: string) => { + await c.db.execute("SELECT ? AS trace", correlationToken); + }, sqliteFailure: async (c) => { await c.db.execute("SELECT value FROM missing_trace_test_table"); }, diff --git a/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts b/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts index 65de48ff55..e49792c17b 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts @@ -596,6 +596,33 @@ describe.sequential("native NAPI runtime integration", () => { code: "internal_error", message: "An internal error occurred", }); + + // A scheduled fire keeps the defining invocation's ray and starts a + // fresh trace linked to the defining span. + traceExports.length = 0; + const scheduleToken = crypto.randomUUID(); + expect(await handle.scheduleTrace(scheduleToken)).toBe(scheduleToken); + const scheduleSpans = await waitForInvocationSpans( + traceExports, + ["scheduleTrace", "scheduledTrace"], + 15_000, + ); + const definer = findInvocation(scheduleSpans, "scheduleTrace"); + const scheduled = findInvocation(scheduleSpans, "scheduledTrace"); + expect(definer).toBeDefined(); + expect(scheduled?.attributes["rivet.invocation.type"]).toBe( + "scheduled", + ); + expect(scheduled?.attributes["rivet.ray.id"]).toBe( + definer?.attributes["rivet.ray.id"], + ); + // Without this the assertions below hold for a scheduled fire that threw, + // so a broken action body would still pass. + expect(scheduled?.attributes["error.type"]).toBeUndefined(); + expect(scheduled?.traceId).not.toBe(definer?.traceId); + expect(scheduled?.links).toEqual([ + { traceId: definer?.traceId, spanId: definer?.spanId }, + ]); await client.dispose(); const processId = servicesPid(); From ced9252ec2a55b44298b252af299f98db8b12b91 Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Thu, 3 Sep 2026 13:43:10 +0400 Subject: [PATCH 15/21] test(rivetkit): cover actor invocation tracing end to end --- .../tests/fixtures/napi-runtime-server.ts | 22 + .../rivetkit/tests/fixtures/otlp-collector.ts | 33 +- .../tests/napi-runtime-integration.test.ts | 509 ++++++++++++++++-- 3 files changed, 520 insertions(+), 44 deletions(-) diff --git a/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts b/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts index 64c767fe2e..6d4dda333a 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts @@ -137,6 +137,28 @@ const integrationActor = actor({ kvCount: kvValue ? Number(kvValue) : null, }; }, + // Interleaves awaits, SQLite, a child actor call and a log so two + // overlapping invocations of this action have every chance to observe + // each other's telemetry context. + isolationProbe: async (c, token: string, fail: boolean) => { + await new Promise((resolve) => setTimeout(resolve, 20)); + await c.db.execute("SELECT ? AS probe", token); + c.log.warn({ correlation_token: token }, "isolation probe"); + const client = c.client(); + await client.integrationActor + .getForId(c.actorId, { + params: { userId: "internal-integration-test" }, + }) + .getCount(); + await new Promise((resolve) => setTimeout(resolve, 20)); + await c.db.execute("SELECT ? AS probe2", token); + if (fail) { + throw new UserError("isolation probe failure", { + code: "isolation_probe_failed", + }); + } + return token; + }, getCountViaClient: async (c) => { const client = c.client(); return await client.integrationActor diff --git a/rivetkit-typescript/packages/rivetkit/tests/fixtures/otlp-collector.ts b/rivetkit-typescript/packages/rivetkit/tests/fixtures/otlp-collector.ts index 944b7214c1..798d2af0b9 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/fixtures/otlp-collector.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/fixtures/otlp-collector.ts @@ -6,15 +6,39 @@ export interface OtlpCollector { close(): Promise; } -export async function startOtlpCollector(port: number): Promise { +export interface OtlpCollectorOptions { + /** + * Delay before each export is answered. Models a collector that accepts the + * connection and then stalls, which backs up the exporter's queue rather + * than failing its requests outright. + */ + readonly responseDelayMs?: number; +} + +export async function startOtlpCollector( + port: number, + options: OtlpCollectorOptions = {}, +): Promise { const exports: Buffer[] = []; + const pending = new Set(); const server = createServer((request, response) => { const chunks: Buffer[] = []; request.on("data", (chunk: Buffer) => chunks.push(chunk)); request.on("end", () => { exports.push(Buffer.concat(chunks)); - response.writeHead(200, { "content-type": "application/json" }); - response.end(); + const reply = () => { + response.writeHead(200, { "content-type": "application/json" }); + response.end(); + }; + if (!options.responseDelayMs) { + reply(); + return; + } + const timer = setTimeout(() => { + pending.delete(timer); + reply(); + }, options.responseDelayMs); + pending.add(timer); }); }); @@ -27,6 +51,9 @@ export async function startOtlpCollector(port: number): Promise { spans: () => exports, close: () => new Promise((resolve, reject) => { + for (const timer of pending) clearTimeout(timer); + pending.clear(); + server.closeAllConnections(); server.close((error) => (error ? reject(error) : resolve())); }), }; diff --git a/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts b/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts index e49792c17b..725e7001d6 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts @@ -6,6 +6,10 @@ import { fileURLToPath } from "node:url"; import getPort from "get-port"; import { afterEach, describe, expect, test } from "vitest"; import { createClient } from "../src/client/mod"; +import { + type OtlpCollector, + startOtlpCollector, +} from "./fixtures/otlp-collector"; const TEST_DIR = dirname(fileURLToPath(import.meta.url)); const FIXTURE_PATH = join(TEST_DIR, "fixtures", "napi-runtime-server.ts"); @@ -19,9 +23,13 @@ let runtimeLogs = { let engineEndpoint: string | undefined; let storagePath: string | undefined; +function runtimeOutput(): string { + return [runtimeLogs.stdout, runtimeLogs.stderr].filter(Boolean).join("\n"); +} + function childOutput(child: ChildProcess): string { void child; - return [runtimeLogs.stdout, runtimeLogs.stderr].filter(Boolean).join("\n"); + return runtimeOutput(); } async function engineOutput(): Promise { @@ -434,14 +442,181 @@ async function stopTestEngine(): Promise { } } +interface ExportedSpan { + name: string; + traceId: string; + spanId: string; + parentSpanId?: string; + attributes: Record; + links: Array<{ traceId: string; spanId: string }>; +} + +/** Flattens OTLP/JSON export bodies into the spans they carry. */ +function exportedSpans(exports: Buffer[]): ExportedSpan[] { + type OtlpAttribute = { key: string; value: { stringValue?: string } }; + type OtlpSpan = Omit & { + attributes?: OtlpAttribute[]; + links?: Array<{ traceId: string; spanId: string }>; + }; + type OtlpPayload = { + resourceSpans?: Array<{ scopeSpans?: Array<{ spans?: OtlpSpan[] }> }>; + }; + return exports.flatMap((body) => { + const payload = JSON.parse(body.toString("utf8")) as OtlpPayload; + return (payload.resourceSpans ?? []).flatMap((resource) => + (resource.scopeSpans ?? []).flatMap((scope) => + (scope.spans ?? []).map((span) => ({ + name: span.name, + traceId: span.traceId, + spanId: span.spanId, + parentSpanId: span.parentSpanId || undefined, + attributes: Object.fromEntries( + (span.attributes ?? []).map((attribute) => [ + attribute.key, + attribute.value.stringValue, + ]), + ), + links: (span.links ?? []).map((link) => ({ + traceId: link.traceId, + spanId: link.spanId, + })), + })), + ), + ); + }); +} + +/** + * Polls until the exported spans satisfy `ready`, then returns them. Parent + * and child spans can land in different export batches, so callers that + * assert parentage must wait for both. + */ +async function waitForSpans( + exports: Buffer[], + description: string, + ready: (spans: ExportedSpan[]) => boolean, + timeoutMs: number, +): Promise { + const deadline = Date.now() + timeoutMs; + while (Date.now() < deadline) { + const spans = exportedSpans(exports); + if (ready(spans)) { + return spans; + } + await new Promise((resolve) => setTimeout(resolve, 100)); + } + throw new Error(`timed out waiting for ${description}`); +} + +function isSqliteSpan(span: ExportedSpan): boolean { + return span.name === "rivet.sqlite.execute"; +} + +function isFailedSqliteSpan(span: ExportedSpan): boolean { + return isSqliteSpan(span) && span.attributes["error.type"] !== undefined; +} + +function findInvocation( + spans: ExportedSpan[], + actionName: string, +): ExportedSpan | undefined { + return spans.find( + (span) => + span.attributes["rivet.invocation.type"] !== undefined && + span.attributes["rivet.action.name"] === actionName, + ); +} + +/** Polls until an invocation span has been exported for every named action. */ +async function waitForInvocationSpans( + exports: Buffer[], + actionNames: string[], + timeoutMs: number, +): Promise { + return waitForSpans( + exports, + `invocation spans: ${actionNames.join(", ")}`, + (spans) => actionNames.every((name) => findInvocation(spans, name)), + timeoutMs, + ); +} + +async function waitForRuntimeLog( + correlationToken: string, + timeoutMs: number, +): Promise { + const deadline = Date.now() + timeoutMs; + const marker = `correlation_token=${correlationToken}`; + while (Date.now() < deadline) { + const line = runtimeOutput() + .split("\n") + .find((candidate) => candidate.includes(marker)); + if (line) { + return line; + } + await new Promise((resolve) => setTimeout(resolve, 100)); + } + throw new Error(`timed out waiting for runtime log ${correlationToken}`); +} + +/** + * Starts an engine and a native runtime pointed at one OTLP endpoint, and + * returns the pieces every telemetry test needs. + */ +async function startTracedRuntime( + tracesEndpoint: string, + extraEnv: Record = {}, +): Promise<{ endpoint: string; poolName: string; child: ChildProcess }> { + const poolName = "default"; + const port = await getPort({ host: "127.0.0.1" }); + const endpoint = `http://127.0.0.1:${port}`; + engineEndpoint = endpoint; + storagePath = await mkdtemp(join(tmpdir(), "rivetkit-services-")); + runtimeLogs = { stdout: "", stderr: "" }; + const child = spawn(process.execPath, ["--import", "tsx", FIXTURE_PATH], { + cwd: dirname(TEST_DIR), + env: { + ...process.env, + RIVET_TOKEN: TOKEN, + RIVET_NAMESPACE: NAMESPACE, + RIVET_RUN_ENGINE_HOST: "127.0.0.1", + RIVET_RUN_ENGINE_PORT: String(port), + RIVETKIT_TEST_ENDPOINT: endpoint, + RIVETKIT_TEST_POOL_NAME: poolName, + RIVETKIT_STORAGE_PATH: storagePath, + OTEL_EXPORTER_OTLP_TRACES_ENDPOINT: tracesEndpoint, + OTEL_EXPORTER_OTLP_TRACES_PROTOCOL: "http/json", + OTEL_TRACES_SAMPLER: "always_on", + OTEL_BSP_SCHEDULE_DELAY: "10", + ...extraEnv, + }, + stdio: ["ignore", "pipe", "pipe"], + }); + child.stdout?.on("data", (chunk) => { + runtimeLogs.stdout += chunk.toString(); + }); + child.stderr?.on("data", (chunk) => { + runtimeLogs.stderr += chunk.toString(); + }); + await waitForHealth(child, endpoint, 90_000); + await upsertNormalRunnerConfig(child, endpoint, poolName); + await waitForEnvoy(child, endpoint, poolName, 30_000); + return { endpoint, poolName, child }; +} + describe.sequential("native NAPI runtime integration", () => { let runtime: ChildProcess | undefined; + let collector: OtlpCollector | undefined; afterEach(async () => { if (runtime) { await stopRuntime(runtime); runtime = undefined; } + if (collector) { + await collector.close(); + collector = undefined; + } await stopTestEngine(); if (storagePath) { await rm(storagePath, { recursive: true, force: true }); @@ -451,36 +626,14 @@ describe.sequential("native NAPI runtime integration", () => { }, 30_000); test("runs a TS actor through registry, NAPI, core, envoy, and engine", async () => { - const poolName = "default"; - const port = await getPort({ host: "127.0.0.1" }); - const endpoint = `http://127.0.0.1:${port}`; - engineEndpoint = endpoint; - storagePath = await mkdtemp(join(tmpdir(), "rivetkit-services-")); - runtimeLogs = { stdout: "", stderr: "" }; - runtime = spawn(process.execPath, ["--import", "tsx", FIXTURE_PATH], { - cwd: dirname(TEST_DIR), - env: { - ...process.env, - RIVET_TOKEN: TOKEN, - RIVET_NAMESPACE: NAMESPACE, - RIVET_RUN_ENGINE_HOST: "127.0.0.1", - RIVET_RUN_ENGINE_PORT: String(port), - RIVETKIT_TEST_ENDPOINT: endpoint, - RIVETKIT_TEST_POOL_NAME: poolName, - RIVETKIT_STORAGE_PATH: storagePath, - }, - stdio: ["ignore", "pipe", "pipe"], - }); - runtime.stdout?.on("data", (chunk) => { - runtimeLogs.stdout += chunk.toString(); - }); - runtime.stderr?.on("data", (chunk) => { - runtimeLogs.stderr += chunk.toString(); - }); - - await waitForHealth(runtime, endpoint, 90_000); - await upsertNormalRunnerConfig(runtime, endpoint, poolName); - await waitForEnvoy(runtime, endpoint, poolName, 30_000); + collector = await startOtlpCollector( + await getPort({ host: "127.0.0.1" }), + ); + const traceExports = collector.spans(); + const { endpoint, poolName, child } = await startTracedRuntime( + collector.endpoint, + ); + runtime = child; await waitForEnvoy(runtime, endpoint, SERVICES_POOL_NAME, 30_000); await expectNormalRunnerConfig(endpoint, SERVICES_POOL_NAME); const servicesActorId = await createServicesActor(endpoint); @@ -494,21 +647,42 @@ describe.sequential("native NAPI runtime integration", () => { disableMetadataLookup: true, }) as any; + const actorKey = `napi-runtime-${crypto.randomUUID()}`; const handle = await waitForActorReady( () => - client.integrationActor.create( - [`napi-runtime-${crypto.randomUUID()}`], - { - params: { userId: "integration-test" }, - }, - ), + client.integrationActor.create([actorKey], { + params: { userId: "integration-test" }, + }), 30_000, ); const actorId = await handle.resolve(); + const correlationToken = crypto.randomUUID(); + expect(await handle.logContext(correlationToken)).toBe( + correlationToken, + ); + const actorLog = await waitForRuntimeLog(correlationToken, 10_000); + expect(actorLog).toContain(`actorId=${actorId}`); + expect(actorLog).toContain("actorName=integrationActor"); + expect(actorLog).toContain(actorKey); + expect(actorLog).toMatch(/ rayId=[0-9a-f-]{36}( |$)/); + expect(actorLog).toMatch(/ trace_id=[0-9a-f]{32}( |$)/); + expect(actorLog).toMatch(/ span_id=[0-9a-f]{16}( |$)/); + expect(await waitForActorReady(() => handle.getCount(), 30_000)).toBe( 0, ); + const getCountSpans = await waitForInvocationSpans( + traceExports, + ["getCount"], + 10_000, + ); + expect( + findInvocation(getCountSpans, "getCount")?.attributes, + ).toMatchObject({ + "rivet.invocation.type": "action", + "rivet.actor.name": "integrationActor", + }); expect( await waitForActorReady( () => handle.validatedAction({ amount: 4 }), @@ -561,6 +735,42 @@ describe.sequential("native NAPI runtime integration", () => { count: 2, sqliteValues: [2], }); + // SQLite spans are children of the action that issued them. + const incrementSpans = await waitForSpans( + traceExports, + "increment invocation and sqlite spans", + (spans) => + spans.some(isSqliteSpan) && + findInvocation(spans, "increment") !== undefined, + 10_000, + ); + const incrementSqlite = incrementSpans.find(isSqliteSpan); + expect(incrementSqlite?.attributes).toMatchObject({ + "rivet.operation.system": "sqlite", + "rivet.operation.name": "execute", + }); + expect(incrementSqlite?.parentSpanId).toBe( + findInvocation(incrementSpans, "increment")?.spanId, + ); + traceExports.length = 0; + await expect(handle.sqliteFailure()).rejects.toMatchObject({ + code: expect.any(String), + }); + const failureSpans = await waitForSpans( + traceExports, + "sqliteFailure invocation and failed sqlite spans", + (spans) => + spans.some(isFailedSqliteSpan) && + findInvocation(spans, "sqliteFailure") !== undefined, + 10_000, + ); + const failedSqlite = failureSpans.find(isFailedSqliteSpan); + expect(failedSqlite?.attributes["error.type"]).toMatch( + /^[a-z_]+\.[a-z_]+$/, + ); + expect(failedSqlite?.parentSpanId).toBe( + findInvocation(failureSpans, "sqliteFailure")?.spanId, + ); expect(await handle.snapshot()).toEqual({ count: 2, kvCount: 2, @@ -578,7 +788,22 @@ describe.sequential("native NAPI runtime integration", () => { ).toEqual({ count: 5, }); + // An actor-owned client carries the calling invocation's trace and ray + // across the real Engine boundary, so the callee is its child. + traceExports.length = 0; expect(await handle.getCountViaClient()).toBe(5); + const clientSpans = await waitForInvocationSpans( + traceExports, + ["getCountViaClient", "getCount"], + 10_000, + ); + const caller = findInvocation(clientSpans, "getCountViaClient"); + const callee = findInvocation(clientSpans, "getCount"); + expect(callee?.traceId).toBe(caller?.traceId); + expect(callee?.parentSpanId).toBe(caller?.spanId); + expect(callee?.attributes["rivet.ray.id"]).toBe( + caller?.attributes["rivet.ray.id"], + ); expect(await handle.stateSnapshot()).toEqual({ count: 5, kvCount: 5, @@ -597,8 +822,7 @@ describe.sequential("native NAPI runtime integration", () => { message: "An internal error occurred", }); - // A scheduled fire keeps the defining invocation's ray and starts a - // fresh trace linked to the defining span. + // Scheduled work starts a new trace linked to its origin. traceExports.length = 0; const scheduleToken = crypto.randomUUID(); expect(await handle.scheduleTrace(scheduleToken)).toBe(scheduleToken); @@ -616,8 +840,7 @@ describe.sequential("native NAPI runtime integration", () => { expect(scheduled?.attributes["rivet.ray.id"]).toBe( definer?.attributes["rivet.ray.id"], ); - // Without this the assertions below hold for a scheduled fire that threw, - // so a broken action body would still pass. + // Ensure the scheduled action succeeded before checking its trace. expect(scheduled?.attributes["error.type"]).toBeUndefined(); expect(scheduled?.traceId).not.toBe(definer?.traceId); expect(scheduled?.links).toEqual([ @@ -630,4 +853,208 @@ describe.sequential("native NAPI runtime integration", () => { runtime = undefined; await waitForProcessExit(processId, 5_000); }, 120_000); + + test("keeps overlapping invocations of one actor telemetrically isolated", async () => { + collector = await startOtlpCollector( + await getPort({ host: "127.0.0.1" }), + ); + const traceExports = collector.spans(); + const { endpoint, poolName, child } = await startTracedRuntime( + collector.endpoint, + ); + runtime = child; + + const client = createClient({ + endpoint, + token: TOKEN, + namespace: NAMESPACE, + poolName, + disableMetadataLookup: true, + }) as any; + const handle = await waitForActorReady( + () => + client.integrationActor.create( + [`napi-isolation-${crypto.randomUUID()}`], + { params: { userId: "integration-test" } }, + ), + 30_000, + ); + await waitForActorReady(() => handle.getCount(), 30_000); + + // Use the same actor to exercise isolation between concurrent invocations. + const okToken = crypto.randomUUID(); + const failToken = crypto.randomUUID(); + const [ok, failed] = await Promise.allSettled([ + handle.isolationProbe(okToken, false), + handle.isolationProbe(failToken, true), + ]); + expect(ok.status).toBe("fulfilled"); + expect(failed.status).toBe("rejected"); + + const spans = await waitForSpans( + traceExports, + "both isolation probe invocations and the calls each one made", + (exported) => { + const probes = exported.filter( + (span) => + span.attributes["rivet.action.name"] === + "isolationProbe", + ); + return ( + probes.length >= 2 && + probes.every((probe) => + exported.some( + (span) => + span.attributes["rivet.action.name"] === + "getCount" && + span.traceId === probe.traceId, + ), + ) + ); + }, + 20_000, + ); + + const probes = spans.filter( + (span) => span.attributes["rivet.action.name"] === "isolationProbe", + ); + expect(probes).toHaveLength(2); + + const rays = probes.map((probe) => probe.attributes["rivet.ray.id"]); + expect(new Set(rays).size).toBe(2); + expect(new Set(probes.map((probe) => probe.traceId)).size).toBe(2); + const failedProbes = probes.filter( + (probe) => probe.attributes["error.type"] !== undefined, + ); + expect(failedProbes).toHaveLength(1); + expect(failedProbes[0]?.attributes["error.type"]).toBe( + "user.isolation_probe_failed", + ); + + for (const probe of probes) { + const owned = spans.filter( + (span) => + isSqliteSpan(span) && span.parentSpanId === probe.spanId, + ); + expect(owned.length).toBeGreaterThanOrEqual(2); + for (const span of owned) { + expect(span.traceId).toBe(probe.traceId); + expect(span.attributes["rivet.ray.id"]).toBe( + probe.attributes["rivet.ray.id"], + ); + } + } + + // The outbound call each probe makes while the other is mid-flight + // stays inside its own trace and carries its own ray. + for (const probe of probes) { + const callee = spans.find( + (span) => + span.attributes["rivet.action.name"] === "getCount" && + span.traceId === probe.traceId, + ); + expect(callee).toBeDefined(); + expect(callee?.parentSpanId).toBe(probe.spanId); + expect(callee?.attributes["rivet.ray.id"]).toBe( + probe.attributes["rivet.ray.id"], + ); + } + + const okLog = await waitForRuntimeLog(okToken, 10_000); + const failLog = await waitForRuntimeLog(failToken, 10_000); + const rayOf = (line: string) => / rayId=([0-9a-f-]{36})/.exec(line)?.[1]; + expect(rayOf(okLog)).toBeDefined(); + expect(rayOf(okLog)).not.toBe(rayOf(failLog)); + expect(rays).toContain(rayOf(okLog)); + expect(rays).toContain(rayOf(failLog)); + + await client.dispose(); + }, 120_000); + + test("keeps actor behavior intact when the trace exporter is unavailable", async () => { + // Nothing listens on this port, so every OTLP export attempt fails. + const unavailable = `http://127.0.0.1:${await getPort({ host: "127.0.0.1" })}/v1/traces`; + const { endpoint, poolName, child } = + await startTracedRuntime(unavailable); + runtime = child; + + const client = createClient({ + endpoint, + token: TOKEN, + namespace: NAMESPACE, + poolName, + disableMetadataLookup: true, + }) as any; + const handle = await waitForActorReady( + () => + client.integrationActor.create( + [`napi-telemetry-failure-${crypto.randomUUID()}`], + { params: { userId: "integration-test" } }, + ), + 30_000, + ); + + expect(await waitForActorReady(() => handle.getCount(), 30_000)).toBe( + 0, + ); + expect( + await waitForActorReady( + () => handle.validatedAction({ amount: 4 }), + 30_000, + ), + ).toBe(4); + + await client.dispose(); + }, 120_000); + + test("keeps actor behavior intact when the trace exporter is slow", async () => { + // Stall exports long enough to saturate the small queue below. + collector = await startOtlpCollector( + await getPort({ host: "127.0.0.1" }), + { + responseDelayMs: 120_000, + }, + ); + const { endpoint, poolName, child } = await startTracedRuntime( + collector.endpoint, + { + OTEL_BSP_MAX_QUEUE_SIZE: "8", + OTEL_BSP_MAX_EXPORT_BATCH_SIZE: "4", + }, + ); + runtime = child; + + const client = createClient({ + endpoint, + token: TOKEN, + namespace: NAMESPACE, + poolName, + disableMetadataLookup: true, + }) as any; + const handle = await waitForActorReady( + () => + client.integrationActor.create( + [`napi-telemetry-slow-${crypto.randomUUID()}`], + { params: { userId: "integration-test" } }, + ), + 30_000, + ); + + const started = Date.now(); + for (let index = 1; index <= 12; index += 1) { + expect( + await waitForActorReady(() => handle.increment(1), 30_000), + ).toMatchObject({ count: index }); + } + const elapsed = Date.now() - started; + + // Actions must finish before the stalled collector responds. + expect(elapsed).toBeLessThan(60_000); + + expect(await waitForActorReady(() => handle.getCount(), 30_000)).toBe( + 12, + ); + + await client.dispose(); + }, 180_000); }); From 0980cba774bc26f745340785c1cc3ba5899651eb Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Wed, 2 Sep 2026 02:40:32 +0400 Subject: [PATCH 16/21] feat(rivetkit): forward OpenTelemetry SDK warnings to the JavaScript logger --- .../packages/rivetkit-napi/index.d.ts | 6 ++++++ .../packages/rivetkit-napi/index.js | 3 ++- .../packages/rivetkit-napi/src/lib.rs | 9 ++++++++- .../packages/rivetkit-napi/src/registry.rs | 8 ++++++++ .../packages/rivetkit-napi/src/telemetry.rs | 20 ++++++++++++------- .../rivetkit/src/registry/napi-runtime.ts | 8 ++++++++ .../tests/napi-runtime-integration.test.ts | 15 +++++++++++++- .../rivetkit/tests/runtime-parity.test.ts | 1 + 8 files changed, 60 insertions(+), 10 deletions(-) diff --git a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts index eee15955e2..6896169996 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts +++ b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts @@ -274,6 +274,12 @@ export interface JsServerlessStreamError { code: string message: string } +/** + * Routes the OpenTelemetry SDK's own warnings, such as dropped spans, to the + * JavaScript logger. Call before constructing a registry; later calls are + * ignored because the tracing subscriber initializes once. + */ +export declare function setTelemetryLogSink(callback: (...args: any[]) => any): void export interface JsScheduledEventInfo { id: string action: string diff --git a/rivetkit-typescript/packages/rivetkit-napi/index.js b/rivetkit-typescript/packages/rivetkit-napi/index.js index 6f44128343..2f378421f4 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/index.js +++ b/rivetkit-typescript/packages/rivetkit-napi/index.js @@ -310,7 +310,7 @@ if (!nativeBinding) { throw new Error(`Failed to load native binding`) } -const { ActorContext, decodeInspectorRequest, encodeInspectorResponse, NapiActorFactory, CancellationToken, ConnHandle, JsNativeDatabase, JsSqliteTransaction, JsActorStateTransaction, HttpResponseBodyStream, HttpRequestBodyStream, Kv, Queue, QueueMessage, CoreRegistry, Schedule, WebSocket } = nativeBinding +const { ActorContext, decodeInspectorRequest, encodeInspectorResponse, NapiActorFactory, CancellationToken, ConnHandle, JsNativeDatabase, JsSqliteTransaction, JsActorStateTransaction, HttpResponseBodyStream, HttpRequestBodyStream, Kv, Queue, QueueMessage, CoreRegistry, setTelemetryLogSink, Schedule, WebSocket } = nativeBinding module.exports.ActorContext = ActorContext module.exports.decodeInspectorRequest = decodeInspectorRequest @@ -327,5 +327,6 @@ module.exports.Kv = Kv module.exports.Queue = Queue module.exports.QueueMessage = QueueMessage module.exports.CoreRegistry = CoreRegistry +module.exports.setTelemetryLogSink = setTelemetryLogSink module.exports.Schedule = Schedule module.exports.WebSocket = WebSocket diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/lib.rs b/rivetkit-typescript/packages/rivetkit-napi/src/lib.rs index 30a8760639..236f7a05e1 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/lib.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/lib.rs @@ -125,6 +125,10 @@ pub(crate) fn init_tracing(log_level: Option<&str>) { tracing_subscriber::registry() .with(otel_layer) + .with( + telemetry::sdk_log_bridge::SdkLogLayer + .with_filter(tracing_subscriber::EnvFilter::new("opentelemetry_sdk=warn")), + ) .with(match log_format { LogFormat::Logfmt => Some( tracing_logfmt::builder() @@ -150,7 +154,10 @@ pub(crate) fn init_tracing(log_level: Option<&str>) { .init(); if let Some(error) = otel_error { - tracing::warn!(?error, "OpenTelemetry trace export could not be initialized"); + tracing::warn!( + ?error, + "OpenTelemetry trace export could not be initialized" + ); } }); } diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/registry.rs b/rivetkit-typescript/packages/rivetkit-napi/src/registry.rs index 97d9cd0195..0db676e323 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/registry.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/registry.rs @@ -140,6 +140,14 @@ pub struct CoreRegistry { build_complete: Arc, } +/// Routes the OpenTelemetry SDK's own warnings, such as dropped spans, to the +/// JavaScript logger. Each call replaces the previous sink, so a registry +/// started on a fresh Node worker thread takes over from one that has exited. +#[napi] +pub fn set_telemetry_log_sink(env: Env, callback: napi::JsFunction) -> napi::Result<()> { + crate::telemetry::sdk_log_bridge::install(env, callback) +} + #[napi] impl CoreRegistry { #[napi(constructor)] diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/telemetry.rs b/rivetkit-typescript/packages/rivetkit-napi/src/telemetry.rs index 376b56b78b..8afac9d61b 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/telemetry.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/telemetry.rs @@ -7,10 +7,11 @@ /// This layer hands those events to a JS callback instead, so an operator sees /// them alongside everything else the actor logs. pub(crate) mod sdk_log_bridge { - use std::sync::OnceLock; - use napi::bindgen_prelude::*; use napi::threadsafe_function::{ErrorStrategy, ThreadSafeCallContext, ThreadsafeFunction}; + // Forced-sync: read from inside a tracing layer callback, which is a sync + // context and never spans an await. + use parking_lot::RwLock; use tracing::field::{Field, Visit}; use tracing_subscriber::Layer; use tracing_subscriber::layer::Context; @@ -21,10 +22,14 @@ pub(crate) mod sdk_log_bridge { pub(crate) message: String, } - static SINK: OnceLock> = OnceLock::new(); + /// The most recently installed sink. It is replaceable rather than set + /// once, because a Node worker thread that installed it can exit, after + /// which its callback silently drops every event. The next registry to + /// start, on whichever thread, takes over. + static SINK: RwLock>> = + RwLock::new(None); - /// Installs the JavaScript sink. Only the first call takes effect, matching - /// the one-shot initialization of the tracing subscriber itself. + /// Installs the JavaScript sink, replacing any earlier one. /// /// The threadsafe function is unreferenced. A referenced one counts as live /// work on the Node event loop, so a process that had registered the sink @@ -39,7 +44,7 @@ pub(crate) mod sdk_log_bridge { Ok(vec![object.into_unknown()]) })?; tsfn.unref(&env)?; - let _ = SINK.set(tsfn); + *SINK.write() = Some(tsfn); Ok(()) } @@ -72,7 +77,8 @@ pub(crate) mod sdk_log_bridge { impl Layer for SdkLogLayer { fn on_event(&self, event: &tracing::Event<'_>, _ctx: Context<'_, S>) { - let Some(sink) = SINK.get() else { + let sink = SINK.read(); + let Some(sink) = sink.as_ref() else { return; }; let mut fields = FieldCollector::default(); diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts index 4b3acbc7d0..85cc4d7b19 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts @@ -13,6 +13,7 @@ import { readActiveTraceHeaders, runWithActorInvocationSpan, } from "@/common/otel-context"; +import { logger } from "./log"; import type { ActorContextHandle, ActorFactoryHandle, @@ -316,6 +317,13 @@ export class NapiCoreRuntime implements CoreRuntime { } createRegistry(): RegistryHandle { + // Replace the sink on each registry start because its previous worker may have exited. + this.#bindings.setTelemetryLogSink((event) => { + logger().warn( + { otelEvent: event.name }, + event.message || event.name, + ); + }); return asRegistryHandle(new this.#bindings.CoreRegistry()); } diff --git a/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts b/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts index 725e7001d6..400dd8a724 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts @@ -4,7 +4,7 @@ import { tmpdir } from "node:os"; import { dirname, join } from "node:path"; import { fileURLToPath } from "node:url"; import getPort from "get-port"; -import { afterEach, describe, expect, test } from "vitest"; +import { afterEach, describe, expect, test, vi } from "vitest"; import { createClient } from "../src/client/mod"; import { type OtlpCollector, @@ -1020,6 +1020,8 @@ describe.sequential("native NAPI runtime integration", () => { { OTEL_BSP_MAX_QUEUE_SIZE: "8", OTEL_BSP_MAX_EXPORT_BATCH_SIZE: "4", + // Disable Rust SDK logs so only the JS bridge can satisfy the log assertion. + RUST_LOG: "warn,opentelemetry_sdk=off", }, ); runtime = child; @@ -1051,6 +1053,17 @@ describe.sequential("native NAPI runtime integration", () => { // Actions must finish before the stalled collector responds. expect(elapsed).toBeLessThan(60_000); + // The processor reports dropped spans on its own export cycle, which + // runs after the actions return, so there is nothing to await. + await vi.waitFor( + () => { + expect(runtimeOutput()).toContain( + "BatchSpanProcessor.SpanDroppingStarted", + ); + }, + { timeout: 15_000, interval: 250 }, + ); + expect(await waitForActorReady(() => handle.getCount(), 30_000)).toBe( 12, ); diff --git a/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts b/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts index c1adaad769..867a9a6743 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts @@ -314,6 +314,7 @@ function fakeNapiBindings(scenario: ParityScenario) { NapiActorFactory: FakeActorFactory, CancellationToken: FakeCancellationToken, ActorContext: class {}, + setTelemetryLogSink: () => {}, }; } From e3932dc589969bd2c4eaf53f13b61460816b9f94 Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Mon, 7 Sep 2026 02:25:08 +0400 Subject: [PATCH 17/21] feat(rivetkit): trace actor-to-actor calls --- pnpm-lock.yaml | 76 ++++++++++++++ .../rivetkit-core/src/actor/context.rs | 17 ++++ .../packages/rivetkit-core/src/lib.rs | 3 +- .../packages/rivetkit-core/src/telemetry.rs | 99 ++++++++++++++++++- .../packages/rivetkit-napi/index.d.ts | 30 +++++- .../packages/rivetkit-napi/index.js | 3 +- .../rivetkit-napi/src/actor_context.rs | 58 ++++++++++- .../rivetkit-napi/src/actor_factory.rs | 8 ++ .../packages/rivetkit/package.json | 1 + .../packages/rivetkit/src/actor/errors.ts | 12 +++ .../rivetkit/src/client/actor-handle.ts | 45 +++++++-- .../packages/rivetkit/src/client/client.ts | 17 ++-- .../rivetkit/src/client/outbound-telemetry.ts | 15 ++- .../rivetkit/src/registry/napi-runtime.ts | 17 ++++ .../packages/rivetkit/src/registry/native.ts | 4 + .../packages/rivetkit/src/registry/runtime.ts | 38 ++++++- .../rivetkit/src/registry/wasm-runtime.ts | 11 +++ .../tests/fixtures/napi-runtime-server.ts | 26 +++++ .../tests/napi-runtime-integration.test.ts | 71 +++++++++++-- 19 files changed, 512 insertions(+), 39 deletions(-) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 50dc6369c5..5f4dda7d3c 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -3617,6 +3617,9 @@ importers: '@hono/node-ws': specifier: ^1.1.1 version: 1.3.0(@hono/node-server@1.19.9(hono@4.11.9))(hono@4.11.9) + '@opentelemetry/sdk-trace-node': + specifier: 2.11.0 + version: 2.11.0(@opentelemetry/api@1.9.0) '@rivet-dev/agent-os-common': specifier: '*' version: 0.0.260331072558 @@ -7491,6 +7494,42 @@ packages: resolution: {integrity: sha512-3giAOQvZiH5F9bMlMiv8+GSPMeqg0dbaeo58/0SlA9sxSqZhnUtxzX9/2FzyhS9sWQf5S0GJE0AKBrFqjpeYcg==} engines: {node: '>=8.0.0'} + '@opentelemetry/context-async-hooks@2.11.0': + resolution: {integrity: sha512-Tr79DyWI8itsBdg+jH+opjfrwLzX+erk1/ExkIwhWoAVjVrJIn2y5+cGjTC0Vy8fyNIA/y8wuJPZwr1T3xCZeQ==} + engines: {node: ^18.19.0 || >=20.6.0} + peerDependencies: + '@opentelemetry/api': '>=1.0.0 <1.10.0' + + '@opentelemetry/core@2.11.0': + resolution: {integrity: sha512-7YP44XH0tV6+Mb54x2YGf84i7yi+31MBZlE8JwvozkxyTvXbSp10X7cI7YE49ChJ3shMJoBmCJF3+1QFBJctGA==} + engines: {node: ^18.19.0 || >=20.6.0} + peerDependencies: + '@opentelemetry/api': '>=1.0.0 <1.10.0' + + '@opentelemetry/resources@2.11.0': + resolution: {integrity: sha512-Ie7+8q8MDF4FAEQCKVMTx3ReUvxiIAgIiiW3c9JdmP8+HMcDy20puT+AHjexnExgnbvBxjQ9fjkFDWrikJ2jQA==} + engines: {node: ^18.19.0 || >=20.6.0} + peerDependencies: + '@opentelemetry/api': '>=1.3.0 <1.10.0' + + '@opentelemetry/sdk-trace-base@2.11.0': + resolution: {integrity: sha512-H19x/TX/LZdqiYOjM7fqtSxwlplC5pgelavqbQdHbhdq0q/AI/TGkM2dfGuuynTXmJPeF2HoZVoPDu+TGoW78A==} + engines: {node: ^18.19.0 || >=20.6.0} + peerDependencies: + '@opentelemetry/api': '>=1.3.0 <1.10.0' + + '@opentelemetry/sdk-trace-node@2.11.0': + resolution: {integrity: sha512-CuvCMJmZxswhNLlM2LfuLOW3h3fZujA4hsG4B+Sz4dX2zvaXO8Ng74cnDHWD64gLszTlhiG3c0iNUjj4g+0/sA==} + engines: {node: ^18.19.0 || >=20.6.0} + peerDependencies: + '@opentelemetry/api': '>=1.0.0 <1.10.0' + + '@opentelemetry/sdk-trace@2.11.0': + resolution: {integrity: sha512-fFnTqGm8/G73GQVnxYi7LXa1ZVYEUvgL6XI1LpvV0bPC7WQ/ZGgKxCSl8FnlZBKto9JHHEFTO6s6CUpvvtwFrA==} + engines: {node: ^18.19.0 || >=20.6.0} + peerDependencies: + '@opentelemetry/api': '>=1.3.0 <1.10.0' + '@opentelemetry/semantic-conventions@1.40.0': resolution: {integrity: sha512-cifvXDhcqMwwTlTK04GBNeIe7yyo28Mfby85QXFe1Yk8nmi36Ab/5UQwptOx84SsoGNRg+EVSjwzfSZMy6pmlw==} engines: {node: '>=14'} @@ -23674,6 +23713,43 @@ snapshots: '@opentelemetry/api@1.9.0': {} + '@opentelemetry/context-async-hooks@2.11.0(@opentelemetry/api@1.9.0)': + dependencies: + '@opentelemetry/api': 1.9.0 + + '@opentelemetry/core@2.11.0(@opentelemetry/api@1.9.0)': + dependencies: + '@opentelemetry/api': 1.9.0 + '@opentelemetry/semantic-conventions': 1.40.0 + + '@opentelemetry/resources@2.11.0(@opentelemetry/api@1.9.0)': + dependencies: + '@opentelemetry/api': 1.9.0 + '@opentelemetry/core': 2.11.0(@opentelemetry/api@1.9.0) + '@opentelemetry/semantic-conventions': 1.40.0 + + '@opentelemetry/sdk-trace-base@2.11.0(@opentelemetry/api@1.9.0)': + dependencies: + '@opentelemetry/api': 1.9.0 + '@opentelemetry/core': 2.11.0(@opentelemetry/api@1.9.0) + '@opentelemetry/resources': 2.11.0(@opentelemetry/api@1.9.0) + '@opentelemetry/sdk-trace': 2.11.0(@opentelemetry/api@1.9.0) + '@opentelemetry/semantic-conventions': 1.40.0 + + '@opentelemetry/sdk-trace-node@2.11.0(@opentelemetry/api@1.9.0)': + dependencies: + '@opentelemetry/api': 1.9.0 + '@opentelemetry/context-async-hooks': 2.11.0(@opentelemetry/api@1.9.0) + '@opentelemetry/core': 2.11.0(@opentelemetry/api@1.9.0) + '@opentelemetry/sdk-trace-base': 2.11.0(@opentelemetry/api@1.9.0) + + '@opentelemetry/sdk-trace@2.11.0(@opentelemetry/api@1.9.0)': + dependencies: + '@opentelemetry/api': 1.9.0 + '@opentelemetry/core': 2.11.0(@opentelemetry/api@1.9.0) + '@opentelemetry/resources': 2.11.0(@opentelemetry/api@1.9.0) + '@opentelemetry/semantic-conventions': 1.40.0 + '@opentelemetry/semantic-conventions@1.40.0': {} '@oxc-project/types@0.142.0': {} diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs index 0e257d3558..df5ed8c9ac 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs @@ -281,6 +281,23 @@ impl ActorContext { self.0.sql.clone().with_invocation_telemetry(self.1.clone()) } + /// Opens the span covering one call out to another actor, or nothing when + /// this handle serves no invocation or tracing is disabled. + /// + /// `actor_name` and `action_name` name the callee. Both come from the + /// caller's own registry rather than from a remote peer, so neither is a + /// cardinality surface. + #[doc(hidden)] + pub fn begin_outbound_call( + &self, + actor_name: &str, + action_name: &str, + ) -> Option { + self.1 + .as_ref()? + .start_outbound_call(actor_name, action_name) + } + /// Returns correlation for the invocation this handle serves, absent when /// the handle is not bound to one or tracing is disabled. pub fn invocation_trace_context(&self) -> Option { diff --git a/rivetkit-rust/packages/rivetkit-core/src/lib.rs b/rivetkit-rust/packages/rivetkit-core/src/lib.rs index 2fa505c912..644bc66f4e 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/lib.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/lib.rs @@ -20,7 +20,8 @@ pub mod telemetry; // Internal bridge types consumed by the NAPI and Wasm runtime adapters. #[doc(hidden)] pub use telemetry::{ - ActorInvocationSpanContext, ActorInvocationTelemetry, ActorInvocationTraceContext, TraceOrigin, + ActorInvocationSpanContext, ActorInvocationTelemetry, ActorInvocationTraceContext, + OutboundCallInvocation, TraceOrigin, }; #[cfg(feature = "native-runtime")] pub mod serverless_http; diff --git a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs index 6c934cdce2..9d00c62887 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs @@ -230,6 +230,19 @@ pub(crate) struct SqliteOperationSpan { span: Option, } +/// One call from this invocation out to another actor, held open across a +/// foreign-runtime boundary. +/// +/// The call is made by the host runtime's client, so it is opened and closed by +/// two separate calls rather than by one Rust scope. Dropping this without +/// finishing records the call as cancelled, matching how a dropped SQLite span +/// is treated. +#[doc(hidden)] +pub struct OutboundCallInvocation { + span: Option, + context: Option, +} + impl ActorInvocation { pub(crate) fn start_action( ctx: &ActorContext, @@ -516,11 +529,11 @@ impl ActorInvocationTelemetry { #[doc(hidden)] pub fn trace_context(&self) -> Option { let active = self.active()?; - let span = active.span.lock().clone().and_then(|span| { - let context = span.context(); - let context_span = context.span(); - w3c_span_context(context_span.span_context()) - }); + let span = active + .span + .lock() + .clone() + .and_then(|span| span_context_of(&span)); Some(ActorInvocationTraceContext { ray_id: active.ray_id.clone(), @@ -549,6 +562,47 @@ impl ActorInvocationTelemetry { } } + /// Opens the span covering one call out to another actor. + /// + /// The callee parents to this span rather than to the invocation making the + /// call, so the time spent reaching it, which includes routing and waking a + /// sleeping actor, is attributed to the call instead of falling in the gap + /// between the two invocations. + /// When this handle carries the application span active in the host + /// runtime, the call span parents there instead, which is what puts a + /// callee under the application span that issued the call rather than + /// beside it. + pub(crate) fn start_outbound_call( + &self, + actor_name: &str, + action_name: &str, + ) -> Option { + let invocation_span = self.active()?.span.lock().clone()?; + let span = tracing::info_span!( + target: "rivetkit::telemetry", + parent: &invocation_span, + "rivet.actor.call", + otel.name = %format!("{actor_name}/{action_name}"), + otel.kind = "client", + // Omit rivet.invocation.type: this measures the caller waiting, not the callee running. + rivet.actor.name = %actor_name, + rivet.action.name = %action_name, + rivet.ray.id = self.0.ray_id.as_deref(), + otel.status_code = tracing::field::Empty, + error.type = tracing::field::Empty, + ); + if let Some(application_span) = &self.1 { + span.set_parent( + opentelemetry::Context::new().with_remote_span_context(application_span.clone()), + ); + } + let context = span_context_of(&span); + Some(OutboundCallInvocation { + span: Some(span), + context, + }) + } + pub(crate) fn start_sqlite(&self, operation: SqliteOperation) -> Option { let parent = self.active()?.span.lock().clone()?; let span = tracing::info_span!( @@ -624,6 +678,33 @@ impl Drop for InvocationWorkGuard { } } +impl OutboundCallInvocation { + /// W3C context of this call's span, to send to the callee so it parents + /// here. Absent when the call is not sampled. + pub fn span_context(&self) -> Option { + self.context.clone() + } + + /// Records the call's outcome. `error` is the failure the callee returned, + /// and its group and code become the span's `error.type`. + pub fn finish(mut self, error: Option<&anyhow::Error>) { + let Some(span) = self.span.take() else { + return; + }; + record_outcome(&span, error); + } +} + +impl Drop for OutboundCallInvocation { + fn drop(&mut self) { + let Some(span) = self.span.take() else { + return; + }; + span.record("otel.status_code", "ERROR"); + span.record("error.type", OPERATION_ABANDONED_ERROR_TYPE); + } +} + impl SqliteOperationSpan { pub(crate) fn span(&self) -> tracing::Span { self.span.as_ref().expect("sqlite span is present").clone() @@ -667,6 +748,14 @@ fn w3c_span_context(span_context: &SpanContext) -> Option Option { + let context = span.context(); + let context_span = context.span(); + w3c_span_context(context_span.span_context()) +} + /// Opens the span covering the moment one queue message is handed to the /// actor. It links to the span that sent the message, which is what connects /// the consumer's trace to the sender's. Inside an invocation it sits under diff --git a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts index 6896169996..42764e5de1 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts +++ b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts @@ -276,8 +276,8 @@ export interface JsServerlessStreamError { } /** * Routes the OpenTelemetry SDK's own warnings, such as dropped spans, to the - * JavaScript logger. Call before constructing a registry; later calls are - * ignored because the tracing subscriber initializes once. + * JavaScript logger. Each call replaces the previous sink, so a registry + * started on a fresh Node worker thread takes over from one that has exited. */ export declare function setTelemetryLogSink(callback: (...args: any[]) => any): void export interface JsScheduledEventInfo { @@ -336,6 +336,12 @@ export declare class ActorContext { */ withApplicationSpan(traceparent?: string | undefined | null, tracestate?: string | undefined | null): ActorContext invocationTraceContext(): JsActorInvocationTraceContext | null + /** + * Opens the span covering one call out to another actor. Returns nothing + * when this handle serves no invocation or tracing is disabled, in which + * case the caller sends its own context as before. + */ + beginOutboundCall(actorName: string, actionName: string): OutboundCall | null provisionActorRuntimeSocket(): Promise schedule(): Schedule queue(): Queue @@ -385,6 +391,26 @@ export declare class ActorContext { runtimeState(): object clearRuntimeState(): void } +/** + * One open call out to another actor. + * + * The call spans a request made by the host runtime, so it is opened and closed + * by two separate calls. Letting this be collected without finishing records + * the call as cancelled rather than silently losing it. + */ +export declare class OutboundCall { + /** + * W3C context of this call's span, to send to the callee so it parents to + * the call rather than to the invocation that made it. + */ + spanContext(): JsActorInvocationSpanContext | null + /** + * Records the call's outcome. `error` is the failure as the bridge encodes + * it, so a structured error keeps its group and code while anything else + * stays unstructured for Core to classify. + */ + finish(error?: string | undefined | null): void +} export declare class NapiActorFactory { constructor(callbacks: object, config?: JsActorConfig | undefined | null) } diff --git a/rivetkit-typescript/packages/rivetkit-napi/index.js b/rivetkit-typescript/packages/rivetkit-napi/index.js index 2f378421f4..80a7f44b37 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/index.js +++ b/rivetkit-typescript/packages/rivetkit-napi/index.js @@ -310,11 +310,12 @@ if (!nativeBinding) { throw new Error(`Failed to load native binding`) } -const { ActorContext, decodeInspectorRequest, encodeInspectorResponse, NapiActorFactory, CancellationToken, ConnHandle, JsNativeDatabase, JsSqliteTransaction, JsActorStateTransaction, HttpResponseBodyStream, HttpRequestBodyStream, Kv, Queue, QueueMessage, CoreRegistry, setTelemetryLogSink, Schedule, WebSocket } = nativeBinding +const { ActorContext, decodeInspectorRequest, encodeInspectorResponse, OutboundCall, NapiActorFactory, CancellationToken, ConnHandle, JsNativeDatabase, JsSqliteTransaction, JsActorStateTransaction, HttpResponseBodyStream, HttpRequestBodyStream, Kv, Queue, QueueMessage, CoreRegistry, setTelemetryLogSink, Schedule, WebSocket } = nativeBinding module.exports.ActorContext = ActorContext module.exports.decodeInspectorRequest = decodeInspectorRequest module.exports.encodeInspectorResponse = encodeInspectorResponse +module.exports.OutboundCall = OutboundCall module.exports.NapiActorFactory = NapiActorFactory module.exports.CancellationToken = CancellationToken module.exports.ConnHandle = ConnHandle diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs b/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs index 48d3c6acaa..ddb3d05a3e 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs @@ -19,14 +19,15 @@ use parking_lot::Mutex; use rivetkit_core::types::ActorKeySegment; use rivetkit_core::{ ActorContext as CoreActorContext, ActorInvocationSpanContext, ActorInvocationTraceContext, - ActorWorkKind, ConnHandle as CoreConnHandle, KeepAwakeRegion, Request as CoreRequest, - RequestSaveOpts, StateDelta, WebSocketCallbackRegion, WorkflowKvWrite, + ActorWorkKind, ConnHandle as CoreConnHandle, KeepAwakeRegion, + OutboundCallInvocation as CoreOutboundCallInvocation, Request as CoreRequest, RequestSaveOpts, + StateDelta, WebSocketCallbackRegion, WorkflowKvWrite, }; use scc::HashMap as SccHashMap; use tokio::sync::mpsc::UnboundedSender; use tokio_util::sync::CancellationToken as CoreCancellationToken; -use crate::actor_factory::BridgeRivetErrorContext; +use crate::actor_factory::{BridgeRivetErrorContext, anyhow_error_from_js_reason}; use crate::connection::ConnHandle; use crate::database::{JsActorStateTransaction, JsNativeDatabase, transaction_timeout}; use crate::kv::Kv; @@ -343,6 +344,22 @@ impl ActorContext { self.inner.invocation_trace_context().map(Into::into) } + /// Opens the span covering one call out to another actor. Returns nothing + /// when this handle serves no invocation or tracing is disabled, in which + /// case the caller sends its own context as before. + #[napi] + pub fn begin_outbound_call( + &self, + actor_name: String, + action_name: String, + ) -> Option { + self.inner + .begin_outbound_call(&actor_name, &action_name) + .map(|invocation| OutboundCall { + invocation: Some(invocation), + }) + } + #[napi] pub async fn provision_actor_runtime_socket( &self, @@ -1121,3 +1138,38 @@ fn js_http_request_to_core_request(request: JsHttpRequest) -> napi::Result, +} + +#[napi] +impl OutboundCall { + /// W3C context of this call's span, to send to the callee so it parents to + /// the call rather than to the invocation that made it. + #[napi] + pub fn span_context(&self) -> Option { + self.invocation + .as_ref() + .and_then(CoreOutboundCallInvocation::span_context) + .map(Into::into) + } + + /// Records the call's outcome. `error` is the failure as the bridge encodes + /// it, so a structured error keeps its group and code while anything else + /// stays unstructured for Core to classify. + #[napi] + pub fn finish(&mut self, error: Option) { + let Some(invocation) = self.invocation.take() else { + return; + }; + let error = error.map(anyhow_error_from_js_reason); + invocation.finish(error.as_ref()); + } +} diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/actor_factory.rs b/rivetkit-typescript/packages/rivetkit-napi/src/actor_factory.rs index 025b63ce5d..40d288d836 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/actor_factory.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/actor_factory.rs @@ -987,6 +987,14 @@ fn parse_bridge_rivet_error(reason: &str) -> Option { })) } +/// Rebuilds an error raised in JavaScript from the reason string the bridge +/// carries. A structured error arrives bridge-encoded and keeps its group and +/// code; anything else stays an unstructured message, which is what lets Core +/// classify and sanitize it rather than trusting the JavaScript text. +pub(crate) fn anyhow_error_from_js_reason(reason: String) -> anyhow::Error { + parse_bridge_rivet_error(&reason).unwrap_or_else(|| anyhow::anyhow!(reason)) +} + pub(crate) fn callback_error(callback_name: &str, error: napi::Error) -> anyhow::Error { let reason = error.reason; if let Some(error) = parse_bridge_rivet_error(&reason) { diff --git a/rivetkit-typescript/packages/rivetkit/package.json b/rivetkit-typescript/packages/rivetkit/package.json index 38549ec69b..7f85fe8c21 100644 --- a/rivetkit-typescript/packages/rivetkit/package.json +++ b/rivetkit-typescript/packages/rivetkit/package.json @@ -236,6 +236,7 @@ "@copilotkit/llmock": "^1.6.0", "@hono/node-server": "^1.18.2", "@hono/node-ws": "^1.1.1", + "@opentelemetry/sdk-trace-node": "2.11.0", "@rivet-dev/agent-os-common": "*", "@rivet-dev/agent-os-pi": "^0.1.1", "@standard-schema/spec": "^1.0.0", diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/errors.ts b/rivetkit-typescript/packages/rivetkit/src/actor/errors.ts index 16a2ff8a44..66d2b732c4 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/errors.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/errors.ts @@ -257,6 +257,18 @@ export function encodeBridgeRivetError(error: RivetErrorLike): string { })}`; } +/** + * Encodes an error for the telemetry bridge. A structured error crosses with + * its group and code. Anything else crosses as text so Core classifies it, + * which is the same rule that keeps raw messages out of every span. + */ +export function encodeErrorForBridge(error: unknown): string { + if (error instanceof RivetError) { + return encodeBridgeRivetError(error); + } + return String(error); +} + export function decodeBridgeRivetErrorPayload( value: string, ): BridgeRivetErrorPayload | undefined { diff --git a/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts b/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts index d571828f85..5e7bd2d8a7 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts @@ -1,5 +1,6 @@ import type { AnyActorDefinition } from "@/actor/definition"; -import type { ActorSpecifier } from "@/actor/errors"; +import { type ActorSpecifier, encodeErrorForBridge } from "@/actor/errors"; +import type { ActorInvocationSpanContext } from "@/common/actor-telemetry-context"; import { HEADER_CONN_PARAMS, HEADER_ENCODING, @@ -24,7 +25,10 @@ import { AsyncMutex } from "@/common/database/shared"; import type { Encoding, JsonCompatValue } from "@/common/encoding"; import { deconstructError } from "@/common/utils"; import type { EngineControlClient } from "@/engine-client/driver"; -import type { CurrentActorInvocation } from "@/registry/runtime"; +import type { + BeginOutboundCall, + CurrentActorInvocation, +} from "@/registry/runtime"; import { decodeCborCompat, deserializeWithEncoding, @@ -85,6 +89,7 @@ export class ActorHandleRaw { #resolvingActorId?: Promise; #queueSendMutex = new AsyncMutex(); #currentActorInvocation?: CurrentActorInvocation; + #beginOutboundCall?: BeginOutboundCall; /** * Do not call this directly. @@ -103,6 +108,7 @@ export class ActorHandleRaw { gatewayOptions: ActorGatewayOptions = {}, signal?: AbortSignal, currentActorInvocation?: CurrentActorInvocation, + beginOutboundCall?: BeginOutboundCall, ) { this.#client = client; this.#driver = driver; @@ -113,6 +119,7 @@ export class ActorHandleRaw { this.#getParams = getParams; this.#signal = signal; this.#currentActorInvocation = currentActorInvocation; + this.#beginOutboundCall = beginOutboundCall; } async #resolveConnectionParams(): Promise { @@ -284,20 +291,41 @@ export class ActorHandleRaw { // when no per-call signal is provided. const signal = opts.signal ?? this.#signal; const optsWithSignal = { ...opts, signal }; + // Open before retries so all attempts share one span and parent context. + const call = this.#beginOutboundCall?.( + getActorNameFromQuery(this.#actorResolutionState), + opts.name, + ); const run = async () => - (await this.#sendActionNow(optsWithSignal)) as Response; - if (opts.name === "destroy") { - return await run(); + (await this.#sendActionAttempts( + optsWithSignal, + call?.span, + )) as Response; + const send = async () => { + if (opts.name === "destroy") { + return await run(); + } + return await retryOnLifecycleBoundary(run, { signal }); + }; + if (!call) { + return await send(); + } + try { + const output = await send(); + call.finish(); + return output; + } catch (error) { + call.finish(encodeErrorForBridge(error)); + throw error; } - - return await retryOnLifecycleBoundary(run, { signal }); } - async #sendActionNow( + async #sendActionAttempts( opts: { name: string; args: unknown[]; } & ActorActionOptions, + callSpan?: ActorInvocationSpanContext, ): Promise { const maxAttempts = this.#getDynamicQueryMaxAttempts(); let useQueryTarget = isDynamicActorQuery(this.#actorResolutionState); @@ -333,6 +361,7 @@ export class ActorHandleRaw { [HEADER_ENCODING]: this.#encoding, ...outboundTelemetryHeaders( this.#currentActorInvocation?.(), + callSpan, ), }; if (this.#params !== undefined) { diff --git a/rivetkit-typescript/packages/rivetkit/src/client/client.ts b/rivetkit-typescript/packages/rivetkit/src/client/client.ts index f717bc4751..72c2ead3b2 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/client.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/client.ts @@ -3,7 +3,10 @@ import type { ActorQuery } from "@/client/query"; import type { Encoding } from "@/common/encoding"; import type { EngineControlClient } from "@/engine-client/driver"; import type { Registry } from "@/registry"; -import type { CurrentActorInvocation } from "@/registry/runtime"; +import type { + BeginOutboundCall, + CurrentActorInvocation, +} from "@/registry/runtime"; import type { ActorActionFunction, ActorGatewayOptions } from "./actor-common"; import { type ActorConn, @@ -176,17 +179,14 @@ export interface Region { export const ACTOR_CONNS_SYMBOL = Symbol("actorConns"); export const CREATE_ACTOR_CONN_PROXY = Symbol("createActorConnProxy"); -/** - * Client for managing & connecting to actors. - * - * @template A The actors map type that defines the available actors. - * @see {@link https://rivet.dev/docs/manage|Create & Manage Actors} - */ +/** Options for constructing a raw actor client. */ export interface ClientRawOptions { encoding?: Encoding; gateway?: ActorGatewayOptions; /** Supplies the calling actor's invocation so actor-to-actor clients propagate its trace and ray ID. */ currentActorInvocation?: CurrentActorInvocation; + /** Opens and finishes the Core span covering an actor-to-actor call. */ + beginOutboundCall?: BeginOutboundCall; } export class ClientRaw { @@ -198,6 +198,7 @@ export class ClientRaw { #encodingKind: Encoding; #gatewayOptions: ActorGatewayOptions; #currentActorInvocation?: CurrentActorInvocation; + #beginOutboundCall?: BeginOutboundCall; /** * Creates an instance of Client. @@ -211,6 +212,7 @@ export class ClientRaw { this.#encodingKind = options.encoding ?? "bare"; this.#gatewayOptions = options.gateway ?? {}; this.#currentActorInvocation = options.currentActorInvocation; + this.#beginOutboundCall = options.beginOutboundCall; } /** @@ -413,6 +415,7 @@ export class ClientRaw { this.#gatewayOptions, signal, this.#currentActorInvocation, + this.#beginOutboundCall, ); } diff --git a/rivetkit-typescript/packages/rivetkit/src/client/outbound-telemetry.ts b/rivetkit-typescript/packages/rivetkit/src/client/outbound-telemetry.ts index abe93b056c..f402c891fd 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/outbound-telemetry.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/outbound-telemetry.ts @@ -4,7 +4,11 @@ import { HEADER_TRACESTATE, } from "@/common/actor-router-consts"; import type { ActorInvocationTraceContext } from "@/common/actor-telemetry-context"; -import { readActiveRayId, readActiveTraceHeaders } from "@/common/otel-context"; +import { + type ActiveTraceHeaders, + readActiveRayId, + readActiveTraceHeaders, +} from "@/common/otel-context"; /** * Headers that carry a caller's ray ID and trace context into an actor. One @@ -13,18 +17,21 @@ import { readActiveRayId, readActiveTraceHeaders } from "@/common/otel-context"; * * The ray ID is the calling invocation's when the caller is itself inside an * actor, else the one placed in OpenTelemetry baggage by the surrounding - * request handler. The trace context is the application span active in this - * JavaScript context, else the calling actor's own Core invocation span. + * request handler. The trace context is `callSpan` when the runtime opened a + * span for this call, else the application span active in this JavaScript + * context, else the calling actor's own Core invocation span. */ export function outboundTelemetryHeaders( invocation: ActorInvocationTraceContext | undefined, + callSpan?: ActiveTraceHeaders, ): Record { const headers: Record = {}; const rayId = invocation?.rayId ?? readActiveRayId(); if (rayId) { headers[HEADER_RIVET_RAY_ID] = rayId; } - const traceHeaders = readActiveTraceHeaders() ?? invocation?.span; + const traceHeaders = + callSpan ?? readActiveTraceHeaders() ?? invocation?.span; if (traceHeaders) { headers[HEADER_TRACEPARENT] = traceHeaders.traceparent; if (traceHeaders.tracestate) { diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts index 85cc4d7b19..c71c5fc93d 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts @@ -33,6 +33,7 @@ import type { RuntimeKvEntry, RuntimeKvListOptions, RuntimeListenerConfig, + RuntimeOutboundCall, RuntimeQueueEnqueueAndWaitOptions, RuntimeQueueMessage, RuntimeQueueNextBatchOptions, @@ -625,6 +626,22 @@ export class NapiCoreRuntime implements CoreRuntime { ); } + beginOutboundCall( + ctx: ActorContextHandle, + actorName: string, + actionName: string, + ): RuntimeOutboundCall | undefined { + const call = this.#actorContextForOperation(ctx).beginOutboundCall( + actorName, + actionName, + ); + if (!call) return undefined; + return { + span: call.spanContext() ?? undefined, + finish: (error?: string) => call.finish(error), + }; + } + actorName(ctx: ActorContextHandle): string { return asNativeActorContext(ctx).name(); } diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts index a93b0d06e2..6d7575a5a4 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts @@ -4034,6 +4034,10 @@ export function buildNativeFactory( callNativeSync(() => runtime.actorInvocationTraceContext(ctx), ), + beginOutboundCall: (actorName, actionName) => + callNativeSync(() => + runtime.beginOutboundCall(ctx, actorName, actionName), + ), }, ); const run = getRunFunction(config.run); diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts index dbe68a6364..f3cff43c5d 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts @@ -1,4 +1,7 @@ -import type { ActorInvocationTraceContext } from "@/common/actor-telemetry-context"; +import type { + ActorInvocationSpanContext, + ActorInvocationTraceContext, +} from "@/common/actor-telemetry-context"; import type { SqliteNativeMetrics, SqliteProfilingOptions, @@ -35,6 +38,30 @@ export type CurrentActorInvocation = () => | ActorInvocationTraceContext | undefined; +/** One open call from an actor out to another actor. */ +export interface RuntimeOutboundCall { + /** + * W3C context of the call's own span, to send to the callee so it parents to + * the call. Absent when tracing is disabled. + */ + readonly span?: ActorInvocationSpanContext; + /** + * Records the call's outcome. `error` is the failure encoded the way the + * bridge encodes errors, so a structured error keeps its group and code. + */ + finish(error?: string): void; +} + +/** + * Opens the span covering one call out to another actor, or returns `undefined` + * outside an invocation or on a runtime without invocation telemetry. Callers + * send their own context when it returns nothing. + */ +export type BeginOutboundCall = ( + actorName: string, + actionName: string, +) => RuntimeOutboundCall | undefined; + export interface RuntimeHttpRequest { method: string; uri: string; @@ -561,6 +588,15 @@ export interface CoreRuntime { actorInvocationTraceContext( ctx: ActorContextHandle, ): ActorInvocationTraceContext | undefined; + /** + * Opens the span covering one call this actor makes to another actor. See + * `BeginOutboundCall` for when this returns nothing. + */ + beginOutboundCall( + ctx: ActorContextHandle, + actorName: string, + actionName: string, + ): RuntimeOutboundCall | undefined; actorName(ctx: ActorContextHandle): string; actorKey(ctx: ActorContextHandle): RuntimeActorKeySegment[]; actorRegion(ctx: ActorContextHandle): string; diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts index a89a9d3fd4..7d4dddb2eb 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts @@ -24,6 +24,7 @@ import type { RuntimeKvEntry, RuntimeKvListOptions, RuntimeListenerConfig, + RuntimeOutboundCall, RuntimeQueueEnqueueAndWaitOptions, RuntimeQueueInspectMessage, RuntimeQueueMessage, @@ -550,6 +551,16 @@ export class WasmCoreRuntime implements CoreRuntime { return undefined; } + beginOutboundCall( + _ctx: ActorContextHandle, + _actorName: string, + _actionName: string, + ): RuntimeOutboundCall | undefined { + // Wasm carries no invocation telemetry, so a call goes out untraced + // rather than failing. + return undefined; + } + actorName(ctx: ActorContextHandle): string { return callHandle(asWasmActorContext(ctx), "name"); } diff --git a/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts b/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts index 6d4dda333a..5e957efd91 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts @@ -1,6 +1,8 @@ import { existsSync } from "node:fs"; import { dirname, resolve } from "node:path"; import { fileURLToPath } from "node:url"; +import { trace } from "@opentelemetry/api"; +import { NodeTracerProvider } from "@opentelemetry/sdk-trace-node"; import { getEnginePath } from "@rivetkit/engine-cli"; import { z } from "zod/v4"; import { db } from "../../src/db/mod"; @@ -14,6 +16,10 @@ const repoEngineBinary = resolve( ); const endpoint = process.env.RIVETKIT_TEST_ENDPOINT ?? "http://127.0.0.1:6642"; + +// Register context propagation without exporting application spans; tests read their IDs. +new NodeTracerProvider().register(); +const applicationTracer = trace.getTracer("napi-runtime-fixture"); const connParamsSchema = z.object({ userId: z.string().min(1), }); @@ -159,6 +165,26 @@ const integrationActor = actor({ } return token; }, + // Calls another actor while an application span is active, and returns + // that span's ID so a test can check the call parented to it. + getCountUnderApplicationSpan: async (c) => { + return await applicationTracer.startActiveSpan( + "agent.generate", + async (span) => { + try { + const client = c.client(); + const count = await client.integrationActor + .getForId(c.actorId, { + params: { userId: "internal-integration-test" }, + }) + .getCount(); + return { count, spanId: span.spanContext().spanId }; + } finally { + span.end(); + } + }, + ); + }, getCountViaClient: async (c) => { const client = c.client(); return await client.integrationActor diff --git a/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts b/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts index 400dd8a724..5e0b57d5dd 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts @@ -442,11 +442,15 @@ async function stopTestEngine(): Promise { } } +/** OTLP `SpanKind.CLIENT`. */ +const OTLP_SPAN_KIND_CLIENT = 3; + interface ExportedSpan { name: string; traceId: string; spanId: string; parentSpanId?: string; + kind?: number; attributes: Record; links: Array<{ traceId: string; spanId: string }>; } @@ -456,6 +460,7 @@ function exportedSpans(exports: Buffer[]): ExportedSpan[] { type OtlpAttribute = { key: string; value: { stringValue?: string } }; type OtlpSpan = Omit & { attributes?: OtlpAttribute[]; + kind?: number; links?: Array<{ traceId: string; spanId: string }>; }; type OtlpPayload = { @@ -470,6 +475,7 @@ function exportedSpans(exports: Buffer[]): ExportedSpan[] { traceId: span.traceId, spanId: span.spanId, parentSpanId: span.parentSpanId || undefined, + kind: span.kind, attributes: Object.fromEntries( (span.attributes ?? []).map((attribute) => [ attribute.key, @@ -799,8 +805,21 @@ describe.sequential("native NAPI runtime integration", () => { ); const caller = findInvocation(clientSpans, "getCountViaClient"); const callee = findInvocation(clientSpans, "getCount"); + // The call out to the other actor is its own span sitting between the + // two invocations, so time spent reaching a cold or busy actor belongs + // to something instead of falling in the gap between them. + const hop = clientSpans.find( + (span) => + span.kind === OTLP_SPAN_KIND_CLIENT && + span.attributes["rivet.action.name"] === "getCount", + ); + expect(hop?.parentSpanId).toBe(caller?.spanId); + expect(callee?.parentSpanId).toBe(hop?.spanId); expect(callee?.traceId).toBe(caller?.traceId); - expect(callee?.parentSpanId).toBe(caller?.spanId); + expect(hop?.traceId).toBe(caller?.traceId); + expect(hop?.attributes["rivet.ray.id"]).toBe( + caller?.attributes["rivet.ray.id"], + ); expect(callee?.attributes["rivet.ray.id"]).toBe( caller?.attributes["rivet.ray.id"], ); @@ -945,29 +964,67 @@ describe.sequential("native NAPI runtime integration", () => { } } - // The outbound call each probe makes while the other is mid-flight - // stays inside its own trace and carries its own ray. for (const probe of probes) { + const hop = spans.find( + (span) => + span.kind === OTLP_SPAN_KIND_CLIENT && + span.traceId === probe.traceId, + ); const callee = spans.find( (span) => + span.attributes["rivet.invocation.type"] !== undefined && span.attributes["rivet.action.name"] === "getCount" && span.traceId === probe.traceId, ); + expect(hop).toBeDefined(); expect(callee).toBeDefined(); - expect(callee?.parentSpanId).toBe(probe.spanId); - expect(callee?.attributes["rivet.ray.id"]).toBe( - probe.attributes["rivet.ray.id"], + expect(hop?.attributes["rivet.actor.name"]).toBe( + "integrationActor", ); + expect(hop?.parentSpanId).toBe(probe.spanId); + expect(callee?.parentSpanId).toBe(hop?.spanId); + for (const span of [hop, callee]) { + expect(span?.attributes["rivet.ray.id"]).toBe( + probe.attributes["rivet.ray.id"], + ); + } } const okLog = await waitForRuntimeLog(okToken, 10_000); const failLog = await waitForRuntimeLog(failToken, 10_000); - const rayOf = (line: string) => / rayId=([0-9a-f-]{36})/.exec(line)?.[1]; + const rayOf = (line: string) => + / rayId=([A-Za-z0-9_-]+)/.exec(line)?.[1]; expect(rayOf(okLog)).toBeDefined(); expect(rayOf(okLog)).not.toBe(rayOf(failLog)); expect(rays).toContain(rayOf(okLog)); expect(rays).toContain(rayOf(failLog)); + // SQLite and actor calls inherit the active application span. + const underApp = await handle.getCountUnderApplicationSpan(); + const appHop = await waitForSpans( + traceExports, + "the hop made under an application span", + (exported) => + exported.some( + (span) => + span.kind === OTLP_SPAN_KIND_CLIENT && + span.parentSpanId === underApp.spanId, + ), + 10_000, + ).then((exported) => + exported.find( + (span) => + span.kind === OTLP_SPAN_KIND_CLIENT && + span.parentSpanId === underApp.spanId, + ), + ); + const appCallee = exportedSpans(traceExports).find( + (span) => + span.attributes["rivet.invocation.type"] !== undefined && + span.parentSpanId === appHop?.spanId, + ); + expect(appCallee?.attributes["rivet.action.name"]).toBe("getCount"); + await client.dispose(); }, 120_000); From cf06366c679609321d26d16b91328260fdcb0ddf Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Wed, 9 Sep 2026 22:53:22 +0400 Subject: [PATCH 18/21] test(rivetkit): cover caller-supplied ray IDs and queue trace origins end to end --- .../tests/fixtures/napi-runtime-server.ts | 61 +- .../tests/napi-runtime-integration.test.ts | 588 ++++++++++++------ 2 files changed, 456 insertions(+), 193 deletions(-) diff --git a/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts b/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts index 5e957efd91..4b8d2251e2 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts @@ -59,6 +59,10 @@ const integrationActor = actor({ jobs: queue({ message: jobSchema }), }, onBeforeConnect: async () => {}, + onRequest: async (c, request) => { + await c.db.execute("SELECT ? AS path", new URL(request.url).pathname); + return new Response("ok", { status: 200 }); + }, actions: { ping: async (c) => { return c.conn.params.userId; @@ -66,13 +70,6 @@ const integrationActor = actor({ getCount: async (c) => { return c.state.count; }, - logContext: async (c, correlationToken: string) => { - c.log.warn( - { correlation_token: correlationToken }, - "native actor log context", - ); - return correlationToken; - }, validatedAction: async (_c, payload: { amount: number }) => { return payload.amount; }, @@ -143,20 +140,19 @@ const integrationActor = actor({ kvCount: kvValue ? Number(kvValue) : null, }; }, - // Interleaves awaits, SQLite, a child actor call and a log so two - // overlapping invocations of this action have every chance to observe - // each other's telemetry context. + // Both calls reach the queue before the test releases either one. isolationProbe: async (c, token: string, fail: boolean) => { - await new Promise((resolve) => setTimeout(resolve, 20)); - await c.db.execute("SELECT ? AS probe", token); c.log.warn({ correlation_token: token }, "isolation probe"); + if (!(await c.queue.next({ names: ["jobs"], timeout: 10_000 }))) { + throw new Error("isolation probe was not released"); + } + await c.db.execute("SELECT ? AS probe", token); const client = c.client(); await client.integrationActor .getForId(c.actorId, { params: { userId: "internal-integration-test" }, }) .getCount(); - await new Promise((resolve) => setTimeout(resolve, 20)); await c.db.execute("SELECT ? AS probe2", token); if (fail) { throw new UserError("isolation probe failure", { @@ -165,13 +161,12 @@ const integrationActor = actor({ } return token; }, - // Calls another actor while an application span is active, and returns - // that span's ID so a test can check the call parented to it. getCountUnderApplicationSpan: async (c) => { return await applicationTracer.startActiveSpan( "agent.generate", async (span) => { try { + await c.db.execute("SELECT 1 AS under_span"); const client = c.client(); const count = await client.integrationActor .getForId(c.actorId, { @@ -185,6 +180,26 @@ const integrationActor = actor({ }, ); }, + // The queue gate releases the database work only after the reply. + insertAfterReply: (c, token: string) => { + c.waitUntil( + c.queue + .next({ names: ["jobs"], timeout: 10_000 }) + .then((message) => { + if (!message) + throw new Error("deferred work was not released"); + return c.db.execute("SELECT ? AS deferred", token); + }), + ); + return "replied"; + }, + consumeJob: async (c) => { + const message = await c.queue.next({ + names: ["jobs"], + timeout: 5_000, + }); + return message?.body ?? null; + }, getCountViaClient: async (c) => { const client = c.client(); return await client.integrationActor @@ -211,9 +226,25 @@ const integrationActor = actor({ }, }); +const runConsumerActor = actor({ + state: {}, + queues: { + runJobs: queue({ message: jobSchema }), + }, + run: async (c) => { + while (!c.aborted) { + await c.queue.waitForNames(["runJobs"], { + signal: c.abortSignal, + }); + } + }, + actions: {}, +}); + const registry = setup({ use: { integrationActor, + runConsumerActor, }, endpoint, namespace: process.env.RIVET_NAMESPACE ?? "default", diff --git a/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts b/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts index 5e0b57d5dd..20bc6c4969 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts @@ -3,6 +3,8 @@ import { mkdtemp, readdir, readFile, rm } from "node:fs/promises"; import { tmpdir } from "node:os"; import { dirname, join } from "node:path"; import { fileURLToPath } from "node:url"; +import { context, propagation, trace } from "@opentelemetry/api"; +import { NodeTracerProvider } from "@opentelemetry/sdk-trace-node"; import getPort from "get-port"; import { afterEach, describe, expect, test, vi } from "vitest"; import { createClient } from "../src/client/mod"; @@ -12,6 +14,19 @@ import { } from "./fixtures/otlp-collector"; const TEST_DIR = dirname(fileURLToPath(import.meta.url)); + +// Register a context manager for caller spans and baggage; no exporter is needed. +new NodeTracerProvider().register(); +const testTracer = trace.getTracer("napi-runtime-integration"); + +/** Runs `run` with `rayId` in OpenTelemetry baggage under `rivet.ray.id`. */ +function withRayBaggage(rayId: string, run: () => Promise): Promise { + const baggage = propagation.createBaggage({ + "rivet.ray.id": { value: rayId }, + }); + return context.with(propagation.setBaggage(context.active(), baggage), run); +} + const FIXTURE_PATH = join(TEST_DIR, "fixtures", "napi-runtime-server.ts"); const NAMESPACE = "default"; const TOKEN = "dev"; @@ -23,6 +38,16 @@ let runtimeLogs = { let engineEndpoint: string | undefined; let storagePath: string | undefined; +function createIntegrationClient(endpoint: string, poolName: string) { + return createClient({ + endpoint, + poolName, + token: TOKEN, + namespace: NAMESPACE, + disableMetadataLookup: true, + }) as any; +} + function runtimeOutput(): string { return [runtimeLogs.stdout, runtimeLogs.stderr].filter(Boolean).join("\n"); } @@ -451,15 +476,20 @@ interface ExportedSpan { spanId: string; parentSpanId?: string; kind?: number; + endTimeUnixNano: bigint; attributes: Record; links: Array<{ traceId: string; spanId: string }>; } /** Flattens OTLP/JSON export bodies into the spans they carry. */ function exportedSpans(exports: Buffer[]): ExportedSpan[] { - type OtlpAttribute = { key: string; value: { stringValue?: string } }; - type OtlpSpan = Omit & { + type OtlpAttribute = { + key: string; + value: { stringValue?: string; intValue?: string | number }; + }; + type OtlpSpan = Omit & { attributes?: OtlpAttribute[]; + endTimeUnixNano?: string | number; kind?: number; links?: Array<{ traceId: string; spanId: string }>; }; @@ -476,10 +506,14 @@ function exportedSpans(exports: Buffer[]): ExportedSpan[] { spanId: span.spanId, parentSpanId: span.parentSpanId || undefined, kind: span.kind, + endTimeUnixNano: BigInt(span.endTimeUnixNano ?? 0), attributes: Object.fromEntries( (span.attributes ?? []).map((attribute) => [ attribute.key, - attribute.value.stringValue, + attribute.value.stringValue ?? + (attribute.value.intValue === undefined + ? undefined + : String(attribute.value.intValue)), ]), ), links: (span.links ?? []).map((link) => ({ @@ -511,7 +545,12 @@ async function waitForSpans( } await new Promise((resolve) => setTimeout(resolve, 100)); } - throw new Error(`timed out waiting for ${description}`); + const arrived = exportedSpans(exports) + .map((span) => `${span.name} (${span.spanId} < ${span.parentSpanId})`) + .join(", "); + throw new Error( + `timed out waiting for ${description}; exported: ${arrived}`, + ); } function isSqliteSpan(span: ExportedSpan): boolean { @@ -533,6 +572,46 @@ function findInvocation( ); } +/** The `onRequest` invocation span carrying `rayId`, ignoring its children. */ +function findRequestInvocation( + spans: ExportedSpan[], + rayId: string, +): ExportedSpan | undefined { + return spans.find( + (span) => + span.attributes["rivet.invocation.type"] === "request" && + span.attributes["rivet.ray.id"] === rayId, + ); +} + +/** + * The `queue.receive` span of `actorName` under `parentSpanId`, or its root + * one when `parentSpanId` is undefined. + */ +function findQueueReceive( + spans: ExportedSpan[], + actorName: string, + parentSpanId: string | undefined, +): ExportedSpan | undefined { + return spans.find( + (span) => + span.name === `${actorName}/queue.receive` && + span.parentSpanId === parentSpanId, + ); +} + +/** The `queue.send` invocation span carrying `rayId`, ignoring its children. */ +function findQueueSendInvocation( + spans: ExportedSpan[], + rayId: string, +): ExportedSpan | undefined { + return spans.find( + (span) => + span.attributes["rivet.invocation.type"] === "queue_send" && + span.attributes["rivet.ray.id"] === rayId, + ); +} + /** Polls until an invocation span has been exported for every named action. */ async function waitForInvocationSpans( exports: Buffer[], @@ -570,7 +649,7 @@ async function waitForRuntimeLog( * returns the pieces every telemetry test needs. */ async function startTracedRuntime( - tracesEndpoint: string, + tracesEndpoint?: string, extraEnv: Record = {}, ): Promise<{ endpoint: string; poolName: string; child: ChildProcess }> { const poolName = "default"; @@ -590,10 +669,15 @@ async function startTracedRuntime( RIVETKIT_TEST_ENDPOINT: endpoint, RIVETKIT_TEST_POOL_NAME: poolName, RIVETKIT_STORAGE_PATH: storagePath, - OTEL_EXPORTER_OTLP_TRACES_ENDPOINT: tracesEndpoint, - OTEL_EXPORTER_OTLP_TRACES_PROTOCOL: "http/json", - OTEL_TRACES_SAMPLER: "always_on", - OTEL_BSP_SCHEDULE_DELAY: "10", + // Export spans only for tests that read them. + ...(tracesEndpoint + ? { + OTEL_EXPORTER_OTLP_TRACES_ENDPOINT: tracesEndpoint, + OTEL_EXPORTER_OTLP_TRACES_PROTOCOL: "http/json", + OTEL_TRACES_SAMPLER: "always_on", + OTEL_BSP_SCHEDULE_DELAY: "10", + } + : {}), ...extraEnv, }, stdio: ["ignore", "pipe", "pipe"], @@ -632,26 +716,14 @@ describe.sequential("native NAPI runtime integration", () => { }, 30_000); test("runs a TS actor through registry, NAPI, core, envoy, and engine", async () => { - collector = await startOtlpCollector( - await getPort({ host: "127.0.0.1" }), - ); - const traceExports = collector.spans(); - const { endpoint, poolName, child } = await startTracedRuntime( - collector.endpoint, - ); + const { endpoint, poolName, child } = await startTracedRuntime(); runtime = child; await waitForEnvoy(runtime, endpoint, SERVICES_POOL_NAME, 30_000); await expectNormalRunnerConfig(endpoint, SERVICES_POOL_NAME); const servicesActorId = await createServicesActor(endpoint); await waitForActorStarted(endpoint, servicesActorId, 30_000); - const client = createClient({ - endpoint, - token: TOKEN, - namespace: NAMESPACE, - poolName, - disableMetadataLookup: true, - }) as any; + const client = createIntegrationClient(endpoint, poolName); const actorKey = `napi-runtime-${crypto.randomUUID()}`; const handle = await waitForActorReady( @@ -663,32 +735,9 @@ describe.sequential("native NAPI runtime integration", () => { ); const actorId = await handle.resolve(); - const correlationToken = crypto.randomUUID(); - expect(await handle.logContext(correlationToken)).toBe( - correlationToken, - ); - const actorLog = await waitForRuntimeLog(correlationToken, 10_000); - expect(actorLog).toContain(`actorId=${actorId}`); - expect(actorLog).toContain("actorName=integrationActor"); - expect(actorLog).toContain(actorKey); - expect(actorLog).toMatch(/ rayId=[0-9a-f-]{36}( |$)/); - expect(actorLog).toMatch(/ trace_id=[0-9a-f]{32}( |$)/); - expect(actorLog).toMatch(/ span_id=[0-9a-f]{16}( |$)/); - expect(await waitForActorReady(() => handle.getCount(), 30_000)).toBe( 0, ); - const getCountSpans = await waitForInvocationSpans( - traceExports, - ["getCount"], - 10_000, - ); - expect( - findInvocation(getCountSpans, "getCount")?.attributes, - ).toMatchObject({ - "rivet.invocation.type": "action", - "rivet.actor.name": "integrationActor", - }); expect( await waitForActorReady( () => handle.validatedAction({ amount: 4 }), @@ -741,42 +790,6 @@ describe.sequential("native NAPI runtime integration", () => { count: 2, sqliteValues: [2], }); - // SQLite spans are children of the action that issued them. - const incrementSpans = await waitForSpans( - traceExports, - "increment invocation and sqlite spans", - (spans) => - spans.some(isSqliteSpan) && - findInvocation(spans, "increment") !== undefined, - 10_000, - ); - const incrementSqlite = incrementSpans.find(isSqliteSpan); - expect(incrementSqlite?.attributes).toMatchObject({ - "rivet.operation.system": "sqlite", - "rivet.operation.name": "execute", - }); - expect(incrementSqlite?.parentSpanId).toBe( - findInvocation(incrementSpans, "increment")?.spanId, - ); - traceExports.length = 0; - await expect(handle.sqliteFailure()).rejects.toMatchObject({ - code: expect.any(String), - }); - const failureSpans = await waitForSpans( - traceExports, - "sqliteFailure invocation and failed sqlite spans", - (spans) => - spans.some(isFailedSqliteSpan) && - findInvocation(spans, "sqliteFailure") !== undefined, - 10_000, - ); - const failedSqlite = failureSpans.find(isFailedSqliteSpan); - expect(failedSqlite?.attributes["error.type"]).toMatch( - /^[a-z_]+\.[a-z_]+$/, - ); - expect(failedSqlite?.parentSpanId).toBe( - findInvocation(failureSpans, "sqliteFailure")?.spanId, - ); expect(await handle.snapshot()).toEqual({ count: 2, kvCount: 2, @@ -794,35 +807,7 @@ describe.sequential("native NAPI runtime integration", () => { ).toEqual({ count: 5, }); - // An actor-owned client carries the calling invocation's trace and ray - // across the real Engine boundary, so the callee is its child. - traceExports.length = 0; expect(await handle.getCountViaClient()).toBe(5); - const clientSpans = await waitForInvocationSpans( - traceExports, - ["getCountViaClient", "getCount"], - 10_000, - ); - const caller = findInvocation(clientSpans, "getCountViaClient"); - const callee = findInvocation(clientSpans, "getCount"); - // The call out to the other actor is its own span sitting between the - // two invocations, so time spent reaching a cold or busy actor belongs - // to something instead of falling in the gap between them. - const hop = clientSpans.find( - (span) => - span.kind === OTLP_SPAN_KIND_CLIENT && - span.attributes["rivet.action.name"] === "getCount", - ); - expect(hop?.parentSpanId).toBe(caller?.spanId); - expect(callee?.parentSpanId).toBe(hop?.spanId); - expect(callee?.traceId).toBe(caller?.traceId); - expect(hop?.traceId).toBe(caller?.traceId); - expect(hop?.attributes["rivet.ray.id"]).toBe( - caller?.attributes["rivet.ray.id"], - ); - expect(callee?.attributes["rivet.ray.id"]).toBe( - caller?.attributes["rivet.ray.id"], - ); expect(await handle.stateSnapshot()).toEqual({ count: 5, kvCount: 5, @@ -841,30 +826,6 @@ describe.sequential("native NAPI runtime integration", () => { message: "An internal error occurred", }); - // Scheduled work starts a new trace linked to its origin. - traceExports.length = 0; - const scheduleToken = crypto.randomUUID(); - expect(await handle.scheduleTrace(scheduleToken)).toBe(scheduleToken); - const scheduleSpans = await waitForInvocationSpans( - traceExports, - ["scheduleTrace", "scheduledTrace"], - 15_000, - ); - const definer = findInvocation(scheduleSpans, "scheduleTrace"); - const scheduled = findInvocation(scheduleSpans, "scheduledTrace"); - expect(definer).toBeDefined(); - expect(scheduled?.attributes["rivet.invocation.type"]).toBe( - "scheduled", - ); - expect(scheduled?.attributes["rivet.ray.id"]).toBe( - definer?.attributes["rivet.ray.id"], - ); - // Ensure the scheduled action succeeded before checking its trace. - expect(scheduled?.attributes["error.type"]).toBeUndefined(); - expect(scheduled?.traceId).not.toBe(definer?.traceId); - expect(scheduled?.links).toEqual([ - { traceId: definer?.traceId, spanId: definer?.spanId }, - ]); await client.dispose(); const processId = servicesPid(); @@ -883,13 +844,7 @@ describe.sequential("native NAPI runtime integration", () => { ); runtime = child; - const client = createClient({ - endpoint, - token: TOKEN, - namespace: NAMESPACE, - poolName, - disableMetadataLookup: true, - }) as any; + const client = createIntegrationClient(endpoint, poolName); const handle = await waitForActorReady( () => client.integrationActor.create( @@ -903,10 +858,17 @@ describe.sequential("native NAPI runtime integration", () => { // Use the same actor to exercise isolation between concurrent invocations. const okToken = crypto.randomUUID(); const failToken = crypto.randomUUID(); - const [ok, failed] = await Promise.allSettled([ + const results = Promise.allSettled([ handle.isolationProbe(okToken, false), handle.isolationProbe(failToken, true), ]); + const [okLog, failLog] = await Promise.all([ + waitForRuntimeLog(okToken, 10_000), + waitForRuntimeLog(failToken, 10_000), + ]); + await handle.send("jobs", { id: okToken }); + await handle.send("jobs", { id: failToken }); + const [ok, failed] = await results; expect(ok.status).toBe("fulfilled"); expect(failed.status).toBe("rejected"); @@ -921,14 +883,24 @@ describe.sequential("native NAPI runtime integration", () => { ); return ( probes.length >= 2 && - probes.every((probe) => - exported.some( + probes.every((probe) => { + const hop = exported.find( (span) => - span.attributes["rivet.action.name"] === - "getCount" && + span.kind === OTLP_SPAN_KIND_CLIENT && span.traceId === probe.traceId, - ), - ) + ); + return ( + !!hop && + exported.some( + (span) => span.parentSpanId === hop.spanId, + ) && + exported.filter( + (span) => + isSqliteSpan(span) && + span.parentSpanId === probe.spanId, + ).length >= 2 + ); + }) ); }, 20_000, @@ -990,8 +962,11 @@ describe.sequential("native NAPI runtime integration", () => { } } - const okLog = await waitForRuntimeLog(okToken, 10_000); - const failLog = await waitForRuntimeLog(failToken, 10_000); + for (const line of [okLog, failLog]) { + expect(line).toContain(`actorId=${await handle.resolve()}`); + expect(line).toMatch(/ trace_id=[0-9a-f]{32}( |$)/); + expect(line).toMatch(/ span_id=[0-9a-f]{16}( |$)/); + } const rayOf = (line: string) => / rayId=([A-Za-z0-9_-]+)/.exec(line)?.[1]; expect(rayOf(okLog)).toBeDefined(); @@ -1001,30 +976,299 @@ describe.sequential("native NAPI runtime integration", () => { // SQLite and actor calls inherit the active application span. const underApp = await handle.getCountUnderApplicationSpan(); - const appHop = await waitForSpans( + const applicationSpans = await waitForSpans( traceExports, - "the hop made under an application span", - (exported) => - exported.some( + "SQLite, outgoing call, and callee under the application span", + (exported) => { + const hop = exported.find( (span) => span.kind === OTLP_SPAN_KIND_CLIENT && span.parentSpanId === underApp.spanId, - ), + ); + return ( + !!hop && + exported.some((span) => span.parentSpanId === hop.spanId) && + exported.some( + (span) => + isSqliteSpan(span) && + span.parentSpanId === underApp.spanId, + ) + ); + }, 10_000, - ).then((exported) => - exported.find( + ); + const appHop = applicationSpans.find( + (span) => + span.kind === OTLP_SPAN_KIND_CLIENT && + span.parentSpanId === underApp.spanId, + ); + expect( + applicationSpans.find( + (span) => span.parentSpanId === appHop?.spanId, + )?.attributes["rivet.action.name"], + ).toBe("getCount"); + expect( + applicationSpans.find( (span) => - span.kind === OTLP_SPAN_KIND_CLIENT && - span.parentSpanId === underApp.spanId, - ), + isSqliteSpan(span) && span.parentSpanId === underApp.spanId, + )?.attributes["rivet.operation.name"], + ).toBe("execute"); + + await client.dispose(); + }, 120_000); + + test("carries a caller-supplied ray through requests, queue sends, and work after the reply", async () => { + collector = await startOtlpCollector( + await getPort({ host: "127.0.0.1" }), + ); + const traceExports = collector.spans(); + const { endpoint, poolName, child } = await startTracedRuntime( + collector.endpoint, + ); + runtime = child; + + const client = createIntegrationClient(endpoint, poolName); + const handle = await waitForActorReady( + () => + client.integrationActor.create( + [`napi-caller-ray-${crypto.randomUUID()}`], + { params: { userId: "integration-test" } }, + ), + 30_000, + ); + await waitForActorReady(() => handle.getCount(), 30_000); + + const callerRay = `caller-${crypto.randomUUID()}`; + await withRayBaggage(callerRay, () => handle.getCount()); + const rayedGetCount = await waitForSpans( + traceExports, + "the getCount invocation carrying the caller-supplied ray", + (exported) => + exported.some( + (span) => + span.attributes["rivet.action.name"] === "getCount" && + span.attributes["rivet.ray.id"] === callerRay, + ), + 10_000, + ); + expect( + rayedGetCount.find( + (span) => span.attributes["rivet.ray.id"] === callerRay, + )?.attributes["rivet.invocation.type"], + ).toBe("action"); + + // Fetch inherits the active span and baggage without explicit headers. + const requestRay = `request-${crypto.randomUUID()}`; + const callerSpan = testTracer.startSpan("request.handle"); + const underCallerSpan = (run: () => Promise) => + withRayBaggage(requestRay, () => + context.with(trace.setSpan(context.active(), callerSpan), run), + ); + const response = await underCallerSpan(() => handle.fetch("hello")); + expect(response.status).toBe(200); + const requestSpans = await waitForSpans( + traceExports, + "the onRequest invocation and its sqlite span", + (exported) => { + const request = findRequestInvocation(exported, requestRay); + return ( + request !== undefined && + exported.some( + (span) => + isSqliteSpan(span) && + span.parentSpanId === request.spanId, + ) + ); + }, + 10_000, + ); + const requestSpan = findRequestInvocation(requestSpans, requestRay); + expect(requestSpan?.name).toBe("integrationActor/onRequest"); + expect(requestSpan?.attributes["http.response.status_code"]).toBe( + "200", + ); + expect(requestSpan?.traceId).toBe(callerSpan.spanContext().traceId); + expect(requestSpan?.parentSpanId).toBe(callerSpan.spanContext().spanId); + + // Headers the caller set on the request win over the active context. + const explicitRay = `explicit-${crypto.randomUUID()}`; + const explicitTraceId = crypto.randomUUID().replaceAll("-", ""); + const explicitSpanId = crypto + .randomUUID() + .replaceAll("-", "") + .slice(0, 16); + const explicitResponse = await underCallerSpan(() => + handle.fetch("hello", { + headers: { + "x-rivet-ray-id": explicitRay, + traceparent: `00-${explicitTraceId}-${explicitSpanId}-01`, + }, + }), + ); + callerSpan.end(); + expect(explicitResponse.status).toBe(200); + const explicitSpans = await waitForSpans( + traceExports, + "the onRequest invocation under the caller's own headers", + (exported) => + findRequestInvocation(exported, explicitRay) !== undefined, + 10_000, ); - const appCallee = exportedSpans(traceExports).find( + const explicitSpan = findRequestInvocation(explicitSpans, explicitRay); + expect(explicitSpan?.traceId).toBe(explicitTraceId); + expect(explicitSpan?.parentSpanId).toBe(explicitSpanId); + + // Release deferred work after the reply; its invocation must remain open. + const deferredToken = crypto.randomUUID(); + expect(await handle.insertAfterReply(deferredToken)).toBe("replied"); + await handle.send("jobs", { id: deferredToken }); + const deferredSpans = await waitForSpans( + traceExports, + "the insertAfterReply invocation and its deferred sqlite span", + (exported) => { + const invocation = findInvocation(exported, "insertAfterReply"); + return ( + invocation !== undefined && + exported.some( + (span) => + isSqliteSpan(span) && + span.parentSpanId === invocation.spanId, + ) + ); + }, + 10_000, + ); + const deferredInvocation = findInvocation( + deferredSpans, + "insertAfterReply", + ); + const deferredSqlite = deferredSpans.find( (span) => - span.attributes["rivet.invocation.type"] !== undefined && - span.parentSpanId === appHop?.spanId, + isSqliteSpan(span) && + span.parentSpanId === deferredInvocation?.spanId, + ); + expect(deferredSqlite).toBeDefined(); + expect( + deferredInvocation !== undefined && + deferredSqlite !== undefined && + deferredInvocation.endTimeUnixNano >= + deferredSqlite.endTimeUnixNano, + ).toBe(true); + + // The action receipt links to the send and keeps the consuming action’s ray. + const queueRay = `queue-${crypto.randomUUID()}`; + await withRayBaggage(queueRay, () => + handle.send("jobs", { id: "job-42" }), + ); + expect(await handle.consumeJob()).toEqual({ id: "job-42" }); + const queueSpans = await waitForSpans( + traceExports, + "the queue send invocation, the consuming action, and its receipt", + (exported) => { + const consumer = findInvocation(exported, "consumeJob"); + return ( + findQueueSendInvocation(exported, queueRay) !== undefined && + consumer !== undefined && + findQueueReceive( + exported, + "integrationActor", + consumer.spanId, + ) !== undefined + ); + }, + 10_000, + ); + const queueSend = findQueueSendInvocation(queueSpans, queueRay); + expect(queueSend?.name).toBe("integrationActor/queue.send"); + expect(queueSend?.attributes).toMatchObject({ + "rivet.invocation.type": "queue_send", + "rivet.queue.name": "jobs", + }); + const consumer = findInvocation(queueSpans, "consumeJob"); + const receipt = findQueueReceive( + queueSpans, + "integrationActor", + consumer?.spanId, + ); + expect(receipt?.attributes).toMatchObject({ + "rivet.queue.name": "jobs", + "rivet.ray.id": consumer?.attributes["rivet.ray.id"], + }); + expect(receipt?.links).toEqual([ + { traceId: queueSend?.traceId, spanId: queueSend?.spanId }, + ]); + + // Without an invocation, the receipt is a root span carrying the sender’s ray. + const runRay = `run-${crypto.randomUUID()}`; + const runConsumer = client.runConsumerActor.getOrCreate([ + `napi-run-consumer-${crypto.randomUUID()}`, + ]); + await withRayBaggage(runRay, () => + runConsumer.send("runJobs", { id: "job-run" }), + ); + const runSpans = await waitForSpans( + traceExports, + "the run handler's receipt of a queue message", + (exported) => + findQueueSendInvocation(exported, runRay) !== undefined && + findQueueReceive(exported, "runConsumerActor", undefined) !== + undefined, + 10_000, + ); + const runSend = findQueueSendInvocation(runSpans, runRay); + const runReceipt = findQueueReceive( + runSpans, + "runConsumerActor", + undefined, + ); + expect(runReceipt?.attributes).toMatchObject({ + "rivet.queue.name": "runJobs", + "rivet.ray.id": runRay, + }); + expect(runReceipt?.links).toEqual([ + { traceId: runSend?.traceId, spanId: runSend?.spanId }, + ]); + + traceExports.length = 0; + await expect(handle.sqliteFailure()).rejects.toMatchObject({ + code: expect.any(String), + }); + const failureSpans = await waitForSpans( + traceExports, + "sqliteFailure invocation and failed sqlite spans", + (spans) => + spans.some(isFailedSqliteSpan) && + findInvocation(spans, "sqliteFailure") !== undefined, + 10_000, + ); + const failedSqlite = failureSpans.find(isFailedSqliteSpan); + expect(failedSqlite?.parentSpanId).toBe( + findInvocation(failureSpans, "sqliteFailure")?.spanId, ); - expect(appCallee?.attributes["rivet.action.name"]).toBe("getCount"); + // Scheduled work starts a new trace linked to its origin. + traceExports.length = 0; + const scheduleToken = crypto.randomUUID(); + expect(await handle.scheduleTrace(scheduleToken)).toBe(scheduleToken); + const scheduleSpans = await waitForInvocationSpans( + traceExports, + ["scheduleTrace", "scheduledTrace"], + 15_000, + ); + const definer = findInvocation(scheduleSpans, "scheduleTrace"); + const scheduled = findInvocation(scheduleSpans, "scheduledTrace"); + expect(scheduled?.attributes["rivet.invocation.type"]).toBe( + "scheduled", + ); + expect(scheduled?.attributes["rivet.ray.id"]).toBe( + definer?.attributes["rivet.ray.id"], + ); + // Ensure the scheduled action succeeded before checking its trace. + expect(scheduled?.attributes["error.type"]).toBeUndefined(); + expect(scheduled?.traceId).not.toBe(definer?.traceId); + expect(scheduled?.links).toEqual([ + { traceId: definer?.traceId, spanId: definer?.spanId }, + ]); await client.dispose(); }, 120_000); @@ -1035,13 +1279,7 @@ describe.sequential("native NAPI runtime integration", () => { await startTracedRuntime(unavailable); runtime = child; - const client = createClient({ - endpoint, - token: TOKEN, - namespace: NAMESPACE, - poolName, - disableMetadataLookup: true, - }) as any; + const client = createIntegrationClient(endpoint, poolName); const handle = await waitForActorReady( () => client.integrationActor.create( @@ -1083,13 +1321,7 @@ describe.sequential("native NAPI runtime integration", () => { ); runtime = child; - const client = createClient({ - endpoint, - token: TOKEN, - namespace: NAMESPACE, - poolName, - disableMetadataLookup: true, - }) as any; + const client = createIntegrationClient(endpoint, poolName); const handle = await waitForActorReady( () => client.integrationActor.create( From f04c8d063de893a6bfa4205f3c3ad954eae5d292 Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Mon, 7 Sep 2026 21:08:09 +0400 Subject: [PATCH 19/21] fix(rivetkit): use standard W3C trace context propagation --- .../packages/client-protocol/src/lib.rs | 2 +- .../packages/client-protocol/src/ray_id.rs | 44 +++++ .../client-protocol/src/telemetry_headers.rs | 37 ----- rivetkit-rust/packages/client/Cargo.toml | 3 +- .../packages/client/src/remote_manager.rs | 104 +++++------- .../packages/rivetkit-core/Cargo.toml | 4 +- .../src/actor/internal_storage/queries.rs | 1 + .../packages/rivetkit-core/src/telemetry.rs | 156 ++++++++++-------- .../packages/rivetkit-napi/index.d.ts | 1 - .../rivetkit-napi/src/actor_context.rs | 2 - .../rivetkit/src/client/outbound-telemetry.ts | 13 +- .../src/common/actor-telemetry-context.ts | 11 -- .../rivetkit/src/common/otel-context.ts | 58 ++++--- .../rivetkit/src/registry/napi-runtime.ts | 11 +- .../tests/napi-runtime-integration.test.ts | 90 ++++++++++ 15 files changed, 323 insertions(+), 214 deletions(-) create mode 100644 rivetkit-rust/packages/client-protocol/src/ray_id.rs delete mode 100644 rivetkit-rust/packages/client-protocol/src/telemetry_headers.rs diff --git a/rivetkit-rust/packages/client-protocol/src/lib.rs b/rivetkit-rust/packages/client-protocol/src/lib.rs index faf8453c77..bef5d9a66c 100644 --- a/rivetkit-rust/packages/client-protocol/src/lib.rs +++ b/rivetkit-rust/packages/client-protocol/src/lib.rs @@ -1,5 +1,5 @@ pub mod generated; -pub mod telemetry_headers; +pub mod ray_id; pub mod versioned; // Re-export latest. diff --git a/rivetkit-rust/packages/client-protocol/src/ray_id.rs b/rivetkit-rust/packages/client-protocol/src/ray_id.rs new file mode 100644 index 0000000000..b47a59eab1 --- /dev/null +++ b/rivetkit-rust/packages/client-protocol/src/ray_id.rs @@ -0,0 +1,44 @@ +use std::fmt; + +pub const HEADER_RIVET_RAY_ID: &str = "x-rivet-ray-id"; +pub const RAY_BAGGAGE_KEY: &str = "rivet.ray.id"; + +const MAX_LEN: usize = 128; + +/// A bounded correlation token carried between RivetKit requests. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct RayId(String); + +impl RayId { + pub fn parse(value: impl Into) -> Result { + let value = value.into(); + if value.is_empty() + || value.len() > MAX_LEN + || !value + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-')) + { + return Err(InvalidRayId); + } + Ok(Self(value)) + } + + pub fn as_str(&self) -> &str { + &self.0 + } + + pub fn into_string(self) -> String { + self.0 + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct InvalidRayId; + +impl fmt::Display for InvalidRayId { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("ray ID must be 1 to 128 characters of [A-Za-z0-9_-]") + } +} + +impl std::error::Error for InvalidRayId {} diff --git a/rivetkit-rust/packages/client-protocol/src/telemetry_headers.rs b/rivetkit-rust/packages/client-protocol/src/telemetry_headers.rs deleted file mode 100644 index dc26a1e1cd..0000000000 --- a/rivetkit-rust/packages/client-protocol/src/telemetry_headers.rs +++ /dev/null @@ -1,37 +0,0 @@ -//! Headers that carry a caller's ray ID and W3C trace context into an actor, -//! shared by every Rust peer of the actor HTTP surface: the runtime that reads -//! them and the clients that send them. - -/// Bounded correlation string that follows work across actors, surviving -/// sampling and trace-root boundaries. -pub const HEADER_RIVET_RAY_ID: &str = "x-rivet-ray-id"; -pub const HEADER_TRACEPARENT: &str = "traceparent"; -pub const HEADER_TRACESTATE: &str = "tracestate"; - -/// W3C Baggage key that carries a ray ID through application code, so a request -/// handler that received a ray ID can hand it to the actors it calls. -pub const RAY_BAGGAGE_KEY: &str = "rivet.ray.id"; - -const RAY_ID_MAX_LEN: usize = 128; - -/// Returns `value` when it is a ray ID the runtime accepts: 1 to 128 characters -/// of `[A-Za-z0-9_-]`. The value arrives from a caller, so anything else -/// counts as absent. -pub fn bounded_ray_id(value: &str) -> Option<&str> { - let valid = !value.is_empty() - && value.len() <= RAY_ID_MAX_LEN - && value - .bytes() - .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_')); - if valid { Some(value) } else { None } -} - -/// Formats a version 00 W3C `traceparent` from the parts of a span context, -/// so every Rust peer writes the header the same way. -pub fn format_traceparent( - trace_id: impl std::fmt::Display, - span_id: impl std::fmt::Display, - trace_flags: u8, -) -> String { - format!("00-{trace_id}-{span_id}-{trace_flags:02x}") -} diff --git a/rivetkit-rust/packages/client/Cargo.toml b/rivetkit-rust/packages/client/Cargo.toml index 7dc2cf96f8..33f399d298 100644 --- a/rivetkit-rust/packages/client/Cargo.toml +++ b/rivetkit-rust/packages/client/Cargo.toml @@ -14,6 +14,8 @@ base64 = "0.22.1" bytes = { workspace = true } futures-util = "0.3.31" opentelemetry = { version = "0.28", default-features = false, features = ["trace"] } +opentelemetry-http.workspace = true +opentelemetry_sdk = { version = "0.28", default-features = false, features = ["trace"] } parking_lot.workspace = true reqwest = { version = "0.12.12", default-features = false, features = ["json", "charset", "http2", "macos-system-configuration", "rustls-tls-native-roots", "rustls-tls-webpki-roots"] } rivetkit-client-protocol.workspace = true @@ -32,7 +34,6 @@ vbare = "0.0.4" [dev-dependencies] axum = { workspace = true, features = ["ws"] } -opentelemetry_sdk = { version = "0.28", default-features = false, features = ["trace"] } tracing-subscriber = { version = "0.3.19", features = ["env-filter", "std", "registry"]} tempfile = "3.10.1" tokio-test = "0.4.3" diff --git a/rivetkit-rust/packages/client/src/remote_manager.rs b/rivetkit-rust/packages/client/src/remote_manager.rs index 780078fc5e..af94047a04 100644 --- a/rivetkit-rust/packages/client/src/remote_manager.rs +++ b/rivetkit-rust/packages/client/src/remote_manager.rs @@ -1,16 +1,15 @@ -use anyhow::{anyhow, Context, Result}; -use base64::{engine::general_purpose, engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; +use anyhow::{Context, Result, anyhow}; +use base64::{Engine as _, engine::general_purpose, engine::general_purpose::URL_SAFE_NO_PAD}; use bytes::Bytes; use opentelemetry::baggage::BaggageExt as _; -use opentelemetry::trace::TraceContextExt as _; +use opentelemetry::propagation::TextMapPropagator as _; +use opentelemetry_http::HeaderInjector; +use opentelemetry_sdk::propagation::TraceContextPropagator; use reqwest::{ - header::{HeaderMap, HeaderName, HeaderValue, USER_AGENT}, Method, + header::{HeaderMap, HeaderName, HeaderValue, USER_AGENT}, }; -use rivetkit_client_protocol::telemetry_headers::{ - HEADER_RIVET_RAY_ID, HEADER_TRACEPARENT, HEADER_TRACESTATE, RAY_BAGGAGE_KEY, bounded_ray_id, - format_traceparent, -}; +use rivetkit_client_protocol::ray_id::{HEADER_RIVET_RAY_ID, RAY_BAGGAGE_KEY, RayId}; use serde::{Deserialize, Serialize}; use serde_cbor; use std::{collections::HashMap, str::FromStr, sync::Arc}; @@ -20,15 +19,18 @@ use tracing_opentelemetry::OpenTelemetrySpanExt as _; use crate::{ common::{ - serialize_actor_key, ActorKey, EncodingKind, RawWebSocket, HEADER_RIVET_ACTOR, - HEADER_RIVET_NAMESPACE, HEADER_RIVET_TARGET, HEADER_RIVET_TOKEN, PATH_CONNECT_WEBSOCKET, - PATH_WEBSOCKET_PREFIX, USER_AGENT_VALUE, WS_PROTOCOL_ACTOR, WS_PROTOCOL_CONN_ID, - WS_PROTOCOL_CONN_PARAMS, WS_PROTOCOL_CONN_TOKEN, WS_PROTOCOL_ENCODING, - WS_PROTOCOL_STANDARD, WS_PROTOCOL_TARGET, WS_PROTOCOL_TOKEN, + ActorKey, EncodingKind, HEADER_RIVET_ACTOR, HEADER_RIVET_NAMESPACE, HEADER_RIVET_TARGET, + HEADER_RIVET_TOKEN, PATH_CONNECT_WEBSOCKET, PATH_WEBSOCKET_PREFIX, RawWebSocket, + USER_AGENT_VALUE, WS_PROTOCOL_ACTOR, WS_PROTOCOL_CONN_ID, WS_PROTOCOL_CONN_PARAMS, + WS_PROTOCOL_CONN_TOKEN, WS_PROTOCOL_ENCODING, WS_PROTOCOL_STANDARD, WS_PROTOCOL_TARGET, + WS_PROTOCOL_TOKEN, serialize_actor_key, }, protocol::query::ActorQuery, }; +const HEADER_TRACEPARENT: &str = "traceparent"; +const HEADER_TRACESTATE: &str = "tracestate"; + #[derive(Clone)] pub struct RemoteManager { endpoint: String, @@ -36,7 +38,7 @@ pub struct RemoteManager { namespace: String, pool_name: String, headers: HashMap, - ray_id: Option, + ray_id: Option, max_input_size: usize, disable_metadata_lookup: bool, resolved_config: Arc>, @@ -152,15 +154,15 @@ impl RemoteManager { max_input_size: Option, disable_metadata_lookup: bool, ) -> Self { - let ray_id = ray_id.and_then(|ray_id| { - let bounded = bounded_ray_id(&ray_id).map(str::to_owned); - if bounded.is_none() { + let ray_id = ray_id.and_then(|ray_id| match RayId::parse(ray_id) { + Ok(ray_id) => Some(ray_id), + Err(error) => { tracing::warn!( - len = ray_id.len(), - "dropping configured ray id; it must be 1 to 128 characters of [A-Za-z0-9_-]" + %error, + "dropping invalid configured ray ID" ); + None } - bounded }); Self { endpoint, @@ -523,15 +525,7 @@ impl RemoteManager { // caller passed for this request win over both, matching the // TypeScript client. let mut headers = headers; - let caller_set_trace_context = - headers.contains_key(HEADER_TRACEPARENT) || headers.contains_key(HEADER_TRACESTATE); - for (name, value) in self.telemetry_headers()? { - let is_trace_context = name == HEADER_TRACEPARENT || name == HEADER_TRACESTATE; - if is_trace_context && caller_set_trace_context { - continue; - } - headers.entry(name).or_insert(value); - } + self.add_telemetry_headers(&mut headers)?; req = req.headers(headers); if let Some(body_data) = body { @@ -546,44 +540,30 @@ impl RemoteManager { /// read from the `tracing` span current at the call. Without a registered /// OpenTelemetry layer the span carries no context and only a configured /// ray ID is sent. - fn telemetry_headers(&self) -> Result> { - let mut headers = Vec::with_capacity(3); + fn add_telemetry_headers(&self, headers: &mut HeaderMap) -> Result<()> { let context = tracing::Span::current().context(); - let span = context.span(); - let span_context = span.span_context(); - if span_context.is_valid() { - let traceparent = format_traceparent( - span_context.trace_id(), - span_context.span_id(), - span_context.trace_flags().to_u8(), - ); - headers.push(( - HeaderName::from_static(HEADER_TRACEPARENT), - HeaderValue::from_str(&traceparent).context("format traceparent header")?, - )); - let tracestate = span_context.trace_state().header(); - if !tracestate.is_empty() { - headers.push(( - HeaderName::from_static(HEADER_TRACESTATE), - HeaderValue::from_str(&tracestate).context("format tracestate header")?, - )); + let caller_set_trace_context = + headers.contains_key(HEADER_TRACEPARENT) || headers.contains_key(HEADER_TRACESTATE); + if !caller_set_trace_context { + TraceContextPropagator::new().inject_context(&context, &mut HeaderInjector(headers)); + if headers + .get(HEADER_TRACESTATE) + .is_some_and(HeaderValue::is_empty) + { + headers.remove(HEADER_TRACESTATE); } } - let baggage_ray = context - .baggage() + + let baggage = context.baggage(); + let baggage_ray = baggage .get(RAY_BAGGAGE_KEY) - .map(|value| value.as_str().into_owned()); - let ray_id = baggage_ray - .as_deref() - .and_then(bounded_ray_id) - .or(self.ray_id.as_deref()); - if let Some(ray_id) = ray_id { - headers.push(( - HeaderName::from_static(HEADER_RIVET_RAY_ID), - HeaderValue::from_str(ray_id).context("format ray id header")?, - )); + .and_then(|value| RayId::parse(value.as_str().into_owned()).ok()); + if let Some(ray_id) = baggage_ray.as_ref().or(self.ray_id.as_ref()) { + headers + .entry(HEADER_RIVET_RAY_ID) + .or_insert(HeaderValue::from_str(ray_id.as_str()).context("format ray ID header")?); } - Ok(headers) + Ok(()) } pub fn gateway_url(&self, query: &ActorQuery) -> Result { diff --git a/rivetkit-rust/packages/rivetkit-core/Cargo.toml b/rivetkit-rust/packages/rivetkit-core/Cargo.toml index 09bbec31e2..15fa510386 100644 --- a/rivetkit-rust/packages/rivetkit-core/Cargo.toml +++ b/rivetkit-rust/packages/rivetkit-core/Cargo.toml @@ -16,7 +16,6 @@ native-runtime = [ "dep:nix", "dep:reqwest", "dep:opentelemetry-otlp", - "dep:opentelemetry_sdk", "dep:tracing-subscriber", "dep:rivetkit-engine-process", "dep:axum", @@ -57,8 +56,9 @@ opentelemetry = { version = "0.28", default-features = false, features = [ # Lets the SDK report dropped spans and export failures through tracing. "internal-logs", ] } +opentelemetry-http.workspace = true opentelemetry-otlp = { version = "0.28", default-features = false, optional = true, features = ["trace", "http-json", "http-proto", "grpc-tonic", "reqwest-blocking-client"] } -opentelemetry_sdk = { version = "0.28", default-features = false, optional = true, features = ["trace", "internal-logs"] } +opentelemetry_sdk = { version = "0.28", default-features = false, features = ["trace", "internal-logs"] } rand.workspace = true reqwest = { workspace = true, optional = true } rusqlite = { workspace = true, optional = true } diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/queries.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/queries.rs index a5a94508b0..af94ce3962 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/queries.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/queries.rs @@ -14,6 +14,7 @@ pub(crate) const DELETE_CONN_STATE_SQL: &str = "DELETE FROM _rivet_conn_state WH pub(crate) const DELETE_CONN_SQL: &str = "DELETE FROM _rivet_conns WHERE conn_id = ?"; pub(crate) const RESET_SCHEDULES_FOR_LEGACY_IMPORT_SQL: &str = "DELETE FROM _rivet_schedule_events"; +// `;` immediately follows `:` in ASCII, so this range selects the prefix via the primary-key index. pub(crate) const RESET_SCHEDULE_TRACE_CONTEXTS_SQL: &str = "DELETE FROM _rivet_meta WHERE key >= 'schedule_trace_context:' AND key < 'schedule_trace_context;'"; pub(crate) const INSERT_SCHEDULE_EVENT_SQL: &str = "INSERT INTO _rivet_schedule_events (event_id, trigger_at, action, args, kind, cron_expression, timezone, interval_ms, last_started_at, max_history) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"; pub(crate) const UPSERT_RECURRING_SCHEDULE_SQL: &str = "INSERT INTO _rivet_schedule_events (event_id, trigger_at, action, args, kind, cron_expression, timezone, interval_ms, last_started_at, max_history) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(event_id) DO UPDATE SET trigger_at = excluded.trigger_at, action = excluded.action, args = excluded.args, kind = excluded.kind, cron_expression = excluded.cron_expression, timezone = excluded.timezone, interval_ms = excluded.interval_ms, max_history = excluded.max_history"; diff --git a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs index 9d00c62887..06c87d2a09 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs @@ -3,15 +3,16 @@ #[cfg(feature = "native-runtime")] pub mod export; -use std::str::FromStr as _; use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use opentelemetry::trace::{ - SpanContext, SpanId, TraceContextExt as _, TraceFlags, TraceId, TraceState, -}; +use opentelemetry::Context; +use opentelemetry::propagation::{Extractor, TextMapPropagator as _}; +use opentelemetry::trace::{SpanContext, TraceContextExt as _}; +use opentelemetry_http::{HeaderExtractor, HeaderInjector}; +use opentelemetry_sdk::propagation::TraceContextPropagator; use parking_lot::Mutex; -use rivetkit_client_protocol::telemetry_headers::format_traceparent; +use rivetkit_client_protocol::ray_id::RayId; use tracing_opentelemetry::OpenTelemetrySpanExt as _; use crate::ActorContext; @@ -26,33 +27,24 @@ pub struct IncomingInvocationContext { remote_parent: Option, } -pub(crate) use rivetkit_client_protocol::telemetry_headers::HEADER_RIVET_RAY_ID; +pub(crate) use rivetkit_client_protocol::ray_id::HEADER_RIVET_RAY_ID; -impl IncomingInvocationContext { - pub(crate) fn from_headers( - ray_id: Option, - traceparent: Option<&str>, - tracestate: Option<&str>, - ) -> Self { - Self { - ray_id, - remote_parent: parse_remote_parent(traceparent, tracestate), - } - } +const HEADER_TRACEPARENT: &str = "traceparent"; +const HEADER_TRACESTATE: &str = "tracestate"; +impl IncomingInvocationContext { /// Reads the ray ID and W3C trace context an HTTP request carries. Every /// HTTP entry point into an actor reads them through here, so they all /// apply the same bounds. pub(crate) fn from_http_headers(headers: &http::HeaderMap) -> Self { - Self::from_headers( - invocation_ray_id(headers), - headers - .get("traceparent") - .and_then(|value| value.to_str().ok()), - headers - .get("tracestate") - .and_then(|value| value.to_str().ok()), - ) + Self { + ray_id: invocation_ray_id(headers), + remote_parent: if headers.contains_key(HEADER_TRACEPARENT) { + extract_remote_parent(&HeaderExtractor(headers)) + } else { + None + }, + } } } @@ -61,7 +53,7 @@ impl IncomingInvocationContext { /// absent. fn invocation_ray_id(headers: &http::HeaderMap) -> Option { let value = headers.get(HEADER_RIVET_RAY_ID)?.to_str().ok()?; - rivetkit_client_protocol::telemetry_headers::bounded_ray_id(value).map(str::to_owned) + RayId::parse(value.to_owned()).ok().map(RayId::into_string) } /// Name a request invocation is reported under, in place of an action name. @@ -171,7 +163,6 @@ pub struct ActorInvocationSpanContext { pub trace_id: String, pub span_id: String, pub trace_flags: u8, - pub traceparent: String, pub tracestate: Option, } @@ -547,16 +538,21 @@ impl ActorInvocationTelemetry { /// back to it then points at the code that caused it rather than at the /// whole invocation around that code. pub(crate) fn trace_origin(&self) -> TraceOrigin { - let Some(context) = self.trace_context() else { + let Some(active) = self.active() else { return TraceOrigin::default(); }; - let span = self.1.as_ref().and_then(w3c_span_context).or(context.span); - let (traceparent, tracestate) = match span { - Some(span) => (Some(span.traceparent), span.tracestate), - None => (None, None), - }; + let span = self + .1 + .clone() + .or_else(|| active.span.lock().as_ref().and_then(otel_span_context_of)); + let (traceparent, tracestate) = span + .as_ref() + .and_then(propagation_headers) + .map_or((None, None), |(traceparent, tracestate)| { + (Some(traceparent), tracestate) + }); TraceOrigin { - ray_id: context.ray_id, + ray_id: active.ray_id.clone(), traceparent, tracestate, } @@ -680,7 +676,7 @@ impl Drop for InvocationWorkGuard { impl OutboundCallInvocation { /// W3C context of this call's span, to send to the callee so it parents - /// here. Absent when the call is not sampled. + /// here. Absent when tracing is disabled. pub fn span_context(&self) -> Option { self.context.clone() } @@ -730,7 +726,7 @@ impl Drop for SqliteOperationSpan { /// W3C fields of a span context, or nothing when it is not valid and so /// carries nothing worth propagating. -fn w3c_span_context(span_context: &SpanContext) -> Option { +fn span_context_fields(span_context: &SpanContext) -> Option { if !span_context.is_valid() { return None; } @@ -739,11 +735,6 @@ fn w3c_span_context(span_context: &SpanContext) -> Option Option Option { + otel_span_context_of(span) + .as_ref() + .and_then(span_context_fields) +} + +fn otel_span_context_of(span: &tracing::Span) -> Option { let context = span.context(); let context_span = context.span(); - w3c_span_context(context_span.span_context()) + let span_context = context_span.span_context(); + span_context.is_valid().then(|| span_context.clone()) } /// Opens the span covering the moment one queue message is handed to the @@ -826,31 +824,55 @@ fn record_outcome(span: &tracing::Span, error: Option<&anyhow::Error>) { } fn parse_remote_parent(traceparent: Option<&str>, tracestate: Option<&str>) -> Option { - let mut fields = traceparent?.split('-'); - let version = fields.next()?; - let trace_id = fields.next()?; - let span_id = fields.next()?; - let flags = fields.next()?; - if fields.next().is_some() - || version.len() != 2 - || version.eq_ignore_ascii_case("ff") - || trace_id.len() != 32 - || span_id.len() != 16 - || flags.len() != 2 - { - return None; + traceparent?; + extract_remote_parent(&TraceHeaders { + traceparent, + tracestate, + }) +} + +fn extract_remote_parent(extractor: &dyn Extractor) -> Option { + let context = TraceContextPropagator::new().extract_with_context(&Context::new(), extractor); + let span = context.span(); + let span_context = span.span_context(); + span_context.is_valid().then(|| span_context.clone()) +} + +fn propagation_headers(span_context: &SpanContext) -> Option<(String, Option)> { + let context = Context::new().with_remote_span_context(span_context.clone()); + let mut headers = http::HeaderMap::new(); + TraceContextPropagator::new().inject_context(&context, &mut HeaderInjector(&mut headers)); + let traceparent = headers.get(HEADER_TRACEPARENT)?.to_str().ok()?.to_owned(); + let tracestate = headers + .get(HEADER_TRACESTATE) + .and_then(|value| value.to_str().ok()) + .filter(|value| !value.is_empty()) + .map(str::to_owned); + Some((traceparent, tracestate)) +} + +struct TraceHeaders<'a> { + traceparent: Option<&'a str>, + tracestate: Option<&'a str>, +} + +impl Extractor for TraceHeaders<'_> { + fn get(&self, key: &str) -> Option<&str> { + match key { + key if key.eq_ignore_ascii_case(HEADER_TRACEPARENT) => self.traceparent, + key if key.eq_ignore_ascii_case(HEADER_TRACESTATE) => self.tracestate, + _ => None, + } } - let trace_id = TraceId::from_hex(trace_id).ok()?; - let span_id = SpanId::from_hex(span_id).ok()?; - let flags = u8::from_str_radix(flags, 16).ok()?; - let trace_state = tracestate - .and_then(|value| TraceState::from_str(value).ok()) - .unwrap_or_default(); - let context = SpanContext::new(trace_id, span_id, TraceFlags::new(flags), true, trace_state); - if context.is_valid() { - Some(context) - } else { - None + fn keys(&self) -> Vec<&str> { + let mut keys = Vec::with_capacity(2); + if self.traceparent.is_some() { + keys.push(HEADER_TRACEPARENT); + } + if self.tracestate.is_some() { + keys.push(HEADER_TRACESTATE); + } + keys } } diff --git a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts index 42764e5de1..8e74e93c00 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts +++ b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts @@ -18,7 +18,6 @@ export interface JsActorInvocationSpanContext { traceId: string spanId: string traceFlags: number - traceparent: string tracestate?: string } export interface JsHttpRequest { diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs b/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs index ddb3d05a3e..9bbab7c6d0 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs @@ -94,7 +94,6 @@ pub struct JsActorInvocationSpanContext { pub trace_id: String, pub span_id: String, pub trace_flags: u8, - pub traceparent: String, pub tracestate: Option, } @@ -113,7 +112,6 @@ impl From for JsActorInvocationSpanContext { trace_id: value.trace_id, span_id: value.span_id, trace_flags: value.trace_flags, - traceparent: value.traceparent, tracestate: value.tracestate, } } diff --git a/rivetkit-typescript/packages/rivetkit/src/client/outbound-telemetry.ts b/rivetkit-typescript/packages/rivetkit/src/client/outbound-telemetry.ts index f402c891fd..77da2b14a1 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/outbound-telemetry.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/outbound-telemetry.ts @@ -3,9 +3,12 @@ import { HEADER_TRACEPARENT, HEADER_TRACESTATE, } from "@/common/actor-router-consts"; -import type { ActorInvocationTraceContext } from "@/common/actor-telemetry-context"; +import type { + ActorInvocationSpanContext, + ActorInvocationTraceContext, +} from "@/common/actor-telemetry-context"; import { - type ActiveTraceHeaders, + actorInvocationTraceHeaders, readActiveRayId, readActiveTraceHeaders, } from "@/common/otel-context"; @@ -23,7 +26,7 @@ import { */ export function outboundTelemetryHeaders( invocation: ActorInvocationTraceContext | undefined, - callSpan?: ActiveTraceHeaders, + callSpan?: ActorInvocationSpanContext, ): Record { const headers: Record = {}; const rayId = invocation?.rayId ?? readActiveRayId(); @@ -31,7 +34,9 @@ export function outboundTelemetryHeaders( headers[HEADER_RIVET_RAY_ID] = rayId; } const traceHeaders = - callSpan ?? readActiveTraceHeaders() ?? invocation?.span; + actorInvocationTraceHeaders(callSpan) ?? + readActiveTraceHeaders() ?? + actorInvocationTraceHeaders(invocation?.span); if (traceHeaders) { headers[HEADER_TRACEPARENT] = traceHeaders.traceparent; if (traceHeaders.tracestate) { diff --git a/rivetkit-typescript/packages/rivetkit/src/common/actor-telemetry-context.ts b/rivetkit-typescript/packages/rivetkit/src/common/actor-telemetry-context.ts index fec4e65275..e5c36634fb 100644 --- a/rivetkit-typescript/packages/rivetkit/src/common/actor-telemetry-context.ts +++ b/rivetkit-typescript/packages/rivetkit/src/common/actor-telemetry-context.ts @@ -14,17 +14,6 @@ export interface ActorInvocationSpanContext { readonly spanId: string; /** OpenTelemetry trace flags encoded as an integer. */ readonly traceFlags: number; - /** Serialized W3C Trace Context for the Core invocation. */ - readonly traceparent: string; /** Optional vendor trace state inherited by the Core invocation. */ readonly tracestate?: string; } - -/** Formats a W3C `traceparent` header from its span identifiers. */ -export function formatTraceparent( - traceId: string, - spanId: string, - traceFlags: number, -): string { - return `00-${traceId}-${spanId}-${traceFlags.toString(16).padStart(2, "0")}`; -} diff --git a/rivetkit-typescript/packages/rivetkit/src/common/otel-context.ts b/rivetkit-typescript/packages/rivetkit/src/common/otel-context.ts index b478afcd6f..9392097874 100644 --- a/rivetkit-typescript/packages/rivetkit/src/common/otel-context.ts +++ b/rivetkit-typescript/packages/rivetkit/src/common/otel-context.ts @@ -1,15 +1,12 @@ import { - type Context, context, createTraceState, isSpanContextValid, propagation, + type SpanContext, trace, } from "@opentelemetry/api"; -import { - type ActorInvocationSpanContext, - formatTraceparent, -} from "./actor-telemetry-context"; +import type { ActorInvocationSpanContext } from "./actor-telemetry-context"; /** W3C headers derived from the active JavaScript OTel context. */ export interface ActiveTraceHeaders { @@ -21,18 +18,16 @@ export interface ActiveTraceHeaders { /** Returns the active W3C trace context, when an OTel provider has installed one. */ export function readActiveTraceHeaders(): ActiveTraceHeaders | undefined { - const spanContext = trace.getSpanContext(context.active()); - if (!spanContext || !isSpanContextValid(spanContext)) return undefined; + return traceHeaders(trace.getSpanContext(context.active())); +} - const tracestate = spanContext.traceState?.serialize(); - return { - traceparent: formatTraceparent( - spanContext.traceId, - spanContext.spanId, - spanContext.traceFlags, - ), - ...(tracestate ? { tracestate } : {}), - }; +/** Serializes the structured span context supplied by Core as W3C headers. */ +export function actorInvocationTraceHeaders( + invocation: ActorInvocationSpanContext | undefined, +): ActiveTraceHeaders | undefined { + const spanContext = invocationSpanContext(invocation); + if (!spanContext) return undefined; + return traceHeaders(spanContext); } /** W3C Baggage key that carries a ray ID through application code. */ @@ -65,9 +60,19 @@ export function runWithActorInvocationSpan( ): T { if (!invocation) return run(); - let parent: Context; + const spanContext = invocationSpanContext(invocation); + if (!spanContext) return run(); + const parent = trace.setSpanContext(context.active(), spanContext); + + return context.with(parent, run); +} + +function invocationSpanContext( + invocation: ActorInvocationSpanContext | undefined, +): SpanContext | undefined { + if (!invocation) return undefined; try { - const spanContext = { + const spanContext: SpanContext = { traceId: invocation.traceId, spanId: invocation.spanId, traceFlags: invocation.traceFlags, @@ -76,12 +81,19 @@ export function runWithActorInvocationSpan( : undefined, isRemote: false, }; - if (!isSpanContextValid(spanContext)) return run(); - parent = trace.setSpanContext(context.active(), spanContext); + return isSpanContextValid(spanContext) ? spanContext : undefined; } catch { - // Invalid telemetry must not prevent the action from running. - return run(); + return undefined; } +} - return context.with(parent, run); +function traceHeaders( + spanContext: SpanContext | undefined, +): ActiveTraceHeaders | undefined { + if (!spanContext || !isSpanContextValid(spanContext)) return undefined; + const tracestate = spanContext.traceState?.serialize(); + return { + traceparent: `00-${spanContext.traceId}-${spanContext.spanId}-${(spanContext.traceFlags & 1).toString(16).padStart(2, "0")}`, + ...(tracestate ? { tracestate } : {}), + }; } diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts index c71c5fc93d..1530b1dbf1 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts @@ -10,6 +10,7 @@ import type { } from "@rivetkit/rivetkit-napi"; import type { ActorInvocationTraceContext } from "@/common/actor-telemetry-context"; import { + actorInvocationTraceHeaders, readActiveTraceHeaders, runWithActorInvocationSpan, } from "@/common/otel-context"; @@ -267,7 +268,7 @@ export class NapiCoreRuntime implements CoreRuntime { #bindings: NativeBindings; #sql = new WeakMap(); #invocationContext: AsyncLocalStorage; - // `traceparent` of each invocation's own Core span, so an operation can + // W3C `traceparent` of each invocation's own Core span, so an operation can // tell that span apart from an application span without a native call. #invocationTraceparent = new WeakMap(); @@ -609,8 +610,12 @@ export class NapiCoreRuntime implements CoreRuntime { runWithActorInvocationContext(ctx: ActorContextHandle, run: () => T): T { const nativeCtx = asNativeActorContext(ctx); const span = nativeCtx.invocationTraceContext()?.span; - if (span) { - this.#invocationTraceparent.set(nativeCtx, span.traceparent); + const traceHeaders = actorInvocationTraceHeaders(span); + if (traceHeaders) { + this.#invocationTraceparent.set( + nativeCtx, + traceHeaders.traceparent, + ); } return this.#invocationContext.run(nativeCtx, () => runWithActorInvocationSpan(span, run), diff --git a/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts b/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts index 20bc6c4969..db49223a78 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts @@ -475,6 +475,7 @@ interface ExportedSpan { traceId: string; spanId: string; parentSpanId?: string; + traceState?: string; kind?: number; endTimeUnixNano: bigint; attributes: Record; @@ -505,6 +506,7 @@ function exportedSpans(exports: Buffer[]): ExportedSpan[] { traceId: span.traceId, spanId: span.spanId, parentSpanId: span.parentSpanId || undefined, + traceState: span.traceState, kind: span.kind, endTimeUnixNano: BigInt(span.endTimeUnixNano ?? 0), attributes: Object.fromEntries( @@ -834,6 +836,94 @@ describe.sequential("native NAPI runtime integration", () => { await waitForProcessExit(processId, 5_000); }, 120_000); + test("preserves vendor trace state across actor calls and ignores invalid trace versions", async () => { + collector = await startOtlpCollector( + await getPort({ host: "127.0.0.1" }), + ); + const traceExports = collector.spans(); + const { endpoint, poolName, child } = await startTracedRuntime( + collector.endpoint, + ); + runtime = child; + const traceId = "1234567890abcdef1234567890abcdef"; + const parentSpanId = "1234567890abcdef"; + const traceState = "vendor=opaque-value"; + const client = createIntegrationClient(endpoint, poolName); + const handle = await waitForActorReady( + () => + client.integrationActor.create( + [`trace-context-${crypto.randomUUID()}`], + { params: { userId: "integration-test" } }, + ), + 30_000, + ); + await waitForActorReady(() => handle.getCount(), 30_000); + const actorId = await handle.resolve(); + const url = new URL(await handle.getGatewayUrl()); + url.pathname = `${url.pathname.replace(/\/$/, "")}/action/getCountViaClient`; + // Every version calls the same actor, so each round takes the caller + // span that no earlier round has claimed. + const claimed = new Set(); + const callChain = (spans: ExportedSpan[]) => { + const caller = spans.find( + (span) => + span.attributes["rivet.actor.id"] === actorId && + span.attributes["rivet.action.name"] === + "getCountViaClient" && + !claimed.has(span.spanId), + ); + const hop = + caller && + spans.find( + (span) => + span.kind === OTLP_SPAN_KIND_CLIENT && + span.traceId === caller.traceId, + ); + const callee = + hop && spans.find((span) => span.parentSpanId === hop.spanId); + return { caller, hop, callee }; + }; + for (const version of ["00", "zz", "0A"]) { + const response = await fetch(url, { + method: "POST", + headers: { + "content-type": "application/json", + "x-rivet-encoding": "json", + "x-rivet-token": TOKEN, + "x-rivet-conn-params": JSON.stringify({ + userId: "integration-test", + }), + traceparent: `${version}-${traceId}-${parentSpanId}-01`, + tracestate: traceState, + }, + body: JSON.stringify({ args: [] }), + }); + expect(response.status).toBe(200); + await response.arrayBuffer(); + const { caller, hop, callee } = callChain( + await waitForSpans( + traceExports, + "caller and callee trace contexts", + (spans) => callChain(spans).callee !== undefined, + 10_000, + ), + ); + claimed.add(caller?.spanId ?? ""); + if (version === "00") { + expect(caller?.traceId).toBe(traceId); + expect(caller?.parentSpanId).toBe(parentSpanId); + for (const span of [caller, hop, callee]) + expect(span?.traceState).toBe(traceState); + } else { + expect(caller?.traceId).not.toBe(traceId); + expect(caller?.parentSpanId).toBeUndefined(); + for (const span of [caller, hop, callee]) + expect(span?.traceState || "").toBe(""); + } + } + await client.dispose(); + }, 120_000); + test("keeps overlapping invocations of one actor telemetrically isolated", async () => { collector = await startOtlpCollector( await getPort({ host: "127.0.0.1" }), From 4bd459a3db25e5803c0bd509de6cb99366fee0d1 Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Wed, 2 Sep 2026 02:40:32 +0400 Subject: [PATCH 20/21] docs(rivetkit): document telemetry architecture --- .claude/reference/testing.md | 1 + CLAUDE.md | 1 + docs-internal/engine/rivetkit-telemetry.md | 128 +++++++++++++++++++++ 3 files changed, 130 insertions(+) create mode 100644 docs-internal/engine/rivetkit-telemetry.md diff --git a/.claude/reference/testing.md b/.claude/reference/testing.md index 32a03f3380..8f34506212 100644 --- a/.claude/reference/testing.md +++ b/.claude/reference/testing.md @@ -47,6 +47,7 @@ For RivetKit runtime or parity bugs, use `rivetkit-typescript/packages/rivetkit` - Keep RivetKit test fixtures scoped to the engine-only runtime. - Prefer targeted integration tests under `rivetkit-typescript/packages/rivetkit/tests/` over shared multi-driver matrices. +- A span and its parent can arrive in different OTLP export batches, so a trace test that waits for the child and then asserts its `parentSpanId` is racy. Wait on a predicate over the whole exported span list until both are present, then assert the relationship. ## Frontend testing diff --git a/CLAUDE.md b/CLAUDE.md index a172a63fad..22abc1cc2c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -392,6 +392,7 @@ Load these only when the task touches the topic. - **[SQLite VFS parity](docs-internal/engine/sqlite-vfs.md)** — native Rust VFS ↔ WASM TypeScript VFS 1:1 parity rule, v2 storage keys, chunk layout, delete/truncate strategy. Read before touching either VFS. - **[SQLite optimizations](docs-internal/engine/SQLITE_OPTIMIZATIONS.md)** — brief tracker for SQLite cold-read, VFS, storage, preload, and benchmark optimization ideas. - **[TLS trust roots](docs-internal/engine/tls-trust-roots.md)** — rustls native+webpki union rationale, which clients use which backend. +- **[RivetKit telemetry](docs-internal/engine/rivetkit-telemetry.md)** — Core-owned invocation and SQLite spans, ray semantics, schedule trace origins in `_rivet_meta`, native OTLP export. Read before touching actor tracing, metrics, or log correlation. - **[Sleep sequence](docs-internal/engine/sleep-sequence.md)** — engine lifecycle authority, `keepAwake` vs `waitUntil` semantics, grace deadline shutdown-token abort, `can_arm_sleep_timer` vs `can_finalize_sleep` predicates. Read before touching sleep/destroy lifecycle. ### Agent procedural (`.claude/reference/`) diff --git a/docs-internal/engine/rivetkit-telemetry.md b/docs-internal/engine/rivetkit-telemetry.md new file mode 100644 index 0000000000..7b0c367de8 --- /dev/null +++ b/docs-internal/engine/rivetkit-telemetry.md @@ -0,0 +1,128 @@ +# RivetKit telemetry + +Internal reference for actor traces, invocation metrics, and log correlation. Core owns telemetry behavior. Runtime adapters activate Core's context in the host language. + +See [NAPI bridge](napi-bridge.md) for binding conventions and [Core internals](rivetkit-core-internals.md) for actor dispatch and lifecycle wiring. + +## Ownership + +Telemetry crosses two runtimes but keeps one trace: + +```text +client + headers: ray ID + W3C trace context + → rivetkit-core + invocation spans + metrics + native operation spans + → NAPI + activates the invocation context in TypeScript + → application spans +``` + +- `rivetkit-core::telemetry` owns spans and completion +- `ActorMetrics` owns invocation metrics and bounded labels +- NAPI only translates context between Core and TypeScript +- Rust and TypeScript export through their own OpenTelemetry SDKs + +Core spans use the `rivetkit::telemetry` tracing target. Log layers exclude this target, and the export layer excludes unrelated diagnostic spans. + +## Telemetry surfaces + +| Work | Span | Kind | Relationship | +| --- | --- | --- | --- | +| Action | `{actor}/{action}` | `server` | Child of incoming context | +| Schedule | `{actor}/{action}` | `internal` | New trace linked to its origin | +| Raw HTTP | `{actor}/onRequest` | `server` | Child of incoming context | +| Queue send | `{actor}/queue.send` | `producer` | Child of incoming context | +| Queue receive | `{actor}/queue.receive` | `consumer` | Linked to the send origin | +| Actor call | `{callee}/{action}` | `client` | Child of application or invocation span | +| SQLite | `rivet.sqlite.{operation}` | `internal` | Child of application or invocation span | + +Core records these attributes: + +| Scope | Attributes | +| --- | --- | +| Actor | `rivet.actor.id`, `rivet.actor.name`, `rivet.actor.key` | +| Correlation | `rivet.ray.id` | +| Invocation | `rivet.invocation.type`, `otel.status_code`, `error.type` | +| Action | `rivet.action.name` | +| HTTP | `http.request.method`, `http.response.status_code` | +| Queue | `rivet.queue.name` | +| SQLite | `rivet.operation.system`, `rivet.operation.name` | + +Raw HTTP spans use `onRequest`, never the request path. Handler errors use their `group.code` as `error.type`. A 5xx response uses the status code. Abandoned SQLite and actor-call tracking uses `actor.operation_abandoned` to represent an unknown outcome. + +`ActorMetrics` replaces undeclared invocation action and queue names with `_OTHER`. Outbound actor-call spans retain the target actor and action names. + +## Metrics and logs + +Core records: + +- `rivetkit_actor_invocations_total` +- `rivetkit_actor_invocation_duration_seconds` + +Both use actor name, action name, invocation type, and result labels. Invocation types are `action`, `scheduled`, `request`, and `queue_send`. Duration uses `MICRO_BUCKETS` for work below 5 ms. + +Actor loggers include actor identity and ray ID. A valid span also adds `trace_id` and `span_id`. Do not retain an invocation logger for unrelated work. + +## Context propagation + +```text +incoming headers + x-rivet-ray-id ───────────────────────────────┐ + traceparent + tracestate ──→ invocation span ├─→ actor call / HTTP / queue + │ +application span ──────────────────────────────┘ preferred parent + +schedule or queue send ── stores origin ── later execution links to origin +``` + +Core accepts correlation headers on actions, raw HTTP requests, and queue sends. + +- Ray IDs match `[A-Za-z0-9_-]` and contain 1–128 characters +- RivetKit propagates ray IDs but does not create them +- Invalid W3C context starts a root span without rejecting the request +- Explicit HTTP trace headers override generated headers as one pair +- Per-call context overrides static client telemetry headers +- Application spans take precedence over the invocation span as outbound parents + +The Engine gateway supplies its guard ray ID when the caller sends none. External clients read `rivet.ray.id` from OpenTelemetry baggage. The Rust client uses `ClientConfig::ray_id` only as a fallback. + +Rust owns ray ID validation in `rivetkit-client-protocol::ray_id`. Core and the Rust client use `TraceContextPropagator`. TypeScript handles context in `common/otel-context.ts`. + +Caller trace context provides correlation, not identity or authorization. Strip incoming correlation headers at an untrusted boundary when callers must not select these values. + +## Invocation lifetime + +```text +dispatch + → start span and timer + → run handler with context + → send reply and record metrics + → wait for tracked waitUntil work + → end span +``` + +`ActorInvocation` completes once. Rejected dispatches record an error. Dropped replies use `actor.dropped_reply`. `c.keepAwake` remains part of the handler, while `waitUntil` may extend the span beyond the reply. + +Schedules and queue messages store their ray ID and W3C origin beside the owning record. Each schedule fire starts a new linked trace. Each queue receipt links to the send origin. Writes and deletes update the record and origin in one batch. + +## Native export + +The `native-runtime` feature enables Core's exporter. Hosts attach `telemetry::export::layer()` and call `flush_best_effort()` during shutdown. + +- An OTLP endpoint enables export +- Protocol selection supports `grpc`, `http/protobuf`, and `http/json` +- Export uses a bounded background queue and never fails actor work +- NAPI forwards OpenTelemetry SDK warnings to Pino +- The NAPI binding must provide `setTelemetryLogSink` + +## Data policy + +Record actor identity, invocation type, HTTP method and status, correlation IDs, operation names, and error identity. Do not record arguments, results, connection parameters, SQL text or bindings, actor state, arbitrary headers, or raw error messages. + +## Gaps + +- WebSocket handlers, lifecycle hooks, connection callbacks, KV, and actor-state operations have no dedicated spans +- WebSocket action messages and inspector actions do not inherit caller context +- Actor creation ray IDs do not reach the actor runtime +- Wasm does not export host spans or expose Core invocation context to TypeScript From a49d62abd9ecaa2d02b562c01ec284ac9ff715ca Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Thu, 10 Sep 2026 21:28:02 +0400 Subject: [PATCH 21/21] fix(rivetkit): serialize telemetry exporter initialization --- .../rivetkit-core/src/telemetry/export.rs | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/rivetkit-rust/packages/rivetkit-core/src/telemetry/export.rs b/rivetkit-rust/packages/rivetkit-core/src/telemetry/export.rs index 35c0051d5a..80d926e67e 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/telemetry/export.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/telemetry/export.rs @@ -6,7 +6,6 @@ //! Core never installs a subscriber itself; which log layers surround the span //! layer is the host's decision. -use std::sync::OnceLock; use std::time::Duration; use anyhow::{Context, Result}; @@ -15,6 +14,7 @@ use opentelemetry::trace::TracerProvider as _; use opentelemetry_otlp::{Protocol, SpanExporter, WithExportConfig as _}; use opentelemetry_sdk::Resource; use opentelemetry_sdk::trace::{SdkTracer, SdkTracerProvider}; +use parking_lot::Mutex; use tracing_subscriber::registry::LookupSpan; use tracing_subscriber::{EnvFilter, Layer}; @@ -23,7 +23,7 @@ use tracing_subscriber::{EnvFilter, Layer}; /// process open. const FLUSH_TIMEOUT: Duration = Duration::from_secs(6); -static PROVIDER: OnceLock = OnceLock::new(); +static PROVIDER: Mutex> = Mutex::new(None); /// Builds the span layer when standard OTel environment variables opt in, or /// nothing when they do not. The layer only sees the runtime's own spans, so a @@ -51,7 +51,8 @@ fn initialize_if_configured() -> Result> { if !export_is_configured() { return Ok(None); } - if let Some(provider) = PROVIDER.get() { + let mut stored_provider = PROVIDER.lock(); + if let Some(provider) = stored_provider.as_ref() { return Ok(Some(provider.tracer("rivetkit"))); } @@ -75,10 +76,7 @@ fn initialize_if_configured() -> Result> { .with_batch_exporter(exporter) .build(); let tracer = provider.tracer("rivetkit"); - PROVIDER - .set(provider) - .ok() - .context("tracer provider already initialized")?; + *stored_provider = Some(provider); Ok(Some(tracer)) } @@ -123,7 +121,8 @@ fn export_is_configured() -> bool { /// [`FLUSH_TIMEOUT`]. Export failures are logged and never returned, because a /// telemetry problem must not turn a clean shutdown into a failed one. pub async fn flush_best_effort() { - let Some(provider) = PROVIDER.get().cloned() else { + let provider = PROVIDER.lock().clone(); + let Some(provider) = provider else { return; }; let flush = tokio::task::spawn_blocking(move || provider.force_flush());