From 97ecac4045907941c7b1c1b213b24fd25923ba98 Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Tue, 1 Sep 2026 20:29:44 +0400 Subject: [PATCH 01/15] 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 | 124 +++++++++++++++ .../rivetkit-core/src/telemetry/export.rs | 141 ++++++++++++++++++ .../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 | 92 ++++++++++++ .../rivetkit/tests/fixtures/otlp-collector.ts | 33 ++++ 15 files changed, 460 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 d1d903ce6c..8354855ce4 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", ] @@ -6312,6 +6316,9 @@ dependencies = [ "include_dir", "js-sys", "nix 0.30.1", + "opentelemetry", + "opentelemetry-otlp", + "opentelemetry_sdk", "parking_lot", "portpicker", "rand 0.8.5", @@ -6341,6 +6348,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 4aee22664c..215b7c6d6d 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 2676b4e9cd..c4da1e831e 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 8788280c2e..7afd28bb41 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..3facf015a3 --- /dev/null +++ b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs @@ -0,0 +1,124 @@ +//! 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); + context.is_valid().then_some(context) +} 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..3bdf33b08d --- /dev/null +++ b/rivetkit-rust/packages/rivetkit-core/src/telemetry/export.rs @@ -0,0 +1,141 @@ +//! 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 behaviour 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"))); + } + + // gRPC and HTTP are different builders, so the transport is chosen here + // rather than by passing a protocol into one of them. + 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() + // `service.version` belongs to the application and comes from + // `OTEL_RESOURCE_ATTRIBUTES` like the rest of its resource. The runtime + // version is recorded under its own key so a trace still says which + // RivetKit produced it. + .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(_))) => tracing::warn!("OpenTelemetry trace flush failed"), + Ok(Err(_)) => tracing::warn!("OpenTelemetry trace flush task failed"), + Err(_) => tracing::warn!("OpenTelemetry trace flush timed out"), + } +} diff --git a/rivetkit-rust/packages/rivetkit-core/tests/task.rs b/rivetkit-rust/packages/rivetkit-core/tests/task.rs index 587776925e..1b32c2c730 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/task.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/task.rs @@ -1929,6 +1929,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, }) @@ -2032,6 +2033,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, }) @@ -3734,6 +3736,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, }) @@ -3778,6 +3781,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, }) @@ -4529,6 +4533,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..f6964f3f25 --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit-napi/src/telemetry.rs @@ -0,0 +1,92 @@ +//! 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) level: &'static str, + 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("level", ctx.value.level)?; + 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 { + level: event.metadata().level().as_str(), + 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 4592199ed170de58986275f064cb0af69b3ec896 Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Tue, 1 Sep 2026 20:42:55 +0400 Subject: [PATCH 02/15] feat(rivetkit): trace sqlite operations --- .../errors/actor.operation_abandoned.json | 5 + .../rivetkit-core/src/actor/context.rs | 51 +++- .../src/actor/lifecycle_hooks.rs | 17 +- .../rivetkit-core/src/actor/sqlite/mod.rs | 84 ++++++- .../rivetkit-core/src/actor/sqlite/tx.rs | 30 ++- .../packages/rivetkit-core/src/actor/task.rs | 4 +- .../packages/rivetkit-core/src/error.rs | 9 + .../packages/rivetkit-core/src/lib.rs | 3 + .../packages/rivetkit-core/src/telemetry.rs | 225 +++++++++++++++--- .../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/src/registry/napi-runtime.ts | 54 +++-- .../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 +- 20 files changed, 482 insertions(+), 82 deletions(-) create mode 100644 rivetkit-rust/engine/artifacts/errors/actor.operation_abandoned.json diff --git a/rivetkit-rust/engine/artifacts/errors/actor.operation_abandoned.json b/rivetkit-rust/engine/artifacts/errors/actor.operation_abandoned.json new file mode 100644 index 0000000000..ec1e8234b8 --- /dev/null +++ b/rivetkit-rust/engine/artifacts/errors/actor.operation_abandoned.json @@ -0,0 +1,5 @@ +{ + "code": "operation_abandoned", + "group": "actor", + "message": "Operation tracking ended before a result was recorded." +} \ No newline at end of file diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs index fde8e5e7e7..7612be30fa 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 } @@ -872,6 +909,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); } @@ -1298,7 +1341,7 @@ impl ActorContext { } pub(crate) fn from_weak(weak: &Weak) -> Option { - weak.upgrade().map(Self) + weak.upgrade().map(|inner| Self(inner, None)) } #[doc(hidden)] diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/lifecycle_hooks.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/lifecycle_hooks.rs index 8ef4822b2e..ef6df02bed 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/lifecycle_hooks.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/lifecycle_hooks.rs @@ -7,9 +7,21 @@ use crate::actor::messages::ActorEvent; pub struct Reply { tx: Option>>, + invocation_telemetry: Option, } impl Reply { + #[doc(hidden)] + pub fn with_invocation_telemetry(mut self, telemetry: crate::ActorInvocationTelemetry) -> Self { + self.invocation_telemetry = Some(telemetry); + self + } + + #[doc(hidden)] + pub fn invocation_telemetry(&self) -> Option { + self.invocation_telemetry.clone() + } + pub fn send(mut self, result: Result) { if let Some(tx) = self.tx.take() { let _ = tx.send(result); @@ -35,7 +47,10 @@ impl std::fmt::Debug for Reply { impl From>> for Reply { fn from(tx: oneshot::Sender>) -> Self { - Self { tx: Some(tx) } + Self { + tx: Some(tx), + invocation_telemetry: None, + } } } 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 717a971c23..b0aeb3142c 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..2df8f9a2a3 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,7 +926,8 @@ impl ActorTask { args, conn: Some(conn), scheduled_fire: None, - reply: Reply::from(tracked_reply_tx), + reply: Reply::from(tracked_reply_tx) + .with_invocation_telemetry(invocation_telemetry), }, ) { Ok(()) => { diff --git a/rivetkit-rust/packages/rivetkit-core/src/error.rs b/rivetkit-rust/packages/rivetkit-core/src/error.rs index 195d7ff837..5b40ab87a0 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/error.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/error.rs @@ -221,6 +221,15 @@ pub enum ActorRuntime { #[error("missing_input", "Actor input is missing.")] MissingInput, + /// Telemetry for an operation was dropped before its result was recorded. + /// The underlying work may still have completed, so this says nothing + /// about whether a write landed or a remote call ran. + #[error( + "operation_abandoned", + "Operation tracking ended before a result was recorded." + )] + OperationAbandoned, + #[error( "invalid_operation", "Actor operation is invalid.", diff --git a/rivetkit-rust/packages/rivetkit-core/src/lib.rs b/rivetkit-rust/packages/rivetkit-core/src/lib.rs index c4da1e831e..6d21fdefc8 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 3facf015a3..ddd196c76f 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,86 @@ 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, +} + +/// 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 +126,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"); + let error = crate::error::ActorRuntime::OperationAbandoned.build(); + record_outcome(&span, Some(&error)); + } +} + +/// 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-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..75579f89c9 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs @@ -398,6 +398,7 @@ pub(crate) async fn dispatch_event( scheduled_fire, reply, } => { + let invocation_telemetry = reply.invocation_telemetry(); tracing::info!( actor_id = %ctx.inner().actor_id(), action_name = %name, @@ -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/src/registry/napi-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts index 3d589d5242..fc7b8e60fa 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,39 @@ 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; + } + + // Only the actor-owned handle is cached, because it is the one closed on + // sleep. Handles resolved inside an invocation carry that invocation's + // telemetry and share the same underlying database, so they are created + // on demand and dropped with the invocation. #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 +575,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(); } @@ -870,13 +897,9 @@ export class NapiCoreRuntime implements CoreRuntime { } async actorSqlClose(ctx: ActorContextHandle): Promise { - const nativeCtx = asNativeActorContext(ctx); - const database = this.#sql.get(nativeCtx); - if (!database) { - return; - } - - this.#sql.delete(nativeCtx); + const ownerCtx = asNativeActorContext(ctx); + const database = this.#sql.get(ownerCtx) ?? ownerCtx.sql(); + this.#sql.delete(ownerCtx); await database.close(); } @@ -1173,9 +1196,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 cee301a9e2..8e308f0a2b 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts @@ -5286,21 +5286,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 a27cc2e5e13876d4027b9f01bbca6f1a4680c272 Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Tue, 1 Sep 2026 20:58:41 +0400 Subject: [PATCH 03/15] feat(rivetkit): propagate actor trace context --- pnpm-lock.yaml | 3 + .../rivetkit-core/src/actor/context.rs | 6 + .../packages/rivetkit-core/src/lib.rs | 4 +- .../rivetkit-core/src/registry/http.rs | 24 ++- .../packages/rivetkit-core/src/telemetry.rs | 137 +++++++++++++----- .../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 | 26 +++- .../packages/rivetkit/src/registry/native.ts | 16 +- .../packages/rivetkit/src/registry/runtime.ts | 19 +++ .../rivetkit/src/registry/wasm-runtime.ts | 7 + 18 files changed, 377 insertions(+), 66 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 7612be30fa..890fa5b9a6 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() } diff --git a/rivetkit-rust/packages/rivetkit-core/src/lib.rs b/rivetkit-rust/packages/rivetkit-core/src/lib.rs index 6d21fdefc8..9d56a3ac4c 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..ef9e1d8317 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs @@ -8,6 +8,7 @@ use ::http; const HEADER_RIVET_ACTOR: &str = "x-rivet-actor"; const HEADER_RIVET_ACTOR_GENERATION: &str = "x-rivet-actor-generation"; const HEADER_RIVET_ACTOR_KEY: &str = "x-rivet-actor-key"; +const HEADER_RIVETKIT_RAY_ID: &str = "x-rivetkit-ray-id"; struct RequestCancellationGuard { token: Option, @@ -256,10 +257,7 @@ impl RegistryDispatcher { 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)), + invocation_ray_id(request.headers()), request .headers() .get("traceparent") @@ -452,6 +450,24 @@ impl RegistryDispatcher { } } +/// 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) +} + enum RegistryHttpRoute { Framework(FrameworkHttpRoute), UserRawRequest, diff --git a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs index ddd196c76f..2b4491d315 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs @@ -10,6 +10,7 @@ use std::sync::atomic::{AtomicBool, Ordering}; use opentelemetry::trace::{ SpanContext, SpanId, TraceContextExt as _, TraceFlags, TraceId, TraceState, }; +use parking_lot::Mutex; use tracing_opentelemetry::OpenTelemetrySpanExt as _; use crate::ActorContext; @@ -55,16 +56,40 @@ 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. Only the span slot is mutable: whichever of the +/// finish and drop paths runs first takes it, which both records the terminal +/// status once and drops the span, and dropping the span is what exports it. +/// `finished` marks the invocation closed even when tracing is off and there +/// is no span to take. #[derive(Debug)] struct InvocationInner { - span: Option, + ray_id: String, + span: Mutex>, finished: AtomicBool, identity: Arc, } +/// Active actor invocation fields exposed to foreign-runtime adapters. +#[doc(hidden)] +#[derive(Clone, Debug)] +pub struct ActorInvocationTraceContext { + pub ray_id: String, + /// 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 @@ -126,33 +151,35 @@ impl ActionInvocationSpan { action_name: &str, incoming: IncomingInvocationContext, ) -> Self { + let ray_id = incoming + .ray_id + .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); 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 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, + ); + 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), + telemetry: ActorInvocationTelemetry::new(ray_id, span, identity), } } @@ -173,26 +200,61 @@ impl Drop for ActionInvocationSpan { impl ActorInvocationTelemetry { fn new( + ray_id: String, span: Option, identity: Arc, ) -> Self { Self(Arc::new(InvocationInner { - span, + ray_id, + span: Mutex::new(span), finished: AtomicBool::new(false), identity, })) } + /// 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: (!tracestate.is_empty()).then_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, rivet.actor.id = %self.0.identity.actor_id, rivet.actor.name = %self.0.identity.actor_name, rivet.actor.key = %self.0.identity.actor_key, @@ -206,7 +268,7 @@ impl ActorInvocationTelemetry { let Some(span) = self.take_span() else { return; }; - record_outcome(span, error); + record_outcome(&span, error); } fn finish_dropped(&self) { @@ -217,13 +279,20 @@ impl ActorInvocationTelemetry { span.record("error.type", "actor.dropped_reply"); } + /// Borrows the invocation while it is still open. A finished invocation + /// yields nothing, so late SQLite work and retained handles cannot attach + /// to a span that has already recorded its status. + fn active(&self) -> Option<&InvocationInner> { + (!self.0.finished.load(Ordering::Acquire)).then_some(&*self.0) + } + /// 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> { + fn take_span(&self) -> Option { if self.0.finished.swap(true, Ordering::AcqRel) { return None; } - self.0.span.as_ref() + self.0.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..34d760d58e 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..fc0229efe1 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: String, + 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 02f5e76bbc..42a3ac62ba 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 5e84d7ae39..7e54b105b1 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, @@ -81,6 +85,7 @@ export class ActorHandleRaw { #resolvedActorId?: string; #resolvingActorId?: Promise; #queueSendMutex = new AsyncMutex(); + #currentActorInvocation?: CurrentActorInvocation; /** * Do not call this directly. @@ -97,6 +102,7 @@ export class ActorHandleRaw { encoding: Encoding, actorResolutionState: ActorResolutionState, gatewayOptions: ActorGatewayOptions = {}, + currentActorInvocation?: CurrentActorInvocation, ) { this.#client = client; this.#driver = driver; @@ -105,6 +111,7 @@ export class ActorHandleRaw { this.#gatewayOptions = gatewayOptions; this.#params = params; this.#getParams = getParams; + this.#currentActorInvocation = currentActorInvocation; } async #resolveConnectionParams(): Promise { @@ -312,6 +319,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, @@ -322,16 +347,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 33262b0096..623a4d12d1 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-owned clients propagate its trace and ray. */ + 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; } /** @@ -397,6 +406,7 @@ export class ClientRaw { this.#encodingKind, actorQuery, this.#gatewayOptions, + this.#currentActorInvocation, ); } @@ -453,9 +463,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..9c0d770e74 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_RIVETKIT_RAY_ID = "x-rivetkit-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..5ae12caf2a --- /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 { + /** Rivet request correlation identifier for the current invocation. */ + 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..54fa5cd2b8 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_RIVETKIT_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_RIVETKIT_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 fc7b8e60fa..3068abc50f 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, @@ -576,7 +578,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 { @@ -897,9 +915,9 @@ export class NapiCoreRuntime implements CoreRuntime { } async actorSqlClose(ctx: ActorContextHandle): Promise { - const ownerCtx = asNativeActorContext(ctx); - const database = this.#sql.get(ownerCtx) ?? ownerCtx.sql(); - this.#sql.delete(ownerCtx); + const nativeCtx = asNativeActorContext(ctx); + const database = this.#sql.get(nativeCtx) ?? nativeCtx.sql(); + this.#sql.delete(nativeCtx); await database.close(); } diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts index 8e308f0a2b..92daa77879 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts @@ -3969,12 +3969,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 = @@ -4018,7 +4024,7 @@ export function buildNativeFactory( new ActorContextHandleAdapter( runtime, ctx, - createClient, + () => createClient(ctx), schemaConfig, databaseProvider, request, @@ -4037,7 +4043,7 @@ export function buildNativeFactory( runtime, ctx, conn, - createClient, + () => createClient(ctx), schemaConfig, databaseProvider, request, @@ -5347,7 +5353,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..889c381a03 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts @@ -2,6 +2,7 @@ import type { SqliteNativeMetrics, SqliteProfilingOptions, } from "@/common/database/config"; +import type { ActorInvocationTraceContext } from "@/common/actor-telemetry-context"; import { stringifyError } from "@/common/utils"; import type { RegistryConfig } from "./config"; import { logger } from "./log"; @@ -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..fee9501673 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts @@ -4,6 +4,7 @@ import type { WasmRuntimeConfig, WasmRuntimeInitInput, } from "./config"; +import type { ActorInvocationTraceContext } from "@/common/actor-telemetry-context"; import type { ActorContextHandle, ActorFactoryHandle, @@ -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 2881fec0cd49ab69fb03692e9adcefa692097c92 Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Tue, 1 Sep 2026 21:18:42 +0400 Subject: [PATCH 04/15] 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 | 57 ++++++++- 4 files changed, 208 insertions(+), 15 deletions(-) diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs index 890fa5b9a6..be69760e6c 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()); @@ -1763,6 +1766,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(), @@ -1776,6 +1781,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"), @@ -1785,6 +1791,7 @@ impl ActorContext { ); } Err(error) => { + invocation_status = InvocationStatus::Dropped; dispatch_error = Some(error.into()); tracing::error!( error = ?dispatch_error.as_ref().expect("just assigned"), @@ -1795,6 +1802,7 @@ impl ActorContext { } }, Err(error) => { + action_ran = false; dispatch_error = Some(error); tracing::error!( error = ?dispatch_error.as_ref().expect("just assigned"), @@ -1804,6 +1812,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 2df8f9a2a3..aea216328f 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 1b32c2c730..5377fcb083 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; @@ -1892,7 +1893,11 @@ pub(crate) mod moved_tests { .lock() .expect("action log lock poisoned") .push(conn.as_ref().map(|conn| conn.id().to_owned())); - reply.send(Ok(name.into_bytes())); + if name == "failed-action" { + reply.send(Err(anyhow!("expected action failure"))); + } else { + reply.send(Ok(name.into_bytes())); + } } ActorEvent::BeginSleep => {} ActorEvent::FinalizeSleep { reply } | ActorEvent::Destroy { reply } => { @@ -1941,6 +1946,21 @@ pub(crate) mod moved_tests { .expect("client action should succeed"), b"client-action".to_vec(), ); + let (failed_reply_tx, failed_reply_rx) = oneshot::channel(); + task.handle_dispatch(DispatchCommand::Action { + name: "failed-action".to_owned(), + args: Vec::new(), + incoming: crate::telemetry::IncomingInvocationContext::default(), + conn: ConnHandle::new("conn-failed", Vec::new(), Vec::new(), false), + reply: failed_reply_tx, + }) + .await; + assert!( + failed_reply_rx + .await + .expect("failed action reply should send") + .is_err() + ); task.ctx .at(0, "alarm-action", &[]) @@ -1951,7 +1971,7 @@ pub(crate) mod moved_tests { .await .expect("scheduled actions should drain"); for _ in 0..50 { - if seen_conns.lock().expect("action log lock poisoned").len() >= 2 { + if seen_conns.lock().expect("action log lock poisoned").len() >= 3 { break; } sleep(Duration::from_millis(10)).await; @@ -1959,8 +1979,39 @@ pub(crate) mod moved_tests { assert_eq!( seen_conns.lock().expect("action log lock poisoned").clone(), - vec![Some("conn-client".to_owned()), None], + vec![ + Some("conn-client".to_owned()), + Some("conn-failed".to_owned()), + None, + ], ); + let mut rendered_metrics = String::new(); + for _ in 0..50 { + rendered_metrics = String::from_utf8( + crate::metrics_endpoint::render_prometheus_metrics() + .expect("render invocation metrics") + .body, + ) + .expect("prometheus metrics should be utf-8"); + if rendered_metrics.contains( + "actor_invocations_total{action_name=\"_OTHER\",actor_name=\"task-action\",invocation_type=\"scheduled\",result=\"ok\"} 1", + ) { + break; + } + sleep(Duration::from_millis(10)).await; + } + assert!(rendered_metrics.contains( + "actor_invocations_total{action_name=\"_OTHER\",actor_name=\"task-action\",invocation_type=\"action\",result=\"ok\"} 1", + )); + assert!(rendered_metrics.contains( + "actor_invocations_total{action_name=\"_OTHER\",actor_name=\"task-action\",invocation_type=\"scheduled\",result=\"ok\"} 1", + )); + assert!(rendered_metrics.contains( + "actor_invocations_total{action_name=\"_OTHER\",actor_name=\"task-action\",invocation_type=\"action\",result=\"error\"} 1", + )); + assert!(rendered_metrics.contains( + "actor_invocation_duration_seconds_count{action_name=\"_OTHER\",actor_name=\"task-action\",invocation_type=\"action\",result=\"ok\"} 1", + )); task.handle_stop(ShutdownKind::Destroy) .await From 530995f6d217509d2d997098dce10dc4874eee29 Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Wed, 2 Sep 2026 02:58:45 +0400 Subject: [PATCH 05/15] 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 92daa77879..a8929a7b70 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, @@ -101,6 +102,7 @@ import { validateQueueComplete, } from "./native-validation"; import { RunHandlerCoordinator } from "./run-handler-coordinator"; +import type { ActorInvocationTraceContext } from "@/common/actor-telemetry-context"; import type { ActorContextHandle, ActorFactoryHandle, @@ -2694,6 +2696,7 @@ export class ActorContextHandleAdapter { #db?: unknown; #dispatchCancelToken?: CancellationTokenHandle; #kv?: NativeKvAdapter; + #log?: ActorLogger; #queue?: NativeQueueAdapter; #request?: Request; #schedule?: NativeScheduleAdapter; @@ -2930,8 +2933,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: 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 5da77cb20b1eda1c9b43b3ef9c1b9d1f49a77cd8 Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Wed, 2 Sep 2026 00:56:54 +0400 Subject: [PATCH 06/15] feat(rivetkit): pass outbound trace context from client calls --- .../rivetkit/src/client/actor-handle.ts | 17 ++++++----- .../src/common/actor-telemetry-context.ts | 9 ++++++ .../rivetkit/src/common/otel-context.ts | 29 ++++++++++++++++++- 3 files changed, 47 insertions(+), 8 deletions(-) diff --git a/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts b/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts index 7e54b105b1..7c26c663f4 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts @@ -57,6 +57,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 { readActiveTraceHeaders } from "@/common/otel-context"; import { createQueueSender, type QueueSendNoWaitOptions, @@ -328,13 +329,15 @@ export class ActorHandleRaw { } 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; - } + } + // An application span active in this JavaScript context wins, + // then the calling actor's own Core invocation span. + const traceHeaders = + readActiveTraceHeaders() ?? invocation?.span; + if (traceHeaders) { + headers[HEADER_TRACEPARENT] = traceHeaders.traceparent; + if (traceHeaders.tracestate) { + headers[HEADER_TRACESTATE] = traceHeaders.tracestate; } } const output = await sendHttpRequest< 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 5ae12caf2a..32781f7d65 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..3b09d1e088 100644 --- a/rivetkit-typescript/packages/rivetkit/src/common/otel-context.ts +++ b/rivetkit-typescript/packages/rivetkit/src/common/otel-context.ts @@ -5,7 +5,34 @@ import { isSpanContextValid, 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 } : {}), + }; +} /** * Runs `run` with the Core invocation span as the active OpenTelemetry span, From 7137858c207c0b8b83987db17d1a602e97026ee5 Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Wed, 2 Sep 2026 09:40:49 +0400 Subject: [PATCH 07/15] feat(rivetkit): trace db and schedule calls under the current action --- .../rivetkit/src/registry/napi-runtime.ts | 43 ++++++++++--------- .../rivetkit/tests/runtime-parity.test.ts | 8 ++++ 2 files changed, 30 insertions(+), 21 deletions(-) diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts index 3068abc50f..f2fda407f5 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts @@ -579,22 +579,19 @@ 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; 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 { @@ -1043,7 +1040,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)); } @@ -1054,13 +1051,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( @@ -1068,12 +1065,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( @@ -1085,7 +1083,7 @@ export class NapiCoreRuntime implements CoreRuntime { args: RuntimeBytes, maxHistory: number | undefined, ) { - await asNativeActorContext(ctx) + await this.#actorContextForOperation(ctx) .schedule() .cronSet( name, @@ -1105,7 +1103,7 @@ export class NapiCoreRuntime implements CoreRuntime { args: RuntimeBytes, maxHistory: number | undefined, ) { - await asNativeActorContext(ctx) + await this.#actorContextForOperation(ctx) .schedule() .cronEvery( name, @@ -1120,20 +1118,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( @@ -1141,7 +1142,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 cc99ba94c78b5a99d4cfd8651c7b2c719721204b Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Wed, 2 Sep 2026 09:50:04 +0400 Subject: [PATCH 08/15] feat(rivetkit-core): trace scheduled invocations --- .../rivetkit-core/src/actor/context.rs | 36 +++-- .../rivetkit-core/src/actor/metrics.rs | 34 +++-- .../packages/rivetkit-core/src/actor/task.rs | 21 +-- .../packages/rivetkit-core/src/telemetry.rs | 124 +++++++++++++----- 4 files changed, 140 insertions(+), 75 deletions(-) diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs index be69760e6c..417e1fcde3 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}; @@ -1761,27 +1761,26 @@ 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), - reply: Reply::from(reply_tx), + reply: Reply::from(reply_tx).with_invocation_telemetry(invocation_telemetry), }, "scheduled_action", ) { 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"), @@ -1790,9 +1789,13 @@ impl ActorContext { "scheduled event execution failed" ); } - Err(error) => { - invocation_status = InvocationStatus::Dropped; - dispatch_error = Some(error.into()); + Err(_) => { + // The receiver is gone, so report the canonical dropped + // reply instead of the raw channel error. This one value + // reaches the invocation span, the invocation metric, and + // the persisted schedule history, so all three agree on + // `actor.dropped_reply`. + dispatch_error = Some(ActorLifecycleError::DroppedReply.build()); tracing::error!( error = ?dispatch_error.as_ref().expect("just assigned"), event_id, @@ -1802,7 +1805,6 @@ impl ActorContext { } }, Err(error) => { - action_ran = false; dispatch_error = Some(error); tracing::error!( error = ?dispatch_error.as_ref().expect("just assigned"), @@ -1812,14 +1814,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; @@ -1837,7 +1832,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 aea216328f..a1b8f15711 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 2b4491d315..2811c1fa2d 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)] @@ -35,10 +37,14 @@ impl IncomingInvocationContext { } } -/// 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. @@ -145,24 +151,55 @@ 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 { - let ray_id = incoming - .ray_id - .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + 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 ray_id = ray_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); let identity = ctx.telemetry_identity(); + // Tracing backends group and chart by span name, so the name carries the + // actor and action, and `rivet.invocation.type` still identifies + // invocation spans for filtering. That also makes the name a cardinality + // surface, so an undeclared action name is folded to a bounded + // placeholder here and the same name is used for the span and the metric. + let action_name = ctx.metrics().label_action_name(action_name).to_owned(); 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", + 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, @@ -172,7 +209,7 @@ impl ActionInvocationSpan { error.type = tracing::field::Empty, ); span.record("rivet.ray.id", &ray_id); - if let Some(parent) = incoming.remote_parent { + if let Some(parent) = parent { span.set_parent(opentelemetry::Context::new().with_remote_span_context(parent)); } span @@ -180,6 +217,10 @@ impl ActionInvocationSpan { Self { telemetry: ActorInvocationTelemetry::new(ray_id, span, identity), + metrics: ctx.metrics().clone(), + action_name, + invocation_type, + started_at: Instant::now(), } } @@ -187,14 +228,50 @@ 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.take_active() 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 `take_active`. + 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); + } } } -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.take_active() else { + return; + }; + let error = crate::error::ActorLifecycle::DroppedReply.build(); + self.record_finished(span, InvocationStatus::Dropped, Some(&error)); } } @@ -264,21 +341,6 @@ impl ActorInvocationTelemetry { 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"); - } - /// Borrows the invocation while it is still open. A finished invocation /// yields nothing, so late SQLite work and retained handles cannot attach /// to a span that has already recorded its status. @@ -288,11 +350,11 @@ impl ActorInvocationTelemetry { /// Claims the terminal record, so the finish and drop paths cannot both /// record a status for the same invocation. - fn take_span(&self) -> Option { + fn take_active(&self) -> Option> { if self.0.finished.swap(true, Ordering::AcqRel) { return None; } - self.0.span.lock().take() + Some(self.0.span.lock().take()) } } From ad1c6b3a64b1910ef6c75989c5f9b3349f81a3d7 Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Wed, 2 Sep 2026 09:51:28 +0400 Subject: [PATCH 09/15] 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 | 243 +++++++++++++----- .../packages/rivetkit-core/src/telemetry.rs | 38 ++- .../rivetkit-core/tests/sql_efficiency.rs | 36 ++- 8 files changed, 328 insertions(+), 65 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 417e1fcde3..790868de7f 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs @@ -1763,7 +1763,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..299b06e84c 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, ScheduleTraceOrigin}; 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: ScheduleTraceOrigin, } #[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) -> ScheduleTraceOrigin { + self.invocation_telemetry() + .map(ActorInvocationTelemetry::schedule_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: ScheduleTraceOrigin, ) -> 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,80 @@ fn read_stored_schedule(row: &[ColumnValue]) -> Result { } } +fn read_due_schedule(row: &[ColumnValue]) -> Result<(StoredSchedule, ScheduleTraceOrigin)> { + 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(ScheduleTraceOrigin::default, |context| { + ScheduleTraceOrigin { + ray_id: context.ray_id, + traceparent: context.traceparent, + tracestate: context.tracestate, + } + }); + Ok((event, origin)) +} + +/// 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: ScheduleTraceOrigin, +) -> Result<()> { + let key = schedule_trace_context_key(event_id); + if origin.ray_id.is_none() && origin.traceparent.is_none() && origin.tracestate.is_none() { + 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 { + ray_id: origin.ray_id, + traceparent: origin.traceparent, + tracestate: origin.tracestate, + }, + 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 2811c1fa2d..8a5689cbe3 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs @@ -75,6 +75,13 @@ struct InvocationInner { identity: Arc, } +#[derive(Clone, Debug, Default)] +pub(crate) struct ScheduleTraceOrigin { + pub(crate) ray_id: Option, + pub(crate) traceparent: Option, + pub(crate) tracestate: Option, +} + /// Active actor invocation fields exposed to foreign-runtime adapters. #[doc(hidden)] #[derive(Clone, Debug)] @@ -163,16 +170,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: ScheduleTraceOrigin, + ) -> 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, ) } @@ -182,6 +197,7 @@ impl ActorInvocation { invocation_type: InvocationType, ray_id: Option, parent: Option, + link: Option, ) -> Self { let ray_id = ray_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); let identity = ctx.telemetry_identity(); @@ -212,6 +228,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); + } span }); @@ -321,6 +340,21 @@ impl ActorInvocationTelemetry { }) } + pub(crate) fn schedule_trace_origin(&self) -> ScheduleTraceOrigin { + self.trace_context() + .map_or_else(ScheduleTraceOrigin::default, |context| { + let (traceparent, tracestate) = match context.span { + Some(span) => (Some(span.traceparent), span.tracestate), + None => (None, None), + }; + ScheduleTraceOrigin { + ray_id: Some(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!( 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 d686f8e541ed0985cdfbe2d990a150c0e64b1dd8 Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Wed, 2 Sep 2026 09:56:53 +0400 Subject: [PATCH 10/15] 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 d16304e2058bd26b3fa2a9fb25a7bd3d80e1838c Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Thu, 3 Sep 2026 13:43:10 +0400 Subject: [PATCH 11/15] test(rivetkit): cover actor tracing end to end --- .../tests/fixtures/napi-runtime-server.ts | 22 + .../rivetkit/tests/fixtures/otlp-collector.ts | 33 +- .../tests/napi-runtime-integration.test.ts | 518 ++++++++++++++++-- 3 files changed, 533 insertions(+), 40 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..cb261e2b38 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,43 @@ 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), + }); + // A failed statement records its error identity as group.code, never the message. + 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 +789,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, @@ -630,4 +856,222 @@ 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); + + // Both calls run against the same actor, so `sameActorInstance` is true + // for both and AsyncLocalStorage is the only thing keeping them apart. + 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); + + // Each invocation owns a distinct ray and trace, and the failing one + // must not mark the invocation running beside it. + 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", + ); + + // Every SQLite span belongs to exactly one probe and carries that + // probe's ray, not the ray of the invocation running beside it. + 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"], + ); + } + + // Logs written from inside each invocation carry that invocation's ray. + 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 () => { + // The collector accepts every export and then stalls for longer than the + // whole test, so the exporter's queue fills instead of failing fast. + collector = await startOtlpCollector( + await getPort({ host: "127.0.0.1" }), + { + responseDelayMs: 120_000, + }, + ); + // A queue this small saturates within a handful of actions, so the test + // reaches the drop path rather than merely filling a buffer. + 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, + ); + + // Each increment emits an invocation span plus several SQLite spans, so + // this run produces far more spans than the queue can hold while the + // collector is stalled. + 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; + + // A blocking exporter would stall each action behind the collector's + // 120s delay. Twelve actions finishing well inside one delay window is + // what proves the queue drops instead of applying backpressure. + expect(elapsed).toBeLessThan(60_000); + + expect(await waitForActorReady(() => handle.getCount(), 30_000)).toBe( + 12, + ); + + await client.dispose(); + }, 180_000); }); From 705d251acb23eb778c0ee23533e0380e97b311c5 Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Wed, 2 Sep 2026 02:40:32 +0400 Subject: [PATCH 12/15] feat(rivetkit): send opentelemetry sdk warnings to the actor logger --- .../packages/rivetkit-napi/index.d.ts | 6 +++++ .../packages/rivetkit-napi/index.js | 3 ++- .../packages/rivetkit-napi/src/lib.rs | 12 ++++++++- .../packages/rivetkit-napi/src/registry.rs | 8 ++++++ .../packages/rivetkit-napi/src/telemetry.rs | 20 +++++++++------ .../rivetkit/src/registry/napi-runtime.ts | 25 +++++++++++++++++++ .../tests/napi-runtime-integration.test.ts | 18 ++++++++++++- .../rivetkit/tests/runtime-parity.test.ts | 3 +++ 8 files changed, 85 insertions(+), 10 deletions(-) diff --git a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts index 34d760d58e..dc1f28e7db 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts +++ b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts @@ -270,6 +270,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..7f1ce911c1 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/lib.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/lib.rs @@ -125,6 +125,13 @@ pub(crate) fn init_tracing(log_level: Option<&str>) { tracing_subscriber::registry() .with(otel_layer) + // The SDK reports dropped spans and export failures on its own + // target. Forward those to the JavaScript logger so they appear + // with the rest of the actor's logs. + .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 +157,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 f6964f3f25..51c26ef364 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; @@ -22,10 +23,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 @@ -41,7 +46,7 @@ pub(crate) mod sdk_log_bridge { Ok(vec![object.into_unknown()]) })?; tsfn.unref(&env)?; - let _ = SINK.set(tsfn); + *SINK.write() = Some(tsfn); Ok(()) } @@ -74,7 +79,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 f2fda407f5..b85a7b9389 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 { runWithActorInvocationSpan } from "@/common/otel-context"; +import { logger } from "./log"; import type { ActorContextHandle, ActorFactoryHandle, @@ -296,6 +297,30 @@ export class NapiCoreRuntime implements CoreRuntime { } createRegistry(): RegistryHandle { + // The OpenTelemetry SDK reports dropped spans and export failures + // through Rust tracing, which writes to stdout in a different format + // from the actor logs. Route them into the same logger instead. + // + // The most recent registry to start owns the sink. A Node worker + // thread that registered it can exit, and its callback would then drop + // every event, so a later registry on any thread takes over. + // + // The addon is published as its own platform package, so it can be + // older than this one and not export the sink at all. Forwarding those + // diagnostics is best effort, and losing them must not stop the actor + // from starting. + if (this.#bindings.setTelemetryLogSink) { + this.#bindings.setTelemetryLogSink((event) => { + logger().warn( + { otelEvent: event.name }, + event.message || event.name, + ); + }); + } else { + logger().warn( + "native addon has no telemetry log sink; OpenTelemetry SDK warnings will not reach the actor logs", + ); + } 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 cb261e2b38..1bad5bc615 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, @@ -1032,6 +1032,11 @@ describe.sequential("native NAPI runtime integration", () => { { OTEL_BSP_MAX_QUEUE_SIZE: "8", OTEL_BSP_MAX_EXPORT_BATCH_SIZE: "4", + // Rust's own log layers admit `opentelemetry_sdk` at warn by + // default, so they print this same message and the assertion + // below would pass with the bridge dead. Silencing that one + // target leaves the JS sink as the only way it reaches stdout. + RUST_LOG: "warn,opentelemetry_sdk=off", }, ); runtime = child; @@ -1068,6 +1073,17 @@ describe.sequential("native NAPI runtime integration", () => { // what proves the queue drops instead of applying backpressure. expect(elapsed).toBeLessThan(60_000); + // The SDK's dropped-span warning is bridged into the actor logger, so + // it appears in the runtime's own log output rather than only in Rust. + 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..b5a06e52b1 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts @@ -314,6 +314,9 @@ function fakeNapiBindings(scenario: ParityScenario) { NapiActorFactory: FakeActorFactory, CancellationToken: FakeCancellationToken, ActorContext: class {}, + // The native module exports this, and `createRegistry` calls it on + // every registry, so the fake has to carry it too. + setTelemetryLogSink: () => {}, }; } From 48bf14ab42512646c5b3249ee28f5f6348251a1f Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Wed, 2 Sep 2026 02:40:32 +0400 Subject: [PATCH 13/15] docs(rivetkit): document telemetry architecture --- .claude/reference/testing.md | 1 + CLAUDE.md | 1 + docs-internal/engine/rivetkit-telemetry.md | 133 +++++++++++++++++++++ 3 files changed, 135 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..dedc9978f0 --- /dev/null +++ b/docs-internal/engine/rivetkit-telemetry.md @@ -0,0 +1,133 @@ +# RivetKit telemetry + +Architecture and operational invariants for RivetKit traces, invocation metrics, and log correlation. Core owns invocation telemetry; adapters only bridge the current Core context to their host runtime. Pair with `napi-bridge.md` for the adapter boundary and `rivetkit-core-internals.md` for surrounding lifecycle context. + +## What it produces + +Three span shapes, all on the `rivetkit::telemetry` tracing target: + +| Span | When | Kind | Parent | +| --- | --- | --- | --- | +| `{actor}/{action}` | an action runs | `server` | remote `traceparent` if valid, else root | +| `{actor}/{action}` | a schedule or cron fires | `internal` | fresh root, plus one link to the defining invocation | +| `rivet.sqlite.{operation}` | any `c.db` call | `internal` | the current invocation | + +Invocation spans carry `rivet.invocation.type`, `rivet.actor.id`, `rivet.actor.name`, `rivet.actor.key`, `rivet.action.name`, `rivet.ray.id`, `otel.status_code`, and `error.type` on failure. SQLite spans carry the same actor identity plus `rivet.operation.system` and `rivet.operation.name`. + +Two metrics, labelled by actor name, action name, invocation type, and status: + +- `rivetkit_actor_invocations_total` +- `rivetkit_actor_invocation_duration_seconds` + +Pino binds actor ID, name, key, and ray ID on every actor log line. Sampled work also binds `trace_id` and `span_id`, in snake case because that is what OTel log-correlation tooling looks for. + +### Why action names are bounded + +Action names arrive from the caller on the URL path and are never validated against the registry before dispatch. An undeclared name becomes `_OTHER`, following the OpenTelemetry semantic convention for unknown caller-supplied values (`http.request.method` uses the same fallback). This applies to the metric label **and** to the span name, because a backend that derives metrics from span names turns an unbounded name into a new series exactly as a label would. `ActorMetrics::label_action_name` is the single source for both. + +### Why SQLite spans parent to the invocation, not the surrounding application span + +Core creates these spans in Rust and cannot see the JavaScript span stack. A `c.db` call made while an application span is active therefore appears as a sibling of that span under the invocation, not as its child. This is the visible consequence of the two-pipeline model below, and it was accepted deliberately: reading the JS context from Rust on every query would mean a context lookup across the NAPI boundary per call. + +## Turning it on + +Set these on the process that loads RivetKit, which is the actor runner. On Rivet Cloud that is the same machine and the same environment block as `RIVET_ENDPOINT`; the dashboard only displays that value for you to copy, so both go wherever your platform keeps environment variables. + +```sh +OTEL_SERVICE_NAME=internal-api +OTEL_EXPORTER_OTLP_TRACES_ENDPOINT=http://collector:4318/v1/traces +OTEL_EXPORTER_OTLP_TRACES_PROTOCOL=http/protobuf +``` + +The endpoint is illustrative; use the team's receiver and check which format and authentication it expects. `grpc` targets the collector's 4317 port rather than 4318. `OTEL_EXPORTER_OTLP_HEADERS` and `OTEL_EXPORTER_OTLP_TRACES_HEADERS` are both honored, which is what makes hosted backends work. + +Application spans come from the application's own OpenTelemetry SDK pointed at the same collector. The two pipelines share only trace context, and that is sufficient: Core invocation and SQLite spans arrive from the Rust exporter, application spans arrive from the JavaScript SDK, and the collector joins them on trace ID. No custom JavaScript exporter or span bridge is involved. Bridging the two into one pipeline was considered and deferred; it buys a single export path and richer nesting, at the cost of RivetKit owning span export on behalf of the application. + +### Two ways this fails silently + +Both look like broken RivetKit tracing rather than application misconfiguration, so they belong in any user-facing doc. + +- **Configuring `NodeSDK` in code and nothing else.** That configures the JavaScript pipeline only. The Rust side reads the environment, so the result is application spans with no RivetKit spans and no error anywhere. Configure both through the same environment variables; `NodeSDK` reads them too. +- **No context manager registered.** `runWithActorInvocationSpan` activates the Core span with `context.with(...)` from `@opentelemetry/api`, which delegates to the globally registered context manager. With none registered the API uses a no-op manager that invokes the callback and stores nothing, so `context.active()` inside the action returns the root context and every application span becomes its own trace root. `NodeSDK.start()` and `NodeTracerProvider.register()` register one; a bare `BasicTracerProvider` plus `trace.setGlobalTracerProvider(...)` does not. + +A hand-built JavaScript provider also needs its own resource. `OTEL_SERVICE_NAME` configures the Rust exporter only, so a provider without `service.name` reports `unknown_service:node` and files the two halves of one trace under different services. + +### Why RivetKit does not register a context manager itself + +Doing so would take ownership of global context away from the application, and it would not generalize beyond Node. + +## Overhead characteristics + +Absolute figures belong in a benchmark artifact, not here; they change with every build and host. Three properties of the design do not. + +- **Sampling does not remove the latency cost.** The span is constructed in `tracing` before `tracing-opentelemetry` runs the sampler, so a sampled-out invocation pays nearly the same request latency as a fully traced one. Operators reaching for `OTEL_TRACES_SAMPLER` to cut latency will not get it; sampling reduces export volume, not construction. +- **Export cost lands in CPU, not latency.** The batch processor ships spans from a background thread, so enabling export raises worker CPU per invocation while leaving request latency close to the sampled-out case. +- **Spans are dropped, not queued indefinitely.** The batch processor holds `OTEL_BSP_MAX_QUEUE_SIZE` spans, 2,048 by default, and discards on overflow. Raising it absorbs bursts but only delays overflow when span production sustainably exceeds export throughput, and costs memory. Drops surface as `BatchSpanProcessor.SpanDroppingStarted` in the actor logs through the SDK log bridge, which is the only reason they are visible at all. + +`rivetkit_actor_invocation_duration_seconds` uses `MICRO_BUCKETS`. Invocations land in the hundreds of microseconds, which the Prometheus default buckets, starting at 5 ms, collapse into a single bucket. + +## Core lifecycle + +- `ActorInvocation` owns one `rivet.actor.invoke` span, its timer, its metric labels, and exactly-once completion. Both action and schedule dispatch create it before enqueueing. +- Rejected enqueue attempts count as `status=error`. Unfinished invocations finish as `actor.dropped_reply`. Raw error messages are never recorded, only bounded `group.code` identity. +- A scheduled fire whose reply channel closes records that same `actor.dropped_reply` identity in `_rivet_schedule_history`, so the invocation span, the invocation metric, and what `cronHistory()` returns all agree. +- The span kind is derived from the invocation type rather than passed alongside it, so the two cannot disagree. +- Span and metric completion share `ActorInvocation`, so the action and schedule paths cannot drift when tracing is disabled or sampled out. + +Export layers enable the `rivetkit::telemetry` target and log layers filter it off. That keeps diagnostic Rust spans out of application traces and keeps OTel spans from duplicating log context. + +## Context propagation + +Core accepts `x-rivetkit-ray-id`, `traceparent`, and `tracestate`. Rays must be 1 to 128 characters from `[A-Za-z0-9_-]`; absent or invalid rays become UUIDs. Invalid W3C context fails closed to a root span and never rejects an action. + +Outbound actor calls resolve context at send time, preferring the active `@opentelemetry/api` span and falling back to the current Core invocation for an actor-owned client. An application span is the more specific parent when one is active, so the callee nests under the work that actually issued the call rather than under the whole invocation. Those headers replace static client telemetry headers, so configuration cannot pin stale context. + +- `@opentelemetry/api` is a hard dependency at `^1.1.0`, the lowest minor that exports every symbol used, and it is inert without a provider. It is needed both to read outbound application spans and to activate inbound Core spans; making it optional would silently separate traces. +- Retained databases, clients, and schedule handles resolve the current invocation from `AsyncLocalStorage`, but only when it belongs to the same Core actor generation. Otherwise they use their creation context. Pointer identity through `Arc::ptr_eq`, rather than actor ID equality, is what isolates overlapping calls and restarted generations. +- KV and queue operations do not resolve the invocation, because Core attaches telemetry only to SQLite and schedule work. +- `c.log` stays creation-scoped because a Pino child is immutable correlation metadata, not an operation resolved on every write. Do not retain an action logger for later work. + +### Why accepting caller trace context at an untrusted edge is acceptable + +This follows the W3C Trace Context model: a caller can name any trace and can set the sampled flag, so an actor's traces are only as trustworthy as the callers that can reach it. Trace context carries no identity and no authorization, and every value derived from it is bounded before it reaches a span or a label. An operator who does not trust their callers should strip `traceparent`, `tracestate`, and `x-rivetkit-ray-id` at their own edge. + +## Scheduled work + +- Creating or redefining a one-shot, interval, or cron schedule captures the current ray, `traceparent`, and `tracestate`. Recurring re-registration refreshes that context even when cadence is unchanged; preserving cadence must not preserve the identity of an older definer. +- Schedule context is a versioned BARE value in the existing `_rivet_meta` table, keyed by schedule ID, following the `run_wake_at` precedent. It is written in the same batch as the schedule row and removed in the same batch as the row delete, guarded so a delete that misses the row leaves the context alone. Only the due-schedule query reads it. +- Malformed values are logged and ignored, and their W3C fields fail closed through the same parser used for actions. + +### Why this does not bump the internal schema version + +Older runtimes ignore the metadata, so a schedule redefined by an older runtime keeps the previous definer's context until a newer runtime redefines it again. Telemetry degrades to best effort across that sequence without affecting schedule behavior. + +### Why a scheduled fire starts a new trace + +It keeps the defining invocation's ray, because a ray means the causal path of the work. Its OTel span is still a fresh trace root linked to the precise defining span: rays give broad correlation even when spans are absent or sampled out, while the span link records the standard durable asynchronous relationship. A cron running for a year would otherwise be one trace with millions of spans. Actor calls made by the fire propagate that same causal ray and the fire's new W3C context. + +## Native export + +Core owns the Rust `SdkTracerProvider`, in `rivetkit_core::telemetry::export` behind the `native-runtime` feature. A host adds `export::layer()` to its own subscriber and calls `export::flush_best_effort()` when it stops serving; NAPI does exactly that and adds nothing of its own, so the Rust crate and any future host get the same export by making the same two calls. Native export is enabled by standard `OTEL_EXPORTER_OTLP_TRACES_ENDPOINT` or `OTEL_EXPORTER_OTLP_ENDPOINT`, unless standard SDK disable or exporter controls turn it off. Sampling, resources, service name, and batching all use standard OTel environment variables. There is no RivetKit sampler, rate limiter, registry field, or exporter switch. Shutdown performs a bounded best-effort flush, and export failures cannot fail actor work. + +- **Protocol selection is read here, not left to the exporter.** `opentelemetry-otlp` takes its default from a compile-time constant chosen by the enabled cargo features, and neither of its builders reads `OTEL_EXPORTER_OTLP_TRACES_PROTOCOL` or `OTEL_EXPORTER_OTLP_PROTOCOL`. Enabling `http-json` would otherwise pin every deployment to JSON whatever an operator set. `configured_protocol` reads both variables and accepts all three OTLP values: `grpc`, `http/protobuf`, and `http/json`. Anything else errors naming the value. The default is `http/protobuf`, which the specification lists as a usual SDK default and which most collectors and hosted backends expect. gRPC matters because the Engine's own exporter speaks it, so a deployment running one collector on 4317 does not need a second receiver for RivetKit. +- **The SDK log bridge forwards the SDK's own diagnostics to the JavaScript logger.** A `tracing` layer filtered to `opentelemetry_sdk=warn` hands events to a `ThreadsafeFunction`, so dropped-span warnings appear in Pino alongside everything else the actor logs instead of on stdout in a different format. `internal-logs` must stay enabled on both `opentelemetry` and `opentelemetry_sdk`, or `otel_warn!` compiles to nothing and the bridge goes silent with no error. +- **The sink is replaceable, and the latest registry owns it.** A Node worker thread that installed the callback can exit, after which that callback drops every event, so `setTelemetryLogSink` overwrites rather than sets once. +- **That threadsafe function is unreferenced once created.** A referenced one counts as live work on the Node event loop, and `createRegistry` registers the sink on every registry, so any process that registered it would never exit on its own. Unreferencing keeps warnings flowing without making the sink a reason to keep running. +- **The addon export is called defensively.** `@rivetkit/rivetkit-napi` ships as its own per-platform package and can be older than the JavaScript calling it, so a missing sink logs a warning rather than failing registry construction. A logging hook must not be able to stop an actor from starting. + +## Data policy + +Allowed: actor ID, name, and key; declared action name; invocation type; bounded status; ray, trace, and span IDs; and error identity as `group.code`. + +Never recorded: action arguments or results, connection parameters, SQL text or bindings, actor state, arbitrary headers, or raw error messages. + +Event-category opt-in is a separate application-facing design. It is not inferred from sampling or from registry configuration. + +## Not covered + +- Dedicated spans for raw fetch and WebSocket handlers, lifecycle hooks, connection callbacks, KV, or actor-state operations. +- Connection and inspector actions inheriting the initiating connection trace. +- A public automatic-tracing category opt-in API. +- Engine-owned ray stamping, public invocation tokens, runtime-specific span lifecycles, or custom sampling controls. +- Wasm host span export. The Wasm adapter runs actions without invocation context and reports no trace context; only rays are retained. +- Integration entry points. There is no public API for attaching an application span to an outbound action, so the Effect bridge and `rivetkit/unstable/otel` are not part of this work. From b5bd3df63d5e0c9c93344017561be9915a58346a Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Mon, 7 Sep 2026 02:25:08 +0400 Subject: [PATCH 14/15] feat(rivetkit): span the call out to another actor --- docs-internal/engine/rivetkit-telemetry.md | 5 +- pnpm-lock.yaml | 76 ++++++++++ .../rivetkit-core/src/actor/context.rs | 18 +++ .../packages/rivetkit-core/src/lib.rs | 1 + .../packages/rivetkit-core/src/telemetry.rs | 136 +++++++++++++++--- .../packages/rivetkit-napi/index.d.ts | 26 ++++ .../packages/rivetkit-napi/index.js | 3 +- .../rivetkit-napi/src/actor_context.rs | 63 +++++++- .../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 | 66 +++++++-- .../packages/rivetkit/src/client/client.ts | 9 +- .../rivetkit/src/registry/napi-runtime.ts | 19 +++ .../packages/rivetkit/src/registry/native.ts | 13 ++ .../packages/rivetkit/src/registry/runtime.ts | 40 +++++- .../rivetkit/src/registry/wasm-runtime.ts | 12 ++ .../tests/fixtures/napi-runtime-server.ts | 29 ++++ .../tests/napi-runtime-integration.test.ts | 74 +++++++++- 19 files changed, 566 insertions(+), 45 deletions(-) diff --git a/docs-internal/engine/rivetkit-telemetry.md b/docs-internal/engine/rivetkit-telemetry.md index dedc9978f0..5a9b5fd26d 100644 --- a/docs-internal/engine/rivetkit-telemetry.md +++ b/docs-internal/engine/rivetkit-telemetry.md @@ -4,12 +4,13 @@ Architecture and operational invariants for RivetKit traces, invocation metrics, ## What it produces -Three span shapes, all on the `rivetkit::telemetry` tracing target: +Every span RivetKit produces is on the `rivetkit::telemetry` tracing target: | Span | When | Kind | Parent | | --- | --- | --- | --- | | `{actor}/{action}` | an action runs | `server` | remote `traceparent` if valid, else root | | `{actor}/{action}` | a schedule or cron fires | `internal` | fresh root, plus one link to the defining invocation | +| `{callee}/{action}` | an actor-owned client calls another actor | `client` | the active application span if one is present, else the current invocation | | `rivet.sqlite.{operation}` | any `c.db` call | `internal` | the current invocation | Invocation spans carry `rivet.invocation.type`, `rivet.actor.id`, `rivet.actor.name`, `rivet.actor.key`, `rivet.action.name`, `rivet.ray.id`, `otel.status_code`, and `error.type` on failure. SQLite spans carry the same actor identity plus `rivet.operation.system` and `rivet.operation.name`. @@ -80,7 +81,7 @@ Export layers enable the `rivetkit::telemetry` target and log layers filter it o Core accepts `x-rivetkit-ray-id`, `traceparent`, and `tracestate`. Rays must be 1 to 128 characters from `[A-Za-z0-9_-]`; absent or invalid rays become UUIDs. Invalid W3C context fails closed to a root span and never rejects an action. -Outbound actor calls resolve context at send time, preferring the active `@opentelemetry/api` span and falling back to the current Core invocation for an actor-owned client. An application span is the more specific parent when one is active, so the callee nests under the work that actually issued the call rather than under the whole invocation. Those headers replace static client telemetry headers, so configuration cannot pin stale context. +Outbound actor calls from an actor-owned client open a `rivet.actor.call` span in Core and send that span's context, so the callee nests under the call and the time spent reaching a cold or busy actor belongs to it. The client passes the active `@opentelemetry/api` span's `traceparent` into `beginOutboundCall`, and Core parents the call span there when one is present, because Core cannot see the JavaScript span stack itself. An application span is the more specific parent when one is active, so the callee ends up under the work that actually issued the call rather than beside it. Without a call span, the client falls back to the active application span and then to the current Core invocation. Those headers replace static client telemetry headers, so configuration cannot pin stale context. - `@opentelemetry/api` is a hard dependency at `^1.1.0`, the lowest minor that exports every symbol used, and it is inert without a provider. It is needed both to read outbound application spans and to activate inbound Core spans; making it optional would silently separate traces. - Retained databases, clients, and schedule handles resolve the current invocation from `AsyncLocalStorage`, but only when it belongs to the same Core actor generation. Otherwise they use their creation context. Pointer identity through `Arc::ptr_eq`, rather than actor ID equality, is what isolates overlapping calls and restarted generations. 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 790868de7f..df3113fe8e 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs @@ -263,6 +263,24 @@ 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 the invocation is not sampled. + /// + /// `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, + application_traceparent: Option<&str>, + ) -> Option { + self.1 + .as_ref()? + .start_outbound_call(actor_name, action_name, application_traceparent) + } + /// 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 9d56a3ac4c..70f6f8940b 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/lib.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/lib.rs @@ -21,6 +21,7 @@ pub mod telemetry; #[doc(hidden)] pub use telemetry::{ ActorInvocationSpanContext, ActorInvocationTelemetry, ActorInvocationTraceContext, + OutboundCallInvocation, }; #[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 8a5689cbe3..c15016300a 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs @@ -158,6 +158,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, @@ -312,27 +325,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(); - 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: (!tracestate.is_empty()).then_some(tracestate), - }) - }); + let span = active + .span + .lock() + .clone() + .and_then(|span| span_context_of(&span)); Some(ActorInvocationTraceContext { ray_id: active.ray_id.clone(), @@ -355,6 +352,52 @@ 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. + /// `application_traceparent` is the W3C context of the application span + /// active in the host runtime when the call is made, if any. Core cannot + /// see that span stack itself, so the caller passes it and the call span + /// parents there, which is what puts a callee under the application span + /// that issued the call rather than beside it. Invalid context falls back + /// to the invocation span. + pub(crate) fn start_outbound_call( + &self, + actor_name: &str, + action_name: &str, + application_traceparent: Option<&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", + // No `rivet.invocation.type`. This span is a call, not an + // invocation, and that attribute is what distinguishes spans where + // an actor ran code from spans where one waited on another actor. + rivet.actor.name = %actor_name, + rivet.action.name = %action_name, + rivet.ray.id = %self.0.ray_id, + otel.status_code = tracing::field::Empty, + error.type = tracing::field::Empty, + ); + if let Some(application_parent) = parse_remote_parent(application_traceparent, None) { + span.set_parent( + opentelemetry::Context::new().with_remote_span_context(application_parent), + ); + } + 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!( @@ -392,6 +435,33 @@ impl ActorInvocationTelemetry { } } +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; + }; + let error = crate::error::ActorRuntime::OperationAbandoned.build(); + record_outcome(&span, Some(&error)); + } +} + impl SqliteOperationSpan { pub(crate) fn span(&self) -> tracing::Span { self.span.as_ref().expect("sqlite span is present").clone() @@ -415,6 +485,30 @@ impl Drop for SqliteOperationSpan { } } +/// Reads a span's W3C context, or nothing when the span is not sampled and so +/// carries no valid context to propagate. +fn span_context_of(span: &tracing::Span) -> Option { + 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: (!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-typescript/packages/rivetkit-napi/index.d.ts b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts index dc1f28e7db..2cf9fbc589 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts +++ b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts @@ -327,6 +327,12 @@ export declare class ActorContext { sql(): JsNativeDatabase sameActorInstance(other: ActorContext): boolean invocationTraceContext(): JsActorInvocationTraceContext | null + /** + * Opens the span covering one call out to another actor. Returns nothing + * when this handle serves no invocation or the invocation is not sampled, + * in which case the caller sends its own context as before. + */ + beginOutboundCall(actorName: string, actionName: string, applicationTraceparent?: string | undefined | null): OutboundCall | null provisionActorRuntimeSocket(): Promise schedule(): Schedule queue(): Queue @@ -376,6 +382,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 fc0229efe1..6d8021beee 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; @@ -327,6 +328,27 @@ 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 the invocation is not sampled, + /// in which case the caller sends its own context as before. + #[napi] + pub fn begin_outbound_call( + &self, + actor_name: String, + action_name: String, + application_traceparent: Option, + ) -> Option { + self.inner + .begin_outbound_call( + &actor_name, + &action_name, + application_traceparent.as_deref(), + ) + .map(|invocation| OutboundCall { + invocation: Some(invocation), + }) + } + #[napi] pub async fn provision_actor_runtime_socket( &self, @@ -1105,3 +1127,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 93a828007f..8f22b037f1 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/actor_factory.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/actor_factory.rs @@ -976,6 +976,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 42a3ac62ba..1c86ca44fa 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 7c26c663f4..bfc2696a08 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, @@ -27,7 +28,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, @@ -87,6 +91,7 @@ export class ActorHandleRaw { #resolvingActorId?: Promise; #queueSendMutex = new AsyncMutex(); #currentActorInvocation?: CurrentActorInvocation; + #beginOutboundCall?: BeginOutboundCall; /** * Do not call this directly. @@ -104,6 +109,7 @@ export class ActorHandleRaw { actorResolutionState: ActorResolutionState, gatewayOptions: ActorGatewayOptions = {}, currentActorInvocation?: CurrentActorInvocation, + beginOutboundCall?: BeginOutboundCall, ) { this.#client = client; this.#driver = driver; @@ -113,6 +119,7 @@ export class ActorHandleRaw { this.#params = params; this.#getParams = getParams; this.#currentActorInvocation = currentActorInvocation; + this.#beginOutboundCall = beginOutboundCall; } async #resolveConnectionParams(): Promise { @@ -276,19 +283,56 @@ export class ActorHandleRaw { `Invalid action call: expected an options object { name, args }, got ${typeof opts}. Use handle.actionName(...args) for the shorthand API.`, ); } - const run = async () => (await this.#sendActionNow(opts)) as Response; - if (opts.name === "destroy") { - return await run(); + // One call from the caller is one client span, whatever it takes to + // complete it. The span opens above the lifecycle retry, so an actor + // that is replaced mid-call does not split the trace into two spans, + // and its duration answers how long the call took including retries. + // It has to open before any attempt reads trace context, so the callee + // parents to the call instead of to the invocation that made it. An + // application span active here becomes the call's parent, so the callee + // ends up under the work that issued the call. + const call = this.#beginOutboundCall?.( + this.#targetActorName(), + opts.name, + readActiveTraceHeaders()?.traceparent, + ); + const run = async () => + (await this.#sendActionAttempts(opts, call?.span)) as Response; + const send = async () => { + if (opts.name === "destroy") { + return await run(); + } + return await retryOnLifecycleBoundary(run, { signal: opts.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: opts.signal }); + #targetActorName(): string { + try { + return getActorNameFromQuery(this.#actorResolutionState); + } catch { + // A malformed query fails later with its own error. Naming a span + // must not be what surfaces it. + return "unknown"; + } } - async #sendActionNow( + async #sendActionAttempts( opts: { name: string; args: unknown[]; } & ActorActionOptions, + callSpan?: ActorInvocationSpanContext, ): Promise { const maxAttempts = this.#getDynamicQueryMaxAttempts(); let useQueryTarget = isDynamicActorQuery(this.#actorResolutionState); @@ -330,10 +374,12 @@ export class ActorHandleRaw { if (invocation) { headers[HEADER_RIVETKIT_RAY_ID] = invocation.rayId; } - // An application span active in this JavaScript context wins, - // then the calling actor's own Core invocation span. + // This call's own span wins, and it already descends from any + // application span that was active when the call opened. Without + // a call span, an application span active in this JavaScript + // context, then the calling actor's own Core invocation span. const traceHeaders = - readActiveTraceHeaders() ?? invocation?.span; + callSpan ?? readActiveTraceHeaders() ?? invocation?.span; if (traceHeaders) { headers[HEADER_TRACEPARENT] = traceHeaders.traceparent; if (traceHeaders.tracestate) { diff --git a/rivetkit-typescript/packages/rivetkit/src/client/client.ts b/rivetkit-typescript/packages/rivetkit/src/client/client.ts index 623a4d12d1..333d0b8015 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, @@ -187,6 +190,7 @@ export interface ClientRawOptions { gateway?: ActorGatewayOptions; /** Supplies the calling actor's invocation so actor-owned clients propagate its trace and ray. */ currentActorInvocation?: CurrentActorInvocation; + beginOutboundCall?: BeginOutboundCall; } export class ClientRaw { @@ -198,6 +202,7 @@ export class ClientRaw { #encodingKind: Encoding; #gatewayOptions: ActorGatewayOptions; #currentActorInvocation?: CurrentActorInvocation; + #beginOutboundCall?: BeginOutboundCall; /** * Creates an instance of Client. @@ -211,6 +216,7 @@ export class ClientRaw { this.#encodingKind = options.encoding ?? "bare"; this.#gatewayOptions = options.gateway ?? {}; this.#currentActorInvocation = options.currentActorInvocation; + this.#beginOutboundCall = options.beginOutboundCall; } /** @@ -407,6 +413,7 @@ export class ClientRaw { actorQuery, this.#gatewayOptions, this.#currentActorInvocation, + this.#beginOutboundCall, ); } diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts index b85a7b9389..3d4ab54173 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts @@ -20,6 +20,7 @@ import type { CoreRuntime, RegistryHandle, RuntimeActorConfig, + RuntimeOutboundCall, RuntimeApplicationFetch, RuntimeApplicationListenerConfig, RuntimeBytes, @@ -619,6 +620,24 @@ export class NapiCoreRuntime implements CoreRuntime { ); } + beginOutboundCall( + ctx: ActorContextHandle, + actorName: string, + actionName: string, + applicationTraceparent?: string, + ): RuntimeOutboundCall | undefined { + const call = this.#actorContextForOperation(ctx).beginOutboundCall( + actorName, + actionName, + applicationTraceparent ?? null, + ); + 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 a8929a7b70..b13495c769 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts @@ -4005,6 +4005,19 @@ export function buildNativeFactory( callNativeSync(() => runtime.actorInvocationTraceContext(ctx), ), + beginOutboundCall: ( + actorName, + actionName, + applicationTraceparent, + ) => + callNativeSync(() => + runtime.beginOutboundCall( + ctx, + actorName, + actionName, + applicationTraceparent, + ), + ), }, ); 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 889c381a03..c83cbf4250 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts @@ -2,7 +2,10 @@ import type { SqliteNativeMetrics, SqliteProfilingOptions, } from "@/common/database/config"; -import type { ActorInvocationTraceContext } from "@/common/actor-telemetry-context"; +import type { + ActorInvocationSpanContext, + ActorInvocationTraceContext, +} from "@/common/actor-telemetry-context"; import { stringifyError } from "@/common/utils"; import type { RegistryConfig } from "./config"; import { logger } from "./log"; @@ -35,6 +38,31 @@ 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 the call is not sampled. + */ + 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, on a runtime without invocation telemetry, or when the + * call is not sampled. Callers send their own context when it returns nothing. + */ +export type BeginOutboundCall = ( + actorName: string, + actionName: string, + applicationTraceparent?: string, +) => RuntimeOutboundCall | undefined; + export interface RuntimeHttpRequest { method: string; uri: string; @@ -560,6 +588,16 @@ 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, + applicationTraceparent?: 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 fee9501673..756dd52f5f 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts @@ -14,6 +14,7 @@ import type { CoreRuntime, RegistryHandle, RuntimeActorConfig, + RuntimeOutboundCall, RuntimeActorKeySegment, RuntimeApplicationListenerConfig, RuntimeBytes, @@ -550,6 +551,17 @@ export class WasmCoreRuntime implements CoreRuntime { return undefined; } + beginOutboundCall( + _ctx: ActorContextHandle, + _actorName: string, + _actionName: string, + _applicationTraceparent?: 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..97d76e8c31 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,13 @@ const repoEngineBinary = resolve( ); const endpoint = process.env.RIVETKIT_TEST_ENDPOINT ?? "http://127.0.0.1:6642"; + +// The application's own OpenTelemetry provider, the way a user would set one +// up. It registers a context manager, which is what lets a span opened inside +// an action become the parent of the actor calls made while it is active. +// Nothing exports these spans; the test reads their IDs off the action result. +new NodeTracerProvider().register(); +const applicationTracer = trace.getTracer("napi-runtime-fixture"); const connParamsSchema = z.object({ userId: z.string().min(1), }); @@ -159,6 +168,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 1bad5bc615..ff442eafeb 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, @@ -800,8 +806,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"], ); @@ -954,29 +973,72 @@ 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. + // stays inside its own trace and carries its own ray, through both the + // call span and the invocation it reaches. 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"], + ); + } } // Logs written from inside each invocation carry that invocation's ray. 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=([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)); + // A call made while an application span is active parents to that + // span, not to the invocation, so the callee sits under the work that + // issued the call. + 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 b7db93fe30d6b44adf0fa205f0402c599dcb61ab Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Mon, 7 Sep 2026 21:08:09 +0400 Subject: [PATCH 15/15] fix(rivetkit): preserve outbound trace state and validate trace versions --- docs-internal/engine/rivetkit-telemetry.md | 10 +- .../rivetkit-core/src/actor/context.rs | 10 +- .../src/actor/internal_storage/queries.rs | 1 + .../packages/rivetkit-core/src/telemetry.rs | 8 +- .../packages/rivetkit-napi/index.d.ts | 2 +- .../rivetkit-napi/src/actor_context.rs | 2 + .../rivetkit/src/client/actor-handle.ts | 4 +- .../rivetkit/src/registry/napi-runtime.ts | 2 + .../packages/rivetkit/src/registry/native.ts | 2 + .../packages/rivetkit/src/registry/runtime.ts | 2 + .../rivetkit/src/registry/wasm-runtime.ts | 1 + .../tests/napi-runtime-integration.test.ts | 107 ++++++++++++++++++ 12 files changed, 140 insertions(+), 11 deletions(-) diff --git a/docs-internal/engine/rivetkit-telemetry.md b/docs-internal/engine/rivetkit-telemetry.md index 5a9b5fd26d..3cc244503c 100644 --- a/docs-internal/engine/rivetkit-telemetry.md +++ b/docs-internal/engine/rivetkit-telemetry.md @@ -59,10 +59,10 @@ Doing so would take ownership of global context away from the application, and i ## Overhead characteristics -Absolute figures belong in a benchmark artifact, not here; they change with every build and host. Three properties of the design do not. +Absolute figures belong in a benchmark artifact; costs depend on the build, workload, and host. -- **Sampling does not remove the latency cost.** The span is constructed in `tracing` before `tracing-opentelemetry` runs the sampler, so a sampled-out invocation pays nearly the same request latency as a fully traced one. Operators reaching for `OTEL_TRACES_SAMPLER` to cut latency will not get it; sampling reduces export volume, not construction. -- **Export cost lands in CPU, not latency.** The batch processor ships spans from a background thread, so enabling export raises worker CPU per invocation while leaving request latency close to the sampled-out case. +- Sampling reduces exported spans but leaves span construction and context propagation work. Its effect on latency must be measured. +- Export runs on a background thread. It consumes CPU and can affect request latency through contention. - **Spans are dropped, not queued indefinitely.** The batch processor holds `OTEL_BSP_MAX_QUEUE_SIZE` spans, 2,048 by default, and discards on overflow. Raising it absorbs bursts but only delays overflow when span production sustainably exceeds export throughput, and costs memory. Drops surface as `BatchSpanProcessor.SpanDroppingStarted` in the actor logs through the SDK log bridge, which is the only reason they are visible at all. `rivetkit_actor_invocation_duration_seconds` uses `MICRO_BUCKETS`. Invocations land in the hundreds of microseconds, which the Prometheus default buckets, starting at 5 ms, collapse into a single bucket. @@ -81,7 +81,7 @@ Export layers enable the `rivetkit::telemetry` target and log layers filter it o Core accepts `x-rivetkit-ray-id`, `traceparent`, and `tracestate`. Rays must be 1 to 128 characters from `[A-Za-z0-9_-]`; absent or invalid rays become UUIDs. Invalid W3C context fails closed to a root span and never rejects an action. -Outbound actor calls from an actor-owned client open a `rivet.actor.call` span in Core and send that span's context, so the callee nests under the call and the time spent reaching a cold or busy actor belongs to it. The client passes the active `@opentelemetry/api` span's `traceparent` into `beginOutboundCall`, and Core parents the call span there when one is present, because Core cannot see the JavaScript span stack itself. An application span is the more specific parent when one is active, so the callee ends up under the work that actually issued the call rather than beside it. Without a call span, the client falls back to the active application span and then to the current Core invocation. Those headers replace static client telemetry headers, so configuration cannot pin stale context. +Outbound actor calls from an actor-owned client open a `rivet.actor.call` span in Core and send that span's context, so the callee nests under the call and the time spent reaching a cold or busy actor belongs to it. The client passes the active `@opentelemetry/api` span's `traceparent` and `tracestate` into `beginOutboundCall`, and Core parents the call span there when one is present, because Core cannot see the JavaScript span stack itself. An application span is the more specific parent when one is active, so the callee ends up under the work that actually issued the call rather than beside it. Without a call span, the client falls back to the active application span and then to the current Core invocation. Those headers replace static client telemetry headers, so configuration cannot pin stale context. - `@opentelemetry/api` is a hard dependency at `^1.1.0`, the lowest minor that exports every symbol used, and it is inert without a provider. It is needed both to read outbound application spans and to activate inbound Core spans; making it optional would silently separate traces. - Retained databases, clients, and schedule handles resolve the current invocation from `AsyncLocalStorage`, but only when it belongs to the same Core actor generation. Otherwise they use their creation context. Pointer identity through `Arc::ptr_eq`, rather than actor ID equality, is what isolates overlapping calls and restarted generations. @@ -108,7 +108,7 @@ It keeps the defining invocation's ray, because a ray means the causal path of t ## Native export -Core owns the Rust `SdkTracerProvider`, in `rivetkit_core::telemetry::export` behind the `native-runtime` feature. A host adds `export::layer()` to its own subscriber and calls `export::flush_best_effort()` when it stops serving; NAPI does exactly that and adds nothing of its own, so the Rust crate and any future host get the same export by making the same two calls. Native export is enabled by standard `OTEL_EXPORTER_OTLP_TRACES_ENDPOINT` or `OTEL_EXPORTER_OTLP_ENDPOINT`, unless standard SDK disable or exporter controls turn it off. Sampling, resources, service name, and batching all use standard OTel environment variables. There is no RivetKit sampler, rate limiter, registry field, or exporter switch. Shutdown performs a bounded best-effort flush, and export failures cannot fail actor work. +Core owns the Rust `SdkTracerProvider`, in `rivetkit_core::telemetry::export` behind the `native-runtime` feature. A host adds `export::layer()` to its own subscriber and calls `export::flush_best_effort()` when it stops serving; NAPI does exactly that and adds nothing of its own, so the Rust crate and any future host get the same export by making the same two calls. Native export is enabled by standard `OTEL_EXPORTER_OTLP_TRACES_ENDPOINT` or `OTEL_EXPORTER_OTLP_ENDPOINT`, unless standard SDK disable or exporter controls turn it off. Sampling, resources, service name, and batching all use standard OTel environment variables. There is no RivetKit sampler, rate limiter, registry field, or exporter switch. Batches export during normal execution. Clean shutdown attempts a bounded flush of remaining spans; a hard kill can lose queued spans. Export failures cannot fail actor work. - **Protocol selection is read here, not left to the exporter.** `opentelemetry-otlp` takes its default from a compile-time constant chosen by the enabled cargo features, and neither of its builders reads `OTEL_EXPORTER_OTLP_TRACES_PROTOCOL` or `OTEL_EXPORTER_OTLP_PROTOCOL`. Enabling `http-json` would otherwise pin every deployment to JSON whatever an operator set. `configured_protocol` reads both variables and accepts all three OTLP values: `grpc`, `http/protobuf`, and `http/json`. Anything else errors naming the value. The default is `http/protobuf`, which the specification lists as a usual SDK default and which most collectors and hosted backends expect. gRPC matters because the Engine's own exporter speaks it, so a deployment running one collector on 4317 does not need a second receiver for RivetKit. - **The SDK log bridge forwards the SDK's own diagnostics to the JavaScript logger.** A `tracing` layer filtered to `opentelemetry_sdk=warn` hands events to a `ThreadsafeFunction`, so dropped-span warnings appear in Pino alongside everything else the actor logs instead of on stdout in a different format. `internal-logs` must stay enabled on both `opentelemetry` and `opentelemetry_sdk`, or `otel_warn!` compiles to nothing and the bridge goes silent with no error. diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs index df3113fe8e..013d7bce4f 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs @@ -275,10 +275,14 @@ impl ActorContext { actor_name: &str, action_name: &str, application_traceparent: Option<&str>, + application_tracestate: Option<&str>, ) -> Option { - self.1 - .as_ref()? - .start_outbound_call(actor_name, action_name, application_traceparent) + self.1.as_ref()?.start_outbound_call( + actor_name, + action_name, + application_traceparent, + application_tracestate, + ) } /// Returns correlation for the invocation this handle serves, absent when 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..986e4092f4 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 c15016300a..7c6ea1c3e5 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs @@ -369,6 +369,7 @@ impl ActorInvocationTelemetry { actor_name: &str, action_name: &str, application_traceparent: Option<&str>, + application_tracestate: Option<&str>, ) -> Option { let invocation_span = self.active()?.span.lock().clone()?; let span = tracing::info_span!( @@ -386,7 +387,9 @@ impl ActorInvocationTelemetry { otel.status_code = tracing::field::Empty, error.type = tracing::field::Empty, ); - if let Some(application_parent) = parse_remote_parent(application_traceparent, None) { + if let Some(application_parent) = + parse_remote_parent(application_traceparent, application_tracestate) + { span.set_parent( opentelemetry::Context::new().with_remote_span_context(application_parent), ); @@ -529,6 +532,9 @@ fn parse_remote_parent(traceparent: Option<&str>, tracestate: Option<&str>) -> O let flags = fields.next()?; if fields.next().is_some() || version.len() != 2 + || !version + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) || version.eq_ignore_ascii_case("ff") || trace_id.len() != 32 || span_id.len() != 16 diff --git a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts index 2cf9fbc589..979aec0483 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts +++ b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts @@ -332,7 +332,7 @@ export declare class ActorContext { * when this handle serves no invocation or the invocation is not sampled, * in which case the caller sends its own context as before. */ - beginOutboundCall(actorName: string, actionName: string, applicationTraceparent?: string | undefined | null): OutboundCall | null + beginOutboundCall(actorName: string, actionName: string, applicationTraceparent?: string | undefined | null, applicationTracestate?: string | undefined | null): OutboundCall | 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 6d8021beee..dd45f2b880 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs @@ -337,12 +337,14 @@ impl ActorContext { actor_name: String, action_name: String, application_traceparent: Option, + application_tracestate: Option, ) -> Option { self.inner .begin_outbound_call( &actor_name, &action_name, application_traceparent.as_deref(), + application_tracestate.as_deref(), ) .map(|invocation| OutboundCall { invocation: Some(invocation), diff --git a/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts b/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts index bfc2696a08..4c1c6bbafe 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts @@ -291,10 +291,12 @@ export class ActorHandleRaw { // parents to the call instead of to the invocation that made it. An // application span active here becomes the call's parent, so the callee // ends up under the work that issued the call. + const applicationContext = readActiveTraceHeaders(); const call = this.#beginOutboundCall?.( this.#targetActorName(), opts.name, - readActiveTraceHeaders()?.traceparent, + applicationContext?.traceparent, + applicationContext?.tracestate, ); const run = async () => (await this.#sendActionAttempts(opts, call?.span)) as Response; diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts index 3d4ab54173..800faa7c13 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts @@ -625,11 +625,13 @@ export class NapiCoreRuntime implements CoreRuntime { actorName: string, actionName: string, applicationTraceparent?: string, + applicationTracestate?: string, ): RuntimeOutboundCall | undefined { const call = this.#actorContextForOperation(ctx).beginOutboundCall( actorName, actionName, applicationTraceparent ?? null, + applicationTracestate ?? null, ); if (!call) return undefined; return { diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts index b13495c769..8f8dc4df38 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts @@ -4009,6 +4009,7 @@ export function buildNativeFactory( actorName, actionName, applicationTraceparent, + applicationTracestate, ) => callNativeSync(() => runtime.beginOutboundCall( @@ -4016,6 +4017,7 @@ export function buildNativeFactory( actorName, actionName, applicationTraceparent, + applicationTracestate, ), ), }, diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts index c83cbf4250..53a1da75bf 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts @@ -61,6 +61,7 @@ export type BeginOutboundCall = ( actorName: string, actionName: string, applicationTraceparent?: string, + applicationTracestate?: string, ) => RuntimeOutboundCall | undefined; export interface RuntimeHttpRequest { @@ -597,6 +598,7 @@ export interface CoreRuntime { actorName: string, actionName: string, applicationTraceparent?: string, + applicationTracestate?: string, ): RuntimeOutboundCall | undefined; actorName(ctx: ActorContextHandle): string; actorKey(ctx: ActorContextHandle): RuntimeActorKeySegment[]; diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts index 756dd52f5f..e19864b667 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts @@ -556,6 +556,7 @@ export class WasmCoreRuntime implements CoreRuntime { _actorName: string, _actionName: string, _applicationTraceparent?: string, + _applicationTracestate?: string, ): RuntimeOutboundCall | undefined { // Wasm carries no invocation telemetry, so a call goes out untraced // rather than failing. 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 ff442eafeb..b0028c0857 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts @@ -450,6 +450,7 @@ interface ExportedSpan { traceId: string; spanId: string; parentSpanId?: string; + traceState?: string; kind?: number; attributes: Record; links: Array<{ traceId: string; spanId: string }>; @@ -475,6 +476,7 @@ function exportedSpans(exports: Buffer[]): ExportedSpan[] { traceId: span.traceId, spanId: span.spanId, parentSpanId: span.parentSpanId || undefined, + traceState: span.traceState, kind: span.kind, attributes: Object.fromEntries( (span.attributes ?? []).map((attribute) => [ @@ -876,6 +878,111 @@ 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"; + for (const version of ["00", "zz", "0A"]) { + const client = createClient({ + endpoint, + token: TOKEN, + namespace: NAMESPACE, + poolName, + disableMetadataLookup: true, + }) as any; + try { + 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`; + 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 spans = await waitForSpans( + traceExports, + "caller and callee trace contexts", + (exported) => { + const caller = exported.find( + (span) => + span.attributes["rivet.actor.id"] === actorId && + span.attributes["rivet.action.name"] === + "getCountViaClient", + ); + return ( + !!caller && + exported.some( + (span) => + span.kind === OTLP_SPAN_KIND_CLIENT && + span.traceId === caller.traceId, + ) + ); + }, + 10_000, + ); + const caller = spans.find( + (span) => + span.attributes["rivet.actor.id"] === actorId && + span.attributes["rivet.action.name"] === + "getCountViaClient", + ); + expect(caller).toBeDefined(); + const hop = spans.find( + (span) => + span.kind === OTLP_SPAN_KIND_CLIENT && + span.traceId === caller?.traceId, + ); + const callee = spans.find( + (span) => span.parentSpanId === hop?.spanId, + ); + expect(callee).toBeDefined(); + 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(""); + } + } finally { + 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" }),