Skip to main content

rivetkit_core/actor/
queue.rs

1use std::collections::BTreeSet;
2use std::fmt;
3use std::future::pending;
4use std::pin::Pin;
5use std::sync::Arc;
6use std::sync::atomic::Ordering;
7use std::time::Duration;
8
9use crate::time::{Instant, SystemTime, UNIX_EPOCH, sleep};
10
11use anyhow::{Context, Result};
12use rivet_error::RivetError;
13use rivetkit_actor_persist::{generated::v4 as persist_v4, versioned as persist_versioned};
14use serde::{Deserialize, Serialize};
15#[cfg(not(target_arch = "wasm32"))]
16use tokio::runtime::{Builder, Handle};
17use tokio::sync::{futures::Notified, oneshot};
18use tokio_util::sync::CancellationToken;
19
20use crate::actor::config::ActorConfig;
21use crate::actor::context::ActorContext;
22use crate::actor::internal_storage;
23use crate::actor::persist::{
24	decode_latest_with_embedded_version, encode_latest_with_embedded_version,
25};
26use crate::actor::task_types::UserTaskKind;
27#[cfg(target_arch = "wasm32")]
28use crate::error::ActorRuntime;
29use crate::telemetry::{self, ActorInvocationTelemetry, IncomingTraceContext};
30
31#[derive(Clone, Debug, Default)]
32pub struct QueueNextOpts {
33	pub names: Option<Vec<String>>,
34	pub timeout: Option<Duration>,
35	pub signal: Option<CancellationToken>,
36	pub completable: bool,
37}
38
39#[derive(Clone, Debug, Default)]
40pub struct QueueWaitOpts {
41	pub timeout: Option<Duration>,
42	pub signal: Option<CancellationToken>,
43	pub completable: bool,
44}
45
46#[derive(Clone, Debug, Default)]
47pub struct EnqueueAndWaitOpts {
48	pub timeout: Option<Duration>,
49	pub signal: Option<CancellationToken>,
50}
51
52#[derive(Clone, Debug)]
53pub struct QueueNextBatchOpts {
54	pub names: Option<Vec<String>>,
55	pub count: u32,
56	pub timeout: Option<Duration>,
57	pub signal: Option<CancellationToken>,
58	pub completable: bool,
59}
60
61impl Default for QueueNextBatchOpts {
62	fn default() -> Self {
63		Self {
64			names: None,
65			count: 1,
66			timeout: None,
67			signal: None,
68			completable: false,
69		}
70	}
71}
72
73#[derive(Clone, Debug, Default)]
74pub struct QueueTryNextOpts {
75	pub names: Option<Vec<String>>,
76	pub completable: bool,
77}
78
79#[derive(Clone, Debug)]
80pub struct QueueTryNextBatchOpts {
81	pub names: Option<Vec<String>>,
82	pub count: u32,
83	pub completable: bool,
84}
85
86impl Default for QueueTryNextBatchOpts {
87	fn default() -> Self {
88		Self {
89			names: None,
90			count: 1,
91			completable: false,
92		}
93	}
94}
95
96pub(super) type QueueWaitActivityCallback = Arc<dyn Fn() + Send + Sync>;
97pub(super) type QueueInspectorUpdateCallback = Arc<dyn Fn(u32) + Send + Sync>;
98
99#[derive(Clone, Debug)]
100pub struct QueueMessage {
101	pub id: u64,
102	pub name: String,
103	pub body: Vec<u8>,
104	pub created_at: i64,
105	/// Ray ID and span of the invocation that sent the message, which the span
106	/// covering its receipt links back to. Empty when the sender carried no
107	/// trace context.
108	pub trace_context: IncomingTraceContext,
109	completion: Option<CompletionHandle>,
110}
111
112#[derive(Clone, Debug)]
113pub struct CompletableQueueMessage {
114	pub id: u64,
115	pub name: String,
116	pub body: Vec<u8>,
117	pub created_at: i64,
118	pub trace_context: IncomingTraceContext,
119	completion: CompletionHandle,
120}
121
122#[derive(Clone)]
123struct CompletionHandle(Arc<CompletionHandleInner>);
124
125struct CompletionHandleInner {
126	ctx: ActorContext,
127	message_id: u64,
128	completed: std::sync::atomic::AtomicBool,
129}
130
131pub(crate) type QueueMetadata = persist_v4::QueueMetadata;
132pub(crate) type PersistedQueueMessage = persist_v4::QueueMessage;
133
134
135
136pub(crate) fn decode_queue_metadata(payload: &[u8]) -> Result<QueueMetadata> {
137	let metadata = decode_latest_with_embedded_version::<persist_versioned::QueueMetadata>(
138		payload,
139		"queue metadata",
140	)?;
141	Ok(metadata)
142}
143
144pub(crate) fn encode_queue_message(message: &PersistedQueueMessage) -> Result<Vec<u8>> {
145	encode_latest_with_embedded_version::<persist_versioned::QueueMessage>(
146		message.clone(),
147		rivetkit_actor_persist::CURRENT_VERSION,
148		"queue message",
149	)
150}
151
152pub(crate) fn decode_queue_message(payload: &[u8]) -> Result<PersistedQueueMessage> {
153	let message = decode_latest_with_embedded_version::<persist_versioned::QueueMessage>(
154		payload,
155		"queue message",
156	)?;
157	Ok(message)
158}
159
160#[derive(RivetError, Serialize, Deserialize)]
161#[error(
162	"queue",
163	"full",
164	"Queue is full",
165	"Queue is full. Limit is {limit} messages."
166)]
167struct QueueFull {
168	limit: u32,
169}
170
171#[derive(RivetError, Serialize, Deserialize)]
172#[error(
173	"queue",
174	"message_too_large",
175	"Queue message is too large",
176	"Queue message too large ({size} bytes). Limit is {limit} bytes."
177)]
178struct QueueMessageTooLarge {
179	size: usize,
180	limit: u32,
181}
182
183#[derive(RivetError)]
184#[error("queue", "already_completed", "Queue message was already completed")]
185struct QueueAlreadyCompleted;
186
187#[derive(RivetError, Serialize, Deserialize)]
188#[error(
189	"queue",
190	"complete_not_configured",
191	"Queue message does not support completion",
192	"Queue '{name}' does not support completion responses."
193)]
194struct QueueCompleteNotConfigured {
195	name: String,
196}
197#[derive(RivetError, Serialize, Deserialize)]
198#[error(
199	"queue",
200	"message_identity_mismatch",
201	"Queue message identity does not match",
202	"Queue message {message_id} is named '{actual_name}', not '{expected_name}'."
203)]
204struct QueueMessageIdentityMismatch {
205	message_id: u64,
206	expected_name: String,
207	actual_name: String,
208}
209
210#[derive(RivetError)]
211#[error("actor", "aborted", "Actor aborted")]
212struct QueueActorAborted;
213
214#[derive(RivetError, Serialize, Deserialize)]
215#[error(
216	"queue",
217	"timed_out",
218	"Queue wait timed out",
219	"Queue wait timed out after {timeout_ms} ms."
220)]
221struct QueueWaitTimedOut {
222	timeout_ms: u64,
223}
224
225#[derive(RivetError, Serialize, Deserialize)]
226#[error(
227	"queue",
228	"completion_waiter_conflict",
229	"Queue completion waiter conflict",
230	"Queue completion waiter is already registered for message {message_id}."
231)]
232struct QueueCompletionWaiterConflict {
233	message_id: u64,
234}
235
236#[derive(RivetError)]
237#[error(
238	"queue",
239	"completion_waiter_dropped",
240	"Queue completion waiter dropped before response"
241)]
242struct QueueCompletionWaiterDropped;
243
244impl ActorContext {
245	pub async fn send(&self, name: &str, body: &[u8]) -> Result<QueueMessage> {
246		self.enqueue_message(name, body, None).await
247	}
248
249	pub async fn enqueue_and_wait(
250		&self,
251		name: &str,
252		body: &[u8],
253		opts: EnqueueAndWaitOpts,
254	) -> Result<Option<Vec<u8>>> {
255		let (sender, receiver) = oneshot::channel();
256		let message = self.enqueue_message(name, body, Some(sender)).await?;
257		let result = self
258			.wait_for_completion_response(message.id, receiver, opts.timeout, opts.signal.as_ref())
259			.await;
260		self.remove_completion_waiter(message.id).await;
261		result
262	}
263
264	async fn enqueue_message(
265		&self,
266		name: &str,
267		body: &[u8],
268		completion_waiter: Option<oneshot::Sender<Option<Vec<u8>>>>,
269	) -> Result<QueueMessage> {
270		self.ensure_initialized().await?;
271
272		let created_at = current_timestamp_ms()?;
273		let persisted = PersistedQueueMessage {
274			name: name.to_owned().into(),
275			body: body.to_vec().into(),
276			created_at,
277			failure_count: None,
278			available_at: None,
279			in_flight: None,
280			in_flight_at: None,
281		};
282		let encoded_message = encode_queue_message(&persisted).context("encode queue message")?;
283		let trace_context = self
284			.invocation_telemetry()
285			.map(ActorInvocationTelemetry::incoming_trace_context)
286			.unwrap_or_default();
287		let config = self.config();
288		if encoded_message.len() > config.max_queue_message_size as usize {
289			return Err(QueueMessageTooLarge {
290				size: encoded_message.len(),
291				limit: config.max_queue_message_size,
292			}
293			.build());
294		}
295
296		let mut metadata = self.0.queue_metadata.lock().await;
297		if metadata.size >= config.max_queue_size {
298			return Err(QueueFull {
299				limit: config.max_queue_size,
300			}
301			.build());
302		}
303
304		let id = if metadata.next_id == 0 {
305			1
306		} else {
307			metadata.next_id
308		};
309		metadata.next_id = id.saturating_add(1);
310		metadata.size = metadata.size.saturating_add(1);
311		let registered_completion_waiter = if let Some(waiter) = completion_waiter {
312			if self
313				.0
314				.queue_completion_waiters
315				.insert_async(id, waiter)
316				.await
317				.is_err()
318			{
319				metadata.next_id = id;
320				metadata.size = metadata.size.saturating_sub(1);
321				return Err(QueueCompletionWaiterConflict { message_id: id }.build());
322			}
323			true
324		} else {
325			false
326		};
327
328		let persist_result = internal_storage::persist_queue_message(
329			self.sql(),
330			id,
331			metadata.next_id,
332			&persisted,
333			trace_context.clone(),
334		)
335		.await;
336
337		if let Err(error) = persist_result {
338			metadata.next_id = id;
339			metadata.size = metadata.size.saturating_sub(1);
340			if registered_completion_waiter {
341				self.remove_completion_waiter(id).await;
342			}
343			return Err(error).context("persist queue message");
344		}
345
346		let queue_size = metadata.size;
347		drop(metadata);
348		self.0.metrics.add_queue_messages_sent(1);
349		self.0
350			.metrics
351			.set_queue_depth(self.0.queue_metadata.lock().await.size);
352		self.notify_inspector_update(queue_size);
353		self.0.queue_notify.notify_waiters();
354
355		Ok(QueueMessage {
356			id,
357			name: name.to_owned(),
358			body: body.to_vec(),
359			created_at,
360			trace_context,
361			completion: None,
362		})
363	}
364
365	pub async fn next(&self, opts: QueueNextOpts) -> Result<Option<QueueMessage>> {
366		let mut messages = self
367			.next_batch(QueueNextBatchOpts {
368				names: opts.names,
369				count: 1,
370				timeout: opts.timeout,
371				signal: opts.signal,
372				completable: opts.completable,
373			})
374			.await?;
375		Ok(messages.pop())
376	}
377
378	pub async fn next_batch(&self, opts: QueueNextBatchOpts) -> Result<Vec<QueueMessage>> {
379		self.ensure_initialized().await?;
380
381		let count = opts.count.max(1);
382		let deadline = opts.timeout.map(|timeout| Instant::now() + timeout);
383		let names = normalize_names(opts.names);
384
385		loop {
386			let notified = self.0.queue_notify.notified();
387			tokio::pin!(notified);
388			notified.as_mut().enable();
389			let messages = self
390				.try_receive_batch(names.as_ref(), count, opts.completable)
391				.await?;
392			if !messages.is_empty() {
393				return Ok(messages);
394			}
395
396			let remaining_timeout =
397				deadline.map(|deadline| deadline.saturating_duration_since(Instant::now()));
398			if matches!(remaining_timeout, Some(timeout) if timeout.is_zero()) {
399				return Ok(Vec::new());
400			}
401
402			let wait_guard = ActiveQueueWaitGuard::new(self);
403			let result = self
404				.wait_for_message(notified.as_mut(), remaining_timeout, opts.signal.as_ref())
405				.await;
406			drop(wait_guard);
407
408			match result {
409				WaitOutcome::Notified => continue,
410				WaitOutcome::TimedOut => return Ok(Vec::new()),
411				WaitOutcome::Aborted => return Err(QueueActorAborted.build()),
412			}
413		}
414	}
415
416	pub async fn wait_for_names(
417		&self,
418		names: Vec<String>,
419		opts: QueueWaitOpts,
420	) -> Result<QueueMessage> {
421		self.ensure_initialized().await?;
422
423		let deadline = opts.timeout.map(|timeout| Instant::now() + timeout);
424		let names = normalize_names(Some(names));
425
426		loop {
427			let notified = self.0.queue_notify.notified();
428			tokio::pin!(notified);
429			notified.as_mut().enable();
430			if let Some(message) = self
431				.try_receive_batch(names.as_ref(), 1, opts.completable)
432				.await?
433				.into_iter()
434				.next()
435			{
436				return Ok(message);
437			}
438
439			let remaining_timeout =
440				deadline.map(|deadline| deadline.saturating_duration_since(Instant::now()));
441			if let Some(timeout) = remaining_timeout
442				&& timeout.is_zero()
443			{
444				return Err(QueueWaitTimedOut {
445					timeout_ms: opts.timeout.map(duration_ms).unwrap_or(0),
446				}
447				.build());
448			}
449
450			let wait_guard = ActiveQueueWaitGuard::new(self);
451			let result = self
452				.wait_for_message(notified.as_mut(), remaining_timeout, opts.signal.as_ref())
453				.await;
454			drop(wait_guard);
455
456			match result {
457				WaitOutcome::Notified => continue,
458				WaitOutcome::TimedOut => {
459					return Err(QueueWaitTimedOut {
460						timeout_ms: opts.timeout.map(duration_ms).unwrap_or(0),
461					}
462					.build());
463				}
464				WaitOutcome::Aborted => return Err(QueueActorAborted.build()),
465			}
466		}
467	}
468
469	pub async fn wait_for_names_available(
470		&self,
471		names: Vec<String>,
472		opts: QueueWaitOpts,
473	) -> Result<()> {
474		self.ensure_initialized().await?;
475
476		let deadline = opts.timeout.map(|timeout| Instant::now() + timeout);
477		let names = normalize_names(Some(names));
478
479		loop {
480			let notified = self.0.queue_notify.notified();
481			tokio::pin!(notified);
482			notified.as_mut().enable();
483			if internal_storage::has_queue_messages(self.sql(), names.as_ref())
484				.await
485				.context("check for matching sqlite queue messages")?
486			{
487				return Ok(());
488			}
489
490			let remaining_timeout =
491				deadline.map(|deadline| deadline.saturating_duration_since(Instant::now()));
492			if let Some(timeout) = remaining_timeout
493				&& timeout.is_zero()
494			{
495				return Err(QueueWaitTimedOut {
496					timeout_ms: opts.timeout.map(duration_ms).unwrap_or(0),
497				}
498				.build());
499			}
500
501			let wait_guard = ActiveQueueWaitGuard::new(self);
502			let result = self
503				.wait_for_message(notified.as_mut(), remaining_timeout, opts.signal.as_ref())
504				.await;
505			drop(wait_guard);
506
507			match result {
508				WaitOutcome::Notified => continue,
509				WaitOutcome::TimedOut => {
510					return Err(QueueWaitTimedOut {
511						timeout_ms: opts.timeout.map(duration_ms).unwrap_or(0),
512					}
513					.build());
514				}
515				WaitOutcome::Aborted => return Err(QueueActorAborted.build()),
516			}
517		}
518	}
519	/// Completes a persisted message by durable identity. Missing or already
520	/// completed IDs are idempotent; an expected-name mismatch leaves the row
521	/// untouched.
522	pub async fn complete_persisted_message(
523		&self,
524		message_id: u64,
525		expected_name: &str,
526		response: Option<Vec<u8>>,
527	) -> Result<bool> {
528		self.ensure_initialized().await?;
529		let _receive_guard = self.0.queue_receive_lock.lock().await;
530		let Some(_) = self
531			.verify_persisted_message_identity_unlocked(message_id, expected_name)
532			.await?
533		else {
534			return Ok(false);
535		};
536		self.complete_message_by_id_unlocked(message_id, response)
537			.await?;
538		Ok(true)
539	}
540
541	async fn verify_persisted_message_identity_unlocked(
542		&self,
543		message_id: u64,
544		expected_name: &str,
545	) -> Result<Option<String>> {
546		let Some(actual_name) =
547			internal_storage::load_queue_message_name(self.sql(), message_id).await?
548		else {
549			return Ok(None);
550		};
551		if actual_name != expected_name {
552			return Err(QueueMessageIdentityMismatch {
553				message_id,
554				expected_name: expected_name.to_owned(),
555				actual_name: actual_name.clone(),
556			}
557			.build());
558		}
559		Ok(Some(actual_name))
560	}
561
562	pub fn try_next(&self, opts: QueueTryNextOpts) -> Result<Option<QueueMessage>> {
563		let mut messages = self.try_next_batch(QueueTryNextBatchOpts {
564			names: opts.names,
565			count: 1,
566			completable: opts.completable,
567		})?;
568		Ok(messages.pop())
569	}
570
571	pub fn try_next_batch(&self, opts: QueueTryNextBatchOpts) -> Result<Vec<QueueMessage>> {
572		self.block_on(async {
573			self.ensure_initialized().await?;
574			self.try_receive_batch(
575				normalize_names(opts.names).as_ref(),
576				opts.count.max(1),
577				opts.completable,
578			)
579			.await
580		})
581	}
582
583	pub async fn inspect_messages(&self) -> Result<Vec<QueueMessage>> {
584		self.ensure_initialized().await?;
585		self.list_messages().await
586	}
587
588	pub fn max_size(&self) -> u32 {
589		self.config().max_queue_size
590	}
591
592	/// Removes all messages from the queue and resets the size counter.
593	pub async fn reset(&self) -> Result<()> {
594		self.ensure_initialized().await?;
595
596		// Serialize against receivers before touching metadata. Lock order matches
597		// try_receive_batch (receive lock then metadata) so there is no deadlock.
598		let _receive_guard = self.0.queue_receive_lock.lock().await;
599
600		let mut metadata = self.0.queue_metadata.lock().await;
601
602		internal_storage::reset_queue(self.sql())
603			.await
604			.context("delete all sqlite queue messages")?;
605
606		metadata.size = 0;
607
608		self.0.queue_completion_waiters.clear_async().await;
609
610		drop(metadata);
611
612		self.0.metrics.set_queue_depth(0);
613		self.notify_inspector_update(0);
614		self.0.queue_notify.notify_waiters();
615
616		Ok(())
617	}
618
619	pub(crate) fn configure_queue(&self, config: ActorConfig) {
620		*self.0.queue_config.lock() = config;
621	}
622
623	pub(crate) fn set_wait_activity_callback(&self, callback: Option<Arc<dyn Fn() + Send + Sync>>) {
624		*self.0.queue_wait_activity_callback.lock() = callback;
625	}
626
627	pub(crate) fn set_inspector_update_callback(
628		&self,
629		callback: Option<Arc<dyn Fn(u32) + Send + Sync>>,
630	) {
631		*self.0.queue_inspector_update_callback.lock() = callback;
632	}
633
634	async fn ensure_initialized(&self) -> Result<()> {
635		self.0
636			.queue_initialize
637			.get_or_try_init(|| async {
638				let metadata = internal_storage::load_queue_metadata(self.sql())
639					.await
640					.context("load queue metadata from sqlite")?;
641				let mut state = self.0.queue_metadata.lock().await;
642				*state = metadata;
643				self.0.metrics.set_queue_depth(state.size);
644				Ok(())
645			})
646			.await
647			.map(|_| ())
648	}
649
650	async fn try_receive_batch(
651		&self,
652		names: Option<&BTreeSet<String>>,
653		count: u32,
654		completable: bool,
655	) -> Result<Vec<QueueMessage>> {
656		let _receive_guard = self.0.queue_receive_lock.lock().await;
657
658		let selected = self.list_messages_matching(names, count).await?;
659
660		if selected.is_empty() {
661			return Ok(Vec::new());
662		}
663
664		let _receive_spans: Vec<tracing::Span> = selected
665			.iter()
666			.map(|message| telemetry::start_queue_receive(self, message))
667			.collect();
668
669		if completable {
670			let queue_size = self.0.queue_metadata.lock().await.size;
671			self.0
672				.metrics
673				.add_queue_messages_received(selected.len().try_into().unwrap_or(u64::MAX));
674			self.notify_inspector_update(queue_size);
675			return Ok(selected
676				.into_iter()
677				.map(|message| self.attach_completion(message))
678				.collect());
679		}
680
681		self.remove_messages(selected.iter().map(|message| message.id).collect())
682			.await?;
683		self.0
684			.metrics
685			.add_queue_messages_received(selected.len().try_into().unwrap_or(u64::MAX));
686
687		Ok(selected)
688	}
689
690	async fn list_messages(&self) -> Result<Vec<QueueMessage>> {
691		let messages: Vec<QueueMessage> = internal_storage::load_queue_messages(self.sql())
692			.await
693			.context("list sqlite queue messages")?
694			.into_iter()
695			.map(queue_message_from_row)
696			.collect();
697
698		let actual_size = messages.len().try_into().unwrap_or(u32::MAX);
699		let mut metadata = self.0.queue_metadata.lock().await;
700		if metadata.size != actual_size {
701			metadata.size = actual_size;
702		}
703		if metadata.next_id == 0 {
704			metadata.next_id = messages
705				.last()
706				.map(|message| message.id.saturating_add(1))
707				.unwrap_or(1);
708		}
709
710		Ok(messages)
711	}
712
713	async fn list_messages_matching(
714		&self,
715		names: Option<&BTreeSet<String>>,
716		limit: u32,
717	) -> Result<Vec<QueueMessage>> {
718		internal_storage::load_queue_messages_matching(self.sql(), names, limit)
719			.await
720			.context("list matching sqlite queue messages")
721			.map(|rows| rows.into_iter().map(queue_message_from_row).collect())
722	}
723
724	fn attach_completion(&self, mut message: QueueMessage) -> QueueMessage {
725		message.completion = Some(CompletionHandle::new(self.clone(), message.id));
726		message
727	}
728
729	async fn remove_messages(&self, message_ids: Vec<u64>) -> Result<()> {
730		if message_ids.is_empty() {
731			return Ok(());
732		}
733
734		let deleted_count = internal_storage::delete_queue_messages(self.sql(), &message_ids)
735			.await
736			.context("delete sqlite queue messages")?;
737
738		let queue_size = {
739			let mut metadata = self.0.queue_metadata.lock().await;
740			metadata.size = metadata.size.saturating_sub(deleted_count);
741			metadata.size
742		};
743		self.0
744			.metrics
745			.set_queue_depth(self.0.queue_metadata.lock().await.size);
746		self.notify_inspector_update(queue_size);
747		Ok(())
748	}
749
750	async fn complete_message_by_id(
751		&self,
752		message_id: u64,
753		response: Option<Vec<u8>>,
754	) -> Result<()> {
755		let _receive_guard = self.0.queue_receive_lock.lock().await;
756		self.complete_message_by_id_unlocked(message_id, response)
757			.await
758	}
759
760	async fn complete_message_by_id_unlocked(
761		&self,
762		message_id: u64,
763		response: Option<Vec<u8>>,
764	) -> Result<()> {
765		self.remove_messages(vec![message_id]).await?;
766		if let Some(waiter) = self.remove_completion_waiter(message_id).await {
767			let _ = waiter.send(response);
768		}
769		Ok(())
770	}
771
772	async fn remove_completion_waiter(
773		&self,
774		message_id: u64,
775	) -> Option<oneshot::Sender<Option<Vec<u8>>>> {
776		self.0
777			.queue_completion_waiters
778			.remove_async(&message_id)
779			.await
780			.map(|(_, waiter)| waiter)
781	}
782
783	async fn wait_for_message(
784		&self,
785		mut notified: Pin<&mut Notified<'_>>,
786		timeout: Option<Duration>,
787		signal: Option<&CancellationToken>,
788	) -> WaitOutcome {
789		let actor_abort_signal = self.0.queue_abort_signal.lock().clone();
790		if signal.is_some_and(CancellationToken::is_cancelled) {
791			return WaitOutcome::Aborted;
792		}
793		if actor_abort_signal.is_cancelled() {
794			return WaitOutcome::Aborted;
795		}
796
797		let actor_aborted = async {
798			actor_abort_signal.cancelled().await;
799		};
800		let external_aborted = async {
801			if let Some(signal) = signal {
802				signal.cancelled().await;
803			} else {
804				pending::<()>().await;
805			}
806		};
807
808		match timeout {
809			Some(timeout) => {
810				tokio::select! {
811					_ = &mut notified => WaitOutcome::Notified,
812					_ = actor_aborted => WaitOutcome::Aborted,
813					_ = external_aborted => WaitOutcome::Aborted,
814					_ = sleep(timeout) => WaitOutcome::TimedOut,
815				}
816			}
817			None => {
818				tokio::select! {
819					_ = &mut notified => WaitOutcome::Notified,
820					_ = actor_aborted => WaitOutcome::Aborted,
821					_ = external_aborted => WaitOutcome::Aborted,
822				}
823			}
824		}
825	}
826
827	/// TS parity: queue-manager.ts keeps `enqueueAndWait` completion waits
828	/// alive across actor aborts; the surrounding tracked user task owns
829	/// shutdown cancellation.
830	async fn wait_for_completion_response(
831		&self,
832		message_id: u64,
833		mut receiver: oneshot::Receiver<Option<Vec<u8>>>,
834		timeout: Option<Duration>,
835		signal: Option<&CancellationToken>,
836	) -> Result<Option<Vec<u8>>> {
837		if signal.is_some_and(CancellationToken::is_cancelled) {
838			return Err(QueueActorAborted.build());
839		}
840
841		let external_aborted = async {
842			if let Some(signal) = signal {
843				signal.cancelled().await;
844			} else {
845				pending::<()>().await;
846			}
847		};
848
849		let wait_result = match timeout {
850			Some(timeout) => {
851				tokio::select! {
852					response = &mut receiver => CompletionWaitOutcome::Response(response),
853					_ = external_aborted => CompletionWaitOutcome::Aborted,
854					_ = sleep(timeout) => CompletionWaitOutcome::TimedOut,
855				}
856			}
857			None => {
858				tokio::select! {
859					response = &mut receiver => CompletionWaitOutcome::Response(response),
860					_ = external_aborted => CompletionWaitOutcome::Aborted,
861				}
862			}
863		};
864
865		match wait_result {
866			CompletionWaitOutcome::Response(Ok(response)) => Ok(response),
867			CompletionWaitOutcome::Response(Err(_)) => Err(QueueCompletionWaiterDropped.build())
868				.context(format!("wait for queue completion on message {message_id}")),
869			CompletionWaitOutcome::TimedOut => Err(QueueWaitTimedOut {
870				timeout_ms: timeout.map(duration_ms).unwrap_or(0),
871			}
872			.build()),
873			CompletionWaitOutcome::Aborted => Err(QueueActorAborted.build()),
874		}
875	}
876
877	fn block_on<T>(&self, future: impl std::future::Future<Output = Result<T>>) -> Result<T> {
878		#[cfg(not(target_arch = "wasm32"))]
879		{
880			if let Ok(handle) = Handle::try_current() {
881				tokio::task::block_in_place(|| handle.block_on(future))
882			} else {
883				Builder::new_current_thread()
884					.enable_all()
885					.build()
886					.context("build temporary runtime for queue operation")?
887					.block_on(future)
888			}
889		}
890
891		#[cfg(target_arch = "wasm32")]
892		{
893			drop(future);
894			Err(ActorRuntime::InvalidOperation {
895				operation: "queue.try_next_batch".to_owned(),
896				reason: "synchronous queue receive requires native runtime support".to_owned(),
897			}
898			.build())
899		}
900	}
901
902	fn config(&self) -> ActorConfig {
903		self.0.queue_config.lock().clone()
904	}
905
906	
907
908	fn notify_wait_activity(&self) {
909		if let Some(callback) = self.0.queue_wait_activity_callback.lock().clone() {
910			callback();
911		}
912	}
913
914	fn notify_inspector_update(&self, queue_size: u32) {
915		if let Some(callback) = self.0.queue_inspector_update_callback.lock().clone() {
916			callback(queue_size);
917		}
918	}
919}
920
921impl QueueMessage {
922	pub async fn complete(self, response: Option<Vec<u8>>) -> Result<()> {
923		let completable = self.into_completable()?;
924		completable.complete(response).await
925	}
926
927	pub fn into_completable(self) -> Result<CompletableQueueMessage> {
928		let completion = self.completion.clone().ok_or_else(|| {
929			QueueCompleteNotConfigured {
930				name: self.name.clone(),
931			}
932			.build()
933		})?;
934
935		Ok(CompletableQueueMessage {
936			id: self.id,
937			name: self.name,
938			body: self.body,
939			created_at: self.created_at,
940			trace_context: self.trace_context,
941			completion,
942		})
943	}
944
945	pub fn is_completable(&self) -> bool {
946		self.completion.is_some()
947	}
948}
949
950impl CompletableQueueMessage {
951	pub async fn complete(self, response: Option<Vec<u8>>) -> Result<()> {
952		self.completion.complete(response).await
953	}
954
955	pub fn into_message(self) -> QueueMessage {
956		QueueMessage {
957			id: self.id,
958			name: self.name,
959			body: self.body,
960			created_at: self.created_at,
961			trace_context: self.trace_context,
962			completion: Some(self.completion),
963		}
964	}
965}
966
967impl CompletionHandle {
968	fn new(ctx: ActorContext, message_id: u64) -> Self {
969		Self(Arc::new(CompletionHandleInner {
970			ctx,
971			message_id,
972			completed: std::sync::atomic::AtomicBool::new(false),
973		}))
974	}
975
976	async fn complete(&self, response: Option<Vec<u8>>) -> Result<()> {
977		if self.0.completed.swap(true, Ordering::SeqCst) {
978			return Err(QueueAlreadyCompleted.build());
979		}
980
981		if let Err(error) = self
982			.0
983			.ctx
984			.complete_message_by_id(self.0.message_id, response)
985			.await
986		{
987			self.0.completed.store(false, Ordering::SeqCst);
988			return Err(error);
989		}
990
991		Ok(())
992	}
993}
994
995impl fmt::Debug for CompletionHandle {
996	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
997		f.debug_struct("CompletionHandle")
998			.field("message_id", &self.0.message_id)
999			.field("completed", &self.0.completed.load(Ordering::SeqCst))
1000			.finish()
1001	}
1002}
1003
1004struct ActiveQueueWaitGuard<'a> {
1005	ctx: &'a ActorContext,
1006	started_at: Instant,
1007}
1008
1009impl<'a> ActiveQueueWaitGuard<'a> {
1010	fn new(ctx: &'a ActorContext) -> Self {
1011		ctx.0.active_queue_wait_count.fetch_add(1, Ordering::SeqCst);
1012		ctx.0.metrics.begin_user_task(UserTaskKind::QueueWait);
1013		ctx.notify_wait_activity();
1014		Self {
1015			ctx,
1016			started_at: Instant::now(),
1017		}
1018	}
1019}
1020
1021impl Drop for ActiveQueueWaitGuard<'_> {
1022	fn drop(&mut self) {
1023		self.ctx
1024			.0
1025			.metrics
1026			.end_user_task(UserTaskKind::QueueWait, self.started_at.elapsed());
1027		let previous = self
1028			.ctx
1029			.0
1030			.active_queue_wait_count
1031			.fetch_sub(1, Ordering::SeqCst);
1032		if previous == 0 {
1033			self.ctx
1034				.0
1035				.active_queue_wait_count
1036				.store(0, Ordering::SeqCst);
1037		}
1038		self.ctx.notify_wait_activity();
1039	}
1040}
1041
1042enum WaitOutcome {
1043	Notified,
1044	TimedOut,
1045	Aborted,
1046}
1047
1048enum CompletionWaitOutcome {
1049	Response(Result<Option<Vec<u8>>, oneshot::error::RecvError>),
1050	TimedOut,
1051	Aborted,
1052}
1053
1054fn normalize_names(names: Option<Vec<String>>) -> Option<BTreeSet<String>> {
1055	names.and_then(|names| {
1056		let normalized = names.into_iter().collect::<BTreeSet<_>>();
1057		if normalized.is_empty() {
1058			None
1059		} else {
1060			Some(normalized)
1061		}
1062	})
1063}
1064
1065fn queue_message_from_row(row: internal_storage::QueueMessageRow) -> QueueMessage {
1066	QueueMessage {
1067		id: row.id,
1068		name: row.message.name.to_string(),
1069		body: row.message.body.to_vec(),
1070		created_at: row.message.created_at,
1071		trace_context: row.trace_context,
1072		completion: None,
1073	}
1074}
1075
1076fn current_timestamp_ms() -> Result<i64> {
1077	let now = SystemTime::now()
1078		.duration_since(UNIX_EPOCH)
1079		.context("current time is before unix epoch")?;
1080	i64::try_from(now.as_millis()).context("queue timestamp exceeds i64")
1081}
1082
1083fn duration_ms(duration: Duration) -> u64 {
1084	duration.as_millis().try_into().unwrap_or(u64::MAX)
1085}
1086
1087// Test shim keeps moved tests in crate-root tests/ with private-module access.
1088