feat(rivetkit-core): trace queue sends and receives

rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/mod.rs

@@ -385,18 +385,21 @@ pub(crate) async fn persist_queue_message(
 	id: u64,
 	next_id: u64,
 	message: &PersistedQueueMessage,
+	trace_context: IncomingTraceContext,
 ) -> 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")?;
+	let mut params = vec![
+		BindParam::Integer(id),
+		BindParam::Text(message.name.clone()),
+		BindParam::Blob(message.body.clone()),
+		BindParam::Integer(message.created_at),
+	];
+	params.extend(trace_context_params(trace_context));
 	db.execute_batch(vec![
 		SqliteBatchStatement {
 			sql: INSERT_QUEUE_MESSAGE_SQL.to_owned(),
-			params: Some(vec![
-				BindParam::Integer(id),
-				BindParam::Text(message.name.clone()),
-				BindParam::Blob(message.body.clone()),
-				BindParam::Integer(message.created_at),
-			]),
+			params: Some(params),
 		},
 		SqliteBatchStatement {
 			sql: UPSERT_QUEUE_NEXT_ID_SQL.to_owned(),
@@ -421,17 +424,18 @@ pub(crate) async fn persist_queue_messages(
 	for chunk in split_queue_tx_chunks(messages) {
 		let mut statements = Vec::with_capacity(chunk.len());
 		for (id, message) in chunk {
+			let mut params = vec![
+				BindParam::Integer(
+					i64::try_from(*id).context("queue message id exceeds sqlite integer range")?,
+				),
+				BindParam::Text(message.name.clone()),
+				BindParam::Blob(message.body.clone()),
+				BindParam::Integer(message.created_at),
+			];
+			params.extend(trace_context_params(IncomingTraceContext::default()));
 			statements.push(SqliteBatchStatement {
 				sql: INSERT_QUEUE_MESSAGE_SQL.to_owned(),
-				params: Some(vec![
-					BindParam::Integer(
-						i64::try_from(*id)
-							.context("queue message id exceeds sqlite integer range")?,
-					),
-					BindParam::Text(message.name.clone()),
-					BindParam::Blob(message.body.clone()),
-					BindParam::Integer(message.created_at),
-				]),
+				params: Some(params),
 			});
 		}
 		db.execute_batch(statements)
@@ -645,6 +649,7 @@ fn decode_queue_message_rows(rows: &[Vec<ColumnValue>]) -> Result<Vec<QueueMessa
 					in_flight: None,
 					in_flight_at: None,
 				},
+				trace_context: read_trace_context(row, 4, "queue message trace context")?,
 			})
 		})
 		.collect()
