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/Cargo.lock b/Cargo.lock index 5ed497b518..bf7a273262 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", ] @@ -6261,6 +6265,8 @@ dependencies = [ "bytes", "fs_extra", "futures-util", + "opentelemetry", + "opentelemetry_sdk", "parking_lot", "portpicker", "reqwest 0.12.22", @@ -6275,6 +6281,7 @@ dependencies = [ "tokio-test", "tokio-tungstenite", "tracing", + "tracing-opentelemetry", "tracing-subscriber", "tungstenite", "urlencoding", @@ -6313,6 +6320,9 @@ dependencies = [ "include_dir", "js-sys", "nix 0.30.1", + "opentelemetry", + "opentelemetry-otlp", + "opentelemetry_sdk", "parking_lot", "portpicker", "rand 0.8.5", @@ -6342,6 +6352,7 @@ dependencies = [ "tokio-util", "tower-http", "tracing", + "tracing-opentelemetry", "tracing-subscriber", "url", "uuid", diff --git a/docs-internal/engine/rivetkit-telemetry.md b/docs-internal/engine/rivetkit-telemetry.md new file mode 100644 index 0000000000..7b0c367de8 --- /dev/null +++ b/docs-internal/engine/rivetkit-telemetry.md @@ -0,0 +1,128 @@ +# RivetKit telemetry + +Internal reference for actor traces, invocation metrics, and log correlation. Core owns telemetry behavior. Runtime adapters activate Core's context in the host language. + +See [NAPI bridge](napi-bridge.md) for binding conventions and [Core internals](rivetkit-core-internals.md) for actor dispatch and lifecycle wiring. + +## Ownership + +Telemetry crosses two runtimes but keeps one trace: + +```text +client + headers: ray ID + W3C trace context + → rivetkit-core + invocation spans + metrics + native operation spans + → NAPI + activates the invocation context in TypeScript + → application spans +``` + +- `rivetkit-core::telemetry` owns spans and completion +- `ActorMetrics` owns invocation metrics and bounded labels +- NAPI only translates context between Core and TypeScript +- Rust and TypeScript export through their own OpenTelemetry SDKs + +Core spans use the `rivetkit::telemetry` tracing target. Log layers exclude this target, and the export layer excludes unrelated diagnostic spans. + +## Telemetry surfaces + +| Work | Span | Kind | Relationship | +| --- | --- | --- | --- | +| Action | `{actor}/{action}` | `server` | Child of incoming context | +| Schedule | `{actor}/{action}` | `internal` | New trace linked to its origin | +| Raw HTTP | `{actor}/onRequest` | `server` | Child of incoming context | +| Queue send | `{actor}/queue.send` | `producer` | Child of incoming context | +| Queue receive | `{actor}/queue.receive` | `consumer` | Linked to the send origin | +| Actor call | `{callee}/{action}` | `client` | Child of application or invocation span | +| SQLite | `rivet.sqlite.{operation}` | `internal` | Child of application or invocation span | + +Core records these attributes: + +| Scope | Attributes | +| --- | --- | +| Actor | `rivet.actor.id`, `rivet.actor.name`, `rivet.actor.key` | +| Correlation | `rivet.ray.id` | +| Invocation | `rivet.invocation.type`, `otel.status_code`, `error.type` | +| Action | `rivet.action.name` | +| HTTP | `http.request.method`, `http.response.status_code` | +| Queue | `rivet.queue.name` | +| SQLite | `rivet.operation.system`, `rivet.operation.name` | + +Raw HTTP spans use `onRequest`, never the request path. Handler errors use their `group.code` as `error.type`. A 5xx response uses the status code. Abandoned SQLite and actor-call tracking uses `actor.operation_abandoned` to represent an unknown outcome. + +`ActorMetrics` replaces undeclared invocation action and queue names with `_OTHER`. Outbound actor-call spans retain the target actor and action names. + +## Metrics and logs + +Core records: + +- `rivetkit_actor_invocations_total` +- `rivetkit_actor_invocation_duration_seconds` + +Both use actor name, action name, invocation type, and result labels. Invocation types are `action`, `scheduled`, `request`, and `queue_send`. Duration uses `MICRO_BUCKETS` for work below 5 ms. + +Actor loggers include actor identity and ray ID. A valid span also adds `trace_id` and `span_id`. Do not retain an invocation logger for unrelated work. + +## Context propagation + +```text +incoming headers + x-rivet-ray-id ───────────────────────────────┐ + traceparent + tracestate ──→ invocation span ├─→ actor call / HTTP / queue + │ +application span ──────────────────────────────┘ preferred parent + +schedule or queue send ── stores origin ── later execution links to origin +``` + +Core accepts correlation headers on actions, raw HTTP requests, and queue sends. + +- Ray IDs match `[A-Za-z0-9_-]` and contain 1–128 characters +- RivetKit propagates ray IDs but does not create them +- Invalid W3C context starts a root span without rejecting the request +- Explicit HTTP trace headers override generated headers as one pair +- Per-call context overrides static client telemetry headers +- Application spans take precedence over the invocation span as outbound parents + +The Engine gateway supplies its guard ray ID when the caller sends none. External clients read `rivet.ray.id` from OpenTelemetry baggage. The Rust client uses `ClientConfig::ray_id` only as a fallback. + +Rust owns ray ID validation in `rivetkit-client-protocol::ray_id`. Core and the Rust client use `TraceContextPropagator`. TypeScript handles context in `common/otel-context.ts`. + +Caller trace context provides correlation, not identity or authorization. Strip incoming correlation headers at an untrusted boundary when callers must not select these values. + +## Invocation lifetime + +```text +dispatch + → start span and timer + → run handler with context + → send reply and record metrics + → wait for tracked waitUntil work + → end span +``` + +`ActorInvocation` completes once. Rejected dispatches record an error. Dropped replies use `actor.dropped_reply`. `c.keepAwake` remains part of the handler, while `waitUntil` may extend the span beyond the reply. + +Schedules and queue messages store their ray ID and W3C origin beside the owning record. Each schedule fire starts a new linked trace. Each queue receipt links to the send origin. Writes and deletes update the record and origin in one batch. + +## Native export + +The `native-runtime` feature enables Core's exporter. Hosts attach `telemetry::export::layer()` and call `flush_best_effort()` during shutdown. + +- An OTLP endpoint enables export +- Protocol selection supports `grpc`, `http/protobuf`, and `http/json` +- Export uses a bounded background queue and never fails actor work +- NAPI forwards OpenTelemetry SDK warnings to Pino +- The NAPI binding must provide `setTelemetryLogSink` + +## Data policy + +Record actor identity, invocation type, HTTP method and status, correlation IDs, operation names, and error identity. Do not record arguments, results, connection parameters, SQL text or bindings, actor state, arbitrary headers, or raw error messages. + +## Gaps + +- WebSocket handlers, lifecycle hooks, connection callbacks, KV, and actor-state operations have no dedicated spans +- WebSocket action messages and inspector actions do not inherit caller context +- Actor creation ray IDs do not reach the actor runtime +- Wasm does not export host spans or expose Core invocation context to TypeScript diff --git a/engine/packages/guard-core/src/request_context.rs b/engine/packages/guard-core/src/request_context.rs index 3dfe3a75d4..caa4c1d857 100644 --- a/engine/packages/guard-core/src/request_context.rs +++ b/engine/packages/guard-core/src/request_context.rs @@ -2,6 +2,8 @@ use anyhow::{Context, Result}; use hyper::{Method, header::HeaderMap}; use rivet_runner_protocol as protocol; use rivet_util::Id; +use rivet_api_builder::X_RIVET_RAY_ID; +use std::collections::HashMap; use std::{ net::{IpAddr, SocketAddr}, sync::Arc, @@ -90,6 +92,14 @@ impl RequestContext { self.ray_id } + /// Adds this request's ray ID to the headers forwarded to an actor when the + /// caller did not send one, so the actor and gateway use the same ray ID. + pub fn forward_ray(&self, headers: &mut HashMap) { + headers + .entry(X_RIVET_RAY_ID.as_str().to_owned()) + .or_insert_with(|| self.ray_id.to_string()); + } + pub fn req_id(&self) -> Id { self.req_id } diff --git a/engine/packages/pegboard-gateway2/src/lib.rs b/engine/packages/pegboard-gateway2/src/lib.rs index 290ceec819..67ee28c7a9 100644 --- a/engine/packages/pegboard-gateway2/src/lib.rs +++ b/engine/packages/pegboard-gateway2/src/lib.rs @@ -127,7 +127,7 @@ impl PegboardGateway2 { let request_id = req_ctx.in_flight_request_id()?; // Extract request parts - let headers = req + let mut headers = req .headers() .iter() .filter_map(|(name, value)| { @@ -137,6 +137,7 @@ impl PegboardGateway2 { .map(|value_str| (name.to_string(), value_str.to_string())) }) .collect::>(); + req_ctx.forward_ray(&mut headers); // NOTE: Size constraints have already been applied by guard let body_bytes = req @@ -360,6 +361,7 @@ impl PegboardGateway2 { request_headers.insert(name.to_string(), value_str.to_string()); } } + req_ctx.forward_ray(&mut request_headers); let (mut stopped_sub, _) = tokio::try_join!( ctx.subscribe::(("actor_id", self.actor_id)), diff --git a/engine/packages/pegboard-gateway3/src/http_stream/handler.rs b/engine/packages/pegboard-gateway3/src/http_stream/handler.rs index 6bcaa7c014..d9387675c9 100644 --- a/engine/packages/pegboard-gateway3/src/http_stream/handler.rs +++ b/engine/packages/pegboard-gateway3/src/http_stream/handler.rs @@ -6,6 +6,8 @@ use std::{ time::Duration, }; +use std::collections::HashMap; + use anyhow::{Result, anyhow}; use bytes::Bytes; use gas::prelude::*; @@ -71,7 +73,7 @@ impl PegboardGateway3 { req_ctx.request_body_is_end_stream(), ); let request_id = req_ctx.in_flight_request_id()?; - let headers = req_ctx + let mut headers: HashMap = req_ctx .headers() .iter() .filter_map(|(name, value)| { @@ -81,6 +83,7 @@ impl PegboardGateway3 { .map(|value| (name.to_string(), value.to_owned())) }) .collect(); + req_ctx.forward_ray(&mut headers); let (mut stopped_sub, _) = tokio::try_join!( ctx.subscribe::(("actor_id", self.actor_id)), pegboard::utils::ensure_ns_metrics_exporter_for_namespace(ctx, self.namespace_id), diff --git a/engine/packages/pegboard-gateway3/src/lib.rs b/engine/packages/pegboard-gateway3/src/lib.rs index 5ff19be290..2bfa284f1c 100644 --- a/engine/packages/pegboard-gateway3/src/lib.rs +++ b/engine/packages/pegboard-gateway3/src/lib.rs @@ -134,6 +134,7 @@ impl PegboardGateway3 { request_headers.insert(name.to_string(), value_str.to_string()); } } + req_ctx.forward_ray(&mut request_headers); let (mut stopped_sub, _) = tokio::try_join!( ctx.subscribe::(("actor_id", self.actor_id)), diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index e8d9fcae90..5f4dda7d3c 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) @@ -3614,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 @@ -7488,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'} @@ -23671,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/actor-persist/src/versioned.rs b/rivetkit-rust/packages/actor-persist/src/versioned.rs index d64bd2d9e0..a2e2c0c462 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,103 @@ 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() + } +} + +/// Trace origin of a queue message, stored beside its `_rivet_queue` row. It +/// has its own versioned type so its persisted meaning cannot be conflated with +/// `ScheduleTraceContext` even though both currently encode the same fields. +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +pub struct QueueTraceContextData { + pub ray_id: Option, + pub traceparent: Option, + pub tracestate: Option, +} + +pub enum QueueTraceContext { + V1(QueueTraceContextData), +} + +impl OwnedVersionedData for QueueTraceContext { + type Latest = QueueTraceContextData; + + fn wrap_latest(latest: Self::Latest) -> Self { + Self::V1(latest) + } + + fn unwrap_latest(self) -> Result { + match self { + Self::V1(data) => Ok(data), + } + } + + fn deserialize_version(payload: &[u8], version: u16) -> Result { + match version { + 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), + _ => bail!("invalid queue trace context version: {version}"), + } + } + + fn serialize_version(self, version: u16) -> Result> { + match (self, version) { + (Self::V1(data), 1) => serde_bare::to_vec(&data).map_err(Into::into), + (_, version) => bail!("unexpected queue trace context version: {version}"), + } + } + + fn deserialize_converters() -> Vec Result> { + Vec:: Result>::new() + } + + fn serialize_converters() -> Vec Result> { + Vec:: Result>::new() + } +} + impl OwnedVersionedData for RunWakeAt { type Latest = Option; diff --git a/rivetkit-rust/packages/client-protocol/src/lib.rs b/rivetkit-rust/packages/client-protocol/src/lib.rs index 70f35bb01b..bef5d9a66c 100644 --- a/rivetkit-rust/packages/client-protocol/src/lib.rs +++ b/rivetkit-rust/packages/client-protocol/src/lib.rs @@ -1,4 +1,5 @@ pub mod generated; +pub mod ray_id; pub mod versioned; // Re-export latest. diff --git a/rivetkit-rust/packages/client-protocol/src/ray_id.rs b/rivetkit-rust/packages/client-protocol/src/ray_id.rs new file mode 100644 index 0000000000..b47a59eab1 --- /dev/null +++ b/rivetkit-rust/packages/client-protocol/src/ray_id.rs @@ -0,0 +1,44 @@ +use std::fmt; + +pub const HEADER_RIVET_RAY_ID: &str = "x-rivet-ray-id"; +pub const RAY_BAGGAGE_KEY: &str = "rivet.ray.id"; + +const MAX_LEN: usize = 128; + +/// A bounded correlation token carried between RivetKit requests. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct RayId(String); + +impl RayId { + pub fn parse(value: impl Into) -> Result { + let value = value.into(); + if value.is_empty() + || value.len() > MAX_LEN + || !value + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-')) + { + return Err(InvalidRayId); + } + Ok(Self(value)) + } + + pub fn as_str(&self) -> &str { + &self.0 + } + + pub fn into_string(self) -> String { + self.0 + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct InvalidRayId; + +impl fmt::Display for InvalidRayId { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("ray ID must be 1 to 128 characters of [A-Za-z0-9_-]") + } +} + +impl std::error::Error for InvalidRayId {} diff --git a/rivetkit-rust/packages/client/Cargo.toml b/rivetkit-rust/packages/client/Cargo.toml index edcadab797..33f399d298 100644 --- a/rivetkit-rust/packages/client/Cargo.toml +++ b/rivetkit-rust/packages/client/Cargo.toml @@ -13,6 +13,9 @@ anyhow = "1.0" base64 = "0.22.1" bytes = { workspace = true } futures-util = "0.3.31" +opentelemetry = { version = "0.28", default-features = false, features = ["trace"] } +opentelemetry-http.workspace = true +opentelemetry_sdk = { version = "0.28", default-features = false, features = ["trace"] } parking_lot.workspace = true reqwest = { version = "0.12.12", default-features = false, features = ["json", "charset", "http2", "macos-system-configuration", "rustls-tls-native-roots", "rustls-tls-webpki-roots"] } rivetkit-client-protocol.workspace = true @@ -24,6 +27,7 @@ serde_json = "1.0" tokio = { version = "1", features = ["full"] } tokio-tungstenite = { version = "0.26.1", features = ["rustls-tls-native-roots", "rustls-tls-webpki-roots", "handshake"] } tracing = "0.1.41" +tracing-opentelemetry = { version = "0.29", default-features = false } tungstenite = "0.26.2" urlencoding = "2.1.3" vbare = "0.0.4" diff --git a/rivetkit-rust/packages/client/src/client.rs b/rivetkit-rust/packages/client/src/client.rs index 79a8a622b4..7578be94fb 100644 --- a/rivetkit-rust/packages/client/src/client.rs +++ b/rivetkit-rust/packages/client/src/client.rs @@ -46,6 +46,9 @@ pub struct ClientConfig { pub encoding: EncodingKind, pub transport: TransportKind, pub headers: Option>, + /// Ray sent on every actor request made outside an active baggage + /// context, so work this client causes can be found by one string. + pub ray_id: Option, pub max_input_size: Option, pub disable_metadata_lookup: bool, } @@ -60,6 +63,7 @@ impl ClientConfig { encoding: EncodingKind::Bare, transport: TransportKind::WebSocket, headers: None, + ray_id: None, max_input_size: None, disable_metadata_lookup: false, } @@ -107,6 +111,16 @@ impl ClientConfig { self } + /// Sets the ray sent on every actor request. It must be 1 to 128 + /// characters of `[A-Za-z0-9_-]`, the same bound the runtime applies; a + /// value outside it is dropped with a warning when the client is built. + /// A ray carried in the active OpenTelemetry baggage under `rivet.ray.id` + /// takes precedence per call. + pub fn ray_id(mut self, ray_id: impl Into) -> Self { + self.ray_id = Some(ray_id.into()); + self + } + pub fn max_input_size(mut self, max_input_size: usize) -> Self { self.max_input_size = Some(max_input_size); self @@ -153,6 +167,7 @@ impl Client { config.namespace, config.pool_name, config.headers, + config.ray_id, config.max_input_size, config.disable_metadata_lookup, ); diff --git a/rivetkit-rust/packages/client/src/remote_manager.rs b/rivetkit-rust/packages/client/src/remote_manager.rs index adf46e666d..af94047a04 100644 --- a/rivetkit-rust/packages/client/src/remote_manager.rs +++ b/rivetkit-rust/packages/client/src/remote_manager.rs @@ -1,27 +1,36 @@ -use anyhow::{anyhow, Context, Result}; -use base64::{engine::general_purpose, engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; +use anyhow::{Context, Result, anyhow}; +use base64::{Engine as _, engine::general_purpose, engine::general_purpose::URL_SAFE_NO_PAD}; use bytes::Bytes; +use opentelemetry::baggage::BaggageExt as _; +use opentelemetry::propagation::TextMapPropagator as _; +use opentelemetry_http::HeaderInjector; +use opentelemetry_sdk::propagation::TraceContextPropagator; use reqwest::{ - header::{HeaderMap, HeaderName, HeaderValue, USER_AGENT}, Method, + header::{HeaderMap, HeaderName, HeaderValue, USER_AGENT}, }; +use rivetkit_client_protocol::ray_id::{HEADER_RIVET_RAY_ID, RAY_BAGGAGE_KEY, RayId}; use serde::{Deserialize, Serialize}; use serde_cbor; use std::{collections::HashMap, str::FromStr, sync::Arc}; use tokio::sync::OnceCell; use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tracing_opentelemetry::OpenTelemetrySpanExt as _; use crate::{ common::{ - serialize_actor_key, ActorKey, EncodingKind, RawWebSocket, HEADER_RIVET_ACTOR, - HEADER_RIVET_NAMESPACE, HEADER_RIVET_TARGET, HEADER_RIVET_TOKEN, PATH_CONNECT_WEBSOCKET, - PATH_WEBSOCKET_PREFIX, USER_AGENT_VALUE, WS_PROTOCOL_ACTOR, WS_PROTOCOL_CONN_ID, - WS_PROTOCOL_CONN_PARAMS, WS_PROTOCOL_CONN_TOKEN, WS_PROTOCOL_ENCODING, - WS_PROTOCOL_STANDARD, WS_PROTOCOL_TARGET, WS_PROTOCOL_TOKEN, + ActorKey, EncodingKind, HEADER_RIVET_ACTOR, HEADER_RIVET_NAMESPACE, HEADER_RIVET_TARGET, + HEADER_RIVET_TOKEN, PATH_CONNECT_WEBSOCKET, PATH_WEBSOCKET_PREFIX, RawWebSocket, + USER_AGENT_VALUE, WS_PROTOCOL_ACTOR, WS_PROTOCOL_CONN_ID, WS_PROTOCOL_CONN_PARAMS, + WS_PROTOCOL_CONN_TOKEN, WS_PROTOCOL_ENCODING, WS_PROTOCOL_STANDARD, WS_PROTOCOL_TARGET, + WS_PROTOCOL_TOKEN, serialize_actor_key, }, protocol::query::ActorQuery, }; +const HEADER_TRACEPARENT: &str = "traceparent"; +const HEADER_TRACESTATE: &str = "tracestate"; + #[derive(Clone)] pub struct RemoteManager { endpoint: String, @@ -29,6 +38,7 @@ pub struct RemoteManager { namespace: String, pool_name: String, headers: HashMap, + ray_id: Option, max_input_size: usize, disable_metadata_lookup: bool, resolved_config: Arc>, @@ -126,6 +136,7 @@ impl RemoteManager { namespace: default_namespace(), pool_name: default_pool_name(), headers: HashMap::new(), + ray_id: None, max_input_size: default_max_input_size(), disable_metadata_lookup: false, resolved_config: Arc::new(OnceCell::new()), @@ -139,15 +150,27 @@ impl RemoteManager { namespace: Option, pool_name: Option, headers: Option>, + ray_id: Option, max_input_size: Option, disable_metadata_lookup: bool, ) -> Self { + let ray_id = ray_id.and_then(|ray_id| match RayId::parse(ray_id) { + Ok(ray_id) => Some(ray_id), + Err(error) => { + tracing::warn!( + %error, + "dropping invalid configured ray ID" + ); + None + } + }); Self { endpoint, token, namespace: namespace.unwrap_or_else(default_namespace), pool_name: pool_name.unwrap_or_else(default_pool_name), headers: headers.unwrap_or_default(), + ray_id, max_input_size: max_input_size.unwrap_or_else(default_max_input_size), disable_metadata_lookup, resolved_config: Arc::new(OnceCell::new()), @@ -214,6 +237,11 @@ impl RemoteManager { req = req.header(USER_AGENT, USER_AGENT_VALUE); for (key, value) in &self.headers { + // Ray ID and trace context are per call, so a configured value cannot + // pin stale context on every request. Matches the TypeScript client. + if is_telemetry_header(key) { + continue; + } let name = HeaderName::from_str(key) .with_context(|| format!("invalid configured header name `{key}`"))?; let value = HeaderValue::from_str(value) @@ -493,6 +521,11 @@ impl RemoteManager { let mut req = self.apply_common_headers_with(builder, &config)?; + // Per-call context wins over configured headers, and headers the + // caller passed for this request win over both, matching the + // TypeScript client. + let mut headers = headers; + self.add_telemetry_headers(&mut headers)?; req = req.headers(headers); if let Some(body_data) = body { @@ -503,6 +536,36 @@ impl RemoteManager { Ok(res) } + /// Headers that carry the caller's trace context and ray ID into the actor, + /// read from the `tracing` span current at the call. Without a registered + /// OpenTelemetry layer the span carries no context and only a configured + /// ray ID is sent. + fn add_telemetry_headers(&self, headers: &mut HeaderMap) -> Result<()> { + let context = tracing::Span::current().context(); + let caller_set_trace_context = + headers.contains_key(HEADER_TRACEPARENT) || headers.contains_key(HEADER_TRACESTATE); + if !caller_set_trace_context { + TraceContextPropagator::new().inject_context(&context, &mut HeaderInjector(headers)); + if headers + .get(HEADER_TRACESTATE) + .is_some_and(HeaderValue::is_empty) + { + headers.remove(HEADER_TRACESTATE); + } + } + + let baggage = context.baggage(); + let baggage_ray = baggage + .get(RAY_BAGGAGE_KEY) + .and_then(|value| RayId::parse(value.as_str().into_owned()).ok()); + if let Some(ray_id) = baggage_ray.as_ref().or(self.ray_id.as_ref()) { + headers + .entry(HEADER_RIVET_RAY_ID) + .or_insert(HeaderValue::from_str(ray_id.as_str()).context("format ray ID header")?); + } + Ok(()) + } + pub fn gateway_url(&self, query: &ActorQuery) -> Result { let config = self.base_config(); match query { @@ -823,3 +886,9 @@ fn default_pool_name() -> String { fn default_max_input_size() -> usize { 4 * 1024 } + +fn is_telemetry_header(name: &str) -> bool { + [HEADER_RIVET_RAY_ID, HEADER_TRACEPARENT, HEADER_TRACESTATE] + .iter() + .any(|header| header.eq_ignore_ascii_case(name)) +} diff --git a/rivetkit-rust/packages/client/tests/bare.rs b/rivetkit-rust/packages/client/tests/bare.rs index 6900a4b91d..e2f6e949ba 100644 --- a/rivetkit-rust/packages/client/tests/bare.rs +++ b/rivetkit-rust/packages/client/tests/bare.rs @@ -64,6 +64,11 @@ struct ConfigHeaderTestState { saw_raw_websocket: Arc, } +#[derive(Clone)] +struct TelemetryHeaderState { + seen: Arc>>>, +} + #[derive(Clone)] struct MetadataLookupState { saw_metadata: Arc, @@ -547,6 +552,96 @@ async fn config_headers_are_sent_on_http_and_websocket_paths() { server.abort(); } +#[tokio::test] +async fn active_span_and_ray_reach_the_actor_over_configured_headers() { + use opentelemetry::trace::{TraceContextExt as _, TracerProvider as _}; + use opentelemetry::{Context as OtelContext, KeyValue, baggage::BaggageExt as _}; + use tracing::Instrument as _; + use tracing_opentelemetry::OpenTelemetrySpanExt as _; + use tracing_subscriber::layer::SubscriberExt as _; + + // Register the application tracing layer without an exporter. + let provider = opentelemetry_sdk::trace::SdkTracerProvider::builder().build(); + let subscriber = tracing_subscriber::registry() + .with(tracing_opentelemetry::layer().with_tracer(provider.tracer("test"))); + let _subscriber = tracing::subscriber::set_default(subscriber); + + let state = TelemetryHeaderState { + seen: Arc::new(std::sync::Mutex::new(Vec::new())), + }; + let app = Router::new() + .route( + "/gateway/{actor_id}/action/{action}", + post(action_capturing_telemetry_headers), + ) + .with_state(state.clone()); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + + // A configured `traceparent` must not pin stale context, and the + // configured ray applies only when no baggage carries one. + let client = Client::new( + ClientConfig::new(endpoint(addr)) + .disable_metadata_lookup(true) + .header( + "traceparent", + "00-000000000000000000000000000000aa-00000000000000aa-01", + ) + .header("x-rivet-ray-id", "stale-configured-ray") + .ray_id("ray-from-config"), + ); + let actor = client + .get_or_create( + "counter", + vec!["telemetry-headers".to_owned()], + GetOrCreateOptions::default(), + ) + .unwrap(); + + let span = tracing::info_span!("caller"); + let expected_span = span.context().span().span_context().clone(); + let output = async { actor.action("increment", vec![json!(2)]).await.unwrap() } + .instrument(span) + .await; + assert_eq!(output, json!({ "count": 3 })); + + let baggage = + OtelContext::current_with_baggage(vec![KeyValue::new("rivet.ray.id", "ray-from-baggage")]); + let baggage_span = tracing::info_span!(parent: None, "handler"); + baggage_span.set_parent(baggage); + let output = async { actor.action("increment", vec![json!(2)]).await.unwrap() } + .instrument(baggage_span) + .await; + assert_eq!(output, json!({ "count": 3 })); + + let seen = state.seen.lock().unwrap().clone(); + assert_eq!(seen.len(), 2); + assert_eq!( + seen[0].get("traceparent").map(String::as_str), + Some( + format!( + "00-{}-{}-01", + expected_span.trace_id(), + expected_span.span_id() + ) + .as_str() + ) + ); + assert_eq!( + seen[0].get("x-rivet-ray-id").map(String::as_str), + Some("ray-from-config") + ); + assert_eq!( + seen[1].get("x-rivet-ray-id").map(String::as_str), + Some("ray-from-baggage") + ); + + server.abort(); +} + #[tokio::test] async fn max_input_size_checks_raw_query_input_before_base64url_encoding() { let client = Client::new( @@ -1169,6 +1264,34 @@ async fn action_with_config_header( .await } +async fn action_capturing_telemetry_headers( + State(state): State, + Path((actor_id, action_name)): Path<(String, String)>, + headers: HeaderMap, + body: Bytes, +) -> impl IntoResponse { + state.seen.lock().unwrap().push( + headers + .iter() + .filter_map(|(name, value)| { + Some((name.as_str().to_owned(), value.to_str().ok()?.to_owned())) + }) + .collect(), + ); + action( + State(TestState { + saw_bare_action: Arc::new(AtomicBool::new(false)), + saw_bare_queue: Arc::new(AtomicBool::new(false)), + saw_raw_fetch: Arc::new(AtomicBool::new(false)), + saw_raw_websocket: Arc::new(AtomicBool::new(false)), + }), + Path((actor_id, action_name)), + headers, + body, + ) + .await +} + async fn action_for_disable_metadata( Path((actor_id, action_name)): Path<(String, String)>, headers: HeaderMap, diff --git a/rivetkit-rust/packages/rivetkit-core/Cargo.toml b/rivetkit-rust/packages/rivetkit-core/Cargo.toml index b1c5d139f1..15fa510386 100644 --- a/rivetkit-rust/packages/rivetkit-core/Cargo.toml +++ b/rivetkit-rust/packages/rivetkit-core/Cargo.toml @@ -15,6 +15,8 @@ default = ["native-runtime"] native-runtime = [ "dep:nix", "dep:reqwest", + "dep:opentelemetry-otlp", + "dep:tracing-subscriber", "dep:rivetkit-engine-process", "dep:axum", "dep:bytes", @@ -49,6 +51,14 @@ 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-http.workspace = true +opentelemetry-otlp = { version = "0.28", default-features = false, optional = true, features = ["trace", "http-json", "http-proto", "grpc-tonic", "reqwest-blocking-client"] } +opentelemetry_sdk = { version = "0.28", default-features = false, 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/config.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/config.rs index e9519f8824..ae941ce7b7 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/config.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/config.rs @@ -55,6 +55,13 @@ pub struct ActionDefinition { pub name: String, } +/// A queue the actor declares. Names are bounded by this set wherever a queue +/// name becomes a telemetry dimension. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct QueueDefinition { + pub name: String, +} + /// Experimental SQLite profiling configuration. /// /// This entire configuration surface, including every field, is subject to @@ -178,6 +185,7 @@ pub struct ActorConfig { pub max_outgoing_message_size: u32, pub overrides: Option, pub actions: Vec, + pub queues: Vec, /// Author-declared inspector tab entries (custom tabs + built-in /// hides). Validated upstream (Zod / builder). pub inspector_tabs: Vec, @@ -212,6 +220,7 @@ pub struct ActorConfigInput { pub max_incoming_message_size: Option, pub max_outgoing_message_size: Option, pub actions: Option>, + pub queues: Option>, pub inspector_tabs: Option>, } @@ -289,6 +298,9 @@ impl ActorConfig { if let Some(actions) = config.actions { actor_config.actions = actions; } + if let Some(queues) = config.queues { + actor_config.queues = queues; + } if let Some(tabs) = config.inspector_tabs { actor_config.inspector_tabs = tabs; } @@ -364,6 +376,7 @@ impl Default for ActorConfig { max_outgoing_message_size: DEFAULT_MAX_OUTGOING_MESSAGE_SIZE, overrides: None, actions: Vec::new(), + queues: Vec::new(), inspector_tabs: Vec::new(), } } diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs index cfc312a148..df5ed8c9ac 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,72 @@ 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 a handle for the same invocation whose spans parent to the + /// application span the host runtime has active, given as W3C + /// `traceparent` and `tracestate`. A handle that serves no invocation is + /// returned unchanged, because it opens no spans. + #[doc(hidden)] + pub fn with_application_span( + &self, + traceparent: Option<&str>, + tracestate: Option<&str>, + ) -> Self { + Self( + self.0.clone(), + self.1 + .as_ref() + .map(|telemetry| telemetry.with_application_span(traceparent, tracestate)), + ) + } + + /// Returns the SQLite handle bound to this handle's invocation. + pub fn invocation_sql(&self) -> crate::actor::sqlite::SqliteDb { + self.0.sql.clone().with_invocation_telemetry(self.1.clone()) + } + + /// Opens the span covering one call out to another actor, or nothing when + /// this handle serves no invocation or tracing is disabled. + /// + /// `actor_name` and `action_name` name the callee. Both come from the + /// caller's own registry rather than from a remote peer, so neither is a + /// cardinality surface. + #[doc(hidden)] + pub fn begin_outbound_call( + &self, + actor_name: &str, + action_name: &str, + ) -> Option { + self.1 + .as_ref()? + .start_outbound_call(actor_name, action_name) + } + + /// Returns correlation for the invocation this handle serves, absent when + /// the handle is not bound to one or tracing is disabled. + pub fn invocation_trace_context(&self) -> Option { + self.1.as_ref()?.trace_context() + } + + 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, @@ -286,8 +358,12 @@ 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.queues.iter().map(|queue| queue.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()); @@ -300,7 +376,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 +463,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 } @@ -714,9 +796,20 @@ impl ActorContext { false } + /// Runs `future` to completion after the current reply, without blocking + /// it. Work started from an invocation keeps that invocation's span open + /// until it settles, so its SQLite operations and logs stay attributed to + /// the request that started them. #[cfg(not(feature = "wasm-runtime"))] pub fn wait_until(&self, future: impl Future + Send + 'static) { - self.spawn_work(ActorWorkKind::WaitUntil, future); + let invocation = self + .1 + .as_ref() + .map(crate::ActorInvocationTelemetry::hold_open); + self.spawn_work(ActorWorkKind::WaitUntil, async move { + future.await; + drop(invocation); + }); } #[cfg(not(feature = "wasm-runtime"))] @@ -726,7 +819,14 @@ impl ActorContext { #[cfg(feature = "wasm-runtime")] pub fn wait_until(&self, future: impl Future + 'static) { - self.spawn_work(ActorWorkKind::WaitUntil, future); + let invocation = self + .1 + .as_ref() + .map(crate::ActorInvocationTelemetry::hold_open); + self.spawn_work(ActorWorkKind::WaitUntil, async move { + future.await; + drop(invocation); + }); } #[cfg(feature = "wasm-runtime")] @@ -911,6 +1011,12 @@ impl ActorContext { &self.0.metrics } + /// Identity fields shared by every invocation on this actor. Built once so a + /// span does not re-allocate them per action. + pub(crate) fn telemetry_identity(&self) -> Arc { + self.0.telemetry_identity.clone() + } + pub(crate) fn record_user_task_started(&self, kind: UserTaskKind) { self.0.metrics.begin_user_task(kind); } @@ -1337,7 +1443,7 @@ impl ActorContext { } pub(crate) fn from_weak(weak: &Weak) -> Option { - weak.upgrade().map(Self) + weak.upgrade().map(|inner| Self(inner, None)) } #[doc(hidden)] @@ -1498,6 +1604,7 @@ impl ActorContext { } let region = match kind { ActorWorkKind::Action => self.internal_keep_awake_region(), + ActorWorkKind::DispatchReply => self.internal_keep_awake_region(), ActorWorkKind::KeepAwake => self.keep_awake_region_state(), ActorWorkKind::InternalKeepAwake => self.internal_keep_awake_region(), ActorWorkKind::WaitUntil => return None, @@ -1748,8 +1855,14 @@ 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, + dispatch.origin, + ); + let invocation_telemetry = invocation.telemetry(); let (reply_tx, reply_rx) = oneshot::channel(); let mut dispatch_error = None; @@ -1759,6 +1872,7 @@ impl ActorContext { args, conn: None, scheduled_fire: Some(scheduled_fire), + invocation_telemetry: Some(invocation_telemetry), reply: Reply::from(reply_tx), }, "scheduled_action", @@ -1774,8 +1888,9 @@ impl ActorContext { "scheduled event execution failed" ); } - Err(error) => { - dispatch_error = Some(error.into()); + Err(_) => { + // Use the same dropped-reply error for tracing, metrics, and schedule history. + dispatch_error = Some(ActorLifecycleError::DroppedReply.build()); tracing::error!( error = ?dispatch_error.as_ref().expect("just assigned"), event_id, @@ -1794,6 +1909,7 @@ impl ActorContext { ); } } + invocation.finish(dispatch_error.as_ref()); ctx.finish_schedule_dispatch(&event_id, history_id, dispatch_error.as_ref()) .await; @@ -1811,7 +1927,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/internal_storage/mod.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/mod.rs index 78020880a2..cf7223b02b 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/mod.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/mod.rs @@ -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, @@ -377,15 +381,18 @@ pub(crate) async fn load_queue_metadata(db: &SqliteDb) -> Result }) } +/// Writes a queue message and, when the sender carried one, its encoded trace +/// context, in one batch so a message never exists without its origin. pub(crate) async fn persist_queue_message( db: &SqliteDb, id: u64, next_id: u64, message: &PersistedQueueMessage, + trace_context: Option>, ) -> Result<()> { let id = i64::try_from(id).context("queue message id exceeds sqlite integer range")?; let next_id = i64::try_from(next_id).context("queue next id exceeds sqlite integer range")?; - db.execute_batch(vec![ + let mut statements = vec![ SqliteBatchStatement { sql: INSERT_QUEUE_MESSAGE_SQL.to_owned(), params: Some(vec![ @@ -403,12 +410,26 @@ pub(crate) async fn persist_queue_message( BindParam::Integer(next_id), ]), }, - ]) - .await - .context("persist internal queue message")?; + ]; + if let Some(trace_context) = trace_context { + statements.push(SqliteBatchStatement { + sql: UPSERT_QUEUE_TRACE_CONTEXT_SQL.to_owned(), + params: Some(vec![ + BindParam::Text(queue_trace_context_key(id)), + BindParam::Blob(trace_context), + ]), + }); + } + db.execute_batch(statements) + .await + .context("persist internal queue message")?; Ok(()) } +pub(crate) fn queue_trace_context_key(id: i64) -> String { + format!("queue_trace_context:{id}") +} + /// Persists imported queue rows without rewriting `queue_next_id` for every /// message. The importer writes that singleton once after every row is copied. pub(crate) async fn persist_queue_messages( @@ -642,6 +663,7 @@ fn decode_queue_message_rows(rows: &[Vec]) -> Result Result< return Ok(0); } - let mut statements = Vec::with_capacity(ids.len()); + // Each message row is followed by the delete of its trace context, so the + // two go in one batch and the row count below reads every other result. + let mut statements = Vec::with_capacity(ids.len() * 2); for id in ids { + let id = i64::try_from(*id).context("queue message id exceeds sqlite integer range")?; statements.push(SqliteBatchStatement { sql: DELETE_QUEUE_MESSAGE_SQL.to_owned(), - params: Some(vec![BindParam::Integer( - i64::try_from(*id).context("queue message id exceeds sqlite integer range")?, - )]), + params: Some(vec![BindParam::Integer(id)]), + }); + statements.push(SqliteBatchStatement { + sql: DELETE_QUEUE_TRACE_CONTEXT_SQL.to_owned(), + params: Some(vec![BindParam::Text(queue_trace_context_key(id))]), }); } let results = db .execute_batch(statements) .await .context("delete internal queue messages")?; - results.into_iter().try_fold(0u32, |deleted, result| { - let changes = u32::try_from(result.changes) - .context("deleted queue message count is outside u32 range")?; - deleted - .checked_add(changes) - .context("deleted queue message count exceeds u32 range") - }) + results + .into_iter() + .step_by(2) + .try_fold(0u32, |deleted, result| { + let changes = u32::try_from(result.changes) + .context("deleted queue message count is outside u32 range")?; + deleted + .checked_add(changes) + .context("deleted queue message count exceeds u32 range") + }) } pub(crate) async fn reset_queue(db: &SqliteDb) -> Result<()> { - db.execute(RESET_QUEUE_SQL, None) - .await - .context("reset internal queue")?; + db.execute_batch(vec![ + SqliteBatchStatement { + sql: RESET_QUEUE_SQL.to_owned(), + params: None, + }, + SqliteBatchStatement { + sql: RESET_QUEUE_TRACE_CONTEXTS_SQL.to_owned(), + params: None, + }, + ]) + .await + .context("reset internal queue")?; Ok(()) } @@ -685,6 +724,8 @@ pub(crate) async fn reset_queue(db: &SqliteDb) -> Result<()> { pub(crate) struct QueueMessageRow { pub id: u64, pub message: PersistedQueueMessage, + /// Encoded `QueueTraceContext`, when the sender carried one. + pub trace_context: Option>, } pub(crate) async fn user_kv_batch_get( @@ -1163,6 +1204,12 @@ 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")?; + db.execute(RESET_QUEUE_TRACE_CONTEXTS_SQL, None) + .await + .context("clear imported queue trace contexts")?; Ok(()) } diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/queries.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/queries.rs index dad1417c00..af94ce3962 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/queries.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/queries.rs @@ -14,6 +14,8 @@ 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"; pub(crate) const CANCEL_SCHEDULE_SQL: &str = @@ -29,7 +31,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,14 +57,20 @@ 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"; -pub(crate) const LOAD_QUEUE_MESSAGES_SQL: &str = - "SELECT id, name, body, created_at FROM _rivet_queue ORDER BY id"; -pub(crate) const LOAD_QUEUE_MESSAGES_LIMITED_SQL: &str = - "SELECT id, name, body, created_at FROM _rivet_queue ORDER BY id LIMIT ?"; -pub(crate) const LOAD_QUEUE_MESSAGES_FOR_NAME_SQL: &str = "SELECT id, name, body, created_at FROM _rivet_queue INDEXED BY _rivet_queue_name_id WHERE name = ? ORDER BY id LIMIT ?"; +// The trace-origin subquery uses a primary-key lookup per returned message. +pub(crate) const LOAD_QUEUE_MESSAGES_SQL: &str = "SELECT id, name, body, created_at, (SELECT value FROM _rivet_meta WHERE key = 'queue_trace_context:' || id) FROM _rivet_queue ORDER BY id"; +pub(crate) const LOAD_QUEUE_MESSAGES_LIMITED_SQL: &str = "SELECT id, name, body, created_at, (SELECT value FROM _rivet_meta WHERE key = 'queue_trace_context:' || id) FROM _rivet_queue ORDER BY id LIMIT ?"; +pub(crate) const LOAD_QUEUE_MESSAGES_FOR_NAME_SQL: &str = "SELECT id, name, body, created_at, (SELECT value FROM _rivet_meta WHERE key = 'queue_trace_context:' || id) FROM _rivet_queue INDEXED BY _rivet_queue_name_id WHERE name = ? ORDER BY id LIMIT ?"; pub(crate) const HAS_QUEUE_MESSAGES_SQL: &str = "SELECT 1 FROM _rivet_queue LIMIT 1"; pub(crate) const HAS_QUEUE_MESSAGES_FOR_NAME_SQL: &str = "SELECT 1 FROM _rivet_queue INDEXED BY _rivet_queue_name_id WHERE name = ? LIMIT 1"; @@ -71,13 +82,17 @@ pub(crate) fn load_queue_messages_by_ids_sql(id_count: usize) -> String { .collect::>() .join(", "); format!( - "SELECT id, name, body, created_at FROM _rivet_queue WHERE id IN ({placeholders}) ORDER BY id" + "SELECT id, name, body, created_at, (SELECT value FROM _rivet_meta WHERE key = 'queue_trace_context:' || id) FROM _rivet_queue WHERE id IN ({placeholders}) ORDER BY id" ) } pub(crate) const INSERT_QUEUE_MESSAGE_SQL: &str = "INSERT OR REPLACE INTO _rivet_queue (id, name, body, created_at) VALUES (?, ?, ?, ?)"; pub(crate) const DELETE_QUEUE_MESSAGE_SQL: &str = "DELETE FROM _rivet_queue WHERE id = ?"; pub(crate) const RESET_QUEUE_SQL: &str = "DELETE FROM _rivet_queue"; +pub(crate) const UPSERT_QUEUE_TRACE_CONTEXT_SQL: &str = "INSERT INTO _rivet_meta (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value"; +pub(crate) const DELETE_QUEUE_TRACE_CONTEXT_SQL: &str = "DELETE FROM _rivet_meta WHERE key = ?"; +pub(crate) const RESET_QUEUE_TRACE_CONTEXTS_SQL: &str = + "DELETE FROM _rivet_meta WHERE key >= 'queue_trace_context:' AND key < 'queue_trace_context;'"; pub(crate) const DELETE_USER_KV_SQL: &str = "DELETE FROM _rivet_user_kv WHERE key = ?"; pub(crate) const DELETE_USER_KV_RANGE_SQL: &str = diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/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/messages.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs index 702fa8f6ec..a252c2d0cf 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs @@ -313,6 +313,15 @@ pub enum ActorHttpResponse { Stream(StreamingResponse), } +impl ActorHttpResponse { + pub fn status(&self) -> u16 { + match self { + Self::Buffered(response) => response.status().as_u16(), + Self::Stream(response) => response.status, + } + } +} + impl From for ActorHttpResponse { fn from(value: Response) -> Self { Self::Buffered(value) @@ -392,10 +401,16 @@ pub enum ActorEvent { args: Vec, conn: Option, scheduled_fire: Option, + /// Telemetry of the invocation this action runs as, for the host + /// runtime to bind onto the context it hands the action. Absent when + /// the dispatch opened no invocation. + invocation_telemetry: Option, reply: Reply>, }, HttpRequest { request: Request, + /// Telemetry of the invocation this request runs as. See `Action`. + invocation_telemetry: Option, reply: Reply, }, QueueSend { @@ -405,6 +420,8 @@ pub enum ActorEvent { request: Request, wait: bool, timeout_ms: Option, + /// Telemetry of the invocation this send runs as. See `Action`. + invocation_telemetry: Option, reply: Reply, }, WebSocketOpen { diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/metrics.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/metrics.rs index 97892f8734..b01601cea4 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,70 @@ pub(crate) struct StartupTimer { finished: bool, } +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum InvocationType { + Action, + Scheduled, + /// A raw HTTP request served by the actor's `onRequest` handler. + Request, + /// A message sent into one of the actor's queues from outside it. + QueueSend, +} + +impl InvocationType { + pub(crate) fn as_label(self) -> &'static str { + match self { + Self::Action => "action", + Self::Scheduled => "scheduled", + Self::Request => "request", + Self::QueueSend => "queue_send", + } + } + + /// OpenTelemetry span kind for this invocation. An action or a raw HTTP request + /// is entered from outside the actor, a scheduled fire originates inside + /// it, and a queue send produces a message the actor consumes later. + pub(crate) fn otel_kind(self) -> &'static str { + match self { + Self::Action | Self::Request => "server", + Self::Scheduled => "internal", + Self::QueueSend => "producer", + } + } +} + +#[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, + queue_names: BTreeSet, #[cfg(feature = "sqlite-local")] sqlite_profiling: crate::SqliteProfilingConfig, #[cfg(feature = "sqlite-local")] @@ -182,6 +244,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 +1155,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 +1490,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 +1609,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 +1674,32 @@ 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(), + 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(), + std::iter::empty(), + _sqlite_profiling, + ) + } + + pub(crate) fn new_for_actor( + actor_name: impl Into, + action_names: impl IntoIterator, + queue_names: impl IntoIterator, + _sqlite_profiling: crate::SqliteProfilingConfig, ) -> Self { let labels = ActorMetricLabels { actor_name: actor_name.into(), @@ -1610,6 +1720,8 @@ impl ActorMetrics { Self { inner: Arc::new(ActorMetricInner { labels, + action_names: action_names.into_iter().collect(), + queue_names: queue_names.into_iter().collect(), #[cfg(feature = "sqlite-local")] sqlite_profiling: _sqlite_profiling, #[cfg(feature = "sqlite-local")] @@ -1864,6 +1976,54 @@ 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" + } + } + + /// Folds an undeclared queue name the same way `label_action_name` folds + /// an action name, for the same reason: it arrives from the caller. + pub(crate) fn label_queue_name<'a>(&'a self, queue_name: &'a str) -> &'a str { + if self.inner.queue_names.contains(queue_name) { + queue_name + } else { + "_OTHER" + } + } + + /// `action_label` is already bounded: the caller folds action names + /// through `label_action_name`, and a request invocation uses a fixed + /// name, which the fold would otherwise turn into `_OTHER`. + pub(crate) fn record_invocation( + &self, + action_label: &str, + invocation_type: InvocationType, + result: InvocationStatus, + duration: Duration, + ) { + let actor_labels = self.actor_labels(); + let labels = [ + actor_labels[0], + action_label, + 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/queue.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/queue.rs index 68ed09adaa..f549f0f485 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/queue.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/queue.rs @@ -10,6 +10,7 @@ use crate::time::{Instant, SystemTime, UNIX_EPOCH, sleep}; use anyhow::{Context, Result}; use rivet_error::RivetError; +use rivetkit_actor_persist::versioned::{QueueTraceContext, QueueTraceContextData}; use rivetkit_actor_persist::{generated::v4 as persist_v4, versioned as persist_versioned}; use serde::{Deserialize, Serialize}; #[cfg(not(target_arch = "wasm32"))] @@ -26,6 +27,9 @@ use crate::actor::persist::{ use crate::actor::task_types::UserTaskKind; #[cfg(target_arch = "wasm32")] use crate::error::ActorRuntime; +use crate::telemetry::{self, TraceOrigin}; + +const QUEUE_TRACE_CONTEXT_VERSION: u16 = 1; #[derive(Clone, Debug, Default)] pub struct QueueNextOpts { @@ -101,6 +105,10 @@ pub struct QueueMessage { pub name: String, pub body: Vec, pub created_at: i64, + /// Ray ID and span of the invocation that sent the message, which the span + /// covering its receipt links back to. Empty when the sender carried no + /// trace context. + pub trace_origin: TraceOrigin, completion: Option, } @@ -110,9 +118,30 @@ pub struct CompletableQueueMessage { pub name: String, pub body: Vec, pub created_at: i64, + pub trace_origin: TraceOrigin, completion: CompletionHandle, } +impl From for TraceOrigin { + fn from(context: QueueTraceContextData) -> Self { + Self { + ray_id: context.ray_id, + traceparent: context.traceparent, + tracestate: context.tracestate, + } + } +} + +impl From for QueueTraceContextData { + fn from(origin: TraceOrigin) -> Self { + Self { + ray_id: origin.ray_id, + traceparent: origin.traceparent, + tracestate: origin.tracestate, + } + } +} + #[derive(Clone)] struct CompletionHandle(Arc); @@ -281,6 +310,24 @@ impl ActorContext { in_flight_at: None, }; let encoded_message = encode_queue_message(&persisted).context("encode queue message")?; + // A send from inside an invocation records that invocation as the + // message's origin. A send from outside one has nothing to record. + let trace_origin = self + .invocation_telemetry() + .map(crate::ActorInvocationTelemetry::trace_origin) + .unwrap_or_default(); + let encoded_trace_context = if trace_origin.is_empty() { + None + } else { + Some( + encode_latest_with_embedded_version::( + QueueTraceContextData::from(trace_origin.clone()), + QUEUE_TRACE_CONTEXT_VERSION, + "queue trace context", + ) + .context("encode queue trace context")?, + ) + }; let config = self.config(); if encoded_message.len() > config.max_queue_message_size as usize { @@ -323,9 +370,14 @@ impl ActorContext { false }; - let persist_result = - internal_storage::persist_queue_message(self.sql(), id, metadata.next_id, &persisted) - .await; + let persist_result = internal_storage::persist_queue_message( + self.sql(), + id, + metadata.next_id, + &persisted, + encoded_trace_context, + ) + .await; if let Err(error) = persist_result { metadata.next_id = id; @@ -350,6 +402,7 @@ impl ActorContext { name: name.to_owned(), body: body.to_vec(), created_at, + trace_origin, completion: None, }) } @@ -653,6 +706,12 @@ impl ActorContext { return Ok(Vec::new()); } + // Keep receipt spans open until the batch is handed back. + let _receive_spans: Vec = selected + .iter() + .map(|message| telemetry::start_queue_receive(self, message)) + .collect(); + if completable { let queue_size = self.0.queue_metadata.lock().await.size; self.0 @@ -927,6 +986,7 @@ impl QueueMessage { name: self.name, body: self.body, created_at: self.created_at, + trace_origin: self.trace_origin, completion, }) } @@ -947,6 +1007,7 @@ impl CompletableQueueMessage { name: self.name, body: self.body, created_at: self.created_at, + trace_origin: self.trace_origin, completion: Some(self.completion), } } @@ -1051,11 +1112,29 @@ fn normalize_names(names: Option>) -> Option> { } fn queue_message_from_row(row: internal_storage::QueueMessageRow) -> QueueMessage { + // A malformed trace context is a telemetry defect, not a queue defect, so + // the message is still delivered and only its origin is lost. + let trace_origin = row + .trace_context + .as_deref() + .and_then(|payload| { + decode_latest_with_embedded_version::(payload, "queue trace context") + .inspect_err(|error| { + tracing::warn!( + message_id = row.id, + ?error, + "ignoring undecodable queue trace context" + ); + }) + .ok() + }) + .map_or_else(TraceOrigin::default, TraceOrigin::from); QueueMessage { id: row.id, name: row.message.name, body: row.message.body, created_at: row.message.created_at, + trace_origin, completion: None, } } diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/schedule.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/schedule.rs index f54026a1b2..671981070e 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/schedule.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/schedule.rs @@ -11,6 +11,7 @@ use futures::future::BoxFuture; use futures::future::{AbortHandle, Abortable}; use rivet_envoy_client::handle::EnvoyHandle; use rivet_error::RivetError; +use rivetkit_actor_persist::versioned::{ScheduleTraceContext, ScheduleTraceContextData}; use serde::{Deserialize, Serialize}; use tokio::runtime::Handle; use tokio::sync::oneshot; @@ -19,8 +20,12 @@ use uuid::Uuid; use crate::actor::context::ActorContext; use crate::actor::internal_storage::queries::*; +use crate::actor::persist::{ + decode_latest_with_embedded_version, encode_latest_with_embedded_version, +}; use crate::error::{ScheduleRuntimeError, client_error_message, client_error_metadata}; use crate::sqlite::{BindParam, ColumnValue, SqliteBatchStatement}; +use crate::telemetry::{ActorInvocationTelemetry, TraceOrigin}; use crate::time::{SystemTime, UNIX_EPOCH, sleep}; const CRON_ID_PREFIX: &str = "cron:"; @@ -34,6 +39,8 @@ pub const MAX_ACTOR_HISTORY: i64 = 10_000; pub const MIN_INTERVAL_MS: i64 = 5_000; const DEFAULT_HISTORY_LIMIT: i64 = 20; const CLAIM_ONE_SHOT_BATCH_SIZE: usize = 128; +const SCHEDULE_TRACE_CONTEXT_META_PREFIX: &str = "schedule_trace_context:"; +const SCHEDULE_TRACE_CONTEXT_VERSION: u16 = 1; pub(crate) const GLOBAL_HISTORY_PRUNE_INTERVAL: usize = 100; pub(crate) const GLOBAL_HISTORY_RETAINED_ROWS: i64 = MAX_ACTOR_HISTORY - GLOBAL_HISTORY_PRUNE_INTERVAL as i64; @@ -148,6 +155,7 @@ pub(crate) struct DueScheduleDispatch { pub args: Vec, pub fire: ScheduledFireInfo, pub history_id: Option, + pub origin: TraceOrigin, } #[derive(Clone, Debug)] @@ -183,6 +191,14 @@ impl ActorContext { system_now_timestamp_ms() } + /// Trace origin of the invocation defining this schedule, empty when the + /// caller is not inside a traced invocation. + fn schedule_trace_origin(&self) -> TraceOrigin { + self.invocation_telemetry() + .map(ActorInvocationTelemetry::trace_origin) + .unwrap_or_default() + } + pub async fn after( &self, duration: Duration, @@ -195,25 +211,29 @@ impl ActorContext { } pub async fn at(&self, timestamp_ms: i64, action_name: &str, args: &[u8]) -> Result { + let origin = self.schedule_trace_origin(); let _mutation = self.0.schedule_mutation_lock.lock().await; self.ensure_schedule_capacity(false).await?; let event_id = Uuid::new_v4().to_string(); + let schedule_params = vec![ + BindParam::Text(event_id.clone()), + BindParam::Integer(timestamp_ms), + BindParam::Text(action_name.to_owned()), + args_param(args), + BindParam::Integer(ScheduleKind::At.as_i64()), + BindParam::Null, + BindParam::Null, + BindParam::Null, + BindParam::Null, + BindParam::Integer(0), + ]; + let mut statements = vec![SqliteBatchStatement { + sql: INSERT_SCHEDULE_EVENT_SQL.to_owned(), + params: Some(schedule_params), + }]; + append_schedule_trace_context_upsert(&mut statements, &event_id, origin)?; self.sql() - .execute( - INSERT_SCHEDULE_EVENT_SQL, - Some(vec![ - BindParam::Text(event_id.clone()), - BindParam::Integer(timestamp_ms), - BindParam::Text(action_name.to_owned()), - args_param(args), - BindParam::Integer(ScheduleKind::At.as_i64()), - BindParam::Null, - BindParam::Null, - BindParam::Null, - BindParam::Null, - BindParam::Integer(0), - ]), - ) + .execute_batch(statements) .await .context("insert one-shot schedule")?; self.mark_schedule_dirty(); @@ -224,18 +244,21 @@ impl ActorContext { pub async fn cancel_schedule(&self, event_id: &str) -> Result { let _mutation = self.0.schedule_mutation_lock.lock().await; - let result = self + let results = self .sql() - .execute( - CANCEL_SCHEDULE_SQL, - Some(vec![ - BindParam::Text(event_id.to_owned()), - BindParam::Integer(ScheduleKind::At.as_i64()), - ]), - ) + .execute_batch(vec![ + SqliteBatchStatement { + sql: CANCEL_SCHEDULE_SQL.to_owned(), + params: Some(vec![ + BindParam::Text(event_id.to_owned()), + BindParam::Integer(ScheduleKind::At.as_i64()), + ]), + }, + delete_orphan_schedule_trace_context(event_id), + ]) .await .context("cancel one-shot schedule")?; - let removed = result.changes > 0; + let removed = results.first().is_some_and(|result| result.changes > 0); if removed { self.mark_schedule_dirty(); self.record_schedules_updated(); @@ -304,6 +327,7 @@ impl ActorContext { args: &[u8], max_history: Option, ) -> Result<()> { + let origin = self.schedule_trace_origin(); validate_name(name)?; let timezone = timezone.unwrap_or("UTC"); let timezone_parsed = parse_timezone(timezone)?; @@ -334,6 +358,7 @@ impl ActorContext { Some(timezone), None, max_history, + origin, ) .await?; self.prune_schedule_history(&event_id, max_history).await?; @@ -350,6 +375,7 @@ impl ActorContext { args: &[u8], max_history: Option, ) -> Result<()> { + let origin = self.schedule_trace_origin(); validate_name(name)?; if interval_ms < MIN_INTERVAL_MS { return Err(ScheduleRuntimeError::InvalidInterval { @@ -382,6 +408,7 @@ impl ActorContext { None, Some(interval_ms), max_history, + origin, ) .await?; self.prune_schedule_history(&event_id, max_history).await?; @@ -402,23 +429,26 @@ impl ActorContext { timezone: Option<&str>, interval_ms: Option, max_history: i64, + origin: TraceOrigin, ) -> Result<()> { + let mut statements = vec![SqliteBatchStatement { + sql: UPSERT_RECURRING_SCHEDULE_SQL.to_owned(), + params: Some(vec![ + BindParam::Text(event_id.to_owned()), + BindParam::Integer(trigger_at), + BindParam::Text(action_name.to_owned()), + args_param(args), + BindParam::Integer(kind.as_i64()), + optional_text_param(cron_expression), + optional_text_param(timezone), + optional_i64_param(interval_ms), + BindParam::Null, + BindParam::Integer(max_history), + ]), + }]; + append_schedule_trace_context_upsert(&mut statements, event_id, origin)?; self.sql() - .execute( - UPSERT_RECURRING_SCHEDULE_SQL, - Some(vec![ - BindParam::Text(event_id.to_owned()), - BindParam::Integer(trigger_at), - BindParam::Text(action_name.to_owned()), - args_param(args), - BindParam::Integer(kind.as_i64()), - optional_text_param(cron_expression), - optional_text_param(timezone), - optional_i64_param(interval_ms), - BindParam::Null, - BindParam::Integer(max_history), - ]), - ) + .execute_batch(statements) .await .context("upsert recurring schedule")?; Ok(()) @@ -442,10 +472,11 @@ impl ActorContext { SqliteBatchStatement { sql: DELETE_CRON_SQL.to_owned(), params: Some(vec![ - BindParam::Text(event_id), + BindParam::Text(event_id.clone()), BindParam::Integer(ScheduleKind::At.as_i64()), ]), }, + delete_orphan_schedule_trace_context(&event_id), ]) .await .context("delete recurring schedule and history")?; @@ -463,19 +494,23 @@ impl ActorContext { pub(crate) async fn cron_delete_if_action(&self, name: &str, action: &str) -> Result { validate_name(name)?; let _mutation = self.0.schedule_mutation_lock.lock().await; - let result = self + let event_id = cron_event_id(name); + let results = self .sql() - .execute( - DELETE_CRON_IF_ACTION_SQL, - Some(vec![ - BindParam::Text(cron_event_id(name)), - BindParam::Integer(ScheduleKind::At.as_i64()), - BindParam::Text(action.to_owned()), - ]), - ) + .execute_batch(vec![ + SqliteBatchStatement { + sql: DELETE_CRON_IF_ACTION_SQL.to_owned(), + params: Some(vec![ + BindParam::Text(event_id.clone()), + BindParam::Integer(ScheduleKind::At.as_i64()), + BindParam::Text(action.to_owned()), + ]), + }, + delete_orphan_schedule_trace_context(&event_id), + ]) .await .context("delete recurring schedule with matching action")?; - let removed = result.changes > 0; + let removed = results.first().is_some_and(|result| result.changes > 0); if removed { self.mark_schedule_dirty(); self.record_schedules_updated(); @@ -594,24 +629,38 @@ impl ActorContext { let due_schedules = result .rows .iter() - .map(|row| read_stored_schedule(row)) + .map(|row| read_due_schedule(row)) .collect::>>()?; let claim_statements = due_schedules .iter() - .filter(|event| event.kind == ScheduleKind::At) + .filter(|(event, _)| event.kind == ScheduleKind::At) .collect::>() .chunks(CLAIM_ONE_SHOT_BATCH_SIZE) - .map(|events| { + .flat_map(|events| { let mut params = Vec::with_capacity(events.len() * 2 + 1); params.push(BindParam::Integer(ScheduleKind::At.as_i64())); - for event in events { + for (event, _) in events { params.push(BindParam::Text(event.event_id.clone())); params.push(BindParam::Integer(event.trigger_at)); } - SqliteBatchStatement { - sql: claim_one_shots_sql(events.len()), - params: Some(params), - } + let delete_contexts = SqliteBatchStatement { + sql: delete_schedule_trace_contexts_sql(events.len()), + params: Some( + events + .iter() + .map(|(event, _)| { + BindParam::Text(schedule_trace_context_key(&event.event_id)) + }) + .collect(), + ), + }; + [ + delete_contexts, + SqliteBatchStatement { + sql: claim_one_shots_sql(events.len()), + params: Some(params), + }, + ] }) .collect::>(); if !claim_statements.is_empty() { @@ -621,7 +670,7 @@ impl ActorContext { .context("claim due one-shot schedules")?; } let mut dispatches = Vec::new(); - for event in due_schedules { + for (event, origin) in due_schedules { if event.kind == ScheduleKind::At { dispatches.push(DueScheduleDispatch { event_id: event.event_id.clone(), @@ -635,6 +684,7 @@ impl ActorContext { fired_at: now_ms, }, history_id: None, + origin, }); continue; } @@ -708,6 +758,7 @@ impl ActorContext { fired_at: now_ms, }, history_id, + origin, }); } self.mark_schedule_dirty(); @@ -1353,6 +1404,90 @@ fn read_stored_schedule(row: &[ColumnValue]) -> Result { } } +fn read_due_schedule(row: &[ColumnValue]) -> Result<(StoredSchedule, TraceOrigin)> { + let event = read_stored_schedule(row)?; + let origin = read_optional_blob(row, 10, "schedule trace context")? + .and_then(|payload| { + decode_latest_with_embedded_version::( + &payload, + "schedule trace context", + ) + .inspect_err(|error| { + tracing::warn!( + event_id = %event.event_id, + ?error, + "ignoring undecodable schedule trace context" + ); + }) + .ok() + }) + .map_or_else(TraceOrigin::default, TraceOrigin::from); + Ok((event, origin)) +} + +impl From for TraceOrigin { + fn from(context: ScheduleTraceContextData) -> Self { + Self { + ray_id: context.ray_id, + traceparent: context.traceparent, + tracestate: context.tracestate, + } + } +} + +impl From for ScheduleTraceContextData { + fn from(origin: TraceOrigin) -> Self { + Self { + ray_id: origin.ray_id, + traceparent: origin.traceparent, + tracestate: origin.tracestate, + } + } +} + +/// Stores the defining invocation's trace context beside a schedule row, or +/// clears a stale one when the definer carried no context. +fn append_schedule_trace_context_upsert( + statements: &mut Vec, + event_id: &str, + origin: TraceOrigin, +) -> Result<()> { + let key = schedule_trace_context_key(event_id); + if origin.is_empty() { + statements.push(SqliteBatchStatement { + sql: DELETE_SCHEDULE_TRACE_CONTEXT_SQL.to_owned(), + params: Some(vec![BindParam::Text(key)]), + }); + return Ok(()); + } + let payload = encode_latest_with_embedded_version::( + ScheduleTraceContextData::from(origin), + SCHEDULE_TRACE_CONTEXT_VERSION, + "schedule trace context", + )?; + statements.push(SqliteBatchStatement { + sql: UPSERT_SCHEDULE_TRACE_CONTEXT_SQL.to_owned(), + params: Some(vec![BindParam::Text(key), BindParam::Blob(payload)]), + }); + Ok(()) +} + +/// Removes a schedule's trace context once its row is gone. Runs after the +/// row delete in the same batch so a mismatched delete leaves the context alone. +fn delete_orphan_schedule_trace_context(event_id: &str) -> SqliteBatchStatement { + SqliteBatchStatement { + sql: DELETE_ORPHAN_SCHEDULE_TRACE_CONTEXT_SQL.to_owned(), + params: Some(vec![ + BindParam::Text(schedule_trace_context_key(event_id)), + BindParam::Text(event_id.to_owned()), + ]), + } +} + +fn schedule_trace_context_key(event_id: &str) -> String { + format!("{SCHEDULE_TRACE_CONTEXT_META_PREFIX}{event_id}") +} + fn read_cron_fire(row: &[ColumnValue]) -> Result { let error_group = read_optional_text(row, 5, "error_group")?; let error_code = read_optional_text(row, 6, "error_code")?; diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/mod.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/mod.rs index c5cc659d00..9f16a05161 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/mod.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/mod.rs @@ -1,4 +1,5 @@ use std::collections::HashSet; +use std::future::Future; use std::io::Cursor; use std::sync::{ Arc, @@ -22,6 +23,9 @@ use serde_json::{Map as JsonMap, Value as JsonValue}; use tokio::sync::Mutex as AsyncMutex; #[cfg(feature = "sqlite-local")] use tokio::task::JoinHandle; +use tracing::Instrument as _; + +use crate::telemetry::SqliteOperation; #[cfg(feature = "sqlite-local")] mod envoy_sqlite_transport; @@ -234,6 +238,7 @@ pub struct SqliteDb { /// always sets up sqlite storage under the hood, so handle/actor_id are /// not a reliable signal for whether the user opted in; this flag is. enabled: bool, + invocation_telemetry: Option, #[cfg(feature = "sqlite-local")] // Forced-sync: native SQLite handles are used inside spawn_blocking and // synchronous diagnostic accessors. @@ -263,6 +268,7 @@ impl Default for SqliteDb { SqliteBackend::RemoteEnvoy }, enabled: false, + invocation_telemetry: None, #[cfg(feature = "sqlite-local")] db: Default::default(), #[cfg(feature = "sqlite-local")] @@ -299,6 +305,7 @@ impl SqliteDb { generation, backend: select_sqlite_backend(remote_sqlite)?, enabled, + invocation_telemetry: None, #[cfg(feature = "sqlite-local")] db: Default::default(), #[cfg(feature = "sqlite-local")] @@ -351,6 +358,34 @@ impl SqliteDb { self.backend } + #[doc(hidden)] + pub fn with_invocation_telemetry( + mut self, + telemetry: Option, + ) -> Self { + self.invocation_telemetry = telemetry; + self + } + + /// Runs one SQLite operation inside a `rivet.sqlite.*` span when the + /// current invocation is traced. + pub(super) async fn traced( + &self, + operation: SqliteOperation, + future: impl Future>, + ) -> Result { + let Some(mut span) = self + .invocation_telemetry + .as_ref() + .and_then(|telemetry| telemetry.start_sqlite(operation)) + else { + return future.await; + }; + let result = future.instrument(span.span()).await; + span.finish(result.as_ref().err()); + result + } + pub async fn get_pages( &self, request: protocol::SqliteGetPagesRequest, @@ -513,6 +548,11 @@ impl SqliteDb { pub async fn exec(&self, sql: impl Into) -> Result { let sql = sql.into(); + self.traced(SqliteOperation::Exec, self.exec_untraced(sql)) + .await + } + + async fn exec_untraced(&self, sql: String) -> Result { let sql_for_log = sql.clone(); #[cfg(feature = "sqlite-local")] let started_at = self @@ -566,6 +606,15 @@ impl SqliteDb { params: Option>, ) -> Result { let sql = sql.into(); + self.traced(SqliteOperation::Query, self.query_untraced(sql, params)) + .await + } + + async fn query_untraced( + &self, + sql: String, + params: Option>, + ) -> Result { let sql_for_log = sql.clone(); let binding_count = bind_param_count(¶ms); #[cfg(feature = "sqlite-local")] @@ -626,6 +675,15 @@ impl SqliteDb { params: Option>, ) -> Result { let sql = sql.into(); + self.traced(SqliteOperation::Run, self.run_untraced(sql, params)) + .await + } + + async fn run_untraced( + &self, + sql: String, + params: Option>, + ) -> Result { let sql_for_log = sql.clone(); let binding_count = bind_param_count(¶ms); #[cfg(feature = "sqlite-local")] @@ -685,6 +743,15 @@ impl SqliteDb { params: Option>, ) -> Result { let sql = sql.into(); + self.traced(SqliteOperation::Execute, self.execute_untraced(sql, params)) + .await + } + + async fn execute_untraced( + &self, + sql: String, + params: Option>, + ) -> Result { let sql_for_log = sql.clone(); let binding_count = bind_param_count(¶ms); #[cfg(feature = "sqlite-local")] @@ -736,6 +803,17 @@ impl SqliteDb { pub async fn execute_batch( &self, statements: Vec, + ) -> Result> { + self.traced( + SqliteOperation::ExecuteBatch, + self.execute_batch_untraced(statements), + ) + .await + } + + async fn execute_batch_untraced( + &self, + statements: Vec, ) -> Result> { let statement_count = statements.len(); let binding_count = statements @@ -749,7 +827,11 @@ impl SqliteDb { } } else { async { - let transaction = self.begin_transaction(None).await?; + let transaction = self + .clone() + .with_invocation_telemetry(None) + .begin_transaction(None) + .await?; let mut results = Vec::with_capacity(statements.len()); for statement in statements { match transaction.execute(statement.sql, statement.params).await { diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/tx.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/tx.rs index 1f011fcff7..1102395883 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/tx.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/tx.rs @@ -22,6 +22,7 @@ use tokio_util::sync::CancellationToken; #[cfg(not(target_arch = "wasm32"))] use crate::runtime::RuntimeSpawner; +use crate::telemetry::SqliteOperation; #[cfg(feature = "sqlite-local")] use super::profiling::{FINGERPRINT_FORMAT_VERSION, TransactionProfile}; @@ -277,7 +278,7 @@ impl SqliteDb { } let db = self.clone(); - run_detached_transaction_task( + let task = run_detached_transaction_task( async move { db.begin_transaction_profiled_inner( key, @@ -289,8 +290,8 @@ impl SqliteDb { .await }, "sqlite transaction begin task failed", - ) - .await + ); + self.traced(SqliteOperation::TransactionBegin, task).await } #[cfg(test)] @@ -495,11 +496,11 @@ impl SqliteDb { async fn transaction_exec(&self, key: &str, sql: String) -> Result { let db = self.clone(); let key = key.to_owned(); - run_detached_transaction_task( + let task = run_detached_transaction_task( async move { db.transaction_exec_inner(&key, sql).await }, "sqlite transaction exec task failed", - ) - .await + ); + self.traced(SqliteOperation::TransactionExec, task).await } async fn transaction_exec_inner(&self, key: &str, sql: String) -> Result { @@ -547,11 +548,11 @@ impl SqliteDb { ) -> Result { let db = self.clone(); let key = key.to_owned(); - run_detached_transaction_task( + let task = run_detached_transaction_task( async move { db.transaction_execute_inner(&key, sql, params).await }, "sqlite transaction execute task failed", - ) - .await + ); + self.traced(SqliteOperation::TransactionExecute, task).await } async fn transaction_execute_inner( @@ -616,11 +617,16 @@ impl SqliteDb { async fn finish_transaction(&self, key: &str, commit: bool) -> Result<()> { let db = self.clone(); let key = key.to_owned(); - run_detached_transaction_task( + let operation = if commit { + SqliteOperation::TransactionCommit + } else { + SqliteOperation::TransactionRollback + }; + let task = run_detached_transaction_task( async move { db.finish_transaction_inner(&key, commit).await }, "sqlite transaction finish task failed", - ) - .await + ); + self.traced(operation, task).await } async fn finish_transaction_inner(&self, key: &str, commit: bool) -> Result<()> { diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs index c3cdbe3d22..e5bcdbcf41 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs @@ -57,6 +57,7 @@ use crate::actor::task_types::ShutdownKind; use crate::actor::work_registry::ActorWorkKind; use crate::error::{ActorLifecycle as ActorLifecycleError, ActorRuntime}; use crate::runtime::RuntimeSpawner; +use crate::telemetry::{ActorInvocation, IncomingInvocationContext}; #[cfg(test)] use crate::time::sleep; use crate::time::{Instant, sleep_until, timeout}; @@ -201,12 +202,14 @@ pub enum DispatchCommand { Action { name: String, args: Vec, + incoming: crate::telemetry::IncomingInvocationContext, conn: ConnHandle, reply: oneshot::Sender>>, }, QueueSend { name: String, body: Vec, + incoming: crate::telemetry::IncomingInvocationContext, conn: ConnHandle, request: Request, wait: bool, @@ -902,9 +905,12 @@ impl ActorTask { DispatchCommand::Action { name, args, + incoming, conn, reply, } => { + let invocation = ActorInvocation::start_action(&self.ctx, &name, incoming); + let invocation_telemetry = invocation.telemetry(); tracing::info!( actor_id = %self.ctx.actor_id(), action_name = %name, @@ -921,6 +927,7 @@ impl ActorTask { args, conn: Some(conn), scheduled_fire: None, + invocation_telemetry: Some(invocation_telemetry), reply: Reply::from(tracked_reply_tx), }, ) { @@ -933,32 +940,23 @@ impl ActorTask { self.log_dispatch_command_handled(command_kind, "enqueued"); let actor_id = self.ctx.actor_id().to_owned(); let ctx = self.ctx.clone(); - self.ctx.spawn_work(ActorWorkKind::Action, async move { - match tracked_reply_rx.await { - Ok(result) => { - let result = - result.map_err(|error| ctx.attach_actor_to_error(error)); - tracing::info!( - actor_id = %actor_id, - action_name = %action_name_for_log, - ok = result.is_ok(), - "actor task: tracked reply received, forwarding" - ); - let _ = reply.send(result); - } - Err(_) => { - tracing::warn!( - actor_id = %actor_id, - action_name = %action_name_for_log, - "actor task: tracked reply dropped before completion" - ); - let error = ctx.attach_actor_to_error( - ActorLifecycleError::DroppedReply.build(), - ); - let _ = reply.send(Err(error)); - } - } - }); + self.forward_tracked_reply( + ActorWorkKind::Action, + tracked_reply_rx, + reply, + move |result| { + let result = + result.map_err(|error| ctx.attach_actor_to_error(error)); + tracing::info!( + actor_id = %actor_id, + action_name = %action_name_for_log, + ok = result.is_ok(), + "actor task: tracked reply received, forwarding" + ); + invocation.finish(result.as_ref().err()); + result + }, + ); } Err(error) => { tracing::warn!( @@ -967,7 +965,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"); } } @@ -975,42 +975,76 @@ impl ActorTask { DispatchCommand::QueueSend { name, body, + incoming, conn, request, wait, timeout_ms, reply, - } => match self.send_actor_event( - "dispatch_queue_send", - ActorEvent::QueueSend { - name, - body, - conn, - request, - wait, - timeout_ms, - reply: Reply::from(reply), - }, - ) { - Ok(()) => { - self.log_dispatch_command_handled(command_kind, "enqueued"); - } - Err(_error) => { - self.log_dispatch_command_handled(command_kind, "enqueue_failed"); + } => { + let invocation = ActorInvocation::start_queue_send(&self.ctx, &name, incoming); + let invocation_telemetry = invocation.telemetry(); + let (tracked_reply_tx, tracked_reply_rx) = oneshot::channel(); + match self.send_actor_event( + "dispatch_queue_send", + ActorEvent::QueueSend { + name, + body, + conn, + request, + wait, + timeout_ms, + invocation_telemetry: Some(invocation_telemetry), + reply: Reply::from(tracked_reply_tx), + }, + ) { + Ok(()) => { + self.log_dispatch_command_handled(command_kind, "enqueued"); + self.forward_tracked_reply( + ActorWorkKind::DispatchReply, + tracked_reply_rx, + reply, + move |result| { + invocation.finish(result.as_ref().err()); + result + }, + ); + } + Err(error) => { + invocation.finish(Some(&error)); + let _ = reply.send(Err(error)); + self.log_dispatch_command_handled(command_kind, "enqueue_failed"); + } } - }, + } DispatchCommand::Http { request, reply } => { + let incoming = IncomingInvocationContext::from_http_headers(request.headers()); + let invocation = ActorInvocation::start_request(&self.ctx, &request, incoming); + let invocation_telemetry = invocation.telemetry(); + let (tracked_reply_tx, tracked_reply_rx) = oneshot::channel(); match self.send_actor_event( "dispatch_http", ActorEvent::HttpRequest { request, - reply: Reply::from(reply), + invocation_telemetry: Some(invocation_telemetry), + reply: Reply::from(tracked_reply_tx), }, ) { Ok(()) => { self.log_dispatch_command_handled(command_kind, "enqueued"); + self.forward_tracked_reply( + ActorWorkKind::DispatchReply, + tracked_reply_rx, + reply, + move |result| { + invocation.finish_request(result.as_ref()); + result + }, + ); } - Err(_error) => { + Err(error) => { + invocation.finish(Some(&error)); + let _ = reply.send(Err(error)); self.log_dispatch_command_handled(command_kind, "enqueue_failed"); } } @@ -1072,6 +1106,35 @@ impl ActorTask { } } + /// Waits for the runtime adapter's reply to one dispatched invocation and + /// forwards it to the caller. `on_reply` sees the reply first, finishes the + /// invocation, and returns what the caller gets. A reply channel closed + /// without an answer counts as a dropped reply, so every arm reports that + /// case the same way. + fn forward_tracked_reply( + &self, + kind: ActorWorkKind, + tracked_reply_rx: oneshot::Receiver>, + reply: oneshot::Sender>, + on_reply: impl FnOnce(Result) -> Result + Send + 'static, + ) { + let actor_id = self.ctx.actor_id().to_owned(); + self.ctx.spawn_work(kind, async move { + let result = match tracked_reply_rx.await { + Ok(result) => result, + Err(_) => { + tracing::warn!( + actor_id = %actor_id, + "actor task: tracked reply dropped before completion; the runtime adapter dropped its reply handle" + ); + Err(ActorLifecycleError::DroppedReply.build()) + } + }; + let result = on_reply(result); + let _ = reply.send(result); + }); + } + fn log_dispatch_command_handled(&self, command: &'static str, outcome: &'static str) { tracing::debug!( actor_id = %self.ctx.actor_id(), diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/work_registry.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/work_registry.rs index df3bfabf26..028042ef0a 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/work_registry.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/work_registry.rs @@ -17,6 +17,8 @@ use crate::actor::task_types::UserTaskKind; pub enum ActorWorkKind { /// Dispatched action work awaiting its reply from the runtime adapter. Action, + /// Non-action invocation work awaiting its reply from the runtime adapter. + DispatchReply, /// User work that keeps the actor out of idle sleep while it runs. KeepAwake, /// Runtime-owned work that should behave like keep-awake without exposing a user API. @@ -54,6 +56,12 @@ impl ActorWorkKind { aborts_at_shutdown_deadline: true, user_task_kind: Some(UserTaskKind::Action), }, + ActorWorkKind::DispatchReply => ActorWorkPolicy { + blocks_idle_sleep: true, + drains_shutdown_grace: true, + aborts_at_shutdown_deadline: true, + user_task_kind: None, + }, ActorWorkKind::KeepAwake => ActorWorkPolicy { blocks_idle_sleep: true, drains_shutdown_grace: true, @@ -97,6 +105,7 @@ impl ActorWorkKind { pub(crate) fn label(self) -> &'static str { match self { ActorWorkKind::Action => "action", + ActorWorkKind::DispatchReply => "dispatch_reply", ActorWorkKind::KeepAwake => "keep_awake", ActorWorkKind::InternalKeepAwake => "internal_keep_awake", ActorWorkKind::WaitUntil => "wait_until", diff --git a/rivetkit-rust/packages/rivetkit-core/src/lib.rs b/rivetkit-rust/packages/rivetkit-core/src/lib.rs index ae6fa5a94e..644bc66f4e 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/lib.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/lib.rs @@ -16,6 +16,13 @@ pub mod registry; 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::{ + ActorInvocationSpanContext, ActorInvocationTelemetry, ActorInvocationTraceContext, + OutboundCallInvocation, TraceOrigin, +}; #[cfg(feature = "native-runtime")] pub mod serverless_http; #[cfg(feature = "native-runtime")] @@ -124,7 +131,7 @@ pub use actor::{kv, sqlite}; pub use actor::action::ActionDispatchError; pub use actor::config::{ ActionDefinition, ActorConfig, ActorConfigInput, ActorConfigOverrides, CanHibernateWebSocket, - SqliteProfilingConfig, SqliteProfilingConfigInput, + QueueDefinition, SqliteProfilingConfig, SqliteProfilingConfigInput, }; pub use actor::connection::ConnHandle; pub use actor::context::{ diff --git a/rivetkit-rust/packages/rivetkit-core/src/registry/dispatch.rs b/rivetkit-rust/packages/rivetkit-core/src/registry/dispatch.rs index 8a38fcb0d1..7640a23f59 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/registry/dispatch.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/registry/dispatch.rs @@ -7,6 +7,7 @@ pub(super) async fn dispatch_action_through_task( conn: ConnHandle, name: String, args: Vec, + incoming: crate::telemetry::IncomingInvocationContext, ) -> std::result::Result, ActionDispatchError> { let (reply_tx, reply_rx) = oneshot::channel(); try_send_dispatch_command( @@ -14,6 +15,7 @@ pub(super) async fn dispatch_action_through_task( DispatchCommand::Action { name, args, + incoming, conn, reply: reply_tx, }, diff --git a/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs b/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs index 2d0d566bc1..ae1b12664a 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs @@ -255,6 +255,7 @@ impl RegistryDispatcher { conn.clone(), action_name.clone(), args, + crate::telemetry::IncomingInvocationContext::from_http_headers(request.headers()), ), ) .await; @@ -368,12 +369,15 @@ impl RegistryDispatcher { } }; + let incoming = + crate::telemetry::IncomingInvocationContext::from_http_headers(request.headers()); let (reply_tx, reply_rx) = oneshot::channel(); let dispatch_result = try_send_dispatch_command( &instance.dispatch, DispatchCommand::QueueSend { name: queue_name, body: queue_request.body, + incoming, conn: conn.clone(), request, wait: queue_request.wait, diff --git a/rivetkit-rust/packages/rivetkit-core/src/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..06c87d2a09 --- /dev/null +++ b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs @@ -0,0 +1,878 @@ +//! Internal OpenTelemetry spans owned by the actor runtime. + +#[cfg(feature = "native-runtime")] +pub mod export; + +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + +use opentelemetry::Context; +use opentelemetry::propagation::{Extractor, TextMapPropagator as _}; +use opentelemetry::trace::{SpanContext, TraceContextExt as _}; +use opentelemetry_http::{HeaderExtractor, HeaderInjector}; +use opentelemetry_sdk::propagation::TraceContextPropagator; +use parking_lot::Mutex; +use rivetkit_client_protocol::ray_id::RayId; +use tracing_opentelemetry::OpenTelemetrySpanExt as _; + +use crate::ActorContext; +use crate::actor::metrics::{ActorMetrics, InvocationStatus, InvocationType}; +use crate::actor::queue::QueueMessage; +use crate::time::Instant; + +/// Correlation fields accepted at an invocation boundary. +#[derive(Debug, Default)] +pub struct IncomingInvocationContext { + pub(crate) ray_id: Option, + remote_parent: Option, +} + +pub(crate) use rivetkit_client_protocol::ray_id::HEADER_RIVET_RAY_ID; + +const HEADER_TRACEPARENT: &str = "traceparent"; +const HEADER_TRACESTATE: &str = "tracestate"; + +impl IncomingInvocationContext { + /// Reads the ray ID and W3C trace context an HTTP request carries. Every + /// HTTP entry point into an actor reads them through here, so they all + /// apply the same bounds. + pub(crate) fn from_http_headers(headers: &http::HeaderMap) -> Self { + Self { + ray_id: invocation_ray_id(headers), + remote_parent: if headers.contains_key(HEADER_TRACEPARENT) { + extract_remote_parent(&HeaderExtractor(headers)) + } else { + None + }, + } + } +} + +/// Reads the caller's ray ID. The header is untrusted, so it is bounded by +/// the rule shared with the clients that send it; anything else counts as +/// absent. +fn invocation_ray_id(headers: &http::HeaderMap) -> Option { + let value = headers.get(HEADER_RIVET_RAY_ID)?.to_str().ok()?; + RayId::parse(value.to_owned()).ok().map(RayId::into_string) +} + +/// Name a request invocation is reported under, in place of an action name. +/// It cannot collide with an action, because the metric and span carry the +/// invocation type beside it. +const REQUEST_INVOCATION_NAME: &str = "onRequest"; + +/// Name a queue send invocation is reported under. The queue itself is an +/// attribute, so one series covers every queue of an actor. +const QUEUE_SEND_INVOCATION_NAME: &str = "queue.send"; + +/// What an invocation ran, which decides its name and the attributes that +/// identify it on the span. +#[derive(Clone, Copy, Debug)] +enum InvocationSubject<'a> { + Action(&'a str), + Request { method: &'a str }, + QueueSend { queue: &'a str }, +} + +/// Owns the complete lifecycle of one actor invocation. +#[derive(Debug)] +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. +/// +/// The second field is the application span the host runtime had active when +/// it resolved this handle. Core cannot see the host's span stack, so a span +/// Core opens through this handle parents there when it is set and to the +/// invocation span otherwise. Every clone of a handle shares one invocation. +#[doc(hidden)] +#[derive(Clone, Debug)] +pub struct ActorInvocationTelemetry(Arc, Option); + +/// Identity fields that do not change while an actor is alive. Built once per +/// actor and shared by every invocation, so starting one does not re-allocate +/// 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. `finished` lets exactly one of the finish and +/// drop paths record the terminal status, and marks the invocation closed even +/// when tracing is off and there is no span. The span slot is emptied, which +/// is what exports the span, once the status is recorded and no work handed +/// to `wait_until` from this invocation is still running. `pending_work` +/// counts that work, so the invocation stays active for it after the reply. +#[derive(Debug)] +struct InvocationInner { + ray_id: Option, + // Forced-sync: the slot is emptied from `Drop` paths and sync accessors, + // and the guard is never held across an await. + span: Mutex>, + finished: AtomicBool, + pending_work: AtomicUsize, + identity: Arc, +} + +const OPERATION_ABANDONED_ERROR_TYPE: &str = "actor.operation_abandoned"; + +/// Keeps an invocation open while one piece of `wait_until` work runs. The +/// span is released when the last guard drops after the terminal status has +/// been recorded, so work that settles before the reply changes nothing. +pub(crate) struct InvocationWorkGuard(ActorInvocationTelemetry); + +/// Where later work came from: the ray ID of the invocation that caused it and +/// the span that was active there. Persisted beside schedules and queue +/// messages so the work they cause can link back to its origin. +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct TraceOrigin { + pub ray_id: Option, + pub traceparent: Option, + pub tracestate: Option, +} + +impl TraceOrigin { + /// True when there is nothing to persist: the work was caused outside any + /// traced invocation. + pub fn is_empty(&self) -> bool { + self.ray_id.is_none() && self.traceparent.is_none() && self.tracestate.is_none() + } +} + +/// Active actor invocation fields exposed to foreign-runtime adapters. +#[doc(hidden)] +#[derive(Clone, Debug)] +pub struct ActorInvocationTraceContext { + pub ray_id: Option, + /// Present only while the invocation runs inside a valid span. + pub span: Option, +} + +/// W3C span context of the current invocation span. A span context is either +/// complete or absent, so these fields are never optional individually. +#[doc(hidden)] +#[derive(Clone, Debug)] +pub struct ActorInvocationSpanContext { + pub trace_id: String, + pub span_id: String, + pub trace_flags: u8, + pub tracestate: Option, +} + +/// 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, +} + +/// 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, + action_name: &str, + incoming: IncomingInvocationContext, + ) -> Self { + Self::start( + ctx, + InvocationSubject::Action(action_name), + InvocationType::Action, + incoming.ray_id, + incoming.remote_parent, + None, + ) + } + + pub(crate) fn start_scheduled( + ctx: &ActorContext, + action_name: &str, + origin: TraceOrigin, + ) -> Self { + let origin_parent = + parse_remote_parent(origin.traceparent.as_deref(), origin.tracestate.as_deref()); + Self::start( + ctx, + InvocationSubject::Action(action_name), + InvocationType::Scheduled, + origin.ray_id, + None, + origin_parent, + ) + } + + /// Starts the invocation for one message sent into `queue_name` from + /// outside the actor. It ends when the send is acknowledged, or when the + /// sender's wait for a completion ends. + pub(crate) fn start_queue_send( + ctx: &ActorContext, + queue_name: &str, + incoming: IncomingInvocationContext, + ) -> Self { + Self::start( + ctx, + InvocationSubject::QueueSend { queue: queue_name }, + InvocationType::QueueSend, + incoming.ray_id, + incoming.remote_parent, + None, + ) + } + + /// Starts the invocation for one raw HTTP request served by `onRequest`. + /// The span is named after the handler rather than the path, because a + /// path is caller-supplied and would make the name a cardinality surface. + pub(crate) fn start_request( + ctx: &ActorContext, + request: &crate::actor::messages::Request, + incoming: IncomingInvocationContext, + ) -> Self { + Self::start( + ctx, + InvocationSubject::Request { + method: request.method().as_str(), + }, + InvocationType::Request, + incoming.ray_id, + incoming.remote_parent, + None, + ) + } + + fn start( + ctx: &ActorContext, + subject: InvocationSubject<'_>, + invocation_type: InvocationType, + ray_id: Option, + parent: Option, + link: Option, + ) -> Self { + let identity = ctx.telemetry_identity(); + // Use bounded names for both spans and metrics. + let (action_name, http_method, queue_name) = match subject { + InvocationSubject::Action(name) => (ctx.metrics().label_action_name(name), None, None), + InvocationSubject::Request { method } => (REQUEST_INVOCATION_NAME, Some(method), None), + InvocationSubject::QueueSend { queue } => ( + QUEUE_SEND_INVOCATION_NAME, + None, + Some(ctx.metrics().label_queue_name(queue)), + ), + }; + let action_name = action_name.to_owned(); + let span = if tracing::enabled!(target: "rivetkit::telemetry", tracing::Level::INFO) { + let span = tracing::info_span!( + target: "rivetkit::telemetry", + parent: None, + "rivet.actor.invoke", + otel.name = %format!("{}/{}", identity.actor_name, action_name), + otel.kind = invocation_type.otel_kind(), + rivet.invocation.type = invocation_type.as_label(), + rivet.actor.id = %identity.actor_id, + rivet.actor.name = %identity.actor_name, + rivet.actor.key = %identity.actor_key, + rivet.action.name = tracing::field::Empty, + rivet.ray.id = tracing::field::Empty, + http.request.method = tracing::field::Empty, + http.response.status_code = tracing::field::Empty, + rivet.queue.name = tracing::field::Empty, + otel.status_code = tracing::field::Empty, + error.type = tracing::field::Empty, + ); + span.record("rivet.ray.id", ray_id.as_deref()); + match (http_method, queue_name) { + (Some(method), _) => span.record("http.request.method", method), + (None, Some(queue)) => span.record("rivet.queue.name", queue), + (None, None) => span.record("rivet.action.name", &action_name), + }; + if let Some(parent) = parent { + span.set_parent(opentelemetry::Context::new().with_remote_span_context(parent)); + } + if let Some(link) = link { + span.add_link(link); + } + Some(span) + } else { + None + }; + + Self { + telemetry: ActorInvocationTelemetry::new(ray_id, span, identity), + metrics: ctx.metrics().clone(), + action_name, + invocation_type, + started_at: Instant::now(), + } + } + + pub(crate) fn telemetry(&self) -> ActorInvocationTelemetry { + self.telemetry.clone() + } + + pub(crate) fn finish(mut self, error: Option<&anyhow::Error>) { + self.finish_with_status( + error.map_or(InvocationStatus::Ok, InvocationStatus::from_error), + error, + ); + } + + /// Finishes a request invocation with the HTTP status the handler + /// answered, which is recorded on the span beside the outcome. A 5xx + /// answer counts as a failed invocation with the status as its error + /// identity, following the HTTP server span convention. An error means no + /// response was produced, so only the error identity is recorded. + pub(crate) fn finish_request( + mut self, + response: std::result::Result<&crate::actor::messages::ActorHttpResponse, &anyhow::Error>, + ) { + match response { + Ok(response) => { + let status = response.status(); + self.telemetry.record_http_status(status); + if status >= 500 { + self.finish_with_failure(InvocationFailure::HttpStatus(status)); + } else { + self.finish_with_status(InvocationStatus::Ok, None); + } + } + Err(error) => self.finish(Some(error)), + } + } + + fn finish_with_status(&mut self, status: InvocationStatus, error: Option<&anyhow::Error>) { + let Some(span) = self.telemetry.claim_terminal() else { + return; + }; + self.record_finished(span, status, error.map(InvocationFailure::Error)); + } + + fn finish_with_failure(&mut self, failure: InvocationFailure<'_>) { + let Some(span) = self.telemetry.claim_terminal() else { + return; + }; + self.record_finished(span, InvocationStatus::Error, Some(failure)); + } + + /// Records the terminal metric and span status of an invocation whose + /// completion the caller has already claimed through `claim_terminal`. + /// The metric measures what the caller waited for, so it is recorded here + /// even when `wait_until` work keeps the span open past this point. + fn record_finished( + &self, + span: Option, + status: InvocationStatus, + failure: Option>, + ) { + self.metrics.record_invocation( + &self.action_name, + self.invocation_type, + status, + self.started_at.elapsed(), + ); + if let Some(span) = span { + match failure { + Some(InvocationFailure::HttpStatus(status)) => { + span.record("otel.status_code", "ERROR"); + span.record("error.type", status.to_string()); + } + Some(InvocationFailure::Error(error)) => record_outcome(&span, Some(error)), + None => record_outcome(&span, None), + } + self.telemetry.mark_reply_sent(&span); + } + self.telemetry.release_span_if_settled(); + } +} + +/// Why an invocation is recorded as failed: an error crossing the runtime +/// boundary, or a request the handler answered with a server error status. +enum InvocationFailure<'a> { + Error(&'a anyhow::Error), + HttpStatus(u16), +} + +impl Drop for ActorInvocation { + fn drop(&mut self) { + // `finish` consumes the invocation, so this runs on the completed path + // too. Claim the terminal record first, so the dropped-reply error is + // only built for an invocation that really was dropped. + let Some(span) = self.telemetry.claim_terminal() else { + return; + }; + let error = crate::error::ActorLifecycle::DroppedReply.build(); + self.record_finished( + span, + InvocationStatus::Dropped, + Some(InvocationFailure::Error(&error)), + ); + } +} + +impl ActorInvocationTelemetry { + fn new( + ray_id: Option, + span: Option, + identity: Arc, + ) -> Self { + Self( + Arc::new(InvocationInner { + ray_id, + span: Mutex::new(span), + finished: AtomicBool::new(false), + pending_work: AtomicUsize::new(0), + identity, + }), + None, + ) + } + + /// Returns a handle for the same invocation whose spans parent to the + /// application span identified by `traceparent` and `tracestate`. Invalid + /// or absent context yields a handle that parents to the invocation span. + pub(crate) fn with_application_span( + &self, + traceparent: Option<&str>, + tracestate: Option<&str>, + ) -> Self { + Self(self.0.clone(), parse_remote_parent(traceparent, tracestate)) + } + + /// Records the status a request invocation answered with. + fn record_http_status(&self, status: u16) { + if let Some(span) = self.0.span.lock().as_ref() { + span.record("http.response.status_code", status); + } + } + + /// Registers work that outlives the reply, so the invocation span stays + /// open and keeps parenting operations until the returned guard drops. + pub(crate) fn hold_open(&self) -> InvocationWorkGuard { + self.0.pending_work.fetch_add(1, Ordering::SeqCst); + InvocationWorkGuard(self.clone()) + } + + /// Returns correlation fields only while this actor invocation is active. + #[doc(hidden)] + pub fn trace_context(&self) -> Option { + let active = self.active()?; + let span = active + .span + .lock() + .clone() + .and_then(|span| span_context_of(&span)); + + Some(ActorInvocationTraceContext { + ray_id: active.ray_id.clone(), + span, + }) + } + + /// Trace origin work caused by this invocation records: the invocation's + /// ray ID, and the application span active in the host runtime at that + /// moment, or the invocation span when there was none. Work that links + /// back to it then points at the code that caused it rather than at the + /// whole invocation around that code. + pub(crate) fn trace_origin(&self) -> TraceOrigin { + let Some(active) = self.active() else { + return TraceOrigin::default(); + }; + let span = self + .1 + .clone() + .or_else(|| active.span.lock().as_ref().and_then(otel_span_context_of)); + let (traceparent, tracestate) = span + .as_ref() + .and_then(propagation_headers) + .map_or((None, None), |(traceparent, tracestate)| { + (Some(traceparent), tracestate) + }); + TraceOrigin { + ray_id: active.ray_id.clone(), + traceparent, + tracestate, + } + } + + /// Opens the span covering one call out to another actor. + /// + /// The callee parents to this span rather than to the invocation making the + /// call, so the time spent reaching it, which includes routing and waking a + /// sleeping actor, is attributed to the call instead of falling in the gap + /// between the two invocations. + /// When this handle carries the application span active in the host + /// runtime, the call span parents there instead, which is what puts a + /// callee under the application span that issued the call rather than + /// beside it. + pub(crate) fn start_outbound_call( + &self, + actor_name: &str, + action_name: &str, + ) -> Option { + let invocation_span = self.active()?.span.lock().clone()?; + let span = tracing::info_span!( + target: "rivetkit::telemetry", + parent: &invocation_span, + "rivet.actor.call", + otel.name = %format!("{actor_name}/{action_name}"), + otel.kind = "client", + // Omit rivet.invocation.type: this measures the caller waiting, not the callee running. + rivet.actor.name = %actor_name, + rivet.action.name = %action_name, + rivet.ray.id = self.0.ray_id.as_deref(), + otel.status_code = tracing::field::Empty, + error.type = tracing::field::Empty, + ); + if let Some(application_span) = &self.1 { + span.set_parent( + opentelemetry::Context::new().with_remote_span_context(application_span.clone()), + ); + } + let context = span_context_of(&span); + Some(OutboundCallInvocation { + span: Some(span), + context, + }) + } + + pub(crate) fn start_sqlite(&self, operation: SqliteOperation) -> Option { + let parent = self.active()?.span.lock().clone()?; + let span = tracing::info_span!( + target: "rivetkit::telemetry", + parent: &parent, + "rivet.sqlite.operation", + otel.name = operation.span_name(), + otel.kind = "internal", + rivet.operation.system = "sqlite", + rivet.operation.name = operation.as_str(), + rivet.ray.id = self.0.ray_id.as_deref(), + rivet.actor.id = %self.0.identity.actor_id, + rivet.actor.name = %self.0.identity.actor_name, + rivet.actor.key = %self.0.identity.actor_key, + otel.status_code = tracing::field::Empty, + error.type = tracing::field::Empty, + ); + if let Some(application_span) = &self.1 { + span.set_parent( + opentelemetry::Context::new().with_remote_span_context(application_span.clone()), + ); + } + Some(SqliteOperationSpan { span: Some(span) }) + } + + /// Borrows the invocation while it is still open: before its status is + /// recorded, or after it while `wait_until` work from it still runs. A + /// settled invocation yields nothing, so late SQLite work and retained + /// handles cannot attach to a span that has already ended. + fn active(&self) -> Option<&InvocationInner> { + let open = !self.0.finished.load(Ordering::SeqCst) + || self.0.pending_work.load(Ordering::SeqCst) > 0; + if open { Some(&*self.0) } else { None } + } + + /// Claims the terminal record, so the finish and drop paths cannot both + /// record a status for the same invocation. The span stays in its slot + /// until `release_span_if_settled` empties it. + fn claim_terminal(&self) -> Option> { + if self.0.finished.swap(true, Ordering::SeqCst) { + return None; + } + Some(self.0.span.lock().clone()) + } + + /// Marks the moment the caller got its answer when the span will outlive + /// it, so the reply point stays visible inside a span that is still + /// running `wait_until` work. + fn mark_reply_sent(&self, span: &tracing::Span) { + if self.0.pending_work.load(Ordering::SeqCst) > 0 { + tracing::info!(target: "rivetkit::telemetry", parent: span, "reply sent"); + } + } + + /// Ends the span, which exports it, unless `wait_until` work is still + /// holding the invocation open. The last guard to drop ends it instead. + /// Both sides set their flag before reading the other's, so the two + /// cannot each see the other as still pending and leave the span behind. + fn release_span_if_settled(&self) { + if self.0.pending_work.load(Ordering::SeqCst) == 0 { + self.0.span.lock().take(); + } + } +} + +impl Drop for InvocationWorkGuard { + fn drop(&mut self) { + let inner = &self.0.0; + let was_last = inner.pending_work.fetch_sub(1, Ordering::SeqCst) == 1; + if was_last && inner.finished.load(Ordering::SeqCst) { + inner.span.lock().take(); + } + } +} + +impl OutboundCallInvocation { + /// W3C context of this call's span, to send to the callee so it parents + /// here. Absent when tracing is disabled. + pub fn span_context(&self) -> Option { + self.context.clone() + } + + /// Records the call's outcome. `error` is the failure the callee returned, + /// and its group and code become the span's `error.type`. + pub fn finish(mut self, error: Option<&anyhow::Error>) { + let Some(span) = self.span.take() else { + return; + }; + record_outcome(&span, error); + } +} + +impl Drop for OutboundCallInvocation { + fn drop(&mut self) { + let Some(span) = self.span.take() else { + return; + }; + span.record("otel.status_code", "ERROR"); + span.record("error.type", OPERATION_ABANDONED_ERROR_TYPE); + } +} + +impl SqliteOperationSpan { + pub(crate) fn span(&self) -> tracing::Span { + self.span.as_ref().expect("sqlite span is present").clone() + } + + pub(crate) fn finish(&mut self, error: Option<&anyhow::Error>) { + let Some(span) = self.span.take() else { + return; + }; + record_outcome(&span, error); + } +} + +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", OPERATION_ABANDONED_ERROR_TYPE); + } +} + +/// W3C fields of a span context, or nothing when it is not valid and so +/// carries nothing worth propagating. +fn span_context_fields(span_context: &SpanContext) -> Option { + if !span_context.is_valid() { + return None; + } + let tracestate = span_context.trace_state().header(); + Some(ActorInvocationSpanContext { + trace_id: span_context.trace_id().to_string(), + span_id: span_context.span_id().to_string(), + trace_flags: span_context.trace_flags().to_u8(), + tracestate: (!tracestate.is_empty()).then_some(tracestate), + }) +} + +/// Reads a span's W3C context, or nothing when it carries no valid context to +/// propagate. +fn span_context_of(span: &tracing::Span) -> Option { + otel_span_context_of(span) + .as_ref() + .and_then(span_context_fields) +} + +fn otel_span_context_of(span: &tracing::Span) -> Option { + let context = span.context(); + let context_span = context.span(); + let span_context = context_span.span_context(); + span_context.is_valid().then(|| span_context.clone()) +} + +/// Opens the span covering the moment one queue message is handed to the +/// actor. It links to the span that sent the message, which is what connects +/// the consumer's trace to the sender's. Inside an invocation it sits under +/// that invocation's application span, or the invocation span, and carries its +/// ray ID. Outside one, as from the run handler, it is a root span carrying the +/// ray ID the message was sent under. It closes when the caller drops it, which +/// `try_receive_batch` does as it hands the message back. +pub(crate) fn start_queue_receive(ctx: &ActorContext, message: &QueueMessage) -> tracing::Span { + if !tracing::enabled!(target: "rivetkit::telemetry", tracing::Level::INFO) { + return tracing::Span::none(); + } + let identity = ctx.telemetry_identity(); + let invocation = ctx + .1 + .as_ref() + .and_then(|telemetry| telemetry.active().map(|inner| (telemetry, inner))); + let span = tracing::info_span!( + target: "rivetkit::telemetry", + parent: None, + "rivet.queue.receive", + otel.name = %format!("{}/queue.receive", identity.actor_name), + otel.kind = "consumer", + rivet.actor.id = %identity.actor_id, + rivet.actor.name = %identity.actor_name, + rivet.actor.key = %identity.actor_key, + rivet.queue.name = ctx.metrics().label_queue_name(&message.name), + rivet.ray.id = tracing::field::Empty, + ); + match invocation { + Some((telemetry, inner)) => { + span.record("rivet.ray.id", inner.ray_id.as_deref()); + let parent = match &telemetry.1 { + Some(application_span) => Some( + opentelemetry::Context::new() + .with_remote_span_context(application_span.clone()), + ), + None => inner.span.lock().as_ref().map(|span| span.context()), + }; + if let Some(parent) = parent { + span.set_parent(parent); + } + } + None => { + if let Some(ray_id) = &message.trace_origin.ray_id { + span.record("rivet.ray.id", ray_id); + } + } + } + if let Some(link) = parse_remote_parent( + message.trace_origin.traceparent.as_deref(), + message.trace_origin.tracestate.as_deref(), + ) { + span.add_link(link); + } + span +} + +/// Records the terminal status and error identity of a finished span. +fn record_outcome(span: &tracing::Span, error: Option<&anyhow::Error>) { + span.record( + "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())); + } +} + +fn parse_remote_parent(traceparent: Option<&str>, tracestate: Option<&str>) -> Option { + traceparent?; + extract_remote_parent(&TraceHeaders { + traceparent, + tracestate, + }) +} + +fn extract_remote_parent(extractor: &dyn Extractor) -> Option { + let context = TraceContextPropagator::new().extract_with_context(&Context::new(), extractor); + let span = context.span(); + let span_context = span.span_context(); + span_context.is_valid().then(|| span_context.clone()) +} + +fn propagation_headers(span_context: &SpanContext) -> Option<(String, Option)> { + let context = Context::new().with_remote_span_context(span_context.clone()); + let mut headers = http::HeaderMap::new(); + TraceContextPropagator::new().inject_context(&context, &mut HeaderInjector(&mut headers)); + let traceparent = headers.get(HEADER_TRACEPARENT)?.to_str().ok()?.to_owned(); + let tracestate = headers + .get(HEADER_TRACESTATE) + .and_then(|value| value.to_str().ok()) + .filter(|value| !value.is_empty()) + .map(str::to_owned); + Some((traceparent, tracestate)) +} + +struct TraceHeaders<'a> { + traceparent: Option<&'a str>, + tracestate: Option<&'a str>, +} + +impl Extractor for TraceHeaders<'_> { + fn get(&self, key: &str) -> Option<&str> { + match key { + key if key.eq_ignore_ascii_case(HEADER_TRACEPARENT) => self.traceparent, + key if key.eq_ignore_ascii_case(HEADER_TRACESTATE) => self.tracestate, + _ => None, + } + } + + fn keys(&self) -> Vec<&str> { + let mut keys = Vec::with_capacity(2); + if self.traceparent.is_some() { + keys.push(HEADER_TRACEPARENT); + } + if self.tracestate.is_some() { + keys.push(HEADER_TRACESTATE); + } + keys + } +} diff --git a/rivetkit-rust/packages/rivetkit-core/src/telemetry/export.rs b/rivetkit-rust/packages/rivetkit-core/src/telemetry/export.rs new file mode 100644 index 0000000000..80d926e67e --- /dev/null +++ b/rivetkit-rust/packages/rivetkit-core/src/telemetry/export.rs @@ -0,0 +1,144 @@ +//! Native OTLP export of the runtime's spans. +//! +//! Configuration comes entirely from the standard OpenTelemetry environment +//! variables, so every host that embeds core gets the same behavior by adding +//! [`layer`] to its subscriber and calling [`flush_best_effort`] on shutdown. +//! Core never installs a subscriber itself; which log layers surround the span +//! layer is the host's decision. + +use std::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 parking_lot::Mutex; +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: Mutex> = Mutex::new(None); + +/// Builds the span layer when standard OTel environment variables opt in, or +/// nothing when they do not. The layer only sees the runtime's own spans, so a +/// 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); + } + let mut stored_provider = PROVIDER.lock(); + if let Some(provider) = stored_provider.as_ref() { + return Ok(Some(provider.tracer("rivetkit"))); + } + + let exporter = match configured_protocol()? { + Protocol::Grpc => SpanExporter::builder() + .with_tonic() + .build() + .context("build otlp span exporter")?, + protocol @ (Protocol::HttpBinary | Protocol::HttpJson) => SpanExporter::builder() + .with_http() + .with_protocol(protocol) + .build() + .context("build otlp span exporter")?, + }; + let resource = Resource::builder() + // Leave service.version to the application; record the runtime version separately. + .with_attribute(KeyValue::new("rivetkit.version", env!("CARGO_PKG_VERSION"))) + .build(); + let provider = SdkTracerProvider::builder() + .with_resource(resource) + .with_batch_exporter(exporter) + .build(); + let tracer = provider.tracer("rivetkit"); + *stored_provider = Some(provider); + 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 provider = PROVIDER.lock().clone(); + let Some(provider) = provider else { + return; + }; + let flush = tokio::task::spawn_blocking(move || provider.force_flush()); + match tokio::time::timeout(FLUSH_TIMEOUT, flush).await { + Ok(Ok(Ok(()))) => {} + Ok(Ok(Err(error))) => tracing::warn!( + ?error, + "OpenTelemetry trace flush failed; queued spans may be lost" + ), + Ok(Err(error)) => tracing::warn!( + ?error, + "OpenTelemetry trace flush task failed; queued spans may be lost" + ), + Err(_) => tracing::warn!( + timeout = ?FLUSH_TIMEOUT, + "OpenTelemetry trace flush timed out; the collector may be slow or unreachable" + ), + } +} diff --git a/rivetkit-rust/packages/rivetkit-core/tests/context.rs b/rivetkit-rust/packages/rivetkit-core/tests/context.rs index 00d4b28582..6e8e070d12 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/context.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/context.rs @@ -1137,6 +1137,7 @@ mod moved_tests { args, conn, scheduled_fire, + invocation_telemetry: _, reply, } => { assert_eq!(name, "tick"); diff --git a/rivetkit-rust/packages/rivetkit-core/tests/integration/counter.rs b/rivetkit-rust/packages/rivetkit-core/tests/integration/counter.rs index a2c5b1247e..35caea7fd0 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/integration/counter.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/integration/counter.rs @@ -86,7 +86,11 @@ fn counter_factory() -> ActorFactory { ActorEvent::RunGracefulCleanup { reason: _, reply } => { reply.send(Ok(())); } - ActorEvent::HttpRequest { request: _, reply } => { + ActorEvent::HttpRequest { + request: _, + invocation_telemetry: _, + reply, + } => { reply.send(Err(anyhow::anyhow!("http requests are not handled"))); } ActorEvent::QueueSend { @@ -96,6 +100,7 @@ fn counter_factory() -> ActorFactory { request: _, wait: _, timeout_ms: _, + invocation_telemetry: _, reply, } => { reply.send(Err(anyhow::anyhow!("queue sends are not handled"))); diff --git a/rivetkit-rust/packages/rivetkit-core/tests/integration/sqlite_corruption_fuzz.rs b/rivetkit-rust/packages/rivetkit-core/tests/integration/sqlite_corruption_fuzz.rs index c8b9c2c5e6..7ebc9622e1 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/integration/sqlite_corruption_fuzz.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/integration/sqlite_corruption_fuzz.rs @@ -579,7 +579,11 @@ fn sqlite_fuzz_factory() -> ActorFactory { ActorEvent::RunGracefulCleanup { reason: _, reply } => { reply.send(Ok(())); } - ActorEvent::HttpRequest { request: _, reply } => { + ActorEvent::HttpRequest { + request: _, + invocation_telemetry: _, + reply, + } => { reply.send(Err(anyhow::anyhow!("http requests are not handled"))); } ActorEvent::QueueSend { @@ -589,6 +593,7 @@ fn sqlite_fuzz_factory() -> ActorFactory { request: _, wait: _, timeout_ms: _, + invocation_telemetry: _, reply, } => { reply.send(Err(anyhow::anyhow!("queue sends are not handled"))); diff --git a/rivetkit-rust/packages/rivetkit-core/tests/sql_efficiency.rs b/rivetkit-rust/packages/rivetkit-core/tests/sql_efficiency.rs index eedba335de..4aa43b35df 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(), @@ -376,6 +382,24 @@ fn query_catalog() -> Vec { params: vec![1_i64.into()], expectation: indexed(None, &["_rivet_queue"]), }, + QueryCase { + id: "queue.upsert_trace_context", + sql: internal_storage::UPSERT_QUEUE_TRACE_CONTEXT_SQL.into(), + params: vec![text("queue_trace_context:1"), Value::Blob(vec![1])], + expectation: indexed(None, &["_rivet_meta"]), + }, + QueryCase { + id: "queue.delete_trace_context", + sql: internal_storage::DELETE_QUEUE_TRACE_CONTEXT_SQL.into(), + params: vec![text("queue_trace_context:1")], + expectation: indexed(None, &["_rivet_meta"]), + }, + QueryCase { + id: "queue.reset_trace_contexts", + sql: internal_storage::RESET_QUEUE_TRACE_CONTEXTS_SQL.into(), + params: vec![], + expectation: indexed(None, &["_rivet_meta"]), + }, QueryCase { id: "queue.reset", sql: internal_storage::RESET_QUEUE_SQL.into(), @@ -488,6 +512,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 +597,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 +616,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(), diff --git a/rivetkit-rust/packages/rivetkit-core/tests/task.rs b/rivetkit-rust/packages/rivetkit-core/tests/task.rs index 9e8a56b954..f59f835f5e 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/task.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/task.rs @@ -1,4 +1,5 @@ pub(crate) mod moved_tests { + use anyhow::anyhow; use std::collections::{BTreeMap, HashMap}; use std::path::PathBuf; use std::process::Command; @@ -1927,6 +1928,7 @@ pub(crate) mod moved_tests { task.handle_dispatch(DispatchCommand::Action { name: "client-action".to_owned(), args: Vec::new(), + incoming: crate::telemetry::IncomingInvocationContext::default(), conn: client_conn, reply: reply_tx, }) @@ -2030,6 +2032,7 @@ pub(crate) mod moved_tests { task.handle_dispatch(DispatchCommand::Action { name: "slow-action".to_owned(), args: Vec::new(), + incoming: crate::telemetry::IncomingInvocationContext::default(), conn: client_conn, reply: reply_tx, }) @@ -3732,6 +3735,7 @@ pub(crate) mod moved_tests { .send(DispatchCommand::Action { name: "ping".to_owned(), args: Vec::new(), + incoming: crate::telemetry::IncomingInvocationContext::default(), conn: ConnHandle::new("conn-grace", Vec::new(), Vec::new(), false), reply: action_tx, }) @@ -3776,6 +3780,7 @@ pub(crate) mod moved_tests { task.handle_dispatch(DispatchCommand::Action { name: "ping".to_owned(), args: Vec::new(), + incoming: crate::telemetry::IncomingInvocationContext::default(), conn: ConnHandle::new("conn-finalize", Vec::new(), Vec::new(), false), reply: reply_tx, }) @@ -4533,6 +4538,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-rust/packages/rivetkit/src/event.rs b/rivetkit-rust/packages/rivetkit/src/event.rs index 6a0b259e7e..5cc75f8c64 100644 --- a/rivetkit-rust/packages/rivetkit/src/event.rs +++ b/rivetkit-rust/packages/rivetkit/src/event.rs @@ -95,6 +95,9 @@ impl RuntimeEvent { args, conn, scheduled_fire, + // The typed runtime hands actions the actor-wide `Ctx` built + // at start, so it has no per-invocation handle to bind this to. + invocation_telemetry: _, reply, } => Self::Action(ActionCall { name, @@ -103,7 +106,11 @@ impl RuntimeEvent { scheduled_fire, reply: Some(reply), }), - ActorEvent::HttpRequest { request, reply } => Self::Http(HttpCall { + ActorEvent::HttpRequest { + request, + invocation_telemetry: _, + reply, + } => Self::Http(HttpCall { request: Some(request), reply: Some(reply), }), @@ -114,6 +121,7 @@ impl RuntimeEvent { request, wait, timeout_ms, + invocation_telemetry: _, reply, } => Self::QueueSend(QueueSend { name, @@ -1500,6 +1508,7 @@ mod tests { args: Vec::new(), conn: None, scheduled_fire: None, + invocation_telemetry: None, reply: reply_tx.into(), }) .expect("queue action event"); diff --git a/rivetkit-rust/packages/rivetkit/src/start.rs b/rivetkit-rust/packages/rivetkit/src/start.rs index fb229ab7c5..19deea7bee 100644 --- a/rivetkit-rust/packages/rivetkit/src/start.rs +++ b/rivetkit-rust/packages/rivetkit/src/start.rs @@ -385,7 +385,11 @@ async fn handle_actor_event( } } } - ActorEvent::HttpRequest { request, reply } => { + ActorEvent::HttpRequest { + request, + invocation_telemetry: _, + reply, + } => { if let Some(http_pools) = http_pools { let class = A::classify_http_request(&request); let Some(permit) = http_pools.try_acquire(class) else { @@ -1224,6 +1228,7 @@ mod tests { let (reply_tx, reply_rx) = oneshot::channel(); tx.send(ActorEvent::HttpRequest { request: rivetkit_core::Request::default(), + invocation_telemetry: None, reply: reply_tx.into(), }) .expect("send http event"); @@ -1248,6 +1253,7 @@ mod tests { let (reply_tx, reply_rx) = oneshot::channel(); tx.send(ActorEvent::HttpRequest { request: Request::default(), + invocation_telemetry: None, reply: reply_tx.into(), }) .expect("send http event"); @@ -1286,6 +1292,7 @@ mod tests { tx.send(ActorEvent::HttpRequest { request: Request::from_parts("GET", "/standard", Default::default(), Vec::new()) .expect("standard request"), + invocation_telemetry: None, reply: standard_tx.into(), }) .expect("send standard http event"); @@ -1293,6 +1300,7 @@ mod tests { tx.send(ActorEvent::HttpRequest { request: Request::from_parts("GET", "/live", Default::default(), Vec::new()) .expect("live request"), + invocation_telemetry: None, reply: live_tx.into(), }) .expect("send live http event"); @@ -2486,6 +2494,7 @@ mod tests { args: args.to_vec(), conn, scheduled_fire: None, + invocation_telemetry: None, reply: reply_tx.into(), }) .expect("send action event"); @@ -2506,6 +2515,7 @@ mod tests { request: rivetkit_core::Request::default(), wait: true, timeout_ms: None, + invocation_telemetry: None, reply: reply_tx.into(), }) .expect("send queue event"); diff --git a/rivetkit-rust/packages/rivetkit/tests/integration_canned_events.rs b/rivetkit-rust/packages/rivetkit/tests/integration_canned_events.rs index e76aa2e5bf..5061648bf9 100644 --- a/rivetkit-rust/packages/rivetkit/tests/integration_canned_events.rs +++ b/rivetkit-rust/packages/rivetkit/tests/integration_canned_events.rs @@ -119,6 +119,7 @@ async fn send_action(event_tx: &mpsc::UnboundedSender, name: &str) - args: Vec::new(), conn: None, scheduled_fire: None, + invocation_telemetry: None, reply: reply_tx.into(), }) .expect("send action event"); diff --git a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts index 6e980d29c6..8e74e93c00 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts +++ b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts @@ -8,6 +8,18 @@ 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 + tracestate?: string +} export interface JsHttpRequest { method: string uri: string @@ -51,6 +63,9 @@ export interface JsQueueSendResult { export interface JsActionDefinition { name: string } +export interface JsQueueDefinition { + name: string +} /** * One entry in the actor's `inspector.tabs[]` declaration. Either a * custom-tab descriptor (id + label + source dir) or a built-in modifier @@ -120,6 +135,7 @@ export interface JsActorConfig { maxIncomingMessageSize?: number maxOutgoingMessageSize?: number actions?: Array + queues?: Array inspectorTabs?: Array } export interface JsBindParam { @@ -257,6 +273,12 @@ export interface JsServerlessStreamError { code: string message: string } +/** + * 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. + */ +export declare function setTelemetryLogSink(callback: (...args: any[]) => any): void export interface JsScheduledEventInfo { id: string action: string @@ -306,6 +328,19 @@ export declare class ActorContext { endOnStateChange(): void kv(): Kv sql(): JsNativeDatabase + sameActorInstance(other: ActorContext): boolean + /** + * Returns a handle for the same invocation whose Core spans parent to the + * application span active in JavaScript, given as W3C headers. + */ + withApplicationSpan(traceparent?: string | undefined | null, tracestate?: string | undefined | null): ActorContext + invocationTraceContext(): JsActorInvocationTraceContext | null + /** + * Opens the span covering one call out to another actor. Returns nothing + * when this handle serves no invocation or tracing is disabled, in which + * case the caller sends its own context as before. + */ + beginOutboundCall(actorName: string, actionName: string): OutboundCall | null provisionActorRuntimeSocket(): Promise schedule(): Schedule queue(): Queue @@ -355,6 +390,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 6f44128343..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, 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 @@ -327,5 +328,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/actor_context.rs b/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs index 49478854b7..9bbab7c6d0 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs @@ -18,14 +18,16 @@ 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, + 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; @@ -79,6 +81,42 @@ pub struct JsActorKeySegment { pub number_value: Option, } +/// Active actor invocation correlation exposed to the TypeScript runtime adapter. +#[napi(object)] +pub struct JsActorInvocationTraceContext { + pub ray_id: Option, + pub span: Option, +} + +/// W3C span context of the current invocation span, present only when tracing is active. +#[napi(object)] +pub struct JsActorInvocationSpanContext { + pub trace_id: String, + pub span_id: String, + pub trace_flags: u8, + pub 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, + tracestate: value.tracestate, + } + } +} + #[napi(object)] pub struct JsHttpRequest { pub method: String, @@ -273,11 +311,53 @@ 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) + } + + /// Returns a handle for the same invocation whose Core spans parent to the + /// application span active in JavaScript, given as W3C headers. + #[napi] + pub fn with_application_span( + &self, + traceparent: Option, + tracestate: Option, + ) -> ActorContext { + ActorContext { + inner: self + .inner + .with_application_span(traceparent.as_deref(), tracestate.as_deref()), + shared: self.shared.clone(), + } + } + + #[napi] + pub fn invocation_trace_context(&self) -> Option { + self.inner.invocation_trace_context().map(Into::into) + } + + /// Opens the span covering one call out to another actor. Returns nothing + /// when this handle serves no invocation or tracing is disabled, in which + /// case the caller sends its own context as before. + #[napi] + pub fn begin_outbound_call( + &self, + actor_name: String, + action_name: String, + ) -> Option { + self.inner + .begin_outbound_call(&actor_name, &action_name) + .map(|invocation| OutboundCall { + invocation: Some(invocation), + }) + } + #[napi] pub async fn provision_actor_runtime_socket( &self, @@ -1056,3 +1136,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 d01c1c0187..40d288d836 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/actor_factory.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/actor_factory.rs @@ -11,7 +11,7 @@ use rivet_error::{ActorSpecifier, RivetError, RivetErrorKind}; use rivetkit_core::inspector::InspectorTabEntry; use rivetkit_core::{ ActionDefinition, ActorConfig, ActorConfigInput, ActorContext as CoreActorContext, - ActorFactory as CoreActorFactory, ConnHandle as CoreConnHandle, Request, + ActorFactory as CoreActorFactory, ConnHandle as CoreConnHandle, QueueDefinition, Request, SqliteProfilingConfigInput, WebSocket as CoreWebSocket, }; @@ -53,6 +53,12 @@ pub struct JsActionDefinition { pub name: String, } +#[napi(object)] +#[derive(Clone, Default)] +pub struct JsQueueDefinition { + pub name: String, +} + /// One entry in the actor's `inspector.tabs[]` declaration. Either a /// custom-tab descriptor (id + label + source dir) or a built-in modifier /// (id + hidden=true). Validation already happened on the TS side; the @@ -124,6 +130,7 @@ pub struct JsActorConfig { pub max_incoming_message_size: Option, pub max_outgoing_message_size: Option, pub actions: Option>, + pub queues: Option>, pub inspector_tabs: Option>, } @@ -155,6 +162,7 @@ pub(crate) struct MigratePayload { #[derive(Clone)] pub(crate) struct QueueSendPayload { pub(crate) ctx: CoreActorContext, + pub(crate) telemetry: Option, pub(crate) conn: CoreConnHandle, pub(crate) request: Request, pub(crate) name: String, @@ -196,6 +204,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, @@ -804,7 +813,10 @@ fn build_queue_send_payload( payload: QueueSendPayload, ) -> napi::Result> { let mut object = env.create_object()?; - object.set("ctx", ActorContext::new(payload.ctx))?; + object.set( + "ctx", + ActorContext::new(payload.ctx.with_invocation_telemetry(payload.telemetry)), + )?; object.set("conn", ConnHandle::new(payload.conn))?; object.set("request", build_request_object(env, payload.request)?)?; object.set("name", payload.name)?; @@ -871,7 +883,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()?)?, @@ -972,6 +987,14 @@ fn parse_bridge_rivet_error(reason: &str) -> Option { })) } +/// Rebuilds an error raised in JavaScript from the reason string the bridge +/// carries. A structured error arrives bridge-encoded and keeps its group and +/// code; anything else stays an unstructured message, which is what lets Core +/// classify and sanitize it rather than trusting the JavaScript text. +pub(crate) fn anyhow_error_from_js_reason(reason: String) -> anyhow::Error { + parse_bridge_rivet_error(&reason).unwrap_or_else(|| anyhow::anyhow!(reason)) +} + pub(crate) fn callback_error(callback_name: &str, error: napi::Error) -> anyhow::Error { let reason = error.reason; if let Some(error) = parse_bridge_rivet_error(&reason) { @@ -1022,6 +1045,12 @@ impl From for ActorConfigInput { .map(|action| ActionDefinition { name: action.name }) .collect() }), + queues: value.queues.map(|queues| { + queues + .into_iter() + .map(|queue| QueueDefinition { name: queue.name }) + .collect() + }), inspector_tabs: value.inspector_tabs.map(|tabs| { tabs.into_iter() .map(|tab| { diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/http.rs b/rivetkit-typescript/packages/rivetkit-napi/src/http.rs index 80068e10ef..ae89bb04ad 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/http.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/http.rs @@ -30,6 +30,7 @@ pub struct JsHttpResponse { #[derive(Clone)] pub(crate) struct HttpRequestPayload { pub(crate) ctx: CoreActorContext, + pub(crate) telemetry: Option, pub(crate) request: Request, pub(crate) cancel_token: Option, pub(crate) response_stream: Option, @@ -220,7 +221,10 @@ pub(crate) fn build_http_request_payload( payload: HttpRequestPayload, ) -> napi::Result> { let mut object = env.create_object()?; - object.set("ctx", ActorContext::new(payload.ctx))?; + object.set( + "ctx", + ActorContext::new(payload.ctx.with_invocation_telemetry(payload.telemetry)), + )?; object.set("request", build_request_object(env, payload.request)?)?; match payload.response_stream { Some(response_stream) => object.set("responseBodyStream", response_stream)?, diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/lib.rs b/rivetkit-typescript/packages/rivetkit-napi/src/lib.rs index 1c1c1b0a86..236f7a05e1 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,19 @@ 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( + telemetry::sdk_log_bridge::SdkLogLayer + .with_filter(tracing_subscriber::EnvFilter::new("opentelemetry_sdk=warn")), + ) .with(match log_format { LogFormat::Logfmt => Some( tracing_logfmt::builder() @@ -128,7 +138,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 +147,18 @@ 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/napi_actor_events.rs b/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs index 619f8d49b2..d841261f0d 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs @@ -396,6 +396,7 @@ pub(crate) async fn dispatch_event( args, conn, scheduled_fire, + invocation_telemetry, reply, } => { tracing::info!( @@ -430,6 +431,7 @@ pub(crate) async fn dispatch_event( call_action( &callback, &ctx, + invocation_telemetry, conn, name.clone(), args.clone(), @@ -460,7 +462,11 @@ pub(crate) async fn dispatch_event( } }); } - ActorEvent::HttpRequest { request, reply } => { + ActorEvent::HttpRequest { + request, + invocation_telemetry, + reply, + } => { let Some(callback) = bindings.on_request.clone() else { reply.send(Err(missing_callback("onRequest"))); return; @@ -474,7 +480,7 @@ pub(crate) async fn dispatch_event( "Action timed out", None, timeout, - call_http_request(&callback, &ctx, request), + call_http_request(&callback, &ctx, invocation_telemetry, request), ) .await }); @@ -486,6 +492,7 @@ pub(crate) async fn dispatch_event( request, wait, timeout_ms, + invocation_telemetry, reply, } => { let Some(callback) = bindings.on_queue_send.clone() else { @@ -508,6 +515,7 @@ pub(crate) async fn dispatch_event( &callback, QueueSendPayload { ctx: ctx.inner().clone(), + telemetry: invocation_telemetry, conn, request, name, @@ -1177,6 +1185,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 +1198,7 @@ async fn call_action( callback, ActionPayload { ctx: ctx.inner().clone(), + telemetry, conn, name, args, @@ -1222,6 +1232,7 @@ async fn call_on_before_action_response( async fn call_http_request( callback: &crate::actor_factory::CallbackTsfn, ctx: &ActorContext, + telemetry: Option, request: rivetkit_core::Request, ) -> Result { let request_cancel_token = request.cancellation_token(); @@ -1230,6 +1241,7 @@ async fn call_http_request( callback, HttpRequestPayload { ctx: ctx.inner().clone(), + telemetry, request, cancel_token: Some(request_cancel_token), response_stream: None, diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/registry.rs b/rivetkit-typescript/packages/rivetkit-napi/src/registry.rs index c410453cc7..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)] @@ -322,6 +330,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..8afac9d61b --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit-napi/src/telemetry.rs @@ -0,0 +1,95 @@ +//! 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 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; + + /// One SDK diagnostic, flattened for the JavaScript side. + pub(crate) struct SdkLogEvent { + pub(crate) name: String, + pub(crate) message: String, + } + + /// 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, 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 + /// would never exit on its own. Warnings still cross while the application + /// is running; the sink just stops being a reason to keep running. + pub(crate) fn install(env: Env, callback: JsFunction) -> Result<()> { + let mut tsfn = + callback.create_threadsafe_function(0, |ctx: ThreadSafeCallContext| { + let mut object = ctx.env.create_object()?; + object.set("name", ctx.value.name)?; + object.set("message", ctx.value.message)?; + Ok(vec![object.into_unknown()]) + })?; + tsfn.unref(&env)?; + *SINK.write() = Some(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 sink = SINK.read(); + let Some(sink) = sink.as_ref() else { + return; + }; + let mut fields = FieldCollector::default(); + event.record(&mut fields); + sink.call( + SdkLogEvent { + name: fields.name, + message: fields.message, + }, + napi::threadsafe_function::ThreadsafeFunctionCallMode::NonBlocking, + ); + } + } +} diff --git a/rivetkit-typescript/packages/rivetkit-napi/tests/napi_actor_events.rs b/rivetkit-typescript/packages/rivetkit-napi/tests/napi_actor_events.rs index 171d0dc2fb..293810edde 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/tests/napi_actor_events.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/tests/napi_actor_events.rs @@ -168,6 +168,7 @@ mod moved_tests { args: vec![1, 2, 3], conn: None, scheduled_fire: None, + invocation_telemetry: None, reply: tx.into(), }, &bindings, @@ -490,6 +491,7 @@ mod moved_tests { args: Vec::new(), conn: None, scheduled_fire: None, + invocation_telemetry: None, reply: first_tx.into(), }) .expect("first action event should send"); @@ -499,6 +501,7 @@ mod moved_tests { args: Vec::new(), conn: None, scheduled_fire: None, + invocation_telemetry: None, reply: second_tx.into(), }) .expect("second action event should send"); diff --git a/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs b/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs index f1efda5650..4f68f3630b 100644 --- a/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs +++ b/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs @@ -177,6 +177,12 @@ pub struct WasmActionDefinition { pub name: String, } +#[derive(Clone, Default, serde::Deserialize)] +#[serde(default, rename_all = "camelCase")] +pub struct WasmQueueDefinition { + pub name: String, +} + /// Experimental SQLite profiling configuration. This entire configuration /// surface is subject to change without notice. #[derive(Clone, Default, serde::Deserialize)] @@ -229,6 +235,7 @@ pub struct WasmActorConfig { pub max_incoming_message_size: Option, pub max_outgoing_message_size: Option, pub actions: Option>, + pub queues: Option>, } impl From for ActorConfigInput { @@ -265,6 +272,12 @@ impl From for ActorConfigInput { .map(|action| rivetkit_core::ActionDefinition { name: action.name }) .collect() }), + queues: config.queues.map(|queues| { + queues + .into_iter() + .map(|queue| rivetkit_core::QueueDefinition { name: queue.name }) + .collect() + }), // Custom inspector tabs serve assets from a filesystem `root`, which is a // native/server feature that has no meaning in a browser wasm host. inspector_tabs: None, diff --git a/rivetkit-typescript/packages/rivetkit/package.json b/rivetkit-typescript/packages/rivetkit/package.json index ac27eda52e..7f85fe8c21 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", @@ -235,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-conn.ts b/rivetkit-typescript/packages/rivetkit/src/client/actor-conn.ts index db8edb885f..a106156c19 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/actor-conn.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/actor-conn.ts @@ -50,6 +50,7 @@ import { ACTOR_CONNS_SYMBOL, type ClientRaw } from "./client"; import * as errors from "./errors"; import { isRetryableLifecycleReconnectSignal } from "./lifecycle-errors"; import { logger } from "./log"; +import { outboundTelemetryHeaders } from "./outbound-telemetry"; import { createQueueSender, type QueueSendNoWaitOptions, @@ -229,6 +230,7 @@ export class ActorConnRaw { this.#queueSender = createQueueSender({ encoding: this.#encoding, params: this.#params, + telemetryHeaders: () => outboundTelemetryHeaders(undefined), customFetch: async (request: Request) => { return await this.#driver.sendRequest( getGatewayTarget(this.#actorResolutionState), diff --git a/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts b/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts index e41aa6366d..5e7bd2d8a7 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts @@ -1,5 +1,6 @@ import type { AnyActorDefinition } from "@/actor/definition"; -import type { ActorSpecifier } from "@/actor/errors"; +import { type ActorSpecifier, encodeErrorForBridge } from "@/actor/errors"; +import type { ActorInvocationSpanContext } from "@/common/actor-telemetry-context"; import { HEADER_CONN_PARAMS, HEADER_ENCODING, @@ -24,6 +25,10 @@ import { AsyncMutex } from "@/common/database/shared"; import type { Encoding, JsonCompatValue } from "@/common/encoding"; import { deconstructError } from "@/common/utils"; import type { EngineControlClient } from "@/engine-client/driver"; +import type { + BeginOutboundCall, + CurrentActorInvocation, +} from "@/registry/runtime"; import { decodeCborCompat, deserializeWithEncoding, @@ -53,6 +58,7 @@ import { type ClientRaw, CREATE_ACTOR_CONN_PROXY } from "./client"; import { ActorError, isSchedulingError } from "./errors"; import { retryOnLifecycleBoundary } from "./lifecycle-errors"; import { logger } from "./log"; +import { outboundTelemetryHeaders } from "./outbound-telemetry"; import { createQueueSender, type QueueSendNoWaitOptions, @@ -82,6 +88,8 @@ export class ActorHandleRaw { #resolvedActorId?: string; #resolvingActorId?: Promise; #queueSendMutex = new AsyncMutex(); + #currentActorInvocation?: CurrentActorInvocation; + #beginOutboundCall?: BeginOutboundCall; /** * Do not call this directly. @@ -99,6 +107,8 @@ export class ActorHandleRaw { actorResolutionState: ActorResolutionState, gatewayOptions: ActorGatewayOptions = {}, signal?: AbortSignal, + currentActorInvocation?: CurrentActorInvocation, + beginOutboundCall?: BeginOutboundCall, ) { this.#client = client; this.#driver = driver; @@ -108,6 +118,8 @@ export class ActorHandleRaw { this.#params = params; this.#getParams = getParams; this.#signal = signal; + this.#currentActorInvocation = currentActorInvocation; + this.#beginOutboundCall = beginOutboundCall; } async #resolveConnectionParams(): Promise { @@ -166,6 +178,10 @@ export class ActorHandleRaw { return await createQueueSender({ encoding: this.#encoding, params: this.#params, + telemetryHeaders: () => + outboundTelemetryHeaders( + this.#currentActorInvocation?.(), + ), customFetch: async (request: Request) => { return await this.#driver.sendRequest( target, @@ -275,20 +291,41 @@ export class ActorHandleRaw { // when no per-call signal is provided. const signal = opts.signal ?? this.#signal; const optsWithSignal = { ...opts, signal }; + // Open before retries so all attempts share one span and parent context. + const call = this.#beginOutboundCall?.( + getActorNameFromQuery(this.#actorResolutionState), + opts.name, + ); const run = async () => - (await this.#sendActionNow(optsWithSignal)) as Response; - if (opts.name === "destroy") { - return await run(); + (await this.#sendActionAttempts( + optsWithSignal, + call?.span, + )) as Response; + const send = async () => { + if (opts.name === "destroy") { + return await run(); + } + return await retryOnLifecycleBoundary(run, { signal }); + }; + if (!call) { + return await send(); + } + try { + const output = await send(); + call.finish(); + return output; + } catch (error) { + call.finish(encodeErrorForBridge(error)); + throw error; } - - return await retryOnLifecycleBoundary(run, { signal }); } - async #sendActionNow( + async #sendActionAttempts( opts: { name: string; args: unknown[]; } & ActorActionOptions, + callSpan?: ActorInvocationSpanContext, ): Promise { const maxAttempts = this.#getDynamicQueryMaxAttempts(); let useQueryTarget = isDynamicActorQuery(this.#actorResolutionState); @@ -320,6 +357,16 @@ export class ActorHandleRaw { name: opts.name, encoding: this.#encoding, }); + const headers: Record = { + [HEADER_ENCODING]: this.#encoding, + ...outboundTelemetryHeaders( + this.#currentActorInvocation?.(), + callSpan, + ), + }; + if (this.#params !== undefined) { + headers[HEADER_CONN_PARAMS] = JSON.stringify(this.#params); + } const output = await sendHttpRequest< protocol.HttpActionRequest, protocol.HttpActionResponse, @@ -330,16 +377,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) => @@ -680,6 +718,9 @@ export class ActorHandleRaw { skipReadyWait, }, ); + const telemetryHeaders = outboundTelemetryHeaders( + this.#currentActorInvocation?.(), + ); for (let attempt = 0; attempt < maxAttempts; attempt++) { let actorId: string | undefined; @@ -698,6 +739,7 @@ export class ActorHandleRaw { clonesInputBody ? input.clone() : input, requestInit, gatewayOptions, + telemetryHeaders, ); const retry = await this.#shouldRetryRawFetchResponse( response, diff --git a/rivetkit-typescript/packages/rivetkit/src/client/client.ts b/rivetkit-typescript/packages/rivetkit/src/client/client.ts index 439403871e..72c2ead3b2 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/client.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/client.ts @@ -3,6 +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 { + BeginOutboundCall, + CurrentActorInvocation, +} from "@/registry/runtime"; import type { ActorActionFunction, ActorGatewayOptions } from "./actor-common"; import { type ActorConn, @@ -175,12 +179,16 @@ export interface Region { export const ACTOR_CONNS_SYMBOL = Symbol("actorConns"); export const CREATE_ACTOR_CONN_PROXY = Symbol("createActorConnProxy"); -/** - * Client for managing & connecting to actors. - * - * @template A The actors map type that defines the available actors. - * @see {@link https://rivet.dev/docs/manage|Create & Manage Actors} - */ +/** Options for constructing a raw actor client. */ +export interface ClientRawOptions { + encoding?: Encoding; + gateway?: ActorGatewayOptions; + /** Supplies the calling actor's invocation so actor-to-actor clients propagate its trace and ray ID. */ + currentActorInvocation?: CurrentActorInvocation; + /** Opens and finishes the Core span covering an actor-to-actor call. */ + beginOutboundCall?: BeginOutboundCall; +} + export class ClientRaw { #disposed = false; @@ -189,19 +197,22 @@ export class ClientRaw { #driver: EngineControlClient; #encodingKind: Encoding; #gatewayOptions: ActorGatewayOptions; + #currentActorInvocation?: CurrentActorInvocation; + #beginOutboundCall?: BeginOutboundCall; /** * 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; + this.#beginOutboundCall = options.beginOutboundCall; } /** @@ -403,6 +414,8 @@ export class ClientRaw { actorQuery, this.#gatewayOptions, signal, + this.#currentActorInvocation, + this.#beginOutboundCall, ); } @@ -459,9 +472,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/client/outbound-telemetry.ts b/rivetkit-typescript/packages/rivetkit/src/client/outbound-telemetry.ts new file mode 100644 index 0000000000..77da2b14a1 --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/src/client/outbound-telemetry.ts @@ -0,0 +1,69 @@ +import { + HEADER_RIVET_RAY_ID, + HEADER_TRACEPARENT, + HEADER_TRACESTATE, +} from "@/common/actor-router-consts"; +import type { + ActorInvocationSpanContext, + ActorInvocationTraceContext, +} from "@/common/actor-telemetry-context"; +import { + actorInvocationTraceHeaders, + readActiveRayId, + readActiveTraceHeaders, +} from "@/common/otel-context"; + +/** + * Headers that carry a caller's ray ID and trace context into an actor. One + * rule for every kind of outbound call, so an action call and a queue send + * made from the same place land in the same trace under the same ray ID. + * + * The ray ID is the calling invocation's when the caller is itself inside an + * actor, else the one placed in OpenTelemetry baggage by the surrounding + * request handler. The trace context is `callSpan` when the runtime opened a + * span for this call, else the application span active in this JavaScript + * context, else the calling actor's own Core invocation span. + */ +export function outboundTelemetryHeaders( + invocation: ActorInvocationTraceContext | undefined, + callSpan?: ActorInvocationSpanContext, +): Record { + const headers: Record = {}; + const rayId = invocation?.rayId ?? readActiveRayId(); + if (rayId) { + headers[HEADER_RIVET_RAY_ID] = rayId; + } + const traceHeaders = + actorInvocationTraceHeaders(callSpan) ?? + readActiveTraceHeaders() ?? + actorInvocationTraceHeaders(invocation?.span); + if (traceHeaders) { + headers[HEADER_TRACEPARENT] = traceHeaders.traceparent; + if (traceHeaders.tracestate) { + headers[HEADER_TRACESTATE] = traceHeaders.tracestate; + } + } + return headers; +} + +/** + * Adds outbound telemetry headers to a request whose caller may have set + * some already. A header the caller set wins. `traceparent` and `tracestate` + * describe one span between them, so when the caller set either, both stay + * as the caller set them and neither is added. + */ +export function addOutboundTelemetryHeaders( + headers: Headers, + telemetry: Record, +): void { + const callerSetTraceContext = + headers.has(HEADER_TRACEPARENT) || headers.has(HEADER_TRACESTATE); + for (const [name, value] of Object.entries(telemetry)) { + const isTraceContext = + name === HEADER_TRACEPARENT || name === HEADER_TRACESTATE; + if (isTraceContext ? callerSetTraceContext : headers.has(name)) { + continue; + } + headers.set(name, value); + } +} diff --git a/rivetkit-typescript/packages/rivetkit/src/client/queue.ts b/rivetkit-typescript/packages/rivetkit/src/client/queue.ts index a55fc86852..c081a78b0f 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/queue.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/queue.ts @@ -54,6 +54,8 @@ export interface QueueSendResult { interface QueueSenderOptions { encoding: Encoding; params: unknown; + /** Ray ID and trace context headers, read per send. */ + telemetryHeaders: () => Record; customFetch: (request: Request) => Promise; } @@ -90,6 +92,7 @@ export function createQueueSender( method: "POST", headers: { [HEADER_ENCODING]: senderOptions.encoding, + ...senderOptions.telemetryHeaders(), ...(senderOptions.params !== undefined ? { [HEADER_CONN_PARAMS]: JSON.stringify( diff --git a/rivetkit-typescript/packages/rivetkit/src/client/raw-utils.ts b/rivetkit-typescript/packages/rivetkit/src/client/raw-utils.ts index de89ec1b96..b71d45a05f 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/raw-utils.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/raw-utils.ts @@ -12,6 +12,7 @@ import type { } from "@/engine-client/driver"; import { ActorError } from "./errors"; import { logger } from "./log"; +import { addOutboundTelemetryHeaders } from "./outbound-telemetry"; /** Buffer a one-shot ReadableStream body to bytes so every retry attempt can re-send it. */ export async function prepareRetryableInit( @@ -37,6 +38,7 @@ export async function rawHttpFetch( input: string | URL | Request, init?: RequestInit, options: GatewayRequestOptions = {}, + telemetryHeaders: Record = {}, ): Promise { // Extract path and merge init options let path: string; @@ -111,6 +113,7 @@ export async function rawHttpFetch( if (params) { proxyRequestHeaders.set(HEADER_CONN_PARAMS, JSON.stringify(params)); } + addOutboundTelemetryHeaders(proxyRequestHeaders, telemetryHeaders); // Forward the request to the actor const proxyRequest = new Request(url, { diff --git a/rivetkit-typescript/packages/rivetkit/src/common/actor-router-consts.ts b/rivetkit-typescript/packages/rivetkit/src/common/actor-router-consts.ts index edefc6cc8f..22c1efbcd6 100644 --- a/rivetkit-typescript/packages/rivetkit/src/common/actor-router-consts.ts +++ b/rivetkit-typescript/packages/rivetkit/src/common/actor-router-consts.ts @@ -20,6 +20,9 @@ export const HEADER_ACTOR_GENERATION = "x-rivet-actor-generation"; export const HEADER_ACTOR_KEY = "x-rivet-actor-key"; export const HEADER_RIVET_TOKEN = "x-rivet-token"; +export const HEADER_RIVET_RAY_ID = "x-rivet-ray-id"; +export const HEADER_TRACEPARENT = "traceparent"; +export const HEADER_TRACESTATE = "tracestate"; // MARK: Manager Gateway Headers export const HEADER_RIVET_TARGET = "x-rivet-target"; diff --git a/rivetkit-typescript/packages/rivetkit/src/common/actor-telemetry-context.ts b/rivetkit-typescript/packages/rivetkit/src/common/actor-telemetry-context.ts new file mode 100644 index 0000000000..e5c36634fb --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/src/common/actor-telemetry-context.ts @@ -0,0 +1,19 @@ +/** Correlation owned by the currently executing Core actor invocation. */ +export interface ActorInvocationTraceContext { + /** Ray ID the invocation received, absent when nothing upstream supplied one. */ + readonly rayId?: string; + /** Core-owned invocation span context, absent when tracing is disabled. */ + readonly span?: ActorInvocationSpanContext; +} + +/** W3C span context of the Core invocation span. */ +export interface ActorInvocationSpanContext { + /** W3C trace identifier. */ + readonly traceId: string; + /** W3C span identifier for the Core invocation. */ + readonly spanId: string; + /** OpenTelemetry trace flags encoded as an integer. */ + readonly traceFlags: number; + /** 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..9392097874 --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/src/common/otel-context.ts @@ -0,0 +1,99 @@ +import { + context, + createTraceState, + isSpanContextValid, + propagation, + type SpanContext, + trace, +} from "@opentelemetry/api"; +import type { ActorInvocationSpanContext } 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 { + return traceHeaders(trace.getSpanContext(context.active())); +} + +/** Serializes the structured span context supplied by Core as W3C headers. */ +export function actorInvocationTraceHeaders( + invocation: ActorInvocationSpanContext | undefined, +): ActiveTraceHeaders | undefined { + const spanContext = invocationSpanContext(invocation); + if (!spanContext) return undefined; + return traceHeaders(spanContext); +} + +/** W3C Baggage key that carries a ray ID through application code. */ +export const RAY_BAGGAGE_KEY = "rivet.ray.id"; + +const RAY_ID_PATTERN = /^[A-Za-z0-9_-]{1,128}$/; + +/** + * Returns the ray ID carried in the active OpenTelemetry baggage, so a request + * handler that received a ray ID can pass it to the actors it calls. The value is + * bounded by the same rule Core applies at the actor edge: 1 to 128 characters + * of `[A-Za-z0-9_-]`. Anything else counts as absent. + */ +export function readActiveRayId(): string | undefined { + const rayId = propagation + .getBaggage(context.active()) + ?.getEntry(RAY_BAGGAGE_KEY)?.value; + if (rayId === undefined || !RAY_ID_PATTERN.test(rayId)) return undefined; + return rayId; +} + +/** + * Runs `run` with the Core invocation span as the active OpenTelemetry span, + * 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(); + + const spanContext = invocationSpanContext(invocation); + if (!spanContext) return run(); + const parent = trace.setSpanContext(context.active(), spanContext); + + return context.with(parent, run); +} + +function invocationSpanContext( + invocation: ActorInvocationSpanContext | undefined, +): SpanContext | undefined { + if (!invocation) return undefined; + try { + const spanContext: SpanContext = { + traceId: invocation.traceId, + spanId: invocation.spanId, + traceFlags: invocation.traceFlags, + traceState: invocation.tracestate + ? createTraceState(invocation.tracestate) + : undefined, + isRemote: false, + }; + return isSpanContextValid(spanContext) ? spanContext : undefined; + } catch { + return undefined; + } +} + +function traceHeaders( + spanContext: SpanContext | undefined, +): ActiveTraceHeaders | undefined { + if (!spanContext || !isSpanContextValid(spanContext)) return undefined; + const tracestate = spanContext.traceState?.serialize(); + return { + traceparent: `00-${spanContext.traceId}-${spanContext.spanId}-${(spanContext.traceFlags & 1).toString(16).padStart(2, "0")}`, + ...(tracestate ? { tracestate } : {}), + }; +} diff --git a/rivetkit-typescript/packages/rivetkit/src/engine-client/actor-http-client.ts b/rivetkit-typescript/packages/rivetkit/src/engine-client/actor-http-client.ts index 2fc62b41a8..563fd90f8c 100644 --- a/rivetkit-typescript/packages/rivetkit/src/engine-client/actor-http-client.ts +++ b/rivetkit-typescript/packages/rivetkit/src/engine-client/actor-http-client.ts @@ -4,6 +4,9 @@ import { HEADER_RIVET_SKIP_READY_WAIT, HEADER_RIVET_TARGET, HEADER_RIVET_TOKEN, + HEADER_RIVET_RAY_ID, + HEADER_TRACEPARENT, + HEADER_TRACESTATE, } from "@/common/actor-router-consts"; import { type GatewayRequestOptions, shouldSkipReadyWait } from "./driver"; @@ -55,6 +58,19 @@ function buildGuardHeaders( for (const [key, value] of Object.entries(runConfig.headers)) { headers.set(key, value as string); } + // Invocation headers are per action call. Apply the active request last so + // static client configuration cannot retain or override an earlier action. + for (const name of [ + HEADER_RIVET_RAY_ID, + HEADER_TRACEPARENT, + HEADER_TRACESTATE, + ]) { + headers.delete(name); + const value = actorRequest.headers.get(name); + if (value !== null) { + headers.set(name, value); + } + } // Add guard-specific headers if (runConfig.token) { headers.set(HEADER_RIVET_TOKEN, runConfig.token); diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts index 3d589d5242..1530b1dbf1 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, @@ -7,6 +8,13 @@ import type { HttpResponseBodyStream as NativeHttpResponseBodyStream, WebSocket as NativeWebSocket, } from "@rivetkit/rivetkit-napi"; +import type { ActorInvocationTraceContext } from "@/common/actor-telemetry-context"; +import { + actorInvocationTraceHeaders, + readActiveTraceHeaders, + runWithActorInvocationSpan, +} from "@/common/otel-context"; +import { logger } from "./log"; import type { ActorContextHandle, ActorFactoryHandle, @@ -26,6 +34,7 @@ import type { RuntimeKvEntry, RuntimeKvListOptions, RuntimeListenerConfig, + RuntimeOutboundCall, RuntimeQueueEnqueueAndWaitOptions, RuntimeQueueMessage, RuntimeQueueNextBatchOptions, @@ -65,6 +74,9 @@ type NapiSqlTransaction = Awaited< type NapiActorStateTransaction = Awaited< ReturnType >; +type NapiQueueMessage = Awaited< + ReturnType["send"]> +>; function asNativeRegistry(handle: RegistryHandle): NativeCoreRegistry { return handle as unknown as NativeCoreRegistry; @@ -233,7 +245,7 @@ function toNapiKvEntry(entry: RuntimeKvEntry): { }; } -function toNapiQueueMessage(message: RuntimeQueueMessage): RuntimeQueueMessage { +function toNapiQueueMessage(message: NapiQueueMessage): RuntimeQueueMessage { return { id: () => message.id(), name: () => message.name(), @@ -255,22 +267,65 @@ export class NapiCoreRuntime implements CoreRuntime { #bindings: NativeBindings; #sql = new WeakMap(); - - constructor(bindings: NativeBindings) { + #invocationContext: AsyncLocalStorage; + // W3C `traceparent` of each invocation's own Core span, so an operation can + // tell that span apart from an application span without a native call. + #invocationTraceparent = new WeakMap(); + + constructor( + bindings: NativeBindings, + invocationContext: AsyncLocalStorage, + ) { this.#bindings = bindings; + this.#invocationContext = invocationContext; } + // Core cannot read the active JS span; bind it to the operation handle here. + #actorContextForOperation(owner: ActorContextHandle): NativeActorContext { + const ownerCtx = asNativeActorContext(owner); + const active = this.#invocationContext.getStore(); + if (!active?.sameActorInstance(ownerCtx)) { + return ownerCtx; + } + const applicationSpan = readActiveTraceHeaders(); + // Avoid a native call when Core already has the active parent. + if ( + !applicationSpan || + applicationSpan.traceparent === + this.#invocationTraceparent.get(active) + ) { + return active; + } + return active.withApplicationSpan( + applicationSpan.traceparent, + applicationSpan.tracestate ?? null, + ); + } + + // Cache only the actor-owned handle, which is closed on sleep. + // Invocation handles carry per-call context and must not be reused. #actorSql(ctx: ActorContextHandle): NapiSqlDatabase { - const nativeCtx = asNativeActorContext(ctx); - let database = this.#sql.get(nativeCtx); + const ownerCtx = asNativeActorContext(ctx); + const activeCtx = this.#actorContextForOperation(ctx); + if (activeCtx !== ownerCtx) { + return activeCtx.sql(); + } + let database = this.#sql.get(ownerCtx); if (!database) { - database = nativeCtx.sql(); - this.#sql.set(nativeCtx, database); + database = ownerCtx.sql(); + this.#sql.set(ownerCtx, database); } return database; } createRegistry(): RegistryHandle { + // Replace the sink on each registry start because its previous worker may have exited. + this.#bindings.setTelemetryLogSink((event) => { + logger().warn( + { otelEvent: event.name }, + event.message || event.name, + ); + }); return asRegistryHandle(new this.#bindings.CoreRegistry()); } @@ -552,6 +607,46 @@ export class NapiCoreRuntime implements CoreRuntime { return asNativeActorContext(ctx).actorId(); } + runWithActorInvocationContext(ctx: ActorContextHandle, run: () => T): T { + const nativeCtx = asNativeActorContext(ctx); + const span = nativeCtx.invocationTraceContext()?.span; + const traceHeaders = actorInvocationTraceHeaders(span); + if (traceHeaders) { + this.#invocationTraceparent.set( + nativeCtx, + traceHeaders.traceparent, + ); + } + return this.#invocationContext.run(nativeCtx, () => + runWithActorInvocationSpan(span, run), + ); + } + + actorInvocationTraceContext( + ctx: ActorContextHandle, + ): ActorInvocationTraceContext | undefined { + return ( + this.#actorContextForOperation(ctx).invocationTraceContext() ?? + undefined + ); + } + + beginOutboundCall( + ctx: ActorContextHandle, + actorName: string, + actionName: string, + ): RuntimeOutboundCall | undefined { + const call = this.#actorContextForOperation(ctx).beginOutboundCall( + actorName, + actionName, + ); + if (!call) return undefined; + return { + span: call.spanContext() ?? undefined, + finish: (error?: string) => call.finish(error), + }; + } + actorName(ctx: ActorContextHandle): string { return asNativeActorContext(ctx).name(); } @@ -600,7 +695,7 @@ export class NapiCoreRuntime implements CoreRuntime { } actorWaitUntil(ctx: ActorContextHandle, promise: Promise): void { - asNativeActorContext(ctx).waitUntil(promise); + this.#actorContextForOperation(ctx).waitUntil(promise); } async actorWaitForTrackedShutdownWork( @@ -871,11 +966,7 @@ export class NapiCoreRuntime implements CoreRuntime { async actorSqlClose(ctx: ActorContextHandle): Promise { const nativeCtx = asNativeActorContext(ctx); - const database = this.#sql.get(nativeCtx); - if (!database) { - return; - } - + const database = this.#sql.get(nativeCtx) ?? nativeCtx.sql(); this.#sql.delete(nativeCtx); await database.close(); } @@ -890,7 +981,7 @@ export class NapiCoreRuntime implements CoreRuntime { body: RuntimeBytes, ): Promise { return toNapiQueueMessage( - await asNativeActorContext(ctx) + await this.#actorContextForOperation(ctx) .queue() .send(name, toNapiBuffer(body)), ); @@ -901,7 +992,7 @@ export class NapiCoreRuntime implements CoreRuntime { options?: RuntimeQueueNextBatchOptions | undefined | null, signal?: CancellationTokenHandle | undefined | null, ): Promise { - const messages = await asNativeActorContext(ctx) + const messages = await this.#actorContextForOperation(ctx) .queue() .nextBatch( options, @@ -917,7 +1008,7 @@ export class NapiCoreRuntime implements CoreRuntime { signal?: CancellationTokenHandle | undefined | null, ): Promise { return toNapiQueueMessage( - await asNativeActorContext(ctx) + await this.#actorContextForOperation(ctx) .queue() .waitForNames( names, @@ -964,7 +1055,7 @@ export class NapiCoreRuntime implements CoreRuntime { options?: RuntimeQueueEnqueueAndWaitOptions | undefined | null, signal?: CancellationTokenHandle | undefined | null, ): Promise { - return await asNativeActorContext(ctx) + return await this.#actorContextForOperation(ctx) .queue() .enqueueAndWait( name, @@ -1002,7 +1093,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)); } @@ -1013,13 +1104,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( @@ -1027,12 +1118,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( @@ -1044,7 +1136,7 @@ export class NapiCoreRuntime implements CoreRuntime { args: RuntimeBytes, maxHistory: number | undefined, ) { - await asNativeActorContext(ctx) + await this.#actorContextForOperation(ctx) .schedule() .cronSet( name, @@ -1064,7 +1156,7 @@ export class NapiCoreRuntime implements CoreRuntime { args: RuntimeBytes, maxHistory: number | undefined, ) { - await asNativeActorContext(ctx) + await this.#actorContextForOperation(ctx) .schedule() .cronEvery( name, @@ -1079,20 +1171,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( @@ -1100,7 +1195,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[]; } @@ -1173,9 +1268,12 @@ export async function loadNapiRuntime(): Promise<{ // would snapshot the native `.node` addon into the deploy and 413. The // computed specifier keeps it opaque to static analysis so it is never // bundled. Enforced by scripts/ci/check-edge-native-closure.mjs. - const bindings = await import(["@rivetkit", "rivetkit-napi"].join("/")); + const [{ AsyncLocalStorage }, bindings] = await Promise.all([ + import("node:async_hooks"), + import(["@rivetkit", "rivetkit-napi"].join("/")), + ]); return { bindings, - runtime: new NapiCoreRuntime(bindings), + runtime: new NapiCoreRuntime(bindings, new AsyncLocalStorage()), }; } diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts index 4d3a55b45e..6d7575a5a4 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts @@ -8,6 +8,7 @@ import { type ActorCron, type ActorCronEveryOptions, type ActorCronSetOptions, + type ActorLogger, type ActorSchedule, CONN_STATE_MANAGER_SYMBOL, type CronFire, @@ -46,6 +47,7 @@ import { } from "@/client/client"; import { convertRegistryConfigToClientConfig } from "@/client/config"; import { HEADER_CONN_PARAMS } from "@/common/actor-router-consts"; +import type { ActorInvocationTraceContext } from "@/common/actor-telemetry-context"; import type { AnyDatabaseProvider, SqliteProfilingOptions, @@ -2720,6 +2722,7 @@ export class ActorContextHandleAdapter { #db?: unknown; #dispatchCancelToken?: CancellationTokenHandle; #kv?: NativeKvAdapter; + #log?: ActorLogger; #queue?: NativeQueueAdapter; #request?: Request; #schedule?: NativeScheduleAdapter; @@ -2956,8 +2959,30 @@ export class ActorContextHandleAdapter { return this.#connMap; } + #invocationTraceContext(): ActorInvocationTraceContext | undefined { + return callNativeSync(() => + this.#runtime.actorInvocationTraceContext(this.#ctx), + ); + } + get log() { - return logger(); + if (!this.#log) { + // Actor fields follow the camelCase used by the rest of the + // TypeScript logs. trace_id and span_id stay snake_case because + // that is what OpenTelemetry log correlation tooling looks for. + const invocation = this.#invocationTraceContext(); + this.#log = logger().child({ + actorId: this.actorId, + actorName: this.name, + actorKey: this.key, + ...(invocation?.rayId && { rayId: invocation.rayId }), + ...(invocation?.span && { + trace_id: invocation.span.traceId, + span_id: invocation.span.spanId, + }), + }); + } + return this.#log; } get abortSignal(): AbortSignal { @@ -3881,6 +3906,9 @@ function buildActorConfig( actions: Object.keys(flattenActionHandlers(config.actions)) .sort() .map((name) => ({ name })), + queues: Object.keys(config.queues ?? {}) + .sort() + .map((name) => ({ name })), inspectorTabs: buildInspectorTabs(config.inspector, runtimeKind), }; } @@ -3995,12 +4023,22 @@ 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), + ), + beginOutboundCall: (actorName, actionName) => + callNativeSync(() => + runtime.beginOutboundCall(ctx, actorName, actionName), + ), + }, ); const run = getRunFunction(config.run); const runHandlerCoordinator = @@ -4044,7 +4082,7 @@ export function buildNativeFactory( new ActorContextHandleAdapter( runtime, ctx, - createClient, + () => createClient(ctx), schemaConfig, databaseProvider, request, @@ -4063,7 +4101,7 @@ export function buildNativeFactory( runtime, ctx, conn, - createClient, + () => createClient(ctx), schemaConfig, databaseProvider, request, @@ -5041,129 +5079,147 @@ export function buildNativeFactory( try { const { ctx, request, cancelToken, responseBodyStream } = unwrapTsfnPayload(error, payload); - const inspectorResponse = - await maybeHandleNativeInspectorRequest(ctx, request); - if (inspectorResponse) { - await cancelNativeHttpRequestBody(request.bodyStream); - return ( - await convertNativeHttpResponse( - inspectorResponse, - responseBodyStream, - ) - ).response; - } - - if (typeof config.onRequest !== "function") { - await cancelNativeHttpRequestBody(request.bodyStream); - return ( - await convertNativeHttpResponse( - new Response(null, { status: 404 }), - responseBodyStream, - ) - ).response; - } + return await runtime.runWithActorInvocationContext( + ctx, + async () => { + const inspectorResponse = + await maybeHandleNativeInspectorRequest( + ctx, + request, + ); + if (inspectorResponse) { + await cancelNativeHttpRequestBody( + request.bodyStream, + ); + return ( + await convertNativeHttpResponse( + inspectorResponse, + responseBodyStream, + ) + ).response; + } - const requestAbortController = new AbortController(); - const handlerRequest = buildNativeHttpRequest({ - ...request, - abortController: requestAbortController, - }); - const rawConnParams = - handlerRequest.headers.get(HEADER_CONN_PARAMS); - let requestCtx: - | ReturnType - | undefined; - let conn: ConnHandle | undefined; - let removeRequestAbortListener: (() => void) | undefined; - let cleanupDeferredToBody = false; - let cleanedUp = false; - const cleanupRequest = async () => { - if (cleanedUp) return; - cleanedUp = true; - removeRequestAbortListener?.(); - try { - await requestCtx?.dispose(); - } finally { - if (conn) { - await runtime.connDisconnect(conn); + if (typeof config.onRequest !== "function") { + await cancelNativeHttpRequestBody( + request.bodyStream, + ); + return ( + await convertNativeHttpResponse( + new Response(null, { status: 404 }), + responseBodyStream, + ) + ).response; } - } - }; - try { - const connParams = validateConnParams( - schemaConfig.connParamsSchema, - rawConnParams - ? JSON.parse(rawConnParams) - : undefined, - ); - conn = await callNative(() => - runtime.actorConnectConn( - ctx, - encodeValue(connParams), - request, - ), - ); - requestCtx = makeConnCtx( - ctx, - conn, - handlerRequest, - cancelToken, - ); - const ctxAbortSignal = requestCtx.abortSignal; - const abortRequest = () => - requestAbortController.abort(ctxAbortSignal.reason); - if (ctxAbortSignal.aborted) { - abortRequest(); - } else { - ctxAbortSignal.addEventListener( - "abort", - abortRequest, - { once: true }, - ); - removeRequestAbortListener = () => - ctxAbortSignal.removeEventListener( - "abort", - abortRequest, + + const requestAbortController = + new AbortController(); + const handlerRequest = buildNativeHttpRequest({ + ...request, + abortController: requestAbortController, + }); + const rawConnParams = + handlerRequest.headers.get(HEADER_CONN_PARAMS); + let requestCtx: + | ReturnType + | undefined; + let conn: ConnHandle | undefined; + let removeRequestAbortListener: + | (() => void) + | undefined; + let cleanupDeferredToBody = false; + let cleanedUp = false; + const cleanupRequest = async () => { + if (cleanedUp) return; + cleanedUp = true; + removeRequestAbortListener?.(); + try { + await requestCtx?.dispose(); + } finally { + if (conn) { + await runtime.connDisconnect(conn); + } + } + }; + try { + const connParams = validateConnParams( + schemaConfig.connParamsSchema, + rawConnParams + ? JSON.parse(rawConnParams) + : undefined, ); - } - const response = await config.onRequest( - requestCtx, - handlerRequest, - ); - if (!isResponseLike(response)) { - throw new Error( - "onRequest handler must return a Response", - ); - } - const conversion = await convertNativeHttpResponse( - response, - responseBodyStream, - ); - if (conversion.bodyCompletion) { - cleanupDeferredToBody = true; - void conversion.bodyCompletion - .then(cleanupRequest) - .catch((cleanupError) => { - logger().error({ - msg: "failed to clean up native streaming http request", - error: cleanupError, - }); - }); - } - return conversion.response; - } finally { - try { - // Handler completion ends upload ownership even when - // the Web Request body is locked or partly consumed. - await cancelNativeHttpRequestBody( - request.bodyStream, - ); - } finally { - if (!cleanupDeferredToBody) { - await cleanupRequest(); + conn = await callNative(() => + runtime.actorConnectConn( + ctx, + encodeValue(connParams), + request, + ), + ); + requestCtx = makeConnCtx( + ctx, + conn, + handlerRequest, + cancelToken, + ); + const ctxAbortSignal = requestCtx.abortSignal; + const abortRequest = () => + requestAbortController.abort( + ctxAbortSignal.reason, + ); + if (ctxAbortSignal.aborted) { + abortRequest(); + } else { + ctxAbortSignal.addEventListener( + "abort", + abortRequest, + { once: true }, + ); + removeRequestAbortListener = () => + ctxAbortSignal.removeEventListener( + "abort", + abortRequest, + ); + } + const response = await config.onRequest( + requestCtx, + handlerRequest, + ); + if (!isResponseLike(response)) { + throw new Error( + "onRequest handler must return a Response", + ); + } + const conversion = + await convertNativeHttpResponse( + response, + responseBodyStream, + ); + if (conversion.bodyCompletion) { + cleanupDeferredToBody = true; + void conversion.bodyCompletion + .then(cleanupRequest) + .catch((cleanupError) => { + logger().error({ + msg: "failed to clean up native streaming http request", + error: cleanupError, + }); + }); + } + return conversion.response; + } finally { + try { + // Handler completion ends upload ownership even when + // the Web Request body is locked or partly consumed. + await cancelNativeHttpRequestBody( + request.bodyStream, + ); + } finally { + if (!cleanupDeferredToBody) { + await cleanupRequest(); + } + } } - } - } + }, + ); } catch (error) { logger().error({ msg: "native onRequest failed", @@ -5312,21 +5368,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, + ); }, ), ]), @@ -5365,7 +5429,7 @@ export function buildNativeFactory( runtime, ctx, conn, - createClient, + () => createClient(ctx), schemaConfig, databaseProvider, jsRequest, @@ -5374,62 +5438,75 @@ export function buildNativeFactory( cancelToken, run !== undefined, ); - try { - if ( - !schemaConfig.queues || - !hasSchemaConfigKey(schemaConfig.queues, name) - ) { - return { status: "completed" }; - } - - const canPublish = getQueueCanPublish( - schemaConfig.queues, - name, - ); - if (canPublish && !(await canPublish(actorCtx))) { - throw forbiddenError(); - } - - const decodedBody = decodeValue(body); - if (wait) { + return await runtime.runWithActorInvocationContext( + ctx, + async () => { try { - const response = - await actorCtx.queue.enqueueAndWait( - name, - decodedBody, - { - timeout: - timeoutMs === undefined || - timeoutMs === null - ? undefined - : Number(timeoutMs), - }, - ); - return { - status: "completed", - response: - response === undefined - ? undefined - : encodeValue(response), - }; - } catch (error) { if ( - (error as { group?: string; code?: string }) - .group === "queue" && - (error as { group?: string; code?: string }) - .code === "timed_out" + !schemaConfig.queues || + !hasSchemaConfigKey(schemaConfig.queues, name) ) { - return { status: "timedOut" }; + return { status: "completed" }; } - throw error; - } - } - await actorCtx.queue.send(name, decodedBody); - return { status: "completed" }; - } finally { - await actorCtx.dispose(); - } + const canPublish = getQueueCanPublish( + schemaConfig.queues, + name, + ); + if (canPublish && !(await canPublish(actorCtx))) { + throw forbiddenError(); + } + + const decodedBody = decodeValue(body); + if (wait) { + try { + const response = + await actorCtx.queue.enqueueAndWait( + name, + decodedBody, + { + timeout: + timeoutMs === undefined || + timeoutMs === null + ? undefined + : Number(timeoutMs), + }, + ); + return { + status: "completed", + response: + response === undefined + ? undefined + : encodeValue(response), + }; + } catch (error) { + if ( + ( + error as { + group?: string; + code?: string; + } + ).group === "queue" && + ( + error as { + group?: string; + code?: string; + } + ).code === "timed_out" + ) { + return { status: "timedOut" }; + } + throw error; + } + } + + await actorCtx.queue.send(name, decodedBody); + return { status: "completed" }; + } finally { + await actorCtx.dispose(); + } + }, + ); }, ), serializeState: wrapNativeCallback( diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts index 383cec1d1b..f3cff43c5d 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts @@ -1,3 +1,7 @@ +import type { + ActorInvocationSpanContext, + ActorInvocationTraceContext, +} from "@/common/actor-telemetry-context"; import type { SqliteNativeMetrics, SqliteProfilingOptions, @@ -29,6 +33,35 @@ export interface RuntimeActorKeySegment { numberValue?: number; } +/** Resolves correlation at operation time so retained clients cannot freeze stale context. */ +export type CurrentActorInvocation = () => + | ActorInvocationTraceContext + | undefined; + +/** One open call from an actor out to another actor. */ +export interface RuntimeOutboundCall { + /** + * W3C context of the call's own span, to send to the callee so it parents to + * the call. Absent when tracing is disabled. + */ + readonly span?: ActorInvocationSpanContext; + /** + * Records the call's outcome. `error` is the failure encoded the way the + * bridge encodes errors, so a structured error keeps its group and code. + */ + finish(error?: string): void; +} + +/** + * Opens the span covering one call out to another actor, or returns `undefined` + * outside an invocation or on a runtime without invocation telemetry. Callers + * send their own context when it returns nothing. + */ +export type BeginOutboundCall = ( + actorName: string, + actionName: string, +) => RuntimeOutboundCall | undefined; + export interface RuntimeHttpRequest { method: string; uri: string; @@ -318,6 +351,7 @@ export interface RuntimeActorConfig { preloadMaxWorkflowBytes?: number; preloadMaxConnectionsBytes?: number; actions?: Array<{ name: string }>; + queues?: Array<{ name: string }>; inspectorTabs?: Array; } @@ -540,6 +574,29 @@ 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; + /** + * Opens the span covering one call this actor makes to another actor. See + * `BeginOutboundCall` for when this returns nothing. + */ + beginOutboundCall( + ctx: ActorContextHandle, + actorName: string, + actionName: string, + ): RuntimeOutboundCall | undefined; actorName(ctx: ActorContextHandle): string; actorKey(ctx: ActorContextHandle): RuntimeActorKeySegment[]; actorRegion(ctx: ActorContextHandle): string; diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts index b0fb460ccf..7d4dddb2eb 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts @@ -1,4 +1,5 @@ import { decodeBridgeRivetError, RivetError } from "@/actor/errors"; +import type { ActorInvocationTraceContext } from "@/common/actor-telemetry-context"; import type { WasmRuntimeBindings, WasmRuntimeConfig, @@ -23,6 +24,7 @@ import type { RuntimeKvEntry, RuntimeKvListOptions, RuntimeListenerConfig, + RuntimeOutboundCall, RuntimeQueueEnqueueAndWaitOptions, RuntimeQueueInspectMessage, RuntimeQueueMessage, @@ -535,6 +537,30 @@ 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(); + } + + actorInvocationTraceContext( + _ctx: ActorContextHandle, + ): ActorInvocationTraceContext | undefined { + return undefined; + } + + beginOutboundCall( + _ctx: ActorContextHandle, + _actorName: string, + _actionName: string, + ): RuntimeOutboundCall | undefined { + // Wasm carries no invocation telemetry, so a call goes out untraced + // rather than failing. + return undefined; + } + actorName(ctx: ActorContextHandle): string { return callHandle(asWasmActorContext(ctx), "name"); } diff --git a/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts b/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts index 9355ac780a..4b8d2251e2 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/fixtures/napi-runtime-server.ts @@ -1,6 +1,8 @@ import { existsSync } from "node:fs"; import { dirname, resolve } from "node:path"; import { fileURLToPath } from "node:url"; +import { trace } from "@opentelemetry/api"; +import { NodeTracerProvider } from "@opentelemetry/sdk-trace-node"; import { getEnginePath } from "@rivetkit/engine-cli"; import { z } from "zod/v4"; import { db } from "../../src/db/mod"; @@ -14,6 +16,10 @@ const repoEngineBinary = resolve( ); const endpoint = process.env.RIVETKIT_TEST_ENDPOINT ?? "http://127.0.0.1:6642"; + +// Register context propagation without exporting application spans; tests read their IDs. +new NodeTracerProvider().register(); +const applicationTracer = trace.getTracer("napi-runtime-fixture"); const connParamsSchema = z.object({ userId: z.string().min(1), }); @@ -53,6 +59,10 @@ const integrationActor = actor({ jobs: queue({ message: jobSchema }), }, onBeforeConnect: async () => {}, + onRequest: async (c, request) => { + await c.db.execute("SELECT ? AS path", new URL(request.url).pathname); + return new Response("ok", { status: 200 }); + }, actions: { ping: async (c) => { return c.conn.params.userId; @@ -113,6 +123,16 @@ 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"); + }, stateSnapshot: async (c) => { const kvValue = await c.kv.get("count"); return { @@ -120,6 +140,66 @@ const integrationActor = actor({ kvCount: kvValue ? Number(kvValue) : null, }; }, + // Both calls reach the queue before the test releases either one. + isolationProbe: async (c, token: string, fail: boolean) => { + c.log.warn({ correlation_token: token }, "isolation probe"); + if (!(await c.queue.next({ names: ["jobs"], timeout: 10_000 }))) { + throw new Error("isolation probe was not released"); + } + await c.db.execute("SELECT ? AS probe", token); + const client = c.client(); + await client.integrationActor + .getForId(c.actorId, { + params: { userId: "internal-integration-test" }, + }) + .getCount(); + await c.db.execute("SELECT ? AS probe2", token); + if (fail) { + throw new UserError("isolation probe failure", { + code: "isolation_probe_failed", + }); + } + return token; + }, + getCountUnderApplicationSpan: async (c) => { + return await applicationTracer.startActiveSpan( + "agent.generate", + async (span) => { + try { + await c.db.execute("SELECT 1 AS under_span"); + const client = c.client(); + const count = await client.integrationActor + .getForId(c.actorId, { + params: { userId: "internal-integration-test" }, + }) + .getCount(); + return { count, spanId: span.spanContext().spanId }; + } finally { + span.end(); + } + }, + ); + }, + // The queue gate releases the database work only after the reply. + insertAfterReply: (c, token: string) => { + c.waitUntil( + c.queue + .next({ names: ["jobs"], timeout: 10_000 }) + .then((message) => { + if (!message) + throw new Error("deferred work was not released"); + return c.db.execute("SELECT ? AS deferred", token); + }), + ); + return "replied"; + }, + consumeJob: async (c) => { + const message = await c.queue.next({ + names: ["jobs"], + timeout: 5_000, + }); + return message?.body ?? null; + }, getCountViaClient: async (c) => { const client = c.client(); return await client.integrationActor @@ -146,9 +226,25 @@ const integrationActor = actor({ }, }); +const runConsumerActor = actor({ + state: {}, + queues: { + runJobs: queue({ message: jobSchema }), + }, + run: async (c) => { + while (!c.aborted) { + await c.queue.waitForNames(["runJobs"], { + signal: c.abortSignal, + }); + } + }, + actions: {}, +}); + const registry = setup({ use: { integrationActor, + runConsumerActor, }, endpoint, namespace: process.env.RIVET_NAMESPACE ?? "default", diff --git a/rivetkit-typescript/packages/rivetkit/tests/fixtures/otlp-collector.ts b/rivetkit-typescript/packages/rivetkit/tests/fixtures/otlp-collector.ts new file mode 100644 index 0000000000..798d2af0b9 --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/tests/fixtures/otlp-collector.ts @@ -0,0 +1,60 @@ +import { createServer } from "node:http"; + +export interface OtlpCollector { + readonly endpoint: string; + spans(): Buffer[]; + close(): 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)); + 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); + }); + }); + + 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) => { + 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 65de48ff55..db49223a78 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/napi-runtime-integration.test.ts @@ -3,11 +3,30 @@ import { mkdtemp, readdir, readFile, rm } from "node:fs/promises"; import { tmpdir } from "node:os"; import { dirname, join } from "node:path"; import { fileURLToPath } from "node:url"; +import { context, propagation, trace } from "@opentelemetry/api"; +import { NodeTracerProvider } from "@opentelemetry/sdk-trace-node"; import getPort from "get-port"; -import { afterEach, describe, expect, test } from "vitest"; +import { afterEach, describe, expect, test, vi } from "vitest"; import { createClient } from "../src/client/mod"; +import { + type OtlpCollector, + startOtlpCollector, +} from "./fixtures/otlp-collector"; const TEST_DIR = dirname(fileURLToPath(import.meta.url)); + +// Register a context manager for caller spans and baggage; no exporter is needed. +new NodeTracerProvider().register(); +const testTracer = trace.getTracer("napi-runtime-integration"); + +/** Runs `run` with `rayId` in OpenTelemetry baggage under `rivet.ray.id`. */ +function withRayBaggage(rayId: string, run: () => Promise): Promise { + const baggage = propagation.createBaggage({ + "rivet.ray.id": { value: rayId }, + }); + return context.with(propagation.setBaggage(context.active(), baggage), run); +} + const FIXTURE_PATH = join(TEST_DIR, "fixtures", "napi-runtime-server.ts"); const NAMESPACE = "default"; const TOKEN = "dev"; @@ -19,9 +38,23 @@ let runtimeLogs = { let engineEndpoint: string | undefined; let storagePath: string | undefined; +function createIntegrationClient(endpoint: string, poolName: string) { + return createClient({ + endpoint, + poolName, + token: TOKEN, + namespace: NAMESPACE, + disableMetadataLookup: true, + }) as any; +} + +function runtimeOutput(): string { + return [runtimeLogs.stdout, runtimeLogs.stderr].filter(Boolean).join("\n"); +} + function childOutput(child: ChildProcess): string { void child; - return [runtimeLogs.stdout, runtimeLogs.stderr].filter(Boolean).join("\n"); + return runtimeOutput(); } async function engineOutput(): Promise { @@ -434,14 +467,248 @@ async function stopTestEngine(): Promise { } } +/** OTLP `SpanKind.CLIENT`. */ +const OTLP_SPAN_KIND_CLIENT = 3; + +interface ExportedSpan { + name: string; + traceId: string; + spanId: string; + parentSpanId?: string; + traceState?: string; + kind?: number; + endTimeUnixNano: bigint; + attributes: Record; + links: Array<{ traceId: string; spanId: string }>; +} + +/** Flattens OTLP/JSON export bodies into the spans they carry. */ +function exportedSpans(exports: Buffer[]): ExportedSpan[] { + type OtlpAttribute = { + key: string; + value: { stringValue?: string; intValue?: string | number }; + }; + type OtlpSpan = Omit & { + attributes?: OtlpAttribute[]; + endTimeUnixNano?: string | number; + kind?: number; + 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, + traceState: span.traceState, + kind: span.kind, + endTimeUnixNano: BigInt(span.endTimeUnixNano ?? 0), + attributes: Object.fromEntries( + (span.attributes ?? []).map((attribute) => [ + attribute.key, + attribute.value.stringValue ?? + (attribute.value.intValue === undefined + ? undefined + : String(attribute.value.intValue)), + ]), + ), + 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)); + } + const arrived = exportedSpans(exports) + .map((span) => `${span.name} (${span.spanId} < ${span.parentSpanId})`) + .join(", "); + throw new Error( + `timed out waiting for ${description}; exported: ${arrived}`, + ); +} + +function isSqliteSpan(span: ExportedSpan): boolean { + 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, + ); +} + +/** The `onRequest` invocation span carrying `rayId`, ignoring its children. */ +function findRequestInvocation( + spans: ExportedSpan[], + rayId: string, +): ExportedSpan | undefined { + return spans.find( + (span) => + span.attributes["rivet.invocation.type"] === "request" && + span.attributes["rivet.ray.id"] === rayId, + ); +} + +/** + * The `queue.receive` span of `actorName` under `parentSpanId`, or its root + * one when `parentSpanId` is undefined. + */ +function findQueueReceive( + spans: ExportedSpan[], + actorName: string, + parentSpanId: string | undefined, +): ExportedSpan | undefined { + return spans.find( + (span) => + span.name === `${actorName}/queue.receive` && + span.parentSpanId === parentSpanId, + ); +} + +/** The `queue.send` invocation span carrying `rayId`, ignoring its children. */ +function findQueueSendInvocation( + spans: ExportedSpan[], + rayId: string, +): ExportedSpan | undefined { + return spans.find( + (span) => + span.attributes["rivet.invocation.type"] === "queue_send" && + span.attributes["rivet.ray.id"] === rayId, + ); +} + +/** Polls until an invocation span has been exported for every named action. */ +async function waitForInvocationSpans( + exports: Buffer[], + 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, + // Export spans only for tests that read them. + ...(tracesEndpoint + ? { + OTEL_EXPORTER_OTLP_TRACES_ENDPOINT: tracesEndpoint, + OTEL_EXPORTER_OTLP_TRACES_PROTOCOL: "http/json", + OTEL_TRACES_SAMPLER: "always_on", + OTEL_BSP_SCHEDULE_DELAY: "10", + } + : {}), + ...extraEnv, + }, + stdio: ["ignore", "pipe", "pipe"], + }); + 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,57 +718,21 @@ 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); + const { endpoint, poolName, child } = await startTracedRuntime(); + runtime = child; await waitForEnvoy(runtime, endpoint, SERVICES_POOL_NAME, 30_000); await expectNormalRunnerConfig(endpoint, SERVICES_POOL_NAME); const servicesActorId = await createServicesActor(endpoint); await waitForActorStarted(endpoint, servicesActorId, 30_000); - const client = createClient({ - endpoint, - token: TOKEN, - namespace: NAMESPACE, - poolName, - disableMetadataLookup: true, - }) as any; + const client = createIntegrationClient(endpoint, poolName); + const actorKey = `napi-runtime-${crypto.randomUUID()}`; const handle = await waitForActorReady( () => - 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(); @@ -596,6 +827,7 @@ describe.sequential("native NAPI runtime integration", () => { code: "internal_error", message: "An internal error occurred", }); + await client.dispose(); const processId = servicesPid(); @@ -603,4 +835,618 @@ describe.sequential("native NAPI runtime integration", () => { runtime = undefined; await waitForProcessExit(processId, 5_000); }, 120_000); + + test("preserves vendor trace state across actor calls and ignores invalid trace versions", async () => { + collector = await startOtlpCollector( + await getPort({ host: "127.0.0.1" }), + ); + const traceExports = collector.spans(); + const { endpoint, poolName, child } = await startTracedRuntime( + collector.endpoint, + ); + runtime = child; + const traceId = "1234567890abcdef1234567890abcdef"; + const parentSpanId = "1234567890abcdef"; + const traceState = "vendor=opaque-value"; + const client = createIntegrationClient(endpoint, poolName); + const handle = await waitForActorReady( + () => + client.integrationActor.create( + [`trace-context-${crypto.randomUUID()}`], + { params: { userId: "integration-test" } }, + ), + 30_000, + ); + await waitForActorReady(() => handle.getCount(), 30_000); + const actorId = await handle.resolve(); + const url = new URL(await handle.getGatewayUrl()); + url.pathname = `${url.pathname.replace(/\/$/, "")}/action/getCountViaClient`; + // Every version calls the same actor, so each round takes the caller + // span that no earlier round has claimed. + const claimed = new Set(); + const callChain = (spans: ExportedSpan[]) => { + const caller = spans.find( + (span) => + span.attributes["rivet.actor.id"] === actorId && + span.attributes["rivet.action.name"] === + "getCountViaClient" && + !claimed.has(span.spanId), + ); + const hop = + caller && + spans.find( + (span) => + span.kind === OTLP_SPAN_KIND_CLIENT && + span.traceId === caller.traceId, + ); + const callee = + hop && spans.find((span) => span.parentSpanId === hop.spanId); + return { caller, hop, callee }; + }; + for (const version of ["00", "zz", "0A"]) { + const response = await fetch(url, { + method: "POST", + headers: { + "content-type": "application/json", + "x-rivet-encoding": "json", + "x-rivet-token": TOKEN, + "x-rivet-conn-params": JSON.stringify({ + userId: "integration-test", + }), + traceparent: `${version}-${traceId}-${parentSpanId}-01`, + tracestate: traceState, + }, + body: JSON.stringify({ args: [] }), + }); + expect(response.status).toBe(200); + await response.arrayBuffer(); + const { caller, hop, callee } = callChain( + await waitForSpans( + traceExports, + "caller and callee trace contexts", + (spans) => callChain(spans).callee !== undefined, + 10_000, + ), + ); + claimed.add(caller?.spanId ?? ""); + if (version === "00") { + expect(caller?.traceId).toBe(traceId); + expect(caller?.parentSpanId).toBe(parentSpanId); + for (const span of [caller, hop, callee]) + expect(span?.traceState).toBe(traceState); + } else { + expect(caller?.traceId).not.toBe(traceId); + expect(caller?.parentSpanId).toBeUndefined(); + for (const span of [caller, hop, callee]) + expect(span?.traceState || "").toBe(""); + } + } + await client.dispose(); + }, 120_000); + + test("keeps overlapping invocations of one actor telemetrically isolated", async () => { + collector = await startOtlpCollector( + await getPort({ host: "127.0.0.1" }), + ); + const traceExports = collector.spans(); + const { endpoint, poolName, child } = await startTracedRuntime( + collector.endpoint, + ); + runtime = child; + + const client = createIntegrationClient(endpoint, poolName); + const handle = await waitForActorReady( + () => + client.integrationActor.create( + [`napi-isolation-${crypto.randomUUID()}`], + { params: { userId: "integration-test" } }, + ), + 30_000, + ); + await waitForActorReady(() => handle.getCount(), 30_000); + + // Use the same actor to exercise isolation between concurrent invocations. + const okToken = crypto.randomUUID(); + const failToken = crypto.randomUUID(); + const results = Promise.allSettled([ + handle.isolationProbe(okToken, false), + handle.isolationProbe(failToken, true), + ]); + const [okLog, failLog] = await Promise.all([ + waitForRuntimeLog(okToken, 10_000), + waitForRuntimeLog(failToken, 10_000), + ]); + await handle.send("jobs", { id: okToken }); + await handle.send("jobs", { id: failToken }); + const [ok, failed] = await results; + expect(ok.status).toBe("fulfilled"); + expect(failed.status).toBe("rejected"); + + 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) => { + const hop = exported.find( + (span) => + span.kind === OTLP_SPAN_KIND_CLIENT && + span.traceId === probe.traceId, + ); + return ( + !!hop && + exported.some( + (span) => span.parentSpanId === hop.spanId, + ) && + exported.filter( + (span) => + isSqliteSpan(span) && + span.parentSpanId === probe.spanId, + ).length >= 2 + ); + }) + ); + }, + 20_000, + ); + + const probes = spans.filter( + (span) => span.attributes["rivet.action.name"] === "isolationProbe", + ); + expect(probes).toHaveLength(2); + + const rays = probes.map((probe) => probe.attributes["rivet.ray.id"]); + expect(new Set(rays).size).toBe(2); + expect(new Set(probes.map((probe) => probe.traceId)).size).toBe(2); + const failedProbes = probes.filter( + (probe) => probe.attributes["error.type"] !== undefined, + ); + expect(failedProbes).toHaveLength(1); + expect(failedProbes[0]?.attributes["error.type"]).toBe( + "user.isolation_probe_failed", + ); + + for (const probe of probes) { + const owned = spans.filter( + (span) => + isSqliteSpan(span) && span.parentSpanId === probe.spanId, + ); + expect(owned.length).toBeGreaterThanOrEqual(2); + for (const span of owned) { + expect(span.traceId).toBe(probe.traceId); + expect(span.attributes["rivet.ray.id"]).toBe( + probe.attributes["rivet.ray.id"], + ); + } + } + + 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(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"], + ); + } + } + + for (const line of [okLog, failLog]) { + expect(line).toContain(`actorId=${await handle.resolve()}`); + expect(line).toMatch(/ trace_id=[0-9a-f]{32}( |$)/); + expect(line).toMatch(/ span_id=[0-9a-f]{16}( |$)/); + } + const rayOf = (line: string) => + / rayId=([A-Za-z0-9_-]+)/.exec(line)?.[1]; + expect(rayOf(okLog)).toBeDefined(); + expect(rayOf(okLog)).not.toBe(rayOf(failLog)); + expect(rays).toContain(rayOf(okLog)); + expect(rays).toContain(rayOf(failLog)); + + // SQLite and actor calls inherit the active application span. + const underApp = await handle.getCountUnderApplicationSpan(); + const applicationSpans = await waitForSpans( + traceExports, + "SQLite, outgoing call, and callee under the application span", + (exported) => { + const hop = exported.find( + (span) => + span.kind === OTLP_SPAN_KIND_CLIENT && + span.parentSpanId === underApp.spanId, + ); + return ( + !!hop && + exported.some((span) => span.parentSpanId === hop.spanId) && + exported.some( + (span) => + isSqliteSpan(span) && + span.parentSpanId === underApp.spanId, + ) + ); + }, + 10_000, + ); + const appHop = applicationSpans.find( + (span) => + span.kind === OTLP_SPAN_KIND_CLIENT && + span.parentSpanId === underApp.spanId, + ); + expect( + applicationSpans.find( + (span) => span.parentSpanId === appHop?.spanId, + )?.attributes["rivet.action.name"], + ).toBe("getCount"); + expect( + applicationSpans.find( + (span) => + isSqliteSpan(span) && span.parentSpanId === underApp.spanId, + )?.attributes["rivet.operation.name"], + ).toBe("execute"); + + await client.dispose(); + }, 120_000); + + test("carries a caller-supplied ray through requests, queue sends, and work after the reply", async () => { + collector = await startOtlpCollector( + await getPort({ host: "127.0.0.1" }), + ); + const traceExports = collector.spans(); + const { endpoint, poolName, child } = await startTracedRuntime( + collector.endpoint, + ); + runtime = child; + + const client = createIntegrationClient(endpoint, poolName); + const handle = await waitForActorReady( + () => + client.integrationActor.create( + [`napi-caller-ray-${crypto.randomUUID()}`], + { params: { userId: "integration-test" } }, + ), + 30_000, + ); + await waitForActorReady(() => handle.getCount(), 30_000); + + const callerRay = `caller-${crypto.randomUUID()}`; + await withRayBaggage(callerRay, () => handle.getCount()); + const rayedGetCount = await waitForSpans( + traceExports, + "the getCount invocation carrying the caller-supplied ray", + (exported) => + exported.some( + (span) => + span.attributes["rivet.action.name"] === "getCount" && + span.attributes["rivet.ray.id"] === callerRay, + ), + 10_000, + ); + expect( + rayedGetCount.find( + (span) => span.attributes["rivet.ray.id"] === callerRay, + )?.attributes["rivet.invocation.type"], + ).toBe("action"); + + // Fetch inherits the active span and baggage without explicit headers. + const requestRay = `request-${crypto.randomUUID()}`; + const callerSpan = testTracer.startSpan("request.handle"); + const underCallerSpan = (run: () => Promise) => + withRayBaggage(requestRay, () => + context.with(trace.setSpan(context.active(), callerSpan), run), + ); + const response = await underCallerSpan(() => handle.fetch("hello")); + expect(response.status).toBe(200); + const requestSpans = await waitForSpans( + traceExports, + "the onRequest invocation and its sqlite span", + (exported) => { + const request = findRequestInvocation(exported, requestRay); + return ( + request !== undefined && + exported.some( + (span) => + isSqliteSpan(span) && + span.parentSpanId === request.spanId, + ) + ); + }, + 10_000, + ); + const requestSpan = findRequestInvocation(requestSpans, requestRay); + expect(requestSpan?.name).toBe("integrationActor/onRequest"); + expect(requestSpan?.attributes["http.response.status_code"]).toBe( + "200", + ); + expect(requestSpan?.traceId).toBe(callerSpan.spanContext().traceId); + expect(requestSpan?.parentSpanId).toBe(callerSpan.spanContext().spanId); + + // Headers the caller set on the request win over the active context. + const explicitRay = `explicit-${crypto.randomUUID()}`; + const explicitTraceId = crypto.randomUUID().replaceAll("-", ""); + const explicitSpanId = crypto + .randomUUID() + .replaceAll("-", "") + .slice(0, 16); + const explicitResponse = await underCallerSpan(() => + handle.fetch("hello", { + headers: { + "x-rivet-ray-id": explicitRay, + traceparent: `00-${explicitTraceId}-${explicitSpanId}-01`, + }, + }), + ); + callerSpan.end(); + expect(explicitResponse.status).toBe(200); + const explicitSpans = await waitForSpans( + traceExports, + "the onRequest invocation under the caller's own headers", + (exported) => + findRequestInvocation(exported, explicitRay) !== undefined, + 10_000, + ); + const explicitSpan = findRequestInvocation(explicitSpans, explicitRay); + expect(explicitSpan?.traceId).toBe(explicitTraceId); + expect(explicitSpan?.parentSpanId).toBe(explicitSpanId); + + // Release deferred work after the reply; its invocation must remain open. + const deferredToken = crypto.randomUUID(); + expect(await handle.insertAfterReply(deferredToken)).toBe("replied"); + await handle.send("jobs", { id: deferredToken }); + const deferredSpans = await waitForSpans( + traceExports, + "the insertAfterReply invocation and its deferred sqlite span", + (exported) => { + const invocation = findInvocation(exported, "insertAfterReply"); + return ( + invocation !== undefined && + exported.some( + (span) => + isSqliteSpan(span) && + span.parentSpanId === invocation.spanId, + ) + ); + }, + 10_000, + ); + const deferredInvocation = findInvocation( + deferredSpans, + "insertAfterReply", + ); + const deferredSqlite = deferredSpans.find( + (span) => + isSqliteSpan(span) && + span.parentSpanId === deferredInvocation?.spanId, + ); + expect(deferredSqlite).toBeDefined(); + expect( + deferredInvocation !== undefined && + deferredSqlite !== undefined && + deferredInvocation.endTimeUnixNano >= + deferredSqlite.endTimeUnixNano, + ).toBe(true); + + // The action receipt links to the send and keeps the consuming action’s ray. + const queueRay = `queue-${crypto.randomUUID()}`; + await withRayBaggage(queueRay, () => + handle.send("jobs", { id: "job-42" }), + ); + expect(await handle.consumeJob()).toEqual({ id: "job-42" }); + const queueSpans = await waitForSpans( + traceExports, + "the queue send invocation, the consuming action, and its receipt", + (exported) => { + const consumer = findInvocation(exported, "consumeJob"); + return ( + findQueueSendInvocation(exported, queueRay) !== undefined && + consumer !== undefined && + findQueueReceive( + exported, + "integrationActor", + consumer.spanId, + ) !== undefined + ); + }, + 10_000, + ); + const queueSend = findQueueSendInvocation(queueSpans, queueRay); + expect(queueSend?.name).toBe("integrationActor/queue.send"); + expect(queueSend?.attributes).toMatchObject({ + "rivet.invocation.type": "queue_send", + "rivet.queue.name": "jobs", + }); + const consumer = findInvocation(queueSpans, "consumeJob"); + const receipt = findQueueReceive( + queueSpans, + "integrationActor", + consumer?.spanId, + ); + expect(receipt?.attributes).toMatchObject({ + "rivet.queue.name": "jobs", + "rivet.ray.id": consumer?.attributes["rivet.ray.id"], + }); + expect(receipt?.links).toEqual([ + { traceId: queueSend?.traceId, spanId: queueSend?.spanId }, + ]); + + // Without an invocation, the receipt is a root span carrying the sender’s ray. + const runRay = `run-${crypto.randomUUID()}`; + const runConsumer = client.runConsumerActor.getOrCreate([ + `napi-run-consumer-${crypto.randomUUID()}`, + ]); + await withRayBaggage(runRay, () => + runConsumer.send("runJobs", { id: "job-run" }), + ); + const runSpans = await waitForSpans( + traceExports, + "the run handler's receipt of a queue message", + (exported) => + findQueueSendInvocation(exported, runRay) !== undefined && + findQueueReceive(exported, "runConsumerActor", undefined) !== + undefined, + 10_000, + ); + const runSend = findQueueSendInvocation(runSpans, runRay); + const runReceipt = findQueueReceive( + runSpans, + "runConsumerActor", + undefined, + ); + expect(runReceipt?.attributes).toMatchObject({ + "rivet.queue.name": "runJobs", + "rivet.ray.id": runRay, + }); + expect(runReceipt?.links).toEqual([ + { traceId: runSend?.traceId, spanId: runSend?.spanId }, + ]); + + traceExports.length = 0; + await expect(handle.sqliteFailure()).rejects.toMatchObject({ + code: expect.any(String), + }); + const failureSpans = await waitForSpans( + traceExports, + "sqliteFailure invocation and failed sqlite spans", + (spans) => + spans.some(isFailedSqliteSpan) && + findInvocation(spans, "sqliteFailure") !== undefined, + 10_000, + ); + const failedSqlite = failureSpans.find(isFailedSqliteSpan); + expect(failedSqlite?.parentSpanId).toBe( + findInvocation(failureSpans, "sqliteFailure")?.spanId, + ); + + // Scheduled work starts a new trace linked to its origin. + traceExports.length = 0; + const scheduleToken = crypto.randomUUID(); + expect(await handle.scheduleTrace(scheduleToken)).toBe(scheduleToken); + const scheduleSpans = await waitForInvocationSpans( + traceExports, + ["scheduleTrace", "scheduledTrace"], + 15_000, + ); + const definer = findInvocation(scheduleSpans, "scheduleTrace"); + const scheduled = findInvocation(scheduleSpans, "scheduledTrace"); + expect(scheduled?.attributes["rivet.invocation.type"]).toBe( + "scheduled", + ); + expect(scheduled?.attributes["rivet.ray.id"]).toBe( + definer?.attributes["rivet.ray.id"], + ); + // Ensure the scheduled action succeeded before checking its trace. + expect(scheduled?.attributes["error.type"]).toBeUndefined(); + expect(scheduled?.traceId).not.toBe(definer?.traceId); + expect(scheduled?.links).toEqual([ + { traceId: definer?.traceId, spanId: definer?.spanId }, + ]); + await client.dispose(); + }, 120_000); + + 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 = createIntegrationClient(endpoint, poolName); + const handle = await waitForActorReady( + () => + client.integrationActor.create( + [`napi-telemetry-failure-${crypto.randomUUID()}`], + { params: { userId: "integration-test" } }, + ), + 30_000, + ); + + expect(await waitForActorReady(() => handle.getCount(), 30_000)).toBe( + 0, + ); + expect( + await waitForActorReady( + () => handle.validatedAction({ amount: 4 }), + 30_000, + ), + ).toBe(4); + + await client.dispose(); + }, 120_000); + + test("keeps actor behavior intact when the trace exporter is slow", async () => { + // Stall exports long enough to saturate the small queue below. + collector = await startOtlpCollector( + await getPort({ host: "127.0.0.1" }), + { + responseDelayMs: 120_000, + }, + ); + const { endpoint, poolName, child } = await startTracedRuntime( + collector.endpoint, + { + OTEL_BSP_MAX_QUEUE_SIZE: "8", + OTEL_BSP_MAX_EXPORT_BATCH_SIZE: "4", + // Disable Rust SDK logs so only the JS bridge can satisfy the log assertion. + RUST_LOG: "warn,opentelemetry_sdk=off", + }, + ); + runtime = child; + + const client = createIntegrationClient(endpoint, poolName); + const handle = await waitForActorReady( + () => + client.integrationActor.create( + [`napi-telemetry-slow-${crypto.randomUUID()}`], + { params: { userId: "integration-test" } }, + ), + 30_000, + ); + + const started = Date.now(); + for (let index = 1; index <= 12; index += 1) { + expect( + await waitForActorReady(() => handle.increment(1), 30_000), + ).toMatchObject({ count: index }); + } + const elapsed = Date.now() - started; + + // Actions must finish before the stalled collector responds. + expect(elapsed).toBeLessThan(60_000); + + // The processor reports dropped spans on its own export cycle, which + // runs after the actions return, so there is nothing to await. + await vi.waitFor( + () => { + expect(runtimeOutput()).toContain( + "BatchSpanProcessor.SpanDroppingStarted", + ); + }, + { timeout: 15_000, interval: 250 }, + ); + + expect(await waitForActorReady(() => handle.getCount(), 30_000)).toBe( + 12, + ); + + await client.dispose(); + }, 180_000); }); diff --git a/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts b/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts index afe2e1db3a..867a9a6743 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, @@ -157,6 +158,14 @@ class FakeActorContext { return this.runtimeBag; } + invocationTraceContext(): undefined { + return undefined; + } + + sameActorInstance(other: FakeActorContext): boolean { + return this === other; + } + actorId(): string { return "parity-actor"; } @@ -305,6 +314,7 @@ function fakeNapiBindings(scenario: ParityScenario) { NapiActorFactory: FakeActorFactory, CancellationToken: FakeCancellationToken, ActorContext: class {}, + setTelemetryLogSink: () => {}, }; } @@ -335,7 +345,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),