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