@@ -688,6 +693,7 @@ pub(crate) async fn reset_queue(db: &SqliteDb) -> Result<()> {
 pub(crate) struct QueueMessageRow {
 	pub id: u64,
 	pub message: PersistedQueueMessage,
+	pub trace_context: IncomingTraceContext,
 }
 
 pub(crate) async fn user_kv_batch_get(

rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/queries.rs

@@ -55,11 +55,9 @@ pub(crate) fn claim_one_shots_sql(event_count: usize) -> String {
 pub(crate) const LOAD_QUEUE_NEXT_ID_SQL: &str =
 	"SELECT queue_next_id FROM _rivet_runtime WHERE id = 1";
 pub(crate) const LOAD_QUEUE_STATS_SQL: &str = "SELECT COUNT(*), MAX(id) FROM _rivet_queue";
-pub(crate) const LOAD_QUEUE_MESSAGES_SQL: &str =
-	"SELECT id, name, body, created_at FROM _rivet_queue ORDER BY id";
-pub(crate) const LOAD_QUEUE_MESSAGES_LIMITED_SQL: &str =
-	"SELECT id, name, body, created_at FROM _rivet_queue ORDER BY id LIMIT ?";
-pub(crate) const LOAD_QUEUE_MESSAGES_FOR_NAME_SQL: &str = "SELECT id, name, body, created_at FROM _rivet_queue INDEXED BY _rivet_queue_name_id WHERE name = ? ORDER BY id LIMIT ?";
+pub(crate) const LOAD_QUEUE_MESSAGES_SQL: &str = "SELECT id, name, body, created_at, ray_id, traceparent, tracestate FROM _rivet_queue ORDER BY id";
+pub(crate) const LOAD_QUEUE_MESSAGES_LIMITED_SQL: &str = "SELECT id, name, body, created_at, ray_id, traceparent, tracestate FROM _rivet_queue ORDER BY id LIMIT ?";
+pub(crate) const LOAD_QUEUE_MESSAGES_FOR_NAME_SQL: &str = "SELECT id, name, body, created_at, ray_id, traceparent, tracestate 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,11 +69,10 @@ pub(crate) fn load_queue_messages_by_ids_sql(id_count: usize) -> String {
 		.collect::<Vec<_>>()
 		.join(", ");
 	format!(
-		"SELECT id, name, body, created_at FROM _rivet_queue WHERE id IN ({placeholders}) ORDER BY id"
+		"SELECT id, name, body, created_at, ray_id, traceparent, tracestate 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 INSERT_QUEUE_MESSAGE_SQL: &str = "INSERT OR REPLACE INTO _rivet_queue (id, name, body, created_at, ray_id, traceparent, tracestate) 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";
 

rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs

@@ -421,6 +421,8 @@ pub enum ActorEvent {
 		request: Request,
 		wait: bool,
 		timeout_ms: Option<u64>,
+		/// Telemetry of the invocation this send runs as. See `Action`.
+		invocation_telemetry: Option<crate::ActorInvocationTelemetry>,
 		reply: Reply<QueueSendResult>,
 	},
 	WebSocketOpen {

rivetkit-rust/packages/rivetkit-core/src/actor/queue.rs

@@ -26,6 +26,7 @@ use crate::actor::persist::{
 use crate::actor::task_types::UserTaskKind;
 #[cfg(target_arch = "wasm32")]
 use crate::error::ActorRuntime;
+use crate::telemetry::{self, ActorInvocationTelemetry, IncomingTraceContext};
 
 #[derive(Clone, Debug, Default)]
 pub struct QueueNextOpts {
@@ -101,6 +102,10 @@ pub struct QueueMessage {
 	pub name: String,
 	pub body: Vec<u8>,
 	pub created_at: i64,
+	/// Ray ID and span of the invocation that sent the message, which the span
+	/// covering its receipt links back to. Empty when the sender carried no
+	/// trace context.
+	pub trace_context: IncomingTraceContext,
 	completion: Option<CompletionHandle>,
 }
 
@@ -110,6 +115,7 @@ pub struct CompletableQueueMessage {
 	pub name: String,
 	pub body: Vec<u8>,
 	pub created_at: i64,
+	pub trace_context: IncomingTraceContext,
 	completion: CompletionHandle,
 }
 
@@ -281,7 +287,10 @@ impl ActorContext {
 			in_flight_at: None,
 		};
 		let encoded_message = encode_queue_message(&persisted).context("encode queue message")?;
-
+		let trace_context = self
+			.invocation_telemetry()
+			.map(ActorInvocationTelemetry::incoming_trace_context)
+			.unwrap_or_default();
 		let config = self.config();
 		if encoded_message.len() > config.max_queue_message_size as usize {
 			return Err(QueueMessageTooLarge {
@@ -323,9 +332,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,
+			trace_context.clone(),
+		)
+		.await;
 
 		if let Err(error) = persist_result {
 			metadata.next_id = id;
@@ -350,6 +364,7 @@ impl ActorContext {
 			name: name.to_owned(),
 			body: body.to_vec(),
 			created_at,
+			trace_context,
 			completion: None,
 		})
 	}
@@ -653,6 +668,11 @@ impl ActorContext {
 			return Ok(Vec::new());
 		}
 
+		let _receive_spans: Vec<tracing::Span> = 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 +947,7 @@ impl QueueMessage {
 			name: self.name,
 			body: self.body,
 			created_at: self.created_at,
+			trace_context: self.trace_context,
 			completion,
 		})
 	}
@@ -947,6 +968,7 @@ impl CompletableQueueMessage {
 			name: self.name,
 			body: self.body,
 			created_at: self.created_at,
+			trace_context: self.trace_context,
 			completion: Some(self.completion),
 		}
 	}
@@ -1056,6 +1078,7 @@ fn queue_message_from_row(row: internal_storage::QueueMessageRow) -> QueueMessag
 		name: row.message.name,
 		body: row.message.body,
 		created_at: row.message.created_at,
+		trace_context: row.trace_context,
 		completion: None,
 	}
 }

rivetkit-rust/packages/rivetkit-core/src/actor/task.rs

@@ -209,6 +209,7 @@ pub enum DispatchCommand {
 	QueueSend {
 		name: String,
 		body: Vec<u8>,
+		incoming: crate::telemetry::IncomingInvocationContext,
 		conn: ConnHandle,
 		request: Request,
 		wait: bool,
@@ -989,30 +990,36 @@ 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);
+				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(reply)
+							.on_reply(move |result| invocation.finish(result.as_ref().err())),
+					},
+				) {
+					Ok(()) => {
+						self.log_dispatch_command_handled(command_kind, "enqueued");
+					}
+					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);

rivetkit-rust/packages/rivetkit-core/src/lib.rs

@@ -24,6 +24,7 @@ pub mod telemetry;
 #[doc(hidden)]
 pub use telemetry::{
 	ActorInvocationSpanContext, ActorInvocationTelemetry, ActorInvocationTraceContext,
+	IncomingTraceContext,
 };
 #[cfg(any(test, feature = "test-support"))]
 pub mod testing;

rivetkit-rust/packages/rivetkit-core/src/registry/http.rs

@@ -370,12 +370,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,