1#[cfg(feature = "native-runtime")]
4pub mod export;
5#[cfg(feature = "native-runtime")]
6mod sampler;
7
8use std::sync::Arc;
9
10use opentelemetry::Context;
11use opentelemetry::propagation::{Extractor, Injector, TextMapPropagator as _};
12use opentelemetry::trace::{SpanContext, TraceContextExt as _};
13use opentelemetry_sdk::propagation::TraceContextPropagator;
14use parking_lot::Mutex;
15use tracing_opentelemetry::OpenTelemetrySpanExt as _;
16
17use crate::actor::queue::QueueMessage;
18use crate::{ActorContext, ActorTracingConfig};
19
20#[derive(Debug, Default)]
22pub struct IncomingInvocationContext {
23 pub(crate) ray_id: Option<String>,
24 remote_parent: Option<SpanContext>,
25}
26
27pub(crate) use rivetkit_client_protocol::telemetry_headers::HEADER_RIVET_RAY_ID;
28
29impl IncomingInvocationContext {
30 pub(crate) fn from_headers(
31 ray_id: Option<String>,
32 traceparent: Option<&str>,
33 tracestate: Option<&str>,
34 ) -> Self {
35 Self {
36 ray_id,
37 remote_parent: parse_remote_parent(traceparent, tracestate),
38 }
39 }
40
41 pub(crate) fn from_http_headers(headers: &http::HeaderMap) -> Self {
45 Self::from_headers(
46 invocation_ray_id(headers),
47 headers
48 .get("traceparent")
49 .and_then(|value| value.to_str().ok()),
50 headers
51 .get("tracestate")
52 .and_then(|value| value.to_str().ok()),
53 )
54 }
55}
56
57fn invocation_ray_id(headers: &http::HeaderMap) -> Option<String> {
61 let value = headers.get(HEADER_RIVET_RAY_ID)?.to_str().ok()?;
62 rivetkit_client_protocol::telemetry_headers::bounded_ray_id(value).map(str::to_owned)
63}
64
65#[cfg(feature = "native-runtime")]
68pub(crate) const SAMPLE_RATIO_ATTRIBUTE: &str = "rivet.sampling.ratio";
69
70const REQUEST_INVOCATION_NAME: &str = "onRequest";
72
73const QUEUE_SEND_INVOCATION_NAME: &str = "queue.send";
75
76const WORKFLOW_RUN_INVOCATION_NAME: &str = "workflow";
78
79#[derive(Clone, Copy, Debug)]
82enum InvocationSubject<'a> {
83 Action(&'a str),
84 Request { method: &'a str },
85 QueueSend { queue: &'a str },
86 WorkflowRun,
87}
88
89impl<'a> InvocationSubject<'a> {
90 fn name(self) -> &'a str {
91 match self {
92 Self::Action(name) => name,
93 Self::Request { .. } => REQUEST_INVOCATION_NAME,
94 Self::QueueSend { .. } => QUEUE_SEND_INVOCATION_NAME,
95 Self::WorkflowRun => WORKFLOW_RUN_INVOCATION_NAME,
96 }
97 }
98
99 fn sampler(self, config: &ActorTracingConfig) -> Option<f64> {
101 match self {
102 Self::Action(name) => config.sampler_for_action(name),
103 Self::Request { .. } | Self::QueueSend { .. } | Self::WorkflowRun => config.sampler,
104 }
105 }
106
107 fn record_attributes(self, span: &tracing::Span) {
108 match self {
109 Self::Action(name) => span.record("rivet.action.name", name),
110 Self::Request { method } => span.record("http.request.method", method),
111 Self::QueueSend { queue } => span.record("rivet.queue.name", queue),
112 Self::WorkflowRun => span,
113 };
114 }
115}
116
117#[derive(Debug)]
119pub(crate) struct ActorInvocation {
120 telemetry: ActorInvocationTelemetry,
121}
122
123#[derive(Clone, Copy, Debug)]
124enum InvocationType {
125 Action,
126 Scheduled,
127 Request,
128 QueueSend,
129 Workflow,
130}
131
132impl InvocationType {
133 fn as_label(self) -> &'static str {
134 match self {
135 Self::Action => "action",
136 Self::Scheduled => "scheduled",
137 Self::Request => "request",
138 Self::QueueSend => "queue_send",
139 Self::Workflow => "workflow",
140 }
141 }
142}
143
144#[doc(hidden)]
152#[derive(Clone, Debug)]
153pub struct ActorInvocationTelemetry {
154 inner: Arc<InvocationInner>,
155 application_span: Option<SpanContext>,
156 step_span: Option<SpanContext>,
157}
158
159#[derive(Debug)]
163pub(crate) struct ActorTelemetryIdentity {
164 pub(crate) actor_id: String,
165 pub(crate) actor_name: String,
166 pub(crate) actor_key: String,
167}
168
169#[derive(Debug)]
170struct InvocationInner {
171 takes_message_ray: bool,
173 state: Mutex<InvocationState>,
175 identity: Arc<ActorTelemetryIdentity>,
176}
177
178#[derive(Debug)]
179struct InvocationState {
180 ray_id: Option<String>,
181 span: Option<tracing::Span>,
182 finished: bool,
183 pending_work: usize,
184}
185
186const OPERATION_ABANDONED_ERROR_TYPE: &str = "actor.operation_abandoned";
187
188pub(crate) struct InvocationWorkGuard(ActorInvocationTelemetry);
192
193#[derive(Clone, Debug, Default, Eq, PartialEq)]
197pub struct IncomingTraceContext {
198 pub ray_id: Option<String>,
199 pub traceparent: Option<String>,
200 pub tracestate: Option<String>,
201}
202
203#[doc(hidden)]
205#[derive(Clone, Debug)]
206pub struct ActorInvocationTraceContext {
207 pub ray_id: Option<String>,
208 pub span: Option<ActorInvocationSpanContext>,
210}
211
212#[doc(hidden)]
215#[derive(Clone, Debug)]
216pub struct ActorInvocationSpanContext {
217 pub trace_id: String,
218 pub span_id: String,
219 pub trace_flags: u8,
220 pub tracestate: Option<String>,
221}
222
223#[derive(Clone, Copy, Debug)]
229pub(crate) enum SqliteOperation {
230 Exec,
231 Execute,
232 ExecuteBatch,
233 Query,
234 Run,
235 TransactionBegin,
236 TransactionExec,
237 TransactionExecute,
238 TransactionCommit,
239 TransactionRollback,
240}
241
242impl SqliteOperation {
243 fn names(self) -> (&'static str, &'static str) {
244 match self {
245 Self::Exec => ("rivet.sqlite.exec", "exec"),
246 Self::Execute => ("rivet.sqlite.execute", "execute"),
247 Self::ExecuteBatch => ("rivet.sqlite.execute_batch", "execute_batch"),
248 Self::Query => ("rivet.sqlite.query", "query"),
249 Self::Run => ("rivet.sqlite.run", "run"),
250 Self::TransactionBegin => ("rivet.sqlite.transaction.begin", "transaction.begin"),
251 Self::TransactionExec => ("rivet.sqlite.transaction.exec", "transaction.exec"),
252 Self::TransactionExecute => ("rivet.sqlite.transaction.execute", "transaction.execute"),
253 Self::TransactionCommit => ("rivet.sqlite.transaction.commit", "transaction.commit"),
254 Self::TransactionRollback => {
255 ("rivet.sqlite.transaction.rollback", "transaction.rollback")
256 }
257 }
258 }
259}
260
261pub(crate) struct SqliteOperationSpan {
262 span: Option<tracing::Span>,
263}
264
265#[doc(hidden)]
273pub struct OutboundCallInvocation {
274 span: Option<tracing::Span>,
275 context: Option<ActorInvocationSpanContext>,
276}
277
278#[doc(hidden)]
281#[derive(Clone, Copy, Debug, Eq, PartialEq)]
282pub enum WorkflowRunOutcome {
283 Completed,
284 Sleeping,
285 Evicted,
286 Failed,
287 Cancelled,
288}
289
290impl WorkflowRunOutcome {
291 pub fn parse(value: &str) -> anyhow::Result<Self> {
292 match value {
293 "completed" => Ok(Self::Completed),
294 "sleeping" => Ok(Self::Sleeping),
295 "evicted" => Ok(Self::Evicted),
296 "failed" => Ok(Self::Failed),
297 "cancelled" => Ok(Self::Cancelled),
298 other => anyhow::bail!("unknown workflow run outcome `{other}`"),
299 }
300 }
301
302 fn as_label(self) -> &'static str {
303 match self {
304 Self::Completed => "completed",
305 Self::Sleeping => "sleeping",
306 Self::Evicted => "evicted",
307 Self::Failed => "failed",
308 Self::Cancelled => "cancelled",
309 }
310 }
311
312 fn is_error(self) -> bool {
313 match self {
314 Self::Completed | Self::Sleeping | Self::Evicted => false,
315 Self::Failed | Self::Cancelled => true,
316 }
317 }
318}
319
320#[doc(hidden)]
322#[derive(Clone, Copy, Debug, Eq, PartialEq)]
323pub enum WorkflowStepOutcome {
324 Ok,
325 Retry,
326 Failed,
327}
328
329impl WorkflowStepOutcome {
330 pub fn parse(value: &str) -> anyhow::Result<Self> {
331 match value {
332 "ok" => Ok(Self::Ok),
333 "retry" => Ok(Self::Retry),
334 "failed" => Ok(Self::Failed),
335 other => anyhow::bail!("unknown workflow step outcome `{other}`"),
336 }
337 }
338
339 fn as_label(self) -> &'static str {
340 match self {
341 Self::Ok => "ok",
342 Self::Retry => "retry",
343 Self::Failed => "failed",
344 }
345 }
346}
347
348#[doc(hidden)]
351pub struct WorkflowRunInvocation {
352 ctx: ActorContext,
353 invocation: ActorInvocation,
354}
355
356#[doc(hidden)]
359pub struct WorkflowStepSpan {
360 ctx: ActorContext,
361 span: Option<tracing::Span>,
362}
363
364impl ActorInvocation {
365 pub(crate) fn start_action(
366 ctx: &ActorContext,
367 action_name: &str,
368 incoming: IncomingInvocationContext,
369 ) -> Self {
370 Self::start(
371 ctx,
372 InvocationSubject::Action(action_name),
373 InvocationType::Action,
374 incoming.ray_id,
375 incoming.remote_parent,
376 None,
377 )
378 }
379
380 pub(crate) fn start_scheduled(
381 ctx: &ActorContext,
382 action_name: &str,
383 trace_context: IncomingTraceContext,
384 ) -> Self {
385 let remote_parent = parse_remote_parent(
386 trace_context.traceparent.as_deref(),
387 trace_context.tracestate.as_deref(),
388 );
389 Self::start(
390 ctx,
391 InvocationSubject::Action(action_name),
392 InvocationType::Scheduled,
393 trace_context.ray_id,
394 None,
395 remote_parent,
396 )
397 }
398
399 fn start_workflow_run(ctx: &ActorContext, previous: IncomingTraceContext) -> Self {
402 let previous_run = parse_remote_parent(
403 previous.traceparent.as_deref(),
404 previous.tracestate.as_deref(),
405 );
406 Self::start(
407 ctx,
408 InvocationSubject::WorkflowRun,
409 InvocationType::Workflow,
410 previous.ray_id,
411 None,
412 previous_run,
413 )
414 }
415
416 pub(crate) fn start_queue_send(
420 ctx: &ActorContext,
421 queue_name: &str,
422 incoming: IncomingInvocationContext,
423 ) -> Self {
424 Self::start(
425 ctx,
426 InvocationSubject::QueueSend { queue: queue_name },
427 InvocationType::QueueSend,
428 incoming.ray_id,
429 incoming.remote_parent,
430 None,
431 )
432 }
433
434 pub(crate) fn start_request(
438 ctx: &ActorContext,
439 request: &crate::actor::messages::Request,
440 incoming: IncomingInvocationContext,
441 ) -> Self {
442 Self::start(
443 ctx,
444 InvocationSubject::Request {
445 method: request.method().as_str(),
446 },
447 InvocationType::Request,
448 incoming.ray_id,
449 incoming.remote_parent,
450 None,
451 )
452 }
453
454 fn start(
455 ctx: &ActorContext,
456 subject: InvocationSubject<'_>,
457 invocation_type: InvocationType,
458 ray_id: Option<String>,
459 parent: Option<SpanContext>,
460 link: Option<SpanContext>,
461 ) -> Self {
462 let identity = ctx.telemetry_identity();
463 let subject_name = subject.name();
464 let span = if tracing::enabled!(target: "rivetkit::telemetry", tracing::Level::INFO) {
465 let span = tracing::info_span!(
466 target: "rivetkit::telemetry",
467 parent: None,
468 "rivet.actor.invoke",
469 otel.name = %format!("{}/{}", identity.actor_name, subject_name),
470 otel.kind = otel_kind(invocation_type),
471 rivet.invocation.type = invocation_type.as_label(),
472 rivet.actor.id = %identity.actor_id,
473 rivet.actor.name = %identity.actor_name,
474 rivet.actor.key = %identity.actor_key,
475 rivet.action.name = tracing::field::Empty,
476 rivet.ray.id = tracing::field::Empty,
477 http.request.method = tracing::field::Empty,
478 http.response.status_code = tracing::field::Empty,
479 rivet.queue.name = tracing::field::Empty,
480 rivet.workflow.run.outcome = tracing::field::Empty,
481 rivet.sampling.ratio = subject.sampler(ctx.tracing_config()),
482 otel.status_code = tracing::field::Empty,
483 error.type = tracing::field::Empty,
484 );
485 if let Some(ray_id) = ray_id.as_deref() {
486 span.record("rivet.ray.id", ray_id);
487 }
488 subject.record_attributes(&span);
489 if let Some(parent) = parent {
490 span.set_parent(Context::new().with_remote_span_context(parent));
491 }
492 if let Some(link) = link {
493 span.add_link(link);
494 }
495 Some(span)
496 } else {
497 None
498 };
499
500 Self {
501 telemetry: ActorInvocationTelemetry::new(invocation_type, ray_id, span, identity),
502 }
503 }
504
505 pub(crate) fn telemetry(&self) -> ActorInvocationTelemetry {
506 self.telemetry.clone()
507 }
508
509 pub(crate) fn finish(self, error: Option<&anyhow::Error>) {
510 self.telemetry.finish(error);
511 }
512
513 pub(crate) fn finish_request(
514 self,
515 response: std::result::Result<&crate::actor::messages::ActorHttpResponse, &anyhow::Error>,
516 ) {
517 match response {
518 Ok(response) => {
519 let status = response.status();
520 self.telemetry.finish_with(|span| {
521 span.record("http.response.status_code", status);
522 if status >= 500 {
523 span.record("otel.status_code", "ERROR");
524 span.record("error.type", status.to_string());
525 } else {
526 record_outcome(span, None);
527 }
528 });
529 }
530 Err(error) => self.finish(Some(error)),
531 }
532 }
533}
534
535impl WorkflowRunInvocation {
536 pub(crate) fn start(ctx: &ActorContext, previous: IncomingTraceContext) -> Self {
537 let invocation = ActorInvocation::start_workflow_run(ctx, previous);
538 Self {
539 ctx: ctx
540 .clone()
541 .with_invocation_telemetry(Some(invocation.telemetry())),
542 invocation,
543 }
544 }
545
546 pub fn ctx(&self) -> ActorContext {
548 self.ctx.clone()
549 }
550
551 pub(crate) fn finish(self, outcome: WorkflowRunOutcome) -> Option<IncomingTraceContext> {
554 let ray_id = self.invocation.telemetry.ray_id();
555 let mut trace_context = None;
556 self.invocation.telemetry.finish_with(|span| {
557 if let Some(ray_id) = ray_id.as_deref() {
558 span.record("rivet.ray.id", ray_id);
559 }
560 span.record("rivet.workflow.run.outcome", outcome.as_label());
561 span.record(
562 "otel.status_code",
563 if outcome.is_error() { "ERROR" } else { "OK" },
564 );
565 let headers = otel_span_context_of(span)
566 .map(|span_context| w3c_trace_headers(&span_context))
567 .unwrap_or_default();
568 trace_context = Some(IncomingTraceContext {
569 ray_id,
570 traceparent: headers.traceparent,
571 tracestate: headers.tracestate,
572 });
573 });
574 trace_context
575 }
576}
577
578impl Drop for WorkflowRunInvocation {
579 fn drop(&mut self) {
580 let ray_id = self.invocation.telemetry.ray_id();
581 self.invocation.telemetry.finish_with(|span| {
582 if let Some(ray_id) = ray_id.as_deref() {
583 span.record("rivet.ray.id", ray_id);
584 }
585 span.record("otel.status_code", "ERROR");
586 span.record("error.type", OPERATION_ABANDONED_ERROR_TYPE);
587 });
588 }
589}
590
591impl Drop for ActorInvocation {
592 fn drop(&mut self) {
593 self.telemetry.finish_dropped();
594 }
595}
596
597impl ActorInvocationTelemetry {
598 fn new(
599 invocation_type: InvocationType,
600 ray_id: Option<String>,
601 span: Option<tracing::Span>,
602 identity: Arc<ActorTelemetryIdentity>,
603 ) -> Self {
604 Self {
605 inner: Arc::new(InvocationInner {
606 takes_message_ray: matches!(invocation_type, InvocationType::Workflow),
607 state: Mutex::new(InvocationState {
608 ray_id,
609 span,
610 finished: false,
611 pending_work: 0,
612 }),
613 identity,
614 }),
615 application_span: None,
616 step_span: None,
617 }
618 }
619
620 fn ray_id(&self) -> Option<String> {
621 self.inner.state.lock().ray_id.clone()
622 }
623
624 fn record_ray(&self, span: &tracing::Span) {
627 if let Some(ray_id) = self.inner.state.lock().ray_id.as_deref() {
628 span.record("rivet.ray.id", ray_id);
629 }
630 }
631
632 pub(crate) fn with_application_span(
638 &self,
639 traceparent: Option<&str>,
640 tracestate: Option<&str>,
641 ) -> Self {
642 let application_span = parse_remote_parent(traceparent, tracestate).filter(|parent| {
643 self.inner
644 .state
645 .lock()
646 .span
647 .as_ref()
648 .and_then(otel_span_context_of)
649 .map_or(true, |own| own.span_id() != parent.span_id())
650 });
651 Self {
652 inner: self.inner.clone(),
653 application_span,
654 step_span: self.step_span.clone(),
655 }
656 }
657
658 fn parent_context(&self) -> Option<Context> {
662 let state = self.inner.state.lock();
663 if state.finished && state.pending_work == 0 {
664 return None;
665 }
666 match self.application_span.as_ref().or(self.step_span.as_ref()) {
667 Some(span_context) => {
668 Some(Context::new().with_remote_span_context(span_context.clone()))
669 }
670 None => state.span.as_ref().map(tracing::Span::context),
671 }
672 }
673
674 pub(crate) fn hold_open(&self) -> Option<InvocationWorkGuard> {
677 let mut state = self.inner.state.lock();
678 if state.finished {
679 return None;
680 }
681 state.pending_work += 1;
682 Some(InvocationWorkGuard(self.clone()))
683 }
684
685 #[doc(hidden)]
687 pub fn trace_context(&self) -> Option<ActorInvocationTraceContext> {
688 let (ray_id, span) = {
689 let state = self.inner.state.lock();
690 if state.finished && state.pending_work == 0 {
691 return None;
692 }
693 (state.ray_id.clone(), state.span.clone())
694 };
695 let span = match &self.step_span {
696 Some(step_span) => w3c_span_context(step_span),
697 None => span.and_then(|span| span_context_of(&span)),
698 };
699
700 Some(ActorInvocationTraceContext { ray_id, span })
701 }
702
703 pub(crate) fn incoming_trace_context(&self) -> IncomingTraceContext {
710 let headers = self
711 .parent_context()
712 .map(|parent| w3c_trace_headers(parent.span().span_context()))
713 .unwrap_or_default();
714 IncomingTraceContext {
715 ray_id: self.ray_id(),
716 traceparent: headers.traceparent,
717 tracestate: headers.tracestate,
718 }
719 }
720
721 pub(crate) fn start_outbound_call(
732 &self,
733 actor_name: &str,
734 action_name: &str,
735 ) -> Option<OutboundCallInvocation> {
736 let parent = self.parent_context()?;
737 let span = tracing::info_span!(
738 target: "rivetkit::telemetry",
739 parent: None,
740 "rivet.actor.call",
741 otel.name = %format!("{actor_name}/{action_name}"),
742 otel.kind = "client",
743 rivet.actor.name = %actor_name,
744 rivet.action.name = %action_name,
745 rivet.ray.id = tracing::field::Empty,
746 otel.status_code = tracing::field::Empty,
747 error.type = tracing::field::Empty,
748 );
749 self.record_ray(&span);
750 span.set_parent(parent);
751 let context = span_context_of(&span);
752 Some(OutboundCallInvocation {
753 span: Some(span),
754 context,
755 })
756 }
757
758 pub(crate) fn start_workflow_step(
761 &self,
762 step_name: &str,
763 attempt: u32,
764 ) -> Option<(tracing::Span, Self)> {
765 let parent = self.parent_context()?;
766 let span = tracing::info_span!(
767 target: "rivetkit::telemetry",
768 parent: None,
769 "rivet.workflow.step",
770 otel.name = %format!("{}/{}", self.inner.identity.actor_name, step_name),
771 otel.kind = "internal",
772 rivet.workflow.step.name = %step_name,
773 rivet.workflow.step.attempt = attempt,
774 rivet.workflow.step.outcome = tracing::field::Empty,
775 rivet.ray.id = tracing::field::Empty,
776 rivet.actor.id = %self.inner.identity.actor_id,
777 rivet.actor.name = %self.inner.identity.actor_name,
778 rivet.actor.key = %self.inner.identity.actor_key,
779 otel.status_code = tracing::field::Empty,
780 error.type = tracing::field::Empty,
781 );
782 self.record_ray(&span);
783 span.set_parent(parent);
784 let step_telemetry = Self {
785 inner: self.inner.clone(),
786 application_span: None,
787 step_span: otel_span_context_of(&span),
788 };
789 Some((span, step_telemetry))
790 }
791
792 pub(crate) fn start_sqlite(&self, operation: SqliteOperation) -> Option<SqliteOperationSpan> {
793 let parent = self.parent_context()?;
794 let (span_name, operation_name) = operation.names();
795 let span = tracing::info_span!(
796 target: "rivetkit::telemetry",
797 parent: None,
798 "rivet.sqlite.operation",
799 otel.name = span_name,
800 otel.kind = "internal",
801 rivet.operation.system = "sqlite",
802 rivet.operation.name = operation_name,
803 rivet.ray.id = tracing::field::Empty,
804 rivet.actor.id = %self.inner.identity.actor_id,
805 rivet.actor.name = %self.inner.identity.actor_name,
806 rivet.actor.key = %self.inner.identity.actor_key,
807 otel.status_code = tracing::field::Empty,
808 error.type = tracing::field::Empty,
809 );
810 self.record_ray(&span);
811 span.set_parent(parent);
812 Some(SqliteOperationSpan { span: Some(span) })
813 }
814
815 fn finish(&self, error: Option<&anyhow::Error>) {
816 self.finish_with(|span| record_outcome(span, error));
817 }
818
819 fn finish_dropped(&self) {
820 self.finish_with(|span| {
821 span.record("otel.status_code", "ERROR");
822 span.record("error.type", "actor.dropped_reply");
823 });
824 }
825
826 fn finish_with(&self, record: impl FnOnce(&tracing::Span)) {
827 let mut state = self.inner.state.lock();
828 if state.finished {
829 return;
830 }
831 if let Some(span) = state.span.as_ref() {
832 record(span);
833 if state.pending_work > 0 {
834 tracing::info!(target: "rivetkit::telemetry", parent: span, "reply sent");
835 }
836 }
837 state.finished = true;
838 if state.pending_work == 0 {
839 state.span.take();
840 }
841 }
842}
843
844impl Drop for InvocationWorkGuard {
845 fn drop(&mut self) {
846 let mut state = self.0.inner.state.lock();
847 state.pending_work -= 1;
848 if state.pending_work == 0 && state.finished {
849 state.span.take();
850 }
851 }
852}
853
854impl OutboundCallInvocation {
855 pub fn span_context(&self) -> Option<ActorInvocationSpanContext> {
858 self.context.clone()
859 }
860
861 pub fn finish(mut self, error: Option<&anyhow::Error>) {
864 let Some(span) = self.span.take() else {
865 return;
866 };
867 record_outcome(&span, error);
868 }
869}
870
871impl Drop for OutboundCallInvocation {
872 fn drop(&mut self) {
873 let Some(span) = self.span.take() else {
874 return;
875 };
876 span.record("otel.status_code", "ERROR");
877 span.record("error.type", OPERATION_ABANDONED_ERROR_TYPE);
878 }
879}
880
881impl WorkflowStepSpan {
882 pub(crate) fn start(ctx: &ActorContext, step_name: &str, attempt: u32) -> Option<Self> {
883 let (span, step_telemetry) = ctx
884 .invocation_telemetry()?
885 .start_workflow_step(step_name, attempt)?;
886 Some(Self {
887 ctx: ctx.clone().with_invocation_telemetry(Some(step_telemetry)),
888 span: Some(span),
889 })
890 }
891
892 pub fn ctx(&self) -> ActorContext {
894 self.ctx.clone()
895 }
896
897 pub fn finish(mut self, outcome: WorkflowStepOutcome, error: Option<&anyhow::Error>) {
899 let Some(span) = self.span.take() else {
900 return;
901 };
902 span.record("rivet.workflow.step.outcome", outcome.as_label());
903 record_outcome(&span, error);
904 }
905}
906
907impl Drop for WorkflowStepSpan {
908 fn drop(&mut self) {
909 let Some(span) = self.span.take() else {
910 return;
911 };
912 span.record("otel.status_code", "ERROR");
913 span.record("error.type", OPERATION_ABANDONED_ERROR_TYPE);
914 }
915}
916
917impl SqliteOperationSpan {
918 pub(crate) fn span(&self) -> tracing::Span {
919 self.span.as_ref().expect("sqlite span is present").clone()
920 }
921
922 pub(crate) fn finish(&mut self, error: Option<&anyhow::Error>) {
923 let Some(span) = self.span.take() else {
924 return;
925 };
926 record_outcome(&span, error);
927 }
928}
929
930impl Drop for SqliteOperationSpan {
931 fn drop(&mut self) {
932 let Some(span) = self.span.take() else {
933 return;
934 };
935 span.record("otel.status_code", "ERROR");
936 span.record("error.type", OPERATION_ABANDONED_ERROR_TYPE);
937 }
938}
939
940fn otel_span_context_of(span: &tracing::Span) -> Option<SpanContext> {
941 let context = span.context();
942 let context_span = context.span();
943 let span_context = context_span.span_context();
944 if span_context.is_valid() {
945 Some(span_context.clone())
946 } else {
947 None
948 }
949}
950
951fn w3c_span_context(span_context: &SpanContext) -> Option<ActorInvocationSpanContext> {
954 if !span_context.is_valid() {
955 return None;
956 }
957 let tracestate = span_context.trace_state().header();
958 Some(ActorInvocationSpanContext {
959 trace_id: span_context.trace_id().to_string(),
960 span_id: span_context.span_id().to_string(),
961 trace_flags: span_context.trace_flags().to_u8(),
962 tracestate: if tracestate.is_empty() {
963 None
964 } else {
965 Some(tracestate)
966 },
967 })
968}
969
970fn w3c_trace_headers(span_context: &SpanContext) -> OwnedTraceHeaders {
971 if !span_context.is_valid() {
972 return OwnedTraceHeaders::default();
973 }
974 let context = Context::new().with_remote_span_context(span_context.clone());
975 let mut headers = OwnedTraceHeaders::default();
976 TraceContextPropagator::new().inject_context(&context, &mut headers);
977 headers
978}
979
980fn otel_kind(invocation_type: InvocationType) -> &'static str {
984 match invocation_type {
985 InvocationType::Action | InvocationType::Request => "server",
986 InvocationType::Scheduled | InvocationType::Workflow => "internal",
987 InvocationType::QueueSend => "producer",
988 }
989}
990
991fn span_context_of(span: &tracing::Span) -> Option<ActorInvocationSpanContext> {
994 let context = span.context();
995 let context_span = context.span();
996 w3c_span_context(context_span.span_context())
997}
998
999pub(crate) fn start_queue_receive(ctx: &ActorContext, message: &QueueMessage) -> tracing::Span {
1010 let message_ray_id = message.trace_context.ray_id.as_deref();
1011 if let (Some(telemetry), Some(message_ray_id)) = (ctx.invocation_telemetry(), message_ray_id)
1012 && telemetry.inner.takes_message_ray
1013 {
1014 telemetry.inner.state.lock().ray_id = Some(message_ray_id.to_owned());
1015 }
1016 if !tracing::enabled!(target: "rivetkit::telemetry", tracing::Level::INFO) {
1017 return tracing::Span::none();
1018 }
1019 let identity = ctx.telemetry_identity();
1020 let invocation_parent = ctx
1021 .invocation_telemetry()
1022 .and_then(|telemetry| telemetry.parent_context().map(|parent| (telemetry, parent)));
1023 let sampler = match invocation_parent {
1025 Some(_) => None,
1026 None => ctx.tracing_config().sampler,
1027 };
1028 let span = tracing::info_span!(
1029 target: "rivetkit::telemetry",
1030 parent: None,
1031 "rivet.queue.receive",
1032 otel.name = %format!("{}/queue.receive", identity.actor_name),
1033 otel.kind = "consumer",
1034 rivet.actor.id = %identity.actor_id,
1035 rivet.actor.name = %identity.actor_name,
1036 rivet.actor.key = %identity.actor_key,
1037 rivet.queue.name = message.name,
1038 rivet.ray.id = tracing::field::Empty,
1039 rivet.sampling.ratio = sampler,
1040 );
1041 if let Some((telemetry, parent)) = invocation_parent {
1042 let state = telemetry.inner.state.lock();
1043 if let Some(ray_id) = state.ray_id.as_deref().or(message_ray_id) {
1044 span.record("rivet.ray.id", ray_id);
1045 }
1046 drop(state);
1047 span.set_parent(parent);
1048 } else if let Some(ray_id) = message_ray_id {
1049 span.record("rivet.ray.id", ray_id);
1050 }
1051 if let Some(link) = parse_remote_parent(
1052 message.trace_context.traceparent.as_deref(),
1053 message.trace_context.tracestate.as_deref(),
1054 ) {
1055 span.add_link(link);
1056 }
1057 span
1058}
1059
1060fn record_outcome(span: &tracing::Span, error: Option<&anyhow::Error>) {
1062 span.record(
1063 "otel.status_code",
1064 if error.is_none() { "OK" } else { "ERROR" },
1065 );
1066 if let Some(error) = error {
1067 let error = rivet_error::RivetError::extract(error);
1068 span.record("error.type", format!("{}.{}", error.group(), error.code()));
1069 }
1070}
1071
1072fn parse_remote_parent(traceparent: Option<&str>, tracestate: Option<&str>) -> Option<SpanContext> {
1073 traceparent?;
1074 let context = TraceContextPropagator::new().extract(&TraceHeaders {
1075 traceparent,
1076 tracestate,
1077 });
1078 let span = context.span();
1079 if span.span_context().is_valid() {
1080 Some(span.span_context().clone())
1081 } else {
1082 None
1083 }
1084}
1085
1086struct TraceHeaders<'a> {
1087 traceparent: Option<&'a str>,
1088 tracestate: Option<&'a str>,
1089}
1090
1091#[derive(Default)]
1092struct OwnedTraceHeaders {
1093 traceparent: Option<String>,
1094 tracestate: Option<String>,
1095}
1096
1097impl Injector for OwnedTraceHeaders {
1098 fn set(&mut self, key: &str, value: String) {
1099 match key {
1100 "traceparent" => self.traceparent = Some(value),
1101 "tracestate" if !value.is_empty() => self.tracestate = Some(value),
1102 _ => {}
1103 }
1104 }
1105}
1106
1107impl Extractor for TraceHeaders<'_> {
1108 fn get(&self, key: &str) -> Option<&str> {
1109 match key {
1110 "traceparent" => self.traceparent,
1111 "tracestate" => self.tracestate,
1112 _ => None,
1113 }
1114 }
1115
1116 fn keys(&self) -> Vec<&str> {
1117 let mut keys = Vec::with_capacity(2);
1118 if self.traceparent.is_some() {
1119 keys.push("traceparent");
1120 }
1121 if self.tracestate.is_some() {
1122 keys.push("tracestate");
1123 }
1124 keys
1125 }
1126}