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..c666b46e80 --- /dev/null +++ b/docs-internal/engine/rivetkit-telemetry.md @@ -0,0 +1,176 @@ +# RivetKit telemetry + +Internal reference for actor traces, invocation metrics, and log correlation. Core owns telemetry; runtime adapters activate its context in the host language. See [NAPI bridge](napi-bridge.md) for binding conventions and [Core internals](rivetkit-core-internals.md) for dispatch and lifecycle wiring. + +## Ownership + +- `rivetkit-core::telemetry` owns invocation spans, completion, and operation spans. `ActorMetrics` owns metric labels and recording. +- `registry/napi-runtime.ts` resolves the current invocation and passes the active JavaScript span to Core. NAPI bindings only translate types. +- TypeScript and Rust clients send ray and W3C trace headers. An actor-owned client also opens an outbound call span in Core. +- Core exports its spans through the Rust OpenTelemetry (OTel) SDK. Applications export their own spans through their own SDK. Shared trace IDs and parent span IDs connect the two pipelines. + +## Spans + +All Core telemetry spans use the `rivetkit::telemetry` tracing target. Log layers exclude that target; the export layer excludes unrelated diagnostic spans. + +| Span | Kind | Parent or link | +| --- | --- | --- | +| `{actor}/{action}` | `server` | Incoming parent, otherwise root | +| `{actor}/{action}` for a scheduled fire | `internal` | New root linked to the schedule's origin | +| `{actor}/onRequest` | `server` | Incoming parent, otherwise root | +| `{actor}/queue.send` for an external send | `producer` | Incoming parent, otherwise root | +| `{callee}/{action}` for an actor-owned client call | `client` | Active application span, otherwise invocation | +| `rivet.sqlite.{operation}` | `internal` | Active application span, otherwise invocation | +| `{actor}/queue.receive` | `consumer` | Active application span, invocation, or root; linked to the message's origin | + +### Attributes + +| Scope | Attributes | +| --- | --- | +| Actor identity | `rivet.actor.id`, `rivet.actor.name`, `rivet.actor.key` | +| Correlation | `rivet.ray.id` | +| Invocation | `rivet.invocation.type`, `otel.status_code`, `error.type` on failure | +| Action or scheduled invocation | `rivet.action.name` | +| Raw HTTP request | `http.request.method`, `http.response.status_code` | +| Queue send or receipt | `rivet.queue.name` | +| SQLite operation | `rivet.operation.system`, `rivet.operation.name` | + +Raw request spans use the fixed name `onRequest`, not the request path. A 5xx response marks the invocation as failed with the status as `error.type`. A handler error records its `group.code`; Core generates the HTTP response after that invocation ends. + +SQLite and outbound actor-call spans are marked as errors with `actor.operation_abandoned` if tracking ends without a recorded result. This indicates an unknown outcome, not a confirmed operation failure, and is not returned to callers. + +`ActorMetrics::label_action_name` and `label_queue_name` replace undeclared names with `_OTHER`. Use these bounded names for span names and metric labels as well as queue attributes. This follows the OTel fallback convention for unknown caller-supplied values and prevents unbounded metric series. + +## Metrics and logs + +- `rivetkit_actor_invocations_total` +- `rivetkit_actor_invocation_duration_seconds` + +Both metrics use actor name, action name, invocation type, and result labels. Invocation types are `action`, `scheduled`, `request`, and `queue_send`. Request and queue-send invocations use `onRequest` and `queue.send` as their action labels. + +Duration uses `MICRO_BUCKETS` to distinguish invocations shorter than the default Prometheus histogram's first bucket of 5 ms. + +Invocation loggers bind actor ID, name, key, and ray ID. Sampled invocation context also supplies `trace_id` and `span_id`. Actor fields follow existing TypeScript log casing; trace fields follow OTel correlation conventions. A Pino child logger retains the context it was created with; do not retain an action logger for later unrelated work. + +## Context propagation + +Core accepts `x-rivet-ray-id`, `traceparent`, and `tracestate` on actions, raw requests, and queue sends. + +A ray correlates related work across multiple traces and remains available when spans are sampled out. Trace and span IDs describe the individual traces and their parent relationships. + +- Rays accept 1–128 characters from `[A-Za-z0-9_-]`. Missing or invalid rays become UUIDs. +- Invalid W3C trace context starts a root span without rejecting the request. +- An actor-owned client uses the calling invocation's ray. External clients read `rivet.ray.id` from active OTel baggage. The Rust client's configured `ClientConfig::ray_id` is a fallback. +- Outbound parent precedence is the call span, the active application span, then the invocation span. +- TypeScript action calls, `handle.fetch()`, and queue sends propagate context. Explicit raw-request headers win. Setting either `traceparent` or `tracestate` preserves the caller's whole pair rather than mixing contexts. +- Per-call context replaces static client telemetry headers. The Rust client reads context through `tracing-opentelemetry`; without that layer, no active W3C context is available. + +Rust shares names and validation through `rivetkit-client-protocol::telemetry_headers`. TypeScript defines header names in `common/actor-router-consts.ts` and baggage/trace handling in `common/otel-context.ts`. + +### JavaScript context + +Core cannot read the JavaScript span stack. `runWithActorInvocationContext` activates the invocation in `AsyncLocalStorage` and the OTel context manager while the callback runs. Core-bound operations then resolve that invocation and bind a more specific application span when present. + +```text +Incoming request + → Core invocation span + → JavaScript handler with invocation context active + → Application span + → SQLite operation or outbound actor-call span +``` + +- Retained database, client, and schedule handles resolve the active invocation only when it belongs to the same Core actor generation; otherwise they retain their creation context. Core uses pointer identity through `Arc::ptr_eq` for this check. +- The JavaScript context lookup runs inside invocations. The native binding call is skipped when the active span is absent or already matches the invocation span. +- A transaction binds its parent at `beginTransaction`. Statements and commit retain that parent. +- Queue sends, receipts, and `waitUntil` use the same invocation resolution. KV operations have no telemetry to attach. +- `@opentelemetry/api` is a required dependency for reading outbound context and activating inbound context. RivetKit does not register a global context manager; the application owns it. + +Caller trace context provides correlation, not identity or authorization. Operators can strip the three incoming headers when they do not trust caller-selected correlation values. + +## Invocation lifecycle + +`ActorInvocation` owns the span, timer, metric labels, and exactly-once completion. Action, schedule, raw-request, and queue-send dispatch create it before enqueueing work. + +1. The adapter runs the callback with invocation context active. +2. Sending the reply records the outcome and invocation metrics, and adds a `reply sent` event. +3. The span ends after the last tracked `waitUntil` task settles. Deferred database work and logs retain invocation context. + +`c.keepAwake` work is awaited within the callback. Rejected enqueue attempts count as `result=error`. Unfinished invocations record `actor.dropped_reply`; scheduled work also stores that identity in `_rivet_schedule_history` for `cronHistory()`. + +Span kind derives from invocation type. Completion and metrics remain on the same path when tracing is disabled or sampled out. + +## Persisted origins + +Schedules and queue messages capture their sending or defining invocation's ray and W3C context. An active application span is the origin when one exists; otherwise the invocation span is used. + +Origins currently use versioned BARE values in `_rivet_meta`: + +| Origin | Key | +| --- | --- | +| Schedule | `schedule_trace_context:{event_id}` | +| Queue message | `queue_trace_context:{message_id}` | + +Write and delete origins in the same batch as their owning rows. Schedule deletion that misses its row leaves the origin untouched. Malformed origins are logged and ignored without preventing delivery or execution. + +This storage is pending migration into columns on `_rivet_schedule_events` and `_rivet_queue`. It is an exception to the bootstrap-only role described in Core internals, not a precedent for new runtime metadata. + +### Schedules + +- Creation and redefinition capture the current origin. Recurring re-registration refreshes it even when cadence stays unchanged. +- Each fire starts a new trace linked to the defining span and retains the defining ray. This avoids keeping one trace open for the lifetime of a recurring schedule. +- Schedules defined outside an invocation, or before origin storage existed, have no origin; each fire gets a fresh ray. +- Older runtimes ignore origin metadata. Redefining a schedule on an older runtime can leave a stale origin until a newer runtime redefines it. Schedule behavior is unaffected. + +### Queues + +- External sends create `queue.send` invocations. The span ends when the send is acknowledged or the sender's completion wait ends. +- `next`, `nextBatch`, and `waitForNames` create one `queue.receive` span per message in Core. It ends when the message is handed back and links to its stored origin. +- Inside an invocation, the receipt uses the receiving invocation's ray. From `run`, the receipt is a root span with the sender's ray. +- Receipt spans do not cover subsequent application processing. + +## Export configuration + +Set OTel variables on the actor runner process: + +```sh +OTEL_SERVICE_NAME=internal-api +OTEL_EXPORTER_OTLP_TRACES_ENDPOINT=http://collector:4318/v1/traces +OTEL_EXPORTER_OTLP_TRACES_PROTOCOL=http/protobuf +``` + +Core's `telemetry::export` module is behind `native-runtime`. Hosts attach `export::layer()` to their subscriber and call `export::flush_best_effort()` during shutdown. + +- An OTLP endpoint enables native export unless standard SDK-disable or exporter controls disable it. Sampling, resources, service name, and batching use standard OTel environment variables. +- `configured_protocol` reads the traces-specific protocol variable before the general OTLP variable. Supported values are `grpc`, `http/protobuf`, and `http/json`; the default is `http/protobuf`. Invalid values produce an error. Selection is explicit because the exporter builders use feature-dependent defaults rather than reading these variables. +- Both `OTEL_EXPORTER_OTLP_HEADERS` and `OTEL_EXPORTER_OTLP_TRACES_HEADERS` are supported. +- Batches export in the background. Clean shutdown attempts a bounded flush; a hard kill can lose queued spans. Export failures do not fail actor work. +- The batch queue is bounded by `OTEL_BSP_MAX_QUEUE_SIZE`, default 2,048. Overflow drops spans. Sampling reduces exports but does not remove construction or propagation work. Export CPU contention can affect latency; measurements belong in benchmark artifacts. + +### SDK diagnostics + +The NAPI log bridge forwards `opentelemetry_sdk=warn` events to Pino, including `BatchSpanProcessor.SpanDroppingStarted`. + +- Keep `internal-logs` enabled on `opentelemetry` and `opentelemetry_sdk`. +- Each registry replaces the sink because the previous callback's Node worker may have exited. +- Unreference the sink's `ThreadsafeFunction` so it does not keep the event loop alive. +- Older addons may lack `setTelemetryLogSink`. Log a warning rather than preventing actor startup. + +### Configuration failures + +| Symptom | Check | +| --- | --- | +| Application spans appear, Core spans do not | Configuring `NodeSDK` in code does not configure Rust export. Set the runner's OTel environment variables. | +| Application spans start separate traces | Register an OTel context manager. `NodeSDK.start()` and `NodeTracerProvider.register()` do this; setting a bare tracer provider alone does not. | +| Application and Core spans appear under different services | Configure `service.name` on a manually constructed JavaScript provider, or ensure its resource reads the environment. | + +## Data policy + +Record actor identity, bounded action/queue names, invocation type, HTTP method/status, ray/trace/span IDs, operation names, and error identity. Do not record action arguments/results, connection parameters, SQL text/bindings, actor state, arbitrary headers, or raw error messages. + +## Limitations + +- No dedicated spans for WebSocket handlers, lifecycle hooks, connection callbacks, KV, or actor-state operations. +- WebSocket action messages carry no caller ray or trace context. Inspector actions do not inherit caller context. +- Actor creation rays are not delivered to the actor runtime. +- No Wasm host span export; the adapter retains rays but does not propagate invocation trace context. +- No category opt-in API, custom sampling API, dedicated Effect integration, or `rivetkit/unstable/otel`. Standard OTel application context already works for outbound calls. 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/engine/artifacts/errors/actor.operation_abandoned.json b/rivetkit-rust/engine/artifacts/errors/actor.operation_abandoned.json new file mode 100644 index 0000000000..ec1e8234b8 --- /dev/null +++ b/rivetkit-rust/engine/artifacts/errors/actor.operation_abandoned.json @@ -0,0 +1,5 @@ +{ + "code": "operation_abandoned", + "group": "actor", + "message": "Operation tracking ended before a result was recorded." +} \ No newline at end of file diff --git a/rivetkit-rust/packages/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..faf8453c77 100644 --- a/rivetkit-rust/packages/client-protocol/src/lib.rs +++ b/rivetkit-rust/packages/client-protocol/src/lib.rs @@ -1,4 +1,5 @@ pub mod generated; +pub mod telemetry_headers; pub mod versioned; // Re-export latest. diff --git a/rivetkit-rust/packages/client-protocol/src/telemetry_headers.rs b/rivetkit-rust/packages/client-protocol/src/telemetry_headers.rs new file mode 100644 index 0000000000..ce4e14e270 --- /dev/null +++ b/rivetkit-rust/packages/client-protocol/src/telemetry_headers.rs @@ -0,0 +1,37 @@ +//! Headers that carry a caller's ray and W3C trace context into an actor, +//! shared by every Rust peer of the actor HTTP surface: the runtime that reads +//! them and the clients that send them. + +/// Bounded correlation string that follows work across actors, surviving +/// sampling and trace-root boundaries. +pub const HEADER_RIVET_RAY_ID: &str = "x-rivet-ray-id"; +pub const HEADER_TRACEPARENT: &str = "traceparent"; +pub const HEADER_TRACESTATE: &str = "tracestate"; + +/// W3C Baggage key that carries a ray through application code, so a request +/// handler that received a ray can hand it to the actors it calls. +pub const RAY_BAGGAGE_KEY: &str = "rivet.ray.id"; + +const RAY_ID_MAX_LEN: usize = 128; + +/// Returns `value` when it is a ray the runtime accepts: 1 to 128 characters +/// of `[A-Za-z0-9_-]`. The value arrives from a caller, so anything else +/// counts as absent, and the receiving invocation mints a fresh ray instead. +pub fn bounded_ray_id(value: &str) -> Option<&str> { + let valid = !value.is_empty() + && value.len() <= RAY_ID_MAX_LEN + && value + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_')); + if valid { Some(value) } else { None } +} + +/// Formats a version 00 W3C `traceparent` from the parts of a span context, +/// so every Rust peer writes the header the same way. +pub fn format_traceparent( + trace_id: impl std::fmt::Display, + span_id: impl std::fmt::Display, + trace_flags: u8, +) -> String { + format!("00-{trace_id}-{span_id}-{trace_flags:02x}") +} diff --git a/rivetkit-rust/packages/client/Cargo.toml b/rivetkit-rust/packages/client/Cargo.toml index edcadab797..7dc2cf96f8 100644 --- a/rivetkit-rust/packages/client/Cargo.toml +++ b/rivetkit-rust/packages/client/Cargo.toml @@ -13,6 +13,7 @@ anyhow = "1.0" base64 = "0.22.1" bytes = { workspace = true } futures-util = "0.3.31" +opentelemetry = { version = "0.28", default-features = false, features = ["trace"] } parking_lot.workspace = true reqwest = { version = "0.12.12", default-features = false, features = ["json", "charset", "http2", "macos-system-configuration", "rustls-tls-native-roots", "rustls-tls-webpki-roots"] } rivetkit-client-protocol.workspace = true @@ -24,12 +25,14 @@ serde_json = "1.0" tokio = { version = "1", features = ["full"] } tokio-tungstenite = { version = "0.26.1", features = ["rustls-tls-native-roots", "rustls-tls-webpki-roots", "handshake"] } tracing = "0.1.41" +tracing-opentelemetry = { version = "0.29", default-features = false } tungstenite = "0.26.2" urlencoding = "2.1.3" vbare = "0.0.4" [dev-dependencies] axum = { workspace = true, features = ["ws"] } +opentelemetry_sdk = { version = "0.28", default-features = false, features = ["trace"] } tracing-subscriber = { version = "0.3.19", features = ["env-filter", "std", "registry"]} tempfile = "3.10.1" tokio-test = "0.4.3" diff --git a/rivetkit-rust/packages/client/src/client.rs b/rivetkit-rust/packages/client/src/client.rs index 79a8a622b4..7578be94fb 100644 --- a/rivetkit-rust/packages/client/src/client.rs +++ b/rivetkit-rust/packages/client/src/client.rs @@ -46,6 +46,9 @@ pub struct ClientConfig { pub encoding: EncodingKind, pub transport: TransportKind, pub headers: Option>, + /// Ray sent on every actor request made outside an active baggage + /// context, so work this client causes can be found by one string. + pub ray_id: Option, pub max_input_size: Option, pub disable_metadata_lookup: bool, } @@ -60,6 +63,7 @@ impl ClientConfig { encoding: EncodingKind::Bare, transport: TransportKind::WebSocket, headers: None, + ray_id: None, max_input_size: None, disable_metadata_lookup: false, } @@ -107,6 +111,16 @@ impl ClientConfig { self } + /// Sets the ray sent on every actor request. It must be 1 to 128 + /// characters of `[A-Za-z0-9_-]`, the same bound the runtime applies; a + /// value outside it is dropped with a warning when the client is built. + /// A ray carried in the active OpenTelemetry baggage under `rivet.ray.id` + /// takes precedence per call. + pub fn ray_id(mut self, ray_id: impl Into) -> Self { + self.ray_id = Some(ray_id.into()); + self + } + pub fn max_input_size(mut self, max_input_size: usize) -> Self { self.max_input_size = Some(max_input_size); self @@ -153,6 +167,7 @@ impl Client { config.namespace, config.pool_name, config.headers, + config.ray_id, config.max_input_size, config.disable_metadata_lookup, ); diff --git a/rivetkit-rust/packages/client/src/remote_manager.rs b/rivetkit-rust/packages/client/src/remote_manager.rs index adf46e666d..0409553741 100644 --- a/rivetkit-rust/packages/client/src/remote_manager.rs +++ b/rivetkit-rust/packages/client/src/remote_manager.rs @@ -1,23 +1,30 @@ -use anyhow::{anyhow, Context, Result}; -use base64::{engine::general_purpose, engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; +use anyhow::{Context, Result, anyhow}; +use base64::{Engine as _, engine::general_purpose, engine::general_purpose::URL_SAFE_NO_PAD}; use bytes::Bytes; +use opentelemetry::baggage::BaggageExt as _; +use opentelemetry::trace::TraceContextExt as _; use reqwest::{ - header::{HeaderMap, HeaderName, HeaderValue, USER_AGENT}, Method, + header::{HeaderMap, HeaderName, HeaderValue, USER_AGENT}, +}; +use rivetkit_client_protocol::telemetry_headers::{ + HEADER_RIVET_RAY_ID, HEADER_TRACEPARENT, HEADER_TRACESTATE, RAY_BAGGAGE_KEY, bounded_ray_id, + format_traceparent, }; use 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, }; @@ -29,6 +36,7 @@ pub struct RemoteManager { namespace: String, pool_name: String, headers: HashMap, + ray_id: Option, max_input_size: usize, disable_metadata_lookup: bool, resolved_config: Arc>, @@ -126,6 +134,7 @@ impl RemoteManager { namespace: default_namespace(), pool_name: default_pool_name(), headers: HashMap::new(), + ray_id: None, max_input_size: default_max_input_size(), disable_metadata_lookup: false, resolved_config: Arc::new(OnceCell::new()), @@ -139,15 +148,27 @@ impl RemoteManager { namespace: Option, pool_name: Option, headers: Option>, + ray_id: Option, max_input_size: Option, disable_metadata_lookup: bool, ) -> Self { + let ray_id = ray_id.and_then(|ray_id| { + let bounded = bounded_ray_id(&ray_id).map(str::to_owned); + if bounded.is_none() { + tracing::warn!( + len = ray_id.len(), + "dropping configured ray id; it must be 1 to 128 characters of [A-Za-z0-9_-]" + ); + } + bounded + }); Self { endpoint, token, namespace: namespace.unwrap_or_else(default_namespace), pool_name: pool_name.unwrap_or_else(default_pool_name), headers: headers.unwrap_or_default(), + ray_id, max_input_size: max_input_size.unwrap_or_else(default_max_input_size), disable_metadata_lookup, resolved_config: Arc::new(OnceCell::new()), @@ -214,6 +235,11 @@ impl RemoteManager { req = req.header(USER_AGENT, USER_AGENT_VALUE); for (key, value) in &self.headers { + // Ray and trace context are per call, so a configured value cannot + // pin stale context on every request. Matches the TypeScript client. + if is_telemetry_header(key) { + continue; + } let name = HeaderName::from_str(key) .with_context(|| format!("invalid configured header name `{key}`"))?; let value = HeaderValue::from_str(value) @@ -493,6 +519,13 @@ 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; + for (name, value) in self.telemetry_headers()? { + headers.entry(name).or_insert(value); + } req = req.headers(headers); if let Some(body_data) = body { @@ -503,6 +536,50 @@ impl RemoteManager { Ok(res) } + /// Headers that carry the caller's trace context and ray 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 is sent. + fn telemetry_headers(&self) -> Result> { + let mut headers = Vec::with_capacity(3); + let context = tracing::Span::current().context(); + let span = context.span(); + let span_context = span.span_context(); + if span_context.is_valid() { + let traceparent = format_traceparent( + span_context.trace_id(), + span_context.span_id(), + span_context.trace_flags().to_u8(), + ); + headers.push(( + HeaderName::from_static(HEADER_TRACEPARENT), + HeaderValue::from_str(&traceparent).context("format traceparent header")?, + )); + let tracestate = span_context.trace_state().header(); + if !tracestate.is_empty() { + headers.push(( + HeaderName::from_static(HEADER_TRACESTATE), + HeaderValue::from_str(&tracestate).context("format tracestate header")?, + )); + } + } + let baggage_ray = context + .baggage() + .get(RAY_BAGGAGE_KEY) + .map(|value| value.as_str().into_owned()); + let ray_id = baggage_ray + .as_deref() + .and_then(bounded_ray_id) + .or(self.ray_id.as_deref()); + if let Some(ray_id) = ray_id { + headers.push(( + HeaderName::from_static(HEADER_RIVET_RAY_ID), + HeaderValue::from_str(ray_id).context("format ray id header")?, + )); + } + Ok(headers) + } + pub fn gateway_url(&self, query: &ActorQuery) -> Result { let config = self.base_config(); match query { @@ -823,3 +900,13 @@ 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..53d59e4930 100644 --- a/rivetkit-rust/packages/client/tests/bare.rs +++ b/rivetkit-rust/packages/client/tests/bare.rs @@ -2,27 +2,27 @@ use std::{ collections::HashMap, net::SocketAddr, sync::{ - atomic::{AtomicBool, AtomicUsize, Ordering}, Arc, + atomic::{AtomicBool, AtomicUsize, Ordering}, }, time::Duration, }; use axum::{ + Json, Router, body::Bytes, extract::{ - ws::{Message as AxumWsMessage, WebSocket, WebSocketUpgrade}, Path, State, + ws::{Message as AxumWsMessage, WebSocket, WebSocketUpgrade}, }, - http::{header, HeaderMap, Method as AxumMethod, StatusCode, Uri}, + http::{HeaderMap, Method as AxumMethod, StatusCode, Uri, header}, response::IntoResponse, routing::{any, get, post, put}, - Json, Router, }; use futures_util::{SinkExt, StreamExt}; use reqwest::{ - header::{HeaderMap as ReqwestHeaderMap, HeaderValue}, Method, Url, + header::{HeaderMap as ReqwestHeaderMap, HeaderValue}, }; use rivetkit_client::{ Client, ClientConfig, ConnectionStatus, CreateOptions, EncodingKind, GetOptions, @@ -30,10 +30,10 @@ use rivetkit_client::{ }; use rivetkit_client_protocol as wire; use serde::{Deserialize, Serialize}; -use serde_json::{json, Value as JsonValue}; +use serde_json::{Value as JsonValue, json}; use tokio::{ net::TcpListener, - sync::{mpsc, Notify}, + sync::{Notify, mpsc}, time::timeout, }; use vbare::OwnedVersionedData; @@ -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( @@ -669,9 +764,11 @@ fn gateway_url_uses_query_backed_get_or_create_target() { params.get("rvt-token").map(String::as_str), Some("dev-token") ); - assert!(params - .get("rvt-input") - .is_some_and(|value| !value.is_empty())); + assert!( + params + .get("rvt-input") + .is_some_and(|value| !value.is_empty()) + ); } #[tokio::test] @@ -1169,6 +1266,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..09bbec31e2 100644 --- a/rivetkit-rust/packages/rivetkit-core/Cargo.toml +++ b/rivetkit-rust/packages/rivetkit-core/Cargo.toml @@ -15,6 +15,9 @@ default = ["native-runtime"] native-runtime = [ "dep:nix", "dep:reqwest", + "dep:opentelemetry-otlp", + "dep:opentelemetry_sdk", + "dep:tracing-subscriber", "dep:rivetkit-engine-process", "dep:axum", "dep:bytes", @@ -49,6 +52,13 @@ http-body-util = { workspace = true, optional = true } include_dir = { workspace = true } nix = { workspace = true, optional = true, features = ["process"] } parking_lot.workspace = true +opentelemetry = { version = "0.28", default-features = false, features = [ + "trace", + # Lets the SDK report dropped spans and export failures through tracing. + "internal-logs", +] } +opentelemetry-otlp = { version = "0.28", default-features = false, optional = true, features = ["trace", "http-json", "http-proto", "grpc-tonic", "reqwest-blocking-client"] } +opentelemetry_sdk = { version = "0.28", default-features = false, optional = true, features = ["trace", "internal-logs"] } rand.workspace = true reqwest = { workspace = true, optional = true } rusqlite = { workspace = true, optional = true } @@ -74,6 +84,8 @@ tokio-stream = { workspace = true, optional = true } tokio-util.workspace = true tower-http = { workspace = true, optional = true, features = ["fs"] } tracing.workspace = true +tracing-opentelemetry = { version = "0.29", default-features = false } +tracing-subscriber = { workspace = true, optional = true } url.workspace = true vbare.workspace = true diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/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..90cdf500ca 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 the invocation is not sampled. + /// + /// `actor_name` and `action_name` name the callee. Both come from the + /// caller's own registry rather than from a remote peer, so neither is a + /// cardinality surface. + #[doc(hidden)] + pub fn begin_outbound_call( + &self, + actor_name: &str, + action_name: &str, + ) -> 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 database calls 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)] @@ -1748,8 +1854,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 +1871,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 +1887,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 +1908,7 @@ impl ActorContext { ); } } + invocation.finish(dispatch_error.as_ref()); ctx.finish_schedule_dispatch(&event_id, history_id, dispatch_error.as_ref()) .await; @@ -1811,7 +1926,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..f00ad92b54 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 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..fe58092c4f 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 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..159ac2b355 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,31 +940,16 @@ 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(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) => { @@ -967,7 +959,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 +969,66 @@ 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(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(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 +1090,34 @@ 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, + 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(ActorWorkKind::Action, 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/error.rs b/rivetkit-rust/packages/rivetkit-core/src/error.rs index 195d7ff837..5b40ab87a0 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/error.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/error.rs @@ -221,6 +221,15 @@ pub enum ActorRuntime { #[error("missing_input", "Actor input is missing.")] MissingInput, + /// Telemetry for an operation was dropped before its result was recorded. + /// The underlying work may still have completed, so this says nothing + /// about whether a write landed or a remote call ran. + #[error( + "operation_abandoned", + "Operation tracking ended before a result was recorded." + )] + OperationAbandoned, + #[error( "invalid_operation", "Actor operation is invalid.", diff --git a/rivetkit-rust/packages/rivetkit-core/src/lib.rs b/rivetkit-rust/packages/rivetkit-core/src/lib.rs index 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..2da1a7b131 --- /dev/null +++ b/rivetkit-rust/packages/rivetkit-core/src/telemetry.rs @@ -0,0 +1,858 @@ +//! Internal OpenTelemetry spans owned by the actor runtime. + +#[cfg(feature = "native-runtime")] +pub mod export; + +use std::str::FromStr as _; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + +use opentelemetry::trace::{ + SpanContext, SpanId, TraceContextExt as _, TraceFlags, TraceId, TraceState, +}; +use parking_lot::Mutex; +use rivetkit_client_protocol::telemetry_headers::format_traceparent; +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::telemetry_headers::HEADER_RIVET_RAY_ID; + +impl IncomingInvocationContext { + pub(crate) fn from_headers( + ray_id: Option, + traceparent: Option<&str>, + tracestate: Option<&str>, + ) -> Self { + Self { + ray_id, + remote_parent: parse_remote_parent(traceparent, tracestate), + } + } + + /// Reads the ray and W3C trace context an HTTP request carries. Every + /// HTTP entry point into an actor reads them through here, so they all + /// apply the same bounds. + pub(crate) fn from_http_headers(headers: &http::HeaderMap) -> Self { + Self::from_headers( + invocation_ray_id(headers), + headers + .get("traceparent") + .and_then(|value| value.to_str().ok()), + headers + .get("tracestate") + .and_then(|value| value.to_str().ok()), + ) + } +} + +/// Reads the caller's ray id. The header is untrusted, so it is bounded by +/// the rule shared with the clients that send it; anything else counts as +/// absent and the invocation mints a fresh ray instead. +fn invocation_ray_id(headers: &http::HeaderMap) -> Option { + let value = headers.get(HEADER_RIVET_RAY_ID)?.to_str().ok()?; + rivetkit_client_protocol::telemetry_headers::bounded_ray_id(value).map(str::to_owned) +} + +/// 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: String, + // 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, +} + +/// 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 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: String, + /// Present only while the invocation runs inside a valid span. + pub span: Option, +} + +/// W3C span context of the current invocation span. A span context is either +/// complete or absent, so these fields are never optional individually. +#[doc(hidden)] +#[derive(Clone, Debug)] +pub struct ActorInvocationSpanContext { + pub trace_id: String, + pub span_id: String, + pub trace_flags: u8, + pub traceparent: String, + pub tracestate: Option, +} + +/// The closed set of SQLite operations that get a span. +/// +/// Both names are `&'static str`, so starting one of these spans allocates +/// 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 ray_id = ray_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + 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); + 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: String, + 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, and the application span active in the host runtime at that + /// moment, or the invocation span when there was none. Work that links + /// back to it then points at the code that caused it rather than at the + /// whole invocation around that code. + pub(crate) fn trace_origin(&self) -> TraceOrigin { + let Some(context) = self.trace_context() else { + return TraceOrigin::default(); + }; + let span = self.1.as_ref().and_then(w3c_span_context).or(context.span); + let (traceparent, tracestate) = match span { + Some(span) => (Some(span.traceparent), span.tracestate), + None => (None, None), + }; + TraceOrigin { + ray_id: Some(context.ray_id), + 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, + 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, + 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 the call is not sampled. + pub fn span_context(&self) -> Option { + self.context.clone() + } + + /// Records the call's outcome. `error` is the failure the callee returned, + /// and its group and code become the span's `error.type`. + pub fn finish(mut self, error: Option<&anyhow::Error>) { + let Some(span) = self.span.take() else { + return; + }; + record_outcome(&span, error); + } +} + +impl Drop for OutboundCallInvocation { + fn drop(&mut self) { + let Some(span) = self.span.take() else { + return; + }; + let error = crate::error::ActorRuntime::OperationAbandoned.build(); + record_outcome(&span, Some(&error)); + } +} + +impl SqliteOperationSpan { + pub(crate) fn span(&self) -> tracing::Span { + self.span.as_ref().expect("sqlite span is present").clone() + } + + 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; + }; + let error = crate::error::ActorRuntime::OperationAbandoned.build(); + record_outcome(&span, Some(&error)); + } +} + +/// W3C fields of a span context, or nothing when it is not valid and so +/// carries nothing worth propagating. +fn w3c_span_context(span_context: &SpanContext) -> Option { + if !span_context.is_valid() { + return None; + } + let tracestate = span_context.trace_state().header(); + Some(ActorInvocationSpanContext { + trace_id: span_context.trace_id().to_string(), + span_id: span_context.span_id().to_string(), + trace_flags: span_context.trace_flags().to_u8(), + traceparent: format_traceparent( + span_context.trace_id(), + span_context.span_id(), + span_context.trace_flags().to_u8(), + ), + tracestate: (!tracestate.is_empty()).then_some(tracestate), + }) +} + +/// Reads a span's W3C context, or nothing when the span is not sampled and so +/// carries no valid context to propagate. +fn span_context_of(span: &tracing::Span) -> Option { + let context = span.context(); + let context_span = context.span(); + w3c_span_context(context_span.span_context()) +} + +/// Opens the span covering the moment one queue message is handed to the +/// actor. It links to the span that sent the message, which is what connects +/// the consumer's trace to the sender's. Inside an invocation it sits under +/// that invocation's application span, or the invocation span, and carries its +/// ray. Outside one, as from the run handler, it is a root span carrying the +/// ray 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); + 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 { + let mut fields = traceparent?.split('-'); + let version = fields.next()?; + let trace_id = fields.next()?; + let span_id = fields.next()?; + let flags = fields.next()?; + if fields.next().is_some() + || version.len() != 2 + || !version + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) + || version.eq_ignore_ascii_case("ff") + || trace_id.len() != 32 + || span_id.len() != 16 + || flags.len() != 2 + { + return None; + } + + let trace_id = TraceId::from_hex(trace_id).ok()?; + let span_id = SpanId::from_hex(span_id).ok()?; + let flags = u8::from_str_radix(flags, 16).ok()?; + let trace_state = tracestate + .and_then(|value| TraceState::from_str(value).ok()) + .unwrap_or_default(); + let context = SpanContext::new(trace_id, span_id, TraceFlags::new(flags), true, trace_state); + if context.is_valid() { + Some(context) + } else { + None + } +} diff --git a/rivetkit-rust/packages/rivetkit-core/src/telemetry/export.rs b/rivetkit-rust/packages/rivetkit-core/src/telemetry/export.rs new file mode 100644 index 0000000000..77853ab577 --- /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 behaviour by adding +//! [`layer`] to its subscriber and calling [`flush_best_effort`] on shutdown. +//! Core never installs a subscriber itself; which log layers surround the span +//! layer is the host's decision. + +use std::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..aaebdedf42 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts +++ b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts @@ -8,6 +8,19 @@ export interface JsActorKeySegment { stringValue?: string numberValue?: number } +/** Active actor invocation correlation exposed to the TypeScript runtime adapter. */ +export interface JsActorInvocationTraceContext { + rayId: string + span?: JsActorInvocationSpanContext +} +/** W3C span context of the current invocation span, present only when tracing is active. */ +export interface JsActorInvocationSpanContext { + traceId: string + spanId: string + traceFlags: number + traceparent: string + tracestate?: string +} export interface JsHttpRequest { method: string uri: string @@ -51,6 +64,9 @@ export interface JsQueueSendResult { export interface JsActionDefinition { name: string } +export interface JsQueueDefinition { + name: string +} /** * One entry in the actor's `inspector.tabs[]` declaration. Either a * custom-tab descriptor (id + label + source dir) or a built-in modifier @@ -120,6 +136,7 @@ export interface JsActorConfig { maxIncomingMessageSize?: number maxOutgoingMessageSize?: number actions?: Array + queues?: Array inspectorTabs?: Array } export interface JsBindParam { @@ -257,6 +274,12 @@ export interface JsServerlessStreamError { code: string message: string } +/** + * Routes the OpenTelemetry SDK's own warnings, such as dropped spans, to the + * JavaScript logger. 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 +329,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 the invocation is not sampled, + * 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 +391,26 @@ export declare class ActorContext { runtimeState(): object clearRuntimeState(): void } +/** + * One open call out to another actor. + * + * The call spans a request made by the host runtime, so it is opened and closed + * by two separate calls. Letting this be collected without finishing records + * the call as cancelled rather than silently losing it. + */ +export declare class OutboundCall { + /** + * W3C context of this call's span, to send to the callee so it parents to + * the call rather than to the invocation that made it. + */ + spanContext(): JsActorInvocationSpanContext | null + /** + * Records the call's outcome. `error` is the failure as the bridge encodes + * it, so a structured error keeps its group and code while anything else + * stays unstructured for Core to classify. + */ + finish(error?: string | undefined | null): void +} export declare class NapiActorFactory { constructor(callbacks: object, config?: JsActorConfig | undefined | null) } diff --git a/rivetkit-typescript/packages/rivetkit-napi/index.js b/rivetkit-typescript/packages/rivetkit-napi/index.js index 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..85e75ce8d8 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,44 @@ pub struct JsActorKeySegment { pub number_value: Option, } +/// Active actor invocation correlation exposed to the TypeScript runtime adapter. +#[napi(object)] +pub struct JsActorInvocationTraceContext { + pub ray_id: String, + pub span: Option, +} + +/// W3C span context of the current invocation span, present only when tracing is active. +#[napi(object)] +pub struct JsActorInvocationSpanContext { + pub trace_id: String, + pub span_id: String, + pub trace_flags: u8, + pub traceparent: String, + pub tracestate: Option, +} + +impl From for JsActorInvocationTraceContext { + fn from(value: ActorInvocationTraceContext) -> Self { + Self { + ray_id: value.ray_id, + span: value.span.map(JsActorInvocationSpanContext::from), + } + } +} + +impl From for JsActorInvocationSpanContext { + fn from(value: ActorInvocationSpanContext) -> Self { + Self { + trace_id: value.trace_id, + span_id: value.span_id, + trace_flags: value.trace_flags, + traceparent: value.traceparent, + tracestate: value.tracestate, + } + } +} + #[napi(object)] pub struct JsHttpRequest { pub method: String, @@ -273,11 +313,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 the invocation is not sampled, + /// in which case the caller sends its own context as before. + #[napi] + pub fn begin_outbound_call( + &self, + actor_name: String, + action_name: String, + ) -> 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 +1138,38 @@ fn js_http_request_to_core_request(request: JsHttpRequest) -> napi::Result, +} + +#[napi] +impl OutboundCall { + /// W3C context of this call's span, to send to the callee so it parents to + /// the call rather than to the invocation that made it. + #[napi] + pub fn span_context(&self) -> Option { + self.invocation + .as_ref() + .and_then(CoreOutboundCallInvocation::span_context) + .map(Into::into) + } + + /// Records the call's outcome. `error` is the failure as the bridge encodes + /// it, so a structured error keeps its group and code while anything else + /// stays unstructured for Core to classify. + #[napi] + pub fn finish(&mut self, error: Option) { + let Some(invocation) = self.invocation.take() else { + return; + }; + let error = error.map(anyhow_error_from_js_reason); + invocation.finish(error.as_ref()); + } +} diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/actor_factory.rs b/rivetkit-typescript/packages/rivetkit-napi/src/actor_factory.rs index 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..51c26ef364 --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit-napi/src/telemetry.rs @@ -0,0 +1,98 @@ +//! 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) level: &'static str, + 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("level", ctx.value.level)?; + 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 { + level: event.metadata().level().as_str(), + 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 02f5e76bbc..1c86ca44fa 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..c60d85910f 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,51 @@ 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?.( + this.#targetActorName(), + 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 }); + #targetActorName(): string { + try { + return getActorNameFromQuery(this.#actorResolutionState); + } catch { + // A malformed query fails later with its own error. Naming a span + // must not be what surfaces it. + return "unknown"; + } } - async #sendActionNow( + async #sendActionAttempts( opts: { name: string; args: unknown[]; } & ActorActionOptions, + callSpan?: ActorInvocationSpanContext, ): Promise { const maxAttempts = this.#getDynamicQueryMaxAttempts(); let useQueryTarget = isDynamicActorQuery(this.#actorResolutionState); @@ -320,6 +367,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 +387,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 +728,9 @@ export class ActorHandleRaw { skipReadyWait, }, ); + const telemetryHeaders = outboundTelemetryHeaders( + this.#currentActorInvocation?.(), + ); for (let attempt = 0; attempt < maxAttempts; attempt++) { let actorId: string | undefined; @@ -698,6 +749,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..212505f69b 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, @@ -181,6 +185,14 @@ export const CREATE_ACTOR_CONN_PROXY = Symbol("createActorConnProxy"); * @template A The actors map type that defines the available actors. * @see {@link https://rivet.dev/docs/manage|Create & Manage Actors} */ +export interface ClientRawOptions { + encoding?: Encoding; + gateway?: ActorGatewayOptions; + /** Supplies the calling actor's invocation so actor-owned clients propagate its trace and ray. */ + currentActorInvocation?: CurrentActorInvocation; + beginOutboundCall?: BeginOutboundCall; +} + export class ClientRaw { #disposed = false; @@ -189,19 +201,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 +418,8 @@ export class ClientRaw { actorQuery, this.#gatewayOptions, signal, + this.#currentActorInvocation, + this.#beginOutboundCall, ); } @@ -459,9 +476,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..cae6655eea --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/src/client/outbound-telemetry.ts @@ -0,0 +1,64 @@ +import { + HEADER_RIVET_RAY_ID, + HEADER_TRACEPARENT, + HEADER_TRACESTATE, +} from "@/common/actor-router-consts"; +import type { ActorInvocationTraceContext } from "@/common/actor-telemetry-context"; +import { + type ActiveTraceHeaders, + readActiveRayId, + readActiveTraceHeaders, +} from "@/common/otel-context"; + +/** + * Headers that carry a caller's ray 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. + * + * The ray 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?: ActiveTraceHeaders, +): Record { + const headers: Record = {}; + const rayId = invocation?.rayId ?? readActiveRayId(); + if (rayId) { + headers[HEADER_RIVET_RAY_ID] = rayId; + } + const traceHeaders = + callSpan ?? readActiveTraceHeaders() ?? invocation?.span; + if (traceHeaders) { + headers[HEADER_TRACEPARENT] = traceHeaders.traceparent; + if (traceHeaders.tracestate) { + headers[HEADER_TRACESTATE] = traceHeaders.tracestate; + } + } + return headers; +} + +/** + * Adds outbound telemetry headers to a request whose caller may have set + * some already. A header the caller set wins. `traceparent` and `tracestate` + * describe one span between them, so when the caller set either, both stay + * as the caller set them and neither is added. + */ +export function addOutboundTelemetryHeaders( + headers: Headers, + telemetry: Record, +): void { + const callerSetTraceContext = + headers.has(HEADER_TRACEPARENT) || headers.has(HEADER_TRACESTATE); + for (const [name, value] of Object.entries(telemetry)) { + const isTraceContext = + name === HEADER_TRACEPARENT || name === HEADER_TRACESTATE; + if (isTraceContext ? callerSetTraceContext : headers.has(name)) { + continue; + } + headers.set(name, value); + } +} diff --git a/rivetkit-typescript/packages/rivetkit/src/client/queue.ts b/rivetkit-typescript/packages/rivetkit/src/client/queue.ts index a55fc86852..86e8e52fe8 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 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..32781f7d65 --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/src/common/actor-telemetry-context.ts @@ -0,0 +1,30 @@ +/** Correlation owned by the currently executing Core actor invocation. */ +export interface ActorInvocationTraceContext { + /** Rivet request correlation identifier for the current invocation. */ + readonly rayId: string; + /** Core-owned invocation span context, absent when tracing is disabled. */ + readonly span?: ActorInvocationSpanContext; +} + +/** W3C span context of the Core invocation span. */ +export interface ActorInvocationSpanContext { + /** W3C trace identifier. */ + readonly traceId: string; + /** W3C span identifier for the Core invocation. */ + readonly spanId: string; + /** OpenTelemetry trace flags encoded as an integer. */ + readonly traceFlags: number; + /** Serialized W3C Trace Context for the Core invocation. */ + readonly traceparent: string; + /** Optional vendor trace state inherited by the Core invocation. */ + readonly tracestate?: string; +} + +/** Formats a W3C `traceparent` header from its span identifiers. */ +export function formatTraceparent( + traceId: string, + spanId: string, + traceFlags: number, +): string { + return `00-${traceId}-${spanId}-${traceFlags.toString(16).padStart(2, "0")}`; +} diff --git a/rivetkit-typescript/packages/rivetkit/src/common/otel-context.ts b/rivetkit-typescript/packages/rivetkit/src/common/otel-context.ts new file mode 100644 index 0000000000..a4b74687c7 --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/src/common/otel-context.ts @@ -0,0 +1,87 @@ +import { + type Context, + context, + createTraceState, + isSpanContextValid, + propagation, + trace, +} from "@opentelemetry/api"; +import { + type ActorInvocationSpanContext, + formatTraceparent, +} from "./actor-telemetry-context"; + +/** W3C headers derived from the active JavaScript OTel context. */ +export interface ActiveTraceHeaders { + /** W3C Trace Context identifying the active trace and span. */ + readonly traceparent: string; + /** Optional vendor trace state associated with the active span. */ + readonly tracestate?: string; +} + +/** Returns the active W3C trace context, when an OTel provider has installed one. */ +export function readActiveTraceHeaders(): ActiveTraceHeaders | undefined { + const spanContext = trace.getSpanContext(context.active()); + if (!spanContext || !isSpanContextValid(spanContext)) return undefined; + + const tracestate = spanContext.traceState?.serialize(); + return { + traceparent: formatTraceparent( + spanContext.traceId, + spanContext.spanId, + spanContext.traceFlags, + ), + ...(tracestate ? { tracestate } : {}), + }; +} + +/** W3C Baggage key that carries a ray through application code. */ +export const RAY_BAGGAGE_KEY = "rivet.ray.id"; + +const RAY_ID_PATTERN = /^[A-Za-z0-9_-]{1,128}$/; + +/** + * Returns the ray carried in the active OpenTelemetry baggage, so a request + * handler that received a ray 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(); + + let parent: Context; + try { + const spanContext = { + traceId: invocation.traceId, + spanId: invocation.spanId, + traceFlags: invocation.traceFlags, + traceState: invocation.tracestate + ? createTraceState(invocation.tracestate) + : undefined, + isRemote: false, + }; + if (!isSpanContextValid(spanContext)) return run(); + parent = trace.setSpanContext(context.active(), spanContext); + } catch { + // Invalid telemetry must not prevent the action from running. + return run(); + } + + return context.with(parent, run); +} diff --git a/rivetkit-typescript/packages/rivetkit/src/engine-client/actor-http-client.ts b/rivetkit-typescript/packages/rivetkit/src/engine-client/actor-http-client.ts index 2fc62b41a8..563fd90f8c 100644 --- a/rivetkit-typescript/packages/rivetkit/src/engine-client/actor-http-client.ts +++ b/rivetkit-typescript/packages/rivetkit/src/engine-client/actor-http-client.ts @@ -4,6 +4,9 @@ import { HEADER_RIVET_SKIP_READY_WAIT, HEADER_RIVET_TARGET, HEADER_RIVET_TOKEN, + HEADER_RIVET_RAY_ID, + HEADER_TRACEPARENT, + HEADER_TRACESTATE, } from "@/common/actor-router-consts"; import { type GatewayRequestOptions, shouldSkipReadyWait } from "./driver"; @@ -55,6 +58,19 @@ function buildGuardHeaders( for (const [key, value] of Object.entries(runConfig.headers)) { headers.set(key, value as string); } + // Invocation headers are per action call. Apply the active request last so + // static client configuration cannot retain or override an earlier action. + for (const name of [ + HEADER_RIVET_RAY_ID, + HEADER_TRACEPARENT, + HEADER_TRACESTATE, + ]) { + headers.delete(name); + const value = actorRequest.headers.get(name); + if (value !== null) { + headers.set(name, value); + } + } // Add guard-specific headers if (runConfig.token) { headers.set(HEADER_RIVET_TOKEN, runConfig.token); diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts index 3d589d5242..003ddc2ac9 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,12 @@ import type { HttpResponseBodyStream as NativeHttpResponseBodyStream, WebSocket as NativeWebSocket, } from "@rivetkit/rivetkit-napi"; +import type { ActorInvocationTraceContext } from "@/common/actor-telemetry-context"; +import { + readActiveTraceHeaders, + runWithActorInvocationSpan, +} from "@/common/otel-context"; +import { logger } from "./log"; import type { ActorContextHandle, ActorFactoryHandle, @@ -26,6 +33,7 @@ import type { RuntimeKvEntry, RuntimeKvListOptions, RuntimeListenerConfig, + RuntimeOutboundCall, RuntimeQueueEnqueueAndWaitOptions, RuntimeQueueMessage, RuntimeQueueNextBatchOptions, @@ -65,6 +73,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 +244,7 @@ function toNapiKvEntry(entry: RuntimeKvEntry): { }; } -function toNapiQueueMessage(message: RuntimeQueueMessage): RuntimeQueueMessage { +function toNapiQueueMessage(message: NapiQueueMessage): RuntimeQueueMessage { return { id: () => message.id(), name: () => message.name(), @@ -255,22 +266,72 @@ export class NapiCoreRuntime implements CoreRuntime { #bindings: NativeBindings; #sql = new WeakMap(); - - constructor(bindings: NativeBindings) { + #invocationContext: AsyncLocalStorage; + // `traceparent` of each invocation's own Core span, so an operation can + // tell that span apart from an application span without a native call. + #invocationTraceparent = new WeakMap(); + + constructor( + bindings: NativeBindings, + 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. + // Older addons may lack this optional diagnostics binding. + if (this.#bindings.setTelemetryLogSink) { + this.#bindings.setTelemetryLogSink((event) => { + logger().warn( + { otelEvent: event.name }, + event.message || event.name, + ); + }); + } else { + logger().warn( + "native addon has no telemetry log sink; OpenTelemetry SDK warnings will not reach the actor logs", + ); + } return asRegistryHandle(new this.#bindings.CoreRegistry()); } @@ -552,6 +613,42 @@ export class NapiCoreRuntime implements CoreRuntime { return asNativeActorContext(ctx).actorId(); } + runWithActorInvocationContext(ctx: ActorContextHandle, run: () => T): T { + const nativeCtx = asNativeActorContext(ctx); + const span = nativeCtx.invocationTraceContext()?.span; + if (span) { + this.#invocationTraceparent.set(nativeCtx, span.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 +697,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 +968,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 +983,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 +994,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 +1010,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 +1057,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 +1095,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 +1106,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 +1120,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 +1138,7 @@ export class NapiCoreRuntime implements CoreRuntime { args: RuntimeBytes, maxHistory: number | undefined, ) { - await asNativeActorContext(ctx) + await this.#actorContextForOperation(ctx) .schedule() .cronSet( name, @@ -1064,7 +1158,7 @@ export class NapiCoreRuntime implements CoreRuntime { args: RuntimeBytes, maxHistory: number | undefined, ) { - await asNativeActorContext(ctx) + await this.#actorContextForOperation(ctx) .schedule() .cronEvery( name, @@ -1079,20 +1173,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 +1197,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 +1270,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..bd33d86549 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: 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..b8e4b83a11 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 the call is not sampled. + */ + readonly span?: ActorInvocationSpanContext; + /** + * Records the call's outcome. `error` is the failure encoded the way the + * bridge encodes errors, so a structured error keeps its group and code. + */ + finish(error?: string): void; +} + +/** + * Opens the span covering one call out to another actor, or returns `undefined` + * outside an invocation, on a runtime without invocation telemetry, or when the + * call is not sampled. Callers send their own context when it returns nothing. + */ +export type BeginOutboundCall = ( + actorName: string, + actionName: string, +) => 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..b8710aba6e 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,81 @@ 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(); + } + }, + ); + }, + transactionUnderApplicationSpan: async (c) => { + return await applicationTracer.startActiveSpan( + "agent.persist", + async (span) => { + try { + await c.db.transaction(async (tx) => { + await tx.execute("SELECT 1 AS in_transaction"); + }); + return { 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 +241,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..705875bc2e 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,34 @@ 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); +} + +/** OTLP `SpanKind.PRODUCER`. */ +const OTLP_SPAN_KIND_PRODUCER = 4; +/** OTLP `SpanKind.CONSUMER`. */ +const OTLP_SPAN_KIND_CONSUMER = 5; const FIXTURE_PATH = join(TEST_DIR, "fixtures", "napi-runtime-server.ts"); const NAMESPACE = "default"; const TOKEN = "dev"; @@ -19,9 +42,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 +471,249 @@ 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; + events: string[]; + 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< + ExportedSpan, + "attributes" | "events" | "endTimeUnixNano" + > & { + attributes?: OtlpAttribute[]; + events?: Array<{ name: string }>; + 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)), + ]), + ), + events: (span.events ?? []).map((event) => event.name), + 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, + 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 +723,26 @@ describe.sequential("native NAPI runtime integration", () => { }, 30_000); test("runs a TS actor through registry, NAPI, core, envoy, and engine", async () => { - const poolName = "default"; - const port = await getPort({ host: "127.0.0.1" }); - const endpoint = `http://127.0.0.1:${port}`; - engineEndpoint = endpoint; - storagePath = await mkdtemp(join(tmpdir(), "rivetkit-services-")); - runtimeLogs = { stdout: "", stderr: "" }; - runtime = spawn(process.execPath, ["--import", "tsx", FIXTURE_PATH], { - cwd: dirname(TEST_DIR), - env: { - ...process.env, - RIVET_TOKEN: TOKEN, - RIVET_NAMESPACE: NAMESPACE, - RIVET_RUN_ENGINE_HOST: "127.0.0.1", - RIVET_RUN_ENGINE_PORT: String(port), - RIVETKIT_TEST_ENDPOINT: endpoint, - RIVETKIT_TEST_POOL_NAME: poolName, - RIVETKIT_STORAGE_PATH: storagePath, - }, - stdio: ["ignore", "pipe", "pipe"], - }); - runtime.stdout?.on("data", (chunk) => { - runtimeLogs.stdout += chunk.toString(); - }); - runtime.stderr?.on("data", (chunk) => { - runtimeLogs.stderr += chunk.toString(); - }); - - await waitForHealth(runtime, endpoint, 90_000); - await upsertNormalRunnerConfig(runtime, endpoint, poolName); - await waitForEnvoy(runtime, endpoint, poolName, 30_000); + collector = await startOtlpCollector( + await getPort({ host: "127.0.0.1" }), + ); + const { endpoint, poolName, child } = await startTracedRuntime( + collector.endpoint, + ); + runtime = child; await waitForEnvoy(runtime, endpoint, SERVICES_POOL_NAME, 30_000); await expectNormalRunnerConfig(endpoint, SERVICES_POOL_NAME); const servicesActorId = await createServicesActor(endpoint); 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 +837,7 @@ describe.sequential("native NAPI runtime integration", () => { code: "internal_error", message: "An internal error occurred", }); + await client.dispose(); const processId = servicesPid(); @@ -603,4 +845,668 @@ 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"; + for (const version of ["00", "zz", "0A"]) { + const client = createIntegrationClient(endpoint, poolName); + try { + const handle = await waitForActorReady( + () => + client.integrationActor.create( + [`trace-context-${crypto.randomUUID()}`], + { + params: { userId: "integration-test" }, + }, + ), + 30_000, + ); + await waitForActorReady(() => handle.getCount(), 30_000); + const actorId = await handle.resolve(); + const url = new URL(await handle.getGatewayUrl()); + url.pathname = `${url.pathname.replace(/\/$/, "")}/action/getCountViaClient`; + const response = await fetch(url, { + method: "POST", + headers: { + "content-type": "application/json", + "x-rivet-encoding": "json", + "x-rivet-token": TOKEN, + "x-rivet-conn-params": JSON.stringify({ + userId: "integration-test", + }), + traceparent: `${version}-${traceId}-${parentSpanId}-01`, + tracestate: traceState, + }, + body: JSON.stringify({ args: [] }), + }); + expect(response.status).toBe(200); + await response.arrayBuffer(); + const spans = await waitForSpans( + traceExports, + "caller and callee trace contexts", + (exported) => { + const caller = exported.find( + (span) => + span.attributes["rivet.actor.id"] === actorId && + span.attributes["rivet.action.name"] === + "getCountViaClient", + ); + const hop = exported.find( + (span) => + span.kind === OTLP_SPAN_KIND_CLIENT && + span.traceId === caller?.traceId, + ); + return ( + !!hop && + exported.some( + (span) => span.parentSpanId === hop.spanId, + ) + ); + }, + 10_000, + ); + const caller = spans.find( + (span) => + span.attributes["rivet.actor.id"] === actorId && + span.attributes["rivet.action.name"] === + "getCountViaClient", + ); + const hop = spans.find( + (span) => + span.kind === OTLP_SPAN_KIND_CLIENT && + span.traceId === caller?.traceId, + ); + const callee = spans.find( + (span) => span.parentSpanId === hop?.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(""); + } + } finally { + await client.dispose(); + } + } + }, 120_000); + + test("keeps overlapping invocations of one actor telemetrically isolated", async () => { + collector = await startOtlpCollector( + await getPort({ host: "127.0.0.1" }), + ); + 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).toContain("actorName=integrationActor"); + expect(line).toContain("napi-isolation-"); + expect(line).toMatch(/ trace_id=[0-9a-f]{32}( |$)/); + expect(line).toMatch(/ span_id=[0-9a-f]{16}( |$)/); + } + const rayOf = (line: string) => + / rayId=([0-9a-f-]{36})/.exec(line)?.[1]; + expect(rayOf(okLog)).toBeDefined(); + expect(rayOf(okLog)).not.toBe(rayOf(failLog)); + expect(rays).toContain(rayOf(okLog)); + expect(rays).toContain(rayOf(failLog)); + + // 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"); + + // Transactions retain the parent captured at begin. + const underTransaction = await handle.transactionUnderApplicationSpan(); + const transactionSpans = await waitForSpans( + traceExports, + "the transaction statement and commit under the application span", + (exported) => + exported.some( + (span) => + span.name === "rivet.sqlite.transaction.execute" && + span.parentSpanId === underTransaction.spanId, + ) && + exported.some( + (span) => + span.name === "rivet.sqlite.transaction.commit" && + span.parentSpanId === underTransaction.spanId, + ), + 10_000, + ); + expect( + transactionSpans.filter( + (span) => span.parentSpanId === underTransaction.spanId, + ).length, + ).toBeGreaterThanOrEqual(3); + + 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).toMatchObject({ + "rivet.invocation.type": "request", + "http.request.method": "GET", + "http.response.status_code": "200", + }); + expect(requestSpan?.attributes["rivet.action.name"]).toBeUndefined(); + 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(deferredInvocation?.events).toContain("reply sent"); + 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?.kind).toBe(OTLP_SPAN_KIND_PRODUCER); + 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?.kind).toBe(OTLP_SPAN_KIND_CONSUMER); + 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?.kind).toBe(OTLP_SPAN_KIND_CONSUMER); + 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?.attributes).toMatchObject({ + "rivet.operation.system": "sqlite", + "rivet.operation.name": "execute", + }); + expect(failedSqlite?.attributes["error.type"]).toMatch( + /^[a-z_]+\.[a-z_]+$/, + ); + 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); + + 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),