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 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 pub async fn reset(&self) -> Result<()> {
586 self.ensure_initialized().await?;
587
588 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 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#[cfg(test)]
1076#[path = "../../tests/queue.rs"]
1077mod tests;