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 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 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 pub async fn reset(&self) -> Result<()> {
594 self.ensure_initialized().await?;
595
596 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 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