1use crate::admin::{CallbackConfig, CallbackPollResult};
2use crate::dlq::{ListDlqFilter, RetryFromDlqOpts};
3use crate::error::AwaError;
4use crate::insert::prepare_row_raw;
5use crate::{InsertParams, JobRow, JobState};
6use chrono::TimeDelta;
7use chrono::{DateTime, Utc};
8use sqlx::Acquire;
9use sqlx::{PgPool, Postgres, QueryBuilder};
10use std::collections::hash_map::DefaultHasher;
11use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
12use std::hash::{Hash, Hasher};
13use std::sync::atomic::{AtomicU16, AtomicUsize, Ordering};
14use std::sync::Mutex;
15use std::time::Duration;
16use uuid::Uuid;
17
18const DEFAULT_SCHEMA: &str = "awa";
19const DEFAULT_QUEUE_SLOT_COUNT: usize = 16;
20const DEFAULT_LEASE_SLOT_COUNT: usize = 8;
21const DEFAULT_CLAIM_SLOT_COUNT: usize = 8;
22const DEFAULT_QUEUE_STRIPE_COUNT: usize = 1;
23const QUEUE_STRIPE_DELIMITER: &str = "#";
24const COPY_NULL_SENTINEL: &str = "__AWA_NULL__";
25const COPY_CHUNK_TARGET_BYTES: usize = 256 * 1024;
26const TERMINAL_COUNTER_BUCKETS: i16 = 256;
27const RECEIPT_RESCUE_BATCH_LIMIT: i64 = 500;
28const RECEIPT_RESCUE_CURSOR_SCAN_LIMIT: i64 = 10_000;
29const RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT: i64 = 10_000;
30const DEFAULT_PRUNE_LOCK_TIMEOUT: Duration = Duration::from_millis(10);
31
32pub fn ordering_key_hash64(ordering_key: &[u8]) -> u64 {
38 let mut hash: u128 = 14_695_981_039_346_656_037;
39 const PRIME: u128 = 1_099_511_628_211;
40 const MASK: u128 = u64::MAX as u128;
41 for byte in ordering_key {
42 hash = hash.wrapping_mul(PRIME).wrapping_add(*byte as u128) & MASK;
43 }
44 hash as u64
45}
46
47pub fn shard_for_ordering_key(ordering_key: &[u8], shards: i16) -> i16 {
56 if shards <= 1 {
57 return 0;
58 }
59 (ordering_key_hash64(ordering_key) % shards as u64) as i16
60}
61
62fn terminal_counter_bucket(job_id: i64) -> i16 {
63 job_id.rem_euclid(TERMINAL_COUNTER_BUCKETS as i64) as i16
64}
65
66#[derive(Debug, Clone)]
67pub struct QueueStorageConfig {
68 pub schema: String,
69 pub queue_slot_count: usize,
70 pub lease_slot_count: usize,
71 pub claim_slot_count: usize,
78 pub queue_stripe_count: usize,
79 pub lease_claim_receipts: bool,
96}
97
98impl Default for QueueStorageConfig {
99 fn default() -> Self {
100 Self {
101 schema: DEFAULT_SCHEMA.to_string(),
102 queue_slot_count: DEFAULT_QUEUE_SLOT_COUNT,
103 lease_slot_count: DEFAULT_LEASE_SLOT_COUNT,
104 claim_slot_count: DEFAULT_CLAIM_SLOT_COUNT,
105 queue_stripe_count: DEFAULT_QUEUE_STRIPE_COUNT,
106 lease_claim_receipts: true,
107 }
108 }
109}
110
111#[derive(Debug, Clone, PartialEq, Eq, sqlx::FromRow)]
112pub struct ClaimedEntry {
113 pub queue: String,
114 pub priority: i16,
115 pub lane_seq: i64,
116 pub ready_slot: i32,
117 pub ready_generation: i64,
118 pub lease_slot: i32,
119 pub lease_generation: i64,
120 pub claim_slot: i32,
126 pub receipt_id: Option<i64>,
130 pub claim_batch_id: Option<i64>,
133 pub claim_batch_index: Option<i32>,
136 pub lease_claim_receipt: bool,
137 pub enqueue_shard: i16,
143}
144
145#[derive(Debug, Clone)]
146pub struct ClaimedRuntimeJob {
147 pub claim: ClaimedEntry,
148 pub job: JobRow,
149 pub unique_states: Option<String>,
150}
151
152#[derive(Debug, Clone, Copy, PartialEq, Eq, sqlx::FromRow)]
153pub struct QueueClaimerLease {
154 pub claimer_slot: i16,
155 pub lease_epoch: i64,
156}
157
158#[derive(Debug, Clone, Copy, PartialEq, Eq, sqlx::FromRow)]
159struct QueueClaimerLeaseRow {
160 claimer_slot: i16,
161 lease_epoch: i64,
162 last_claimed_at: DateTime<Utc>,
163 expires_at: DateTime<Utc>,
164}
165
166impl QueueClaimerLeaseRow {
167 fn lease(self) -> QueueClaimerLease {
168 QueueClaimerLease {
169 claimer_slot: self.claimer_slot,
170 lease_epoch: self.lease_epoch,
171 }
172 }
173
174 fn needs_refresh(
175 self,
176 now: DateTime<Utc>,
177 lease_ttl: Duration,
178 idle_threshold: Duration,
179 ) -> bool {
180 let Ok(idle_refresh_delta) = TimeDelta::from_std(idle_threshold / 2) else {
181 return true;
182 };
183 let Ok(expiry_refresh_delta) = TimeDelta::from_std(lease_ttl / 2) else {
184 return true;
185 };
186
187 self.last_claimed_at <= now - idle_refresh_delta
188 || self.expires_at <= now + expiry_refresh_delta
189 }
190}
191
192#[derive(Debug, Clone, Copy, PartialEq, Eq, sqlx::FromRow)]
193pub struct QueueClaimerState {
194 pub target_claimers: i16,
195}
196
197impl ClaimedRuntimeJob {
198 fn into_done_row(self, finalized_at: DateTime<Utc>) -> Result<DoneJobRow, AwaError> {
199 let payload = QueueStorage::payload_from_parts(
200 self.job.metadata,
201 self.job.tags,
202 self.job.errors,
203 None,
204 )?;
205
206 Ok(DoneJobRow {
207 ready_slot: self.claim.ready_slot,
208 ready_generation: self.claim.ready_generation,
209 job_id: self.job.id,
210 kind: self.job.kind,
211 queue: self.job.queue,
212 args: self.job.args,
213 state: JobState::Completed,
214 priority: self.claim.priority,
215 attempt: self.job.attempt,
216 run_lease: self.job.run_lease,
217 max_attempts: self.job.max_attempts,
218 lane_seq: self.claim.lane_seq,
219 enqueue_shard: self.claim.enqueue_shard,
220 run_at: self.job.run_at,
221 attempted_at: self.job.attempted_at,
222 finalized_at,
223 created_at: self.job.created_at,
224 unique_key: self.job.unique_key,
225 unique_states: self.unique_states,
226 payload,
227 })
228 }
229}
230
231#[derive(Debug, Clone, Copy, PartialEq, Eq)]
232pub struct QueueCounts {
233 pub available: i64,
234 pub running: i64,
235 pub terminal: i64,
245 pub pruned_failed: i64,
251}
252
253#[derive(Debug, Clone, Copy, PartialEq, Eq)]
265pub(crate) struct AvailableSignal {
266 pub available: i64,
267}
268
269#[derive(Debug, Clone, Copy, PartialEq, Eq)]
270pub enum RotateOutcome {
271 Rotated {
272 slot: i32,
273 generation: i64,
274 },
275 SkippedBusy {
279 slot: i32,
280 busy: BusyCounts,
281 },
282}
283
284#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
290pub struct BusyCounts {
291 pub queue_ready: i64,
293 pub queue_claim_attempt_batches: i64,
295 pub queue_done: i64,
297 pub queue_tombstones: i64,
299 pub queue_ready_segments: i64,
301 pub queue_receipt_completion_batches: i64,
303 pub queue_receipt_completion_tombstones: i64,
305 pub queue_terminal_deltas: i64,
307 pub leases: i64,
309 pub claims: i64,
311 pub closures: i64,
313 pub closure_batches: i64,
315}
316
317#[derive(Debug, Clone, Copy, PartialEq, Eq)]
318pub enum PruneOutcome {
319 Noop,
320 Pruned {
321 slot: i32,
322 carried_failed_rows: u64,
327 },
328 Blocked {
330 slot: i32,
331 },
332 SkippedActive {
335 slot: i32,
336 reason: SkipReason,
337 count: i64,
338 },
339}
340
341#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
342pub struct TerminalDeltaRollupOutcome {
343 pub rolled_slots: usize,
344 pub delta_rows: i64,
345 pub grouped_keys: i64,
346 pub skipped_active_slots: usize,
347 pub blocked_slots: usize,
348 pub skipped_mvcc_pinned: bool,
349}
350
351#[derive(Debug, Clone, Copy, PartialEq, Eq)]
352enum TerminalDeltaSlotRollup {
353 Empty,
354 Rolled { delta_rows: i64, grouped_keys: i64 },
355 SkippedActive,
356 SkippedMvccPinned,
357 Blocked,
358}
359
360#[derive(Debug, Clone, Copy, PartialEq, Eq)]
368pub enum SkipReason {
369 QueueActiveLeases,
371 QueueUnclosedClaimRefs,
373 QueuePendingReady,
375 LeaseCurrent,
377 LeaseActive,
379 ClaimCurrent,
381 ClaimOpen,
383}
384
385impl SkipReason {
386 pub fn as_str(self) -> &'static str {
388 match self {
389 Self::QueueActiveLeases => "queue.active_leases",
390 Self::QueueUnclosedClaimRefs => "queue.unclosed_claim_refs",
391 Self::QueuePendingReady => "queue.pending_ready",
392 Self::LeaseCurrent => "lease.current",
393 Self::LeaseActive => "lease.active",
394 Self::ClaimCurrent => "claim.current",
395 Self::ClaimOpen => "claim.open",
396 }
397 }
398}
399
400fn map_sqlx_error(err: sqlx::Error) -> AwaError {
401 if let sqlx::Error::Database(ref db_err) = err {
402 if db_err.code().as_deref() == Some("23505") {
403 return AwaError::UniqueConflict {
404 constraint: db_err.constraint().map(|c| c.to_string()),
405 };
406 }
407 }
408 AwaError::Database(err)
409}
410
411fn is_lock_contention_error(err: &sqlx::Error) -> bool {
412 matches!(
413 err,
414 sqlx::Error::Database(db_err) if db_err.code().as_deref() == Some("55P03")
415 )
416}
417
418async fn set_prune_lock_timeout_tx(
419 tx: &mut sqlx::Transaction<'_, Postgres>,
420 timeout: Duration,
421) -> Result<(), AwaError> {
422 let millis = timeout.as_millis().max(1).min(i64::MAX as u128);
423 let timeout = format!("{millis}ms");
424 sqlx::query("SELECT set_config('lock_timeout', $1, true)")
425 .bind(timeout)
426 .execute(tx.as_mut())
427 .await
428 .map_err(map_sqlx_error)?;
429 Ok(())
430}
431
432fn validate_ident(ident: &str) -> Result<(), AwaError> {
433 let mut chars = ident.chars();
434 match chars.next() {
435 Some(first) if first.is_ascii_lowercase() || first == '_' => {}
436 _ => {
437 return Err(AwaError::Validation(format!(
438 "invalid SQL identifier: {ident}"
439 )));
440 }
441 }
442
443 if chars.all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_') {
444 Ok(())
445 } else {
446 Err(AwaError::Validation(format!(
447 "invalid SQL identifier: {ident}"
448 )))
449 }
450}
451
452fn ready_child_name(schema: &str, slot: usize) -> String {
453 format!("{schema}.ready_entries_{slot}")
454}
455
456fn ready_claim_attempt_batch_child_name(schema: &str, slot: usize) -> String {
457 format!("{schema}.ready_claim_attempt_batches_{slot}")
458}
459
460fn done_child_name(schema: &str, slot: usize) -> String {
461 format!("{schema}.done_entries_{slot}")
462}
463
464fn ready_tombstone_child_name(schema: &str, slot: usize) -> String {
465 format!("{schema}.ready_tombstones_{slot}")
466}
467
468fn ready_segment_child_name(schema: &str, slot: usize) -> String {
469 format!("{schema}.ready_segments_{slot}")
470}
471
472fn receipt_completion_batch_child_name(schema: &str, slot: usize) -> String {
473 format!("{schema}.receipt_completion_batches_{slot}")
474}
475
476fn receipt_completion_tombstone_child_name(schema: &str, slot: usize) -> String {
477 format!("{schema}.receipt_completion_tombstones_{slot}")
478}
479
480fn terminal_delta_child_name(schema: &str, slot: usize) -> String {
481 format!("{schema}.queue_terminal_count_deltas_{slot}")
482}
483
484fn done_ready_join(schema: &str, done_alias: &str, ready_alias: &str) -> String {
485 format!(
486 r#"
487 LEFT JOIN {schema}.ready_entries AS {ready_alias}
488 ON {ready_alias}.ready_slot = {done_alias}.ready_slot
489 AND {ready_alias}.ready_generation = {done_alias}.ready_generation
490 AND {ready_alias}.queue = {done_alias}.queue
491 AND {ready_alias}.priority = {done_alias}.priority
492 AND {ready_alias}.enqueue_shard = {done_alias}.enqueue_shard
493 AND {ready_alias}.lane_seq = {done_alias}.lane_seq
494 "#
495 )
496}
497
498fn done_row_projection(done_alias: &str, ready_alias: &str) -> String {
499 format!(
500 r#"
501 {done_alias}.ready_slot,
502 {done_alias}.ready_generation,
503 {done_alias}.job_id,
504 {done_alias}.kind,
505 {done_alias}.queue,
506 COALESCE({done_alias}.args, {ready_alias}.args, '{{}}'::jsonb) AS args,
507 {done_alias}.state,
508 {done_alias}.priority,
509 {done_alias}.attempt,
510 {done_alias}.run_lease,
511 COALESCE({done_alias}.max_attempts, {ready_alias}.max_attempts, 25::smallint) AS max_attempts,
512 {done_alias}.lane_seq,
513 {done_alias}.enqueue_shard,
514 COALESCE({done_alias}.run_at, {ready_alias}.run_at, {done_alias}.finalized_at) AS run_at,
515 COALESCE({done_alias}.attempted_at, {ready_alias}.attempted_at) AS attempted_at,
516 {done_alias}.finalized_at,
517 COALESCE({done_alias}.created_at, {ready_alias}.created_at, {done_alias}.finalized_at) AS created_at,
518 COALESCE({done_alias}.unique_key, {ready_alias}.unique_key) AS unique_key,
519 COALESCE({done_alias}.unique_states, {ready_alias}.unique_states) AS unique_states,
520 COALESCE({done_alias}.payload, {ready_alias}.payload, '{{}}'::jsonb) AS payload
521 "#
522 )
523}
524
525fn lease_child_name(schema: &str, slot: usize) -> String {
526 format!("{schema}.leases_{slot}")
527}
528
529fn claim_child_name(schema: &str, slot: usize) -> String {
530 format!("{schema}.lease_claims_{slot}")
531}
532
533fn claim_batch_child_name(schema: &str, slot: usize) -> String {
534 format!("{schema}.lease_claim_batches_{slot}")
535}
536
537fn closure_child_name(schema: &str, slot: usize) -> String {
538 format!("{schema}.lease_claim_closures_{slot}")
539}
540
541fn claim_closure_batch_child_name(schema: &str, slot: usize) -> String {
542 format!("{schema}.lease_claim_closure_batches_{slot}")
543}
544
545fn ring_slot_index(slot: i32, slot_count: usize, ring: &str) -> Result<usize, AwaError> {
546 let index = usize::try_from(slot).map_err(|_| {
547 AwaError::Validation(format!(
548 "invalid {ring} ring slot {slot}: slot must be non-negative"
549 ))
550 })?;
551 if index >= slot_count {
552 return Err(AwaError::Validation(format!(
553 "invalid {ring} ring slot {slot}: configured slot count is {slot_count}"
554 )));
555 }
556 Ok(index)
557}
558
559fn receipt_closed_evidence_sql(
560 schema: &str,
561 closure_rel: &str,
562 closure_batch_rel: &str,
563 claims_alias: &str,
564) -> String {
565 format!(
570 r#"
571 (
572 {claims_alias}.closed_at IS NOT NULL
573 OR EXISTS (
574 SELECT 1 FROM {closure_rel} AS closures
575 WHERE closures.claim_slot = {claims_alias}.claim_slot
576 AND closures.job_id = {claims_alias}.job_id
577 AND closures.run_lease = {claims_alias}.run_lease
578 )
579 OR EXISTS (
580 SELECT 1
581 FROM {closure_batch_rel} AS closure_batches
582 WHERE closure_batches.receipt_ranges @> {claims_alias}.receipt_id
583 )
584 OR EXISTS (
585 SELECT 1 FROM {schema}.done_entries AS done
586 WHERE done.job_id = {claims_alias}.job_id
587 AND done.run_lease = {claims_alias}.run_lease
588 )
589 OR EXISTS (
590 SELECT 1 FROM {schema}.deferred_jobs AS deferred
591 WHERE deferred.job_id = {claims_alias}.job_id
592 AND deferred.run_lease = {claims_alias}.run_lease
593 )
594 OR EXISTS (
595 SELECT 1 FROM {schema}.dlq_entries AS dlq
596 WHERE dlq.job_id = {claims_alias}.job_id
597 AND dlq.run_lease = {claims_alias}.run_lease
598 )
599 )
600 "#
601 )
602}
603
604fn busy_indicator(has_rows: bool) -> i64 {
605 if has_rows {
606 1
607 } else {
608 0
609 }
610}
611
612async fn queue_prune_has_active_leases_tx(
613 tx: &mut sqlx::Transaction<'_, Postgres>,
614 schema: &str,
615 slot: i32,
616 generation: i64,
617) -> Result<bool, AwaError> {
618 sqlx::query_scalar(&format!(
619 r#"
620 SELECT EXISTS (
621 SELECT 1
622 FROM {schema}.leases
623 WHERE ready_slot = $1
624 AND ready_generation = $2
625 LIMIT 1
626 )
627 "#
628 ))
629 .bind(slot)
630 .bind(generation)
631 .fetch_one(tx.as_mut())
632 .await
633 .map_err(map_sqlx_error)
634}
635
636async fn queue_prune_has_pending_ready_tx(
637 tx: &mut sqlx::Transaction<'_, Postgres>,
638 schema: &str,
639 ready_child: &str,
640 generation: i64,
641) -> Result<bool, AwaError> {
642 sqlx::query_scalar(&format!(
643 r#"
644 WITH claim_cursors AS MATERIALIZED (
645 SELECT
646 queue,
647 priority,
648 enqueue_shard,
649 {schema}.sequence_next_value(seq_name) AS claim_seq
650 FROM {schema}.queue_claim_heads
651 )
652 SELECT EXISTS (
653 SELECT 1
654 FROM claim_cursors AS claims
655 CROSS JOIN LATERAL (
656 SELECT 1
657 FROM {ready_child} AS ready
658 WHERE ready.ready_generation = $1
659 AND ready.queue = claims.queue
660 AND ready.priority = claims.priority
661 AND ready.enqueue_shard = claims.enqueue_shard
662 AND ready.lane_seq >= claims.claim_seq
663 LIMIT 1
664 ) AS pending_ready
665 LIMIT 1
666 )
667 "#
668 ))
669 .bind(generation)
670 .fetch_one(tx.as_mut())
671 .await
672 .map_err(map_sqlx_error)
673}
674
675async fn queue_prune_has_unclosed_claim_refs_tx(
676 tx: &mut sqlx::Transaction<'_, Postgres>,
677 schema: &str,
678 slot: i32,
679 generation: i64,
680) -> Result<bool, AwaError> {
681 let count_proves_claim_refs_closed: bool = sqlx::query_scalar(&format!(
682 r#"
683 WITH claim_count AS (
684 SELECT count(*)::bigint AS total
685 FROM {schema}.lease_claims AS claims
686 WHERE claims.ready_slot = $1
687 AND claims.ready_generation = $2
688 ),
689 compact_claim_count AS (
690 SELECT COALESCE(sum(claimed_count), 0)::bigint AS total
691 FROM {schema}.lease_claim_batches AS batches
692 WHERE batches.ready_slot = $1
693 AND batches.ready_generation = $2
694 ),
695 explicit_count AS (
696 SELECT count(*)::bigint AS total
697 FROM {schema}.lease_claims AS claims
698 JOIN {schema}.lease_claim_closures AS closures
699 ON closures.claim_slot = claims.claim_slot
700 AND closures.job_id = claims.job_id
701 AND closures.run_lease = claims.run_lease
702 WHERE claims.ready_slot = $1
703 AND claims.ready_generation = $2
704 ),
705 compact_count AS (
706 SELECT COALESCE(sum(closed_count), 0)::bigint AS total
707 FROM {schema}.lease_claim_closure_batches AS batches
708 WHERE batches.ready_slot = $1
709 AND batches.ready_generation = $2
710 )
711 SELECT claim_count.total + compact_claim_count.total =
712 explicit_count.total + compact_count.total
713 FROM claim_count, compact_claim_count, explicit_count, compact_count
714 "#
715 ))
716 .bind(slot)
717 .bind(generation)
718 .fetch_one(tx.as_mut())
719 .await
720 .map_err(map_sqlx_error)?;
721 if count_proves_claim_refs_closed {
722 return Ok(false);
723 }
724
725 Ok(true)
730}
731
732async fn claim_prune_has_open_claims_tx(
733 tx: &mut sqlx::Transaction<'_, Postgres>,
734 _schema: &str,
735 claim_child: &str,
736 claim_batch_child: &str,
737 closure_child: &str,
738 closure_batch_child: &str,
739) -> Result<bool, AwaError> {
740 let count_proves_claims_closed: bool = sqlx::query_scalar(&format!(
741 r#"
742 WITH claim_count AS (
743 SELECT count(*)::bigint AS total FROM {claim_child}
744 ),
745 compact_claim_count AS (
746 SELECT COALESCE(sum(claimed_count), 0)::bigint AS total
747 FROM {claim_batch_child}
748 ),
749 explicit_count AS (
750 SELECT count(*)::bigint AS total FROM {closure_child}
751 ),
752 compact_count AS (
753 SELECT COALESCE(sum(closed_count), 0)::bigint AS total
754 FROM {closure_batch_child}
755 )
756 SELECT claim_count.total + compact_claim_count.total =
757 explicit_count.total + compact_count.total
758 FROM claim_count, compact_claim_count, explicit_count, compact_count
759 "#
760 ))
761 .fetch_one(tx.as_mut())
762 .await
763 .map_err(map_sqlx_error)?;
764 if count_proves_claims_closed {
765 return Ok(false);
766 }
767
768 Ok(true)
772}
773
774fn oldest_initialized_ring_slot(
775 current_slot: i32,
776 generation: i64,
777 slot_count: i32,
778) -> Option<(i32, i64)> {
779 if slot_count <= 1 {
780 return None;
781 }
782
783 let initialized_slots = (generation + 1).min(slot_count as i64) as i32;
784 if initialized_slots <= 1 {
785 return None;
786 }
787
788 let offset = initialized_slots - 1;
789 let oldest_slot = (current_slot - offset).rem_euclid(slot_count);
790 let oldest_generation = generation - offset as i64;
791 if oldest_generation < 0 {
792 return None;
793 }
794
795 Some((oldest_slot, oldest_generation))
796}
797
798#[cfg(test)]
799mod identifier_tests {
800 use super::{validate_ident, QueueStorage, QueueStorageConfig};
801
802 #[test]
803 fn queue_storage_schema_identifiers_are_lowercase_unquoted_names() {
804 for ident in ["awa", "awa_queue_storage", "_awa123"] {
805 validate_ident(ident).expect("identifier should be accepted");
806 }
807
808 for ident in ["Awa", "awa-queue", "123awa", "awa.queue"] {
809 assert!(
810 validate_ident(ident).is_err(),
811 "identifier should be rejected: {ident}"
812 );
813 }
814 }
815
816 #[test]
817 fn default_queue_storage_schema_requires_default_physical_shape() {
818 for config in [
819 QueueStorageConfig {
820 queue_slot_count: 32,
821 ..Default::default()
822 },
823 QueueStorageConfig {
824 lease_slot_count: 4,
825 ..Default::default()
826 },
827 QueueStorageConfig {
828 claim_slot_count: 4,
829 ..Default::default()
830 },
831 QueueStorageConfig {
832 lease_claim_receipts: false,
833 ..Default::default()
834 },
835 ] {
836 let err = QueueStorage::new(config).expect_err("default awa schema shape must reject");
837 assert!(
838 err.to_string()
839 .contains("default `awa` queue-storage schema"),
840 "unexpected error: {err}"
841 );
842 }
843
844 QueueStorage::new(QueueStorageConfig {
845 schema: "awa_custom".to_string(),
846 queue_slot_count: 4,
847 lease_slot_count: 2,
848 claim_slot_count: 2,
849 lease_claim_receipts: false,
850 ..Default::default()
851 })
852 .expect("custom schema should allow custom physical shape");
853 }
854}
855
856#[cfg(test)]
857mod shard_routing_tests {
858 use super::shard_for_ordering_key;
859 use std::collections::HashSet;
860
861 #[test]
862 fn shards_le_one_collapse_to_zero() {
863 assert_eq!(shard_for_ordering_key(b"customer-42", 1), 0);
864 assert_eq!(shard_for_ordering_key(b"", 1), 0);
865 assert_eq!(shard_for_ordering_key(b"customer-42", 0), 0);
866 }
867
868 #[test]
869 fn same_key_lands_on_same_shard() {
870 let key = b"customer-42";
871 let first = shard_for_ordering_key(key, 8);
872 for _ in 0..100 {
873 assert_eq!(shard_for_ordering_key(key, 8), first);
874 }
875 }
876
877 #[test]
878 fn shard_is_within_range() {
879 for n in 0..256u32 {
880 let key = format!("order-{n}");
881 let shard = shard_for_ordering_key(key.as_bytes(), 8);
882 assert!((0..8).contains(&shard));
883 }
884 }
885
886 #[test]
887 fn distinct_keys_spread_across_shards() {
888 let mut hit: HashSet<i16> = HashSet::new();
889 for n in 0..1024u32 {
890 let key = format!("order-{n}");
891 hit.insert(shard_for_ordering_key(key.as_bytes(), 8));
892 }
893 assert_eq!(hit.len(), 8, "1024 distinct keys should cover all 8 shards");
894 }
895}
896
897#[cfg(test)]
898mod ring_slot_tests {
899 use super::oldest_initialized_ring_slot;
900
901 #[test]
902 fn oldest_initialized_ring_slot_is_none_until_second_slot_exists() {
903 assert_eq!(oldest_initialized_ring_slot(0, 0, 8), None);
904 }
905
906 #[test]
907 fn oldest_initialized_ring_slot_tracks_partial_ring_startup() {
908 assert_eq!(oldest_initialized_ring_slot(1, 1, 8), Some((0, 0)));
909 assert_eq!(oldest_initialized_ring_slot(2, 2, 8), Some((0, 0)));
910 assert_eq!(oldest_initialized_ring_slot(3, 3, 8), Some((0, 0)));
911 }
912
913 #[test]
914 fn oldest_initialized_ring_slot_wraps_after_full_rotation() {
915 assert_eq!(oldest_initialized_ring_slot(7, 7, 8), Some((0, 0)));
916 assert_eq!(oldest_initialized_ring_slot(0, 8, 8), Some((1, 1)));
917 assert_eq!(oldest_initialized_ring_slot(1, 9, 8), Some((2, 2)));
918 }
919}
920
921#[cfg(test)]
922mod claim_cursor_advance_tests {
923 use super::{ClaimCursorAdvance, QueueStorage};
924
925 fn advance(next_seq: i64, only_if_current: Option<i64>) -> ClaimCursorAdvance {
926 ClaimCursorAdvance {
927 queue: "queue".to_string(),
928 priority: 2,
929 enqueue_shard: 0,
930 next_seq,
931 only_if_current,
932 }
933 }
934
935 #[test]
936 fn normalize_claim_cursor_advances_sorts_conditional_lane_updates() {
937 let normalized = QueueStorage::normalize_claim_cursor_advances(&[
938 advance(7, Some(6)),
939 advance(6, Some(5)),
940 advance(8, Some(7)),
941 ]);
942
943 let ordered: Vec<(i64, i64)> = normalized
944 .iter()
945 .map(|advance| (advance.only_if_current.unwrap(), advance.next_seq))
946 .collect();
947 assert_eq!(ordered, vec![(5, 6), (6, 7), (7, 8)]);
948 }
949
950 #[test]
951 fn normalize_claim_cursor_advances_coalesces_unconditional_lane_updates() {
952 let normalized = QueueStorage::normalize_claim_cursor_advances(&[
953 advance(3, None),
954 advance(5, None),
955 advance(4, Some(3)),
956 ]);
957
958 assert_eq!(normalized.len(), 1);
959 assert_eq!(normalized[0].next_seq, 5);
960 assert_eq!(normalized[0].only_if_current, None);
961 }
962}
963
964fn default_payload_metadata() -> serde_json::Value {
965 serde_json::json!({})
966}
967
968fn is_empty_json_object(value: &serde_json::Value) -> bool {
969 value.as_object().is_some_and(serde_json::Map::is_empty)
970}
971
972fn is_compact_receipt_completion_metadata(value: &serde_json::Value) -> bool {
973 let Some(metadata) = value.as_object() else {
974 return false;
975 };
976
977 metadata
978 .keys()
979 .all(|key| key == "_awa_original_priority" || key == "_awa_original_queue")
980}
981
982#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
983struct RuntimePayload {
984 #[serde(
985 default = "default_payload_metadata",
986 skip_serializing_if = "is_empty_json_object"
987 )]
988 metadata: serde_json::Value,
989 #[serde(default, skip_serializing_if = "Vec::is_empty")]
990 tags: Vec<String>,
991 #[serde(default, skip_serializing_if = "Vec::is_empty")]
992 errors: Vec<serde_json::Value>,
993 #[serde(default, skip_serializing_if = "Option::is_none")]
994 progress: Option<serde_json::Value>,
995}
996
997impl Default for RuntimePayload {
998 fn default() -> Self {
999 Self {
1000 metadata: default_payload_metadata(),
1001 tags: Vec::new(),
1002 errors: Vec::new(),
1003 progress: None,
1004 }
1005 }
1006}
1007
1008impl RuntimePayload {
1009 fn from_json(value: serde_json::Value) -> Result<Self, AwaError> {
1010 if value.is_null() {
1011 return Ok(Self::default());
1012 }
1013 let payload: Self = serde_json::from_value(value)?;
1014 if !payload.metadata.is_object() {
1015 return Err(AwaError::Validation(
1016 "queue storage payload metadata must be a JSON object".to_string(),
1017 ));
1018 }
1019 Ok(payload)
1020 }
1021
1022 fn into_json(self) -> serde_json::Value {
1023 serde_json::to_value(self).expect("runtime payload serializes")
1024 }
1025
1026 fn errors_option(&self) -> Option<Vec<serde_json::Value>> {
1027 (!self.errors.is_empty()).then(|| self.errors.clone())
1028 }
1029
1030 fn push_error(&mut self, error: serde_json::Value) {
1031 self.errors.push(error);
1032 }
1033
1034 fn set_progress(&mut self, progress: Option<serde_json::Value>) {
1035 self.progress = progress;
1036 }
1037
1038 fn insert_callback_result(&mut self, payload: Option<serde_json::Value>) {
1039 let metadata = self
1040 .metadata
1041 .as_object_mut()
1042 .expect("runtime payload metadata object");
1043 metadata.insert(
1044 "_awa_callback_result".to_string(),
1045 payload.unwrap_or(serde_json::Value::Null),
1046 );
1047 }
1048}
1049
1050#[cfg(test)]
1051mod runtime_payload_tests {
1052 use super::{
1053 is_compact_receipt_completion_metadata, storage_payload, terminal_storage_payload,
1054 RuntimePayload,
1055 };
1056
1057 #[test]
1058 fn default_runtime_payload_serializes_compactly() {
1059 assert_eq!(
1060 RuntimePayload::default().into_json(),
1061 serde_json::json!({}),
1062 "default payloads should not write empty metadata/tags/errors/progress"
1063 );
1064 assert_eq!(
1065 storage_payload(&RuntimePayload::default().into_json()),
1066 None
1067 );
1068 }
1069
1070 #[test]
1071 fn missing_runtime_payload_fields_round_trip_with_defaults() {
1072 let payload = RuntimePayload::from_json(serde_json::json!({})).unwrap();
1073
1074 assert_eq!(payload.metadata, serde_json::json!({}));
1075 assert!(payload.tags.is_empty());
1076 assert!(payload.errors.is_empty());
1077 assert_eq!(payload.progress, None);
1078 assert_eq!(payload.into_json(), serde_json::json!({}));
1079 }
1080
1081 #[test]
1082 fn null_runtime_payload_round_trips_with_defaults() {
1083 let payload = RuntimePayload::from_json(serde_json::Value::Null).unwrap();
1084
1085 assert_eq!(payload.metadata, serde_json::json!({}));
1086 assert!(payload.tags.is_empty());
1087 assert!(payload.errors.is_empty());
1088 assert_eq!(payload.progress, None);
1089 assert_eq!(storage_payload(&payload.into_json()), None);
1090 }
1091
1092 #[test]
1093 fn compact_receipt_completion_metadata_only_allows_awa_provenance() {
1094 assert!(is_compact_receipt_completion_metadata(&serde_json::json!(
1095 {}
1096 )));
1097 assert!(is_compact_receipt_completion_metadata(
1098 &serde_json::json!({ "_awa_original_priority": 4 })
1099 ));
1100 assert!(is_compact_receipt_completion_metadata(&serde_json::json!({
1101 "_awa_original_queue": "default",
1102 "_awa_original_priority": 4
1103 })));
1104 assert!(!is_compact_receipt_completion_metadata(
1105 &serde_json::json!({ "tenant": "acme" })
1106 ));
1107 assert!(!is_compact_receipt_completion_metadata(&serde_json::json!(
1108 null
1109 )));
1110 }
1111
1112 #[test]
1113 fn legacy_expanded_runtime_payload_round_trips_to_compact_form() {
1114 let payload = RuntimePayload::from_json(serde_json::json!({
1115 "metadata": {},
1116 "tags": [],
1117 "errors": [],
1118 "progress": null
1119 }))
1120 .unwrap();
1121
1122 assert_eq!(payload.metadata, serde_json::json!({}));
1123 assert!(payload.tags.is_empty());
1124 assert!(payload.errors.is_empty());
1125 assert_eq!(payload.progress, None);
1126 assert_eq!(payload.into_json(), serde_json::json!({}));
1127 }
1128
1129 #[test]
1130 fn non_default_runtime_payload_fields_are_preserved() {
1131 let payload = RuntimePayload::from_json(serde_json::json!({
1132 "metadata": { "source": "test" },
1133 "tags": ["fast"],
1134 "errors": [{ "message": "boom" }],
1135 "progress": { "step": 1 }
1136 }))
1137 .unwrap();
1138
1139 assert_eq!(
1140 payload.into_json(),
1141 serde_json::json!({
1142 "metadata": { "source": "test" },
1143 "tags": ["fast"],
1144 "errors": [{ "message": "boom" }],
1145 "progress": { "step": 1 }
1146 })
1147 );
1148 }
1149
1150 #[test]
1151 fn unchanged_terminal_payload_elides_storage_copy() {
1152 let payload = serde_json::json!({
1153 "metadata": { "source": "test" },
1154 "tags": ["fast"]
1155 });
1156
1157 assert_eq!(terminal_storage_payload(&payload, Some(&payload)), None);
1158
1159 let changed = serde_json::json!({
1160 "metadata": { "source": "test" },
1161 "tags": ["fast"],
1162 "errors": [{ "message": "boom" }]
1163 });
1164 assert_eq!(
1165 terminal_storage_payload(&changed, Some(&payload)),
1166 Some(&changed)
1167 );
1168 }
1169}
1170
1171fn unique_state_claims(unique_states: Option<&str>, state: JobState) -> bool {
1172 let Some(bitmask) = unique_states else {
1173 return false;
1174 };
1175 let idx = state.bit_position() as usize;
1176 bitmask.as_bytes().get(idx).is_some_and(|bit| *bit == b'1')
1177}
1178
1179fn write_copy_field(buf: &mut Vec<u8>, value: &str) {
1180 if value.contains(',')
1181 || value.contains('"')
1182 || value.contains('\n')
1183 || value.contains('\r')
1184 || value.contains('\\')
1185 || value == COPY_NULL_SENTINEL
1186 {
1187 buf.push(b'"');
1188 for byte in value.bytes() {
1189 if byte == b'"' {
1190 buf.push(b'"');
1191 }
1192 buf.push(byte);
1193 }
1194 buf.push(b'"');
1195 } else {
1196 buf.extend_from_slice(value.as_bytes());
1197 }
1198}
1199
1200fn write_copy_json(buf: &mut Vec<u8>, value: &serde_json::Value) {
1201 let json = serde_json::to_string(value).expect("JSON serialization should not fail");
1202 write_copy_field(buf, &json);
1203}
1204
1205fn storage_payload(value: &serde_json::Value) -> Option<&serde_json::Value> {
1206 (!is_storage_payload_empty(value)).then_some(value)
1207}
1208
1209fn terminal_storage_payload<'a>(
1210 value: &'a serde_json::Value,
1211 ready_payload: Option<&serde_json::Value>,
1212) -> Option<&'a serde_json::Value> {
1213 if is_storage_payload_empty(value) || ready_payload.is_some_and(|ready| ready == value) {
1214 None
1215 } else {
1216 Some(value)
1217 }
1218}
1219
1220fn is_storage_payload_empty(value: &serde_json::Value) -> bool {
1221 value.is_null() || is_empty_json_object(value)
1222}
1223
1224fn write_copy_storage_payload(buf: &mut Vec<u8>, value: &serde_json::Value) {
1225 match storage_payload(value) {
1226 Some(value) => write_copy_json(buf, value),
1227 None => buf.extend_from_slice(COPY_NULL_SENTINEL.as_bytes()),
1228 }
1229}
1230
1231fn write_copy_datetime(buf: &mut Vec<u8>, value: DateTime<Utc>) {
1232 write_copy_field(buf, &value.to_rfc3339());
1233}
1234
1235fn write_copy_optional_datetime(buf: &mut Vec<u8>, value: Option<DateTime<Utc>>) {
1236 match value {
1237 Some(value) => write_copy_datetime(buf, value),
1238 None => buf.extend_from_slice(COPY_NULL_SENTINEL.as_bytes()),
1239 }
1240}
1241
1242fn write_copy_optional_bytes(buf: &mut Vec<u8>, value: &Option<Vec<u8>>) {
1243 match value {
1244 Some(bytes) => {
1245 let bytea_hex = format!("\\x{}", hex::encode(bytes));
1246 write_copy_field(buf, &bytea_hex);
1247 }
1248 None => buf.extend_from_slice(COPY_NULL_SENTINEL.as_bytes()),
1249 }
1250}
1251
1252fn write_copy_optional_string(buf: &mut Vec<u8>, value: Option<&str>) {
1253 match value {
1254 Some(value) => write_copy_field(buf, value),
1255 None => buf.extend_from_slice(COPY_NULL_SENTINEL.as_bytes()),
1256 }
1257}
1258
1259fn write_ready_copy_row(
1260 buf: &mut Vec<u8>,
1261 ready_slot: i32,
1262 ready_generation: i64,
1263 row: &RuntimeReadyInsert,
1264) {
1265 buf.extend_from_slice(ready_slot.to_string().as_bytes());
1266 buf.push(b',');
1267 buf.extend_from_slice(ready_generation.to_string().as_bytes());
1268 buf.push(b',');
1269 buf.extend_from_slice(row.job_id.to_string().as_bytes());
1270 buf.push(b',');
1271 write_copy_field(buf, &row.kind);
1272 buf.push(b',');
1273 write_copy_field(buf, &row.queue);
1274 buf.push(b',');
1275 write_copy_json(buf, &row.args);
1276 buf.push(b',');
1277 buf.extend_from_slice(row.priority.to_string().as_bytes());
1278 buf.push(b',');
1279 buf.extend_from_slice(row.attempt.to_string().as_bytes());
1280 buf.push(b',');
1281 buf.extend_from_slice(row.run_lease.to_string().as_bytes());
1282 buf.push(b',');
1283 buf.extend_from_slice(row.max_attempts.to_string().as_bytes());
1284 buf.push(b',');
1285 buf.extend_from_slice(row.lane_seq.to_string().as_bytes());
1286 buf.push(b',');
1287 buf.extend_from_slice(row.enqueue_shard.to_string().as_bytes());
1288 buf.push(b',');
1289 write_copy_datetime(buf, row.run_at);
1290 buf.push(b',');
1291 write_copy_optional_datetime(buf, row.attempted_at);
1292 buf.push(b',');
1293 write_copy_datetime(buf, row.created_at);
1294 buf.push(b',');
1295 write_copy_optional_bytes(buf, &row.unique_key);
1296 buf.push(b',');
1297 write_copy_optional_string(buf, row.unique_states.as_deref());
1298 buf.push(b',');
1299 write_copy_storage_payload(buf, &row.payload);
1300 buf.push(b'\n');
1301}
1302
1303fn write_deferred_copy_row(buf: &mut Vec<u8>, row: &DeferredJobRow) {
1304 buf.extend_from_slice(row.job_id.to_string().as_bytes());
1305 buf.push(b',');
1306 write_copy_field(buf, &row.kind);
1307 buf.push(b',');
1308 write_copy_field(buf, &row.queue);
1309 buf.push(b',');
1310 write_copy_json(buf, &row.args);
1311 buf.push(b',');
1312 write_copy_field(buf, &row.state.to_string());
1313 buf.push(b',');
1314 buf.extend_from_slice(row.priority.to_string().as_bytes());
1315 buf.push(b',');
1316 buf.extend_from_slice(row.attempt.to_string().as_bytes());
1317 buf.push(b',');
1318 buf.extend_from_slice(row.run_lease.to_string().as_bytes());
1319 buf.push(b',');
1320 buf.extend_from_slice(row.max_attempts.to_string().as_bytes());
1321 buf.push(b',');
1322 write_copy_datetime(buf, row.run_at);
1323 buf.push(b',');
1324 write_copy_optional_datetime(buf, row.attempted_at);
1325 buf.push(b',');
1326 write_copy_optional_datetime(buf, row.finalized_at);
1327 buf.push(b',');
1328 write_copy_datetime(buf, row.created_at);
1329 buf.push(b',');
1330 write_copy_optional_bytes(buf, &row.unique_key);
1331 buf.push(b',');
1332 write_copy_optional_string(buf, row.unique_states.as_deref());
1333 buf.push(b',');
1334 write_copy_storage_payload(buf, &row.payload);
1335 buf.push(b'\n');
1336}
1337
1338fn lifecycle_error(error: impl Into<String>, attempt: i16, terminal: bool) -> serde_json::Value {
1339 let mut value = serde_json::json!({
1340 "error": error.into(),
1341 "attempt": attempt,
1342 "at": Utc::now().to_rfc3339(),
1343 });
1344 if terminal {
1345 value["terminal"] = serde_json::Value::Bool(true);
1346 }
1347 value
1348}
1349
1350fn transition_timestamp(job: &JobRow) -> DateTime<Utc> {
1351 job.finalized_at
1352 .or(job.heartbeat_at)
1353 .or(job.deadline_at)
1354 .or(job.attempted_at)
1355 .unwrap_or(job.run_at)
1356}
1357
1358fn state_rank(state: JobState) -> u8 {
1359 match state {
1360 JobState::Running | JobState::WaitingExternal => 4,
1361 JobState::Retryable | JobState::Scheduled => 3,
1362 JobState::Available => 2,
1363 JobState::Completed | JobState::Failed | JobState::Cancelled => 1,
1364 }
1365}
1366
1367#[derive(Debug, Clone, sqlx::FromRow)]
1368struct ReadyJobRow {
1369 job_id: i64,
1370 kind: String,
1371 queue: String,
1372 args: serde_json::Value,
1373 priority: i16,
1374 attempt: i16,
1375 run_lease: i64,
1376 max_attempts: i16,
1377 run_at: DateTime<Utc>,
1378 attempted_at: Option<DateTime<Utc>>,
1379 created_at: DateTime<Utc>,
1380 unique_key: Option<Vec<u8>>,
1381 payload: serde_json::Value,
1382}
1383
1384impl ReadyJobRow {
1385 fn into_job_row(self) -> Result<JobRow, AwaError> {
1386 let payload = RuntimePayload::from_json(self.payload)?;
1387 Ok(JobRow {
1388 id: self.job_id,
1389 kind: self.kind,
1390 queue: self.queue,
1391 args: self.args,
1392 state: JobState::Available,
1393 priority: self.priority,
1394 attempt: self.attempt,
1395 run_lease: self.run_lease,
1396 max_attempts: self.max_attempts,
1397 run_at: self.run_at,
1398 heartbeat_at: None,
1399 deadline_at: None,
1400 attempted_at: self.attempted_at,
1401 finalized_at: None,
1402 created_at: self.created_at,
1403 errors: payload.errors_option(),
1404 metadata: payload.metadata,
1405 tags: payload.tags,
1406 unique_key: self.unique_key,
1407 unique_states: None,
1408 callback_id: None,
1409 callback_timeout_at: None,
1410 callback_filter: None,
1411 callback_on_complete: None,
1412 callback_on_fail: None,
1413 callback_transform: None,
1414 progress: payload.progress,
1415 })
1416 }
1417}
1418
1419#[derive(Debug, Clone, sqlx::FromRow)]
1420struct ReadyTransitionRow {
1421 ready_slot: i32,
1422 ready_generation: i64,
1423 job_id: i64,
1424 kind: String,
1425 queue: String,
1426 args: serde_json::Value,
1427 priority: i16,
1428 attempt: i16,
1429 run_lease: i64,
1430 max_attempts: i16,
1431 lane_seq: i64,
1432 enqueue_shard: i16,
1433 run_at: DateTime<Utc>,
1434 attempted_at: Option<DateTime<Utc>>,
1435 created_at: DateTime<Utc>,
1436 unique_key: Option<Vec<u8>>,
1437 unique_states: Option<String>,
1438 payload: serde_json::Value,
1439}
1440
1441#[derive(Debug, Clone)]
1442struct ClaimCursorAdvance {
1443 queue: String,
1444 priority: i16,
1445 enqueue_shard: i16,
1446 next_seq: i64,
1447 only_if_current: Option<i64>,
1448}
1449
1450type ClaimCursorLaneKey = (String, i16, i16);
1451type ConditionalClaimCursorAdvances = BTreeMap<i64, i64>;
1452type GroupedClaimCursorAdvances =
1453 BTreeMap<ClaimCursorLaneKey, (Option<i64>, ConditionalClaimCursorAdvances)>;
1454type TerminalCounterKey = (i32, i64, String, i16, i16, i16);
1455
1456struct CancelJobTxResult {
1457 row: JobRow,
1458 claim_cursor_advance: Option<ClaimCursorAdvance>,
1459}
1460
1461struct ReadyBatchMoveResult {
1462 moved: bool,
1463}
1464
1465impl ReadyTransitionRow {
1466 fn into_existing_ready_row(
1467 self,
1468 queue: String,
1469 priority: i16,
1470 payload: serde_json::Value,
1471 ) -> ExistingReadyRow {
1472 ExistingReadyRow {
1473 job_id: self.job_id,
1474 kind: self.kind,
1475 queue,
1476 args: self.args,
1477 priority,
1478 attempt: self.attempt,
1479 run_lease: self.run_lease,
1480 max_attempts: self.max_attempts,
1481 run_at: self.run_at,
1482 attempted_at: self.attempted_at,
1483 created_at: self.created_at,
1484 unique_key: self.unique_key,
1485 unique_states: self.unique_states,
1486 payload,
1487 }
1488 }
1489
1490 fn into_done_row(
1491 self,
1492 state: JobState,
1493 finalized_at: DateTime<Utc>,
1494 payload: serde_json::Value,
1495 ) -> DoneJobRow {
1496 DoneJobRow {
1497 ready_slot: self.ready_slot,
1498 ready_generation: self.ready_generation,
1499 job_id: self.job_id,
1500 kind: self.kind,
1501 queue: self.queue,
1502 args: self.args,
1503 state,
1504 priority: self.priority,
1505 attempt: self.attempt,
1506 run_lease: self.run_lease,
1507 max_attempts: self.max_attempts,
1508 lane_seq: self.lane_seq,
1509 enqueue_shard: self.enqueue_shard,
1510 run_at: self.run_at,
1511 attempted_at: self.attempted_at,
1512 finalized_at,
1513 created_at: self.created_at,
1514 unique_key: self.unique_key,
1515 unique_states: self.unique_states,
1516 payload,
1517 }
1518 }
1519}
1520
1521#[derive(Debug, Clone, sqlx::FromRow)]
1522struct ReadyJobLeaseRow {
1523 ready_slot: i32,
1524 ready_generation: i64,
1525 lane_seq: i64,
1526 enqueue_shard: i16,
1527 lease_slot: i32,
1528 lease_generation: i64,
1529 claim_slot: i32,
1530 receipt_id: Option<i64>,
1531 claim_batch_id: Option<i64>,
1532 claim_batch_index: Option<i32>,
1533 job_id: i64,
1534 kind: String,
1535 queue: String,
1536 args: serde_json::Value,
1537 lane_priority: i16,
1538 priority: i16,
1539 attempt: i16,
1540 run_lease: i64,
1541 max_attempts: i16,
1542 run_at: DateTime<Utc>,
1543 heartbeat_at: Option<DateTime<Utc>>,
1544 deadline_at: Option<DateTime<Utc>>,
1545 attempted_at: Option<DateTime<Utc>>,
1546 created_at: DateTime<Utc>,
1547 unique_key: Option<Vec<u8>>,
1548 unique_states: Option<String>,
1549 payload: serde_json::Value,
1550}
1551
1552impl ReadyJobLeaseRow {
1553 fn claim_ref(&self, lease_claim_receipt: bool) -> ClaimedEntry {
1554 ClaimedEntry {
1555 queue: self.queue.clone(),
1556 priority: self.lane_priority,
1557 lane_seq: self.lane_seq,
1558 ready_slot: self.ready_slot,
1559 ready_generation: self.ready_generation,
1560 lease_slot: self.lease_slot,
1561 lease_generation: self.lease_generation,
1562 claim_slot: self.claim_slot,
1563 receipt_id: self.receipt_id,
1564 claim_batch_id: self.claim_batch_id,
1565 claim_batch_index: self.claim_batch_index,
1566 lease_claim_receipt,
1567 enqueue_shard: self.enqueue_shard,
1568 }
1569 }
1570
1571 fn into_job_row(self) -> Result<JobRow, AwaError> {
1572 let mut payload = RuntimePayload::from_json(self.payload)?;
1573 if self.priority < self.lane_priority {
1574 let metadata = payload.metadata.as_object_mut().ok_or_else(|| {
1575 AwaError::Validation(
1576 "queue storage payload metadata must be a JSON object".to_string(),
1577 )
1578 })?;
1579 metadata
1580 .entry("_awa_original_priority".to_string())
1581 .or_insert_with(|| serde_json::Value::from(i64::from(self.lane_priority)));
1582 }
1583
1584 Ok(JobRow {
1585 id: self.job_id,
1586 kind: self.kind,
1587 queue: self.queue,
1588 args: self.args,
1589 state: JobState::Running,
1590 priority: self.priority,
1591 attempt: self.attempt,
1592 run_lease: self.run_lease,
1593 max_attempts: self.max_attempts,
1594 run_at: self.run_at,
1595 heartbeat_at: self.heartbeat_at,
1596 deadline_at: self.deadline_at,
1597 attempted_at: self.attempted_at,
1598 finalized_at: None,
1599 created_at: self.created_at,
1600 errors: payload.errors_option(),
1601 metadata: payload.metadata,
1602 tags: payload.tags,
1603 unique_key: self.unique_key,
1604 unique_states: None,
1605 callback_id: None,
1606 callback_timeout_at: None,
1607 callback_filter: None,
1608 callback_on_complete: None,
1609 callback_on_fail: None,
1610 callback_transform: None,
1611 progress: payload.progress,
1612 })
1613 }
1614
1615 fn into_claimed_runtime_job(
1616 self,
1617 lease_claim_receipt: bool,
1618 ) -> Result<ClaimedRuntimeJob, AwaError> {
1619 let claim = self.claim_ref(lease_claim_receipt);
1620 let unique_states = self.unique_states.clone();
1621 let job = self.into_job_row()?;
1622 Ok(ClaimedRuntimeJob {
1623 claim,
1624 job,
1625 unique_states,
1626 })
1627 }
1628}
1629
1630#[derive(Debug, Clone)]
1631struct RuntimeReadyRow {
1632 kind: String,
1633 queue: String,
1634 args: serde_json::Value,
1635 priority: i16,
1636 attempt: i16,
1637 run_lease: i64,
1638 max_attempts: i16,
1639 run_at: DateTime<Utc>,
1640 attempted_at: Option<DateTime<Utc>>,
1641 created_at: DateTime<Utc>,
1642 unique_key: Option<Vec<u8>>,
1643 unique_states: Option<String>,
1644 payload: serde_json::Value,
1645 ordering_key: Option<Vec<u8>>,
1650}
1651
1652#[derive(Debug, Clone)]
1653struct RuntimeReadyInsert {
1654 job_id: i64,
1655 kind: String,
1656 queue: String,
1657 args: serde_json::Value,
1658 priority: i16,
1659 attempt: i16,
1660 run_lease: i64,
1661 max_attempts: i16,
1662 run_at: DateTime<Utc>,
1663 attempted_at: Option<DateTime<Utc>>,
1664 lane_seq: i64,
1665 enqueue_shard: i16,
1669 created_at: DateTime<Utc>,
1670 unique_key: Option<Vec<u8>>,
1671 unique_states: Option<String>,
1672 payload: serde_json::Value,
1673}
1674
1675#[derive(Debug)]
1676struct ReadySegmentInsert {
1677 queue: String,
1678 priority: i16,
1679 enqueue_shard: i16,
1680 first_lane_seq: i64,
1681 next_lane_seq: i64,
1682 first_run_at: DateTime<Utc>,
1686}
1687
1688#[cfg(test)]
1689mod ready_segment_tests {
1690 use super::{QueueStorage, RuntimeReadyInsert};
1691 use chrono::{Duration, TimeZone, Utc};
1692
1693 fn ready_row(lane_seq: i64, run_at: chrono::DateTime<Utc>) -> RuntimeReadyInsert {
1694 RuntimeReadyInsert {
1695 job_id: lane_seq,
1696 kind: "segment_test".to_string(),
1697 queue: "segment-q".to_string(),
1698 args: serde_json::json!({}),
1699 priority: 2,
1700 attempt: 0,
1701 run_lease: 0,
1702 max_attempts: 25,
1703 run_at,
1704 attempted_at: None,
1705 lane_seq,
1706 enqueue_shard: 0,
1707 created_at: run_at,
1708 unique_key: None,
1709 unique_states: None,
1710 payload: serde_json::json!({}),
1711 }
1712 }
1713
1714 #[test]
1715 fn ready_segments_split_on_run_at_boundaries() {
1716 let first = Utc
1717 .with_ymd_and_hms(2026, 6, 14, 12, 0, 0)
1718 .single()
1719 .expect("valid test timestamp");
1720 let second = first + Duration::seconds(1);
1721 let rows = vec![
1722 ready_row(1, first),
1723 ready_row(2, first),
1724 ready_row(3, second),
1725 ready_row(4, first),
1726 ];
1727
1728 let segments = QueueStorage::ready_segments_from_rows(&rows);
1729 let ranges: Vec<_> = segments
1730 .iter()
1731 .map(|segment| {
1732 (
1733 segment.first_lane_seq,
1734 segment.next_lane_seq,
1735 segment.first_run_at,
1736 )
1737 })
1738 .collect();
1739
1740 assert_eq!(ranges, vec![(1, 3, first), (3, 4, second), (4, 5, first)]);
1741 }
1742}
1743
1744#[derive(Debug, Clone, sqlx::FromRow)]
1745struct DoneJobRow {
1746 ready_slot: i32,
1747 ready_generation: i64,
1748 job_id: i64,
1749 kind: String,
1750 queue: String,
1751 args: serde_json::Value,
1752 state: JobState,
1753 priority: i16,
1754 attempt: i16,
1755 run_lease: i64,
1756 max_attempts: i16,
1757 lane_seq: i64,
1758 enqueue_shard: i16,
1763 run_at: DateTime<Utc>,
1764 attempted_at: Option<DateTime<Utc>>,
1765 finalized_at: DateTime<Utc>,
1766 created_at: DateTime<Utc>,
1767 unique_key: Option<Vec<u8>>,
1768 unique_states: Option<String>,
1769 payload: serde_json::Value,
1770}
1771
1772impl DoneJobRow {
1773 fn into_job_row(self) -> Result<JobRow, AwaError> {
1774 let payload = RuntimePayload::from_json(self.payload)?;
1775 Ok(JobRow {
1776 id: self.job_id,
1777 kind: self.kind,
1778 queue: self.queue,
1779 args: self.args,
1780 state: self.state,
1781 priority: self.priority,
1782 attempt: self.attempt,
1783 run_lease: self.run_lease,
1784 max_attempts: self.max_attempts,
1785 run_at: self.run_at,
1786 heartbeat_at: None,
1787 deadline_at: None,
1788 attempted_at: self.attempted_at,
1789 finalized_at: Some(self.finalized_at),
1790 created_at: self.created_at,
1791 errors: payload.errors_option(),
1792 metadata: payload.metadata,
1793 tags: payload.tags,
1794 unique_key: self.unique_key,
1795 unique_states: None,
1796 callback_id: None,
1797 callback_timeout_at: None,
1798 callback_filter: None,
1799 callback_on_complete: None,
1800 callback_on_fail: None,
1801 callback_transform: None,
1802 progress: payload.progress,
1803 })
1804 }
1805
1806 fn into_dlq_row(self, dlq_reason: String, dlq_at: DateTime<Utc>) -> DlqJobRow {
1807 DlqJobRow {
1808 job_id: self.job_id,
1809 kind: self.kind,
1810 queue: self.queue,
1811 args: self.args,
1812 state: self.state,
1813 priority: self.priority,
1814 attempt: self.attempt,
1815 run_lease: self.run_lease,
1816 max_attempts: self.max_attempts,
1817 run_at: self.run_at,
1818 attempted_at: self.attempted_at,
1819 finalized_at: self.finalized_at,
1820 created_at: self.created_at,
1821 unique_key: self.unique_key,
1822 unique_states: self.unique_states,
1823 payload: self.payload,
1824 dlq_reason,
1825 dlq_at,
1826 original_run_lease: self.run_lease,
1827 }
1828 }
1829}
1830
1831#[derive(Debug, Clone, sqlx::FromRow)]
1832struct DlqJobRow {
1833 job_id: i64,
1834 kind: String,
1835 queue: String,
1836 args: serde_json::Value,
1837 state: JobState,
1838 priority: i16,
1839 attempt: i16,
1840 run_lease: i64,
1841 max_attempts: i16,
1842 run_at: DateTime<Utc>,
1843 attempted_at: Option<DateTime<Utc>>,
1844 finalized_at: DateTime<Utc>,
1845 created_at: DateTime<Utc>,
1846 unique_key: Option<Vec<u8>>,
1847 unique_states: Option<String>,
1848 payload: serde_json::Value,
1849 dlq_reason: String,
1850 dlq_at: DateTime<Utc>,
1851 original_run_lease: i64,
1852}
1853
1854impl DlqJobRow {
1855 fn into_job_row(self) -> Result<JobRow, AwaError> {
1856 let payload = RuntimePayload::from_json(self.payload)?;
1857 Ok(JobRow {
1858 id: self.job_id,
1859 kind: self.kind,
1860 queue: self.queue,
1861 args: self.args,
1862 state: self.state,
1863 priority: self.priority,
1864 attempt: self.attempt,
1865 run_lease: self.run_lease,
1866 max_attempts: self.max_attempts,
1867 run_at: self.run_at,
1868 heartbeat_at: None,
1869 deadline_at: None,
1870 attempted_at: self.attempted_at,
1871 finalized_at: Some(self.finalized_at),
1872 created_at: self.created_at,
1873 errors: payload.errors_option(),
1874 metadata: payload.metadata,
1875 tags: payload.tags,
1876 unique_key: self.unique_key,
1877 unique_states: None,
1878 callback_id: None,
1879 callback_timeout_at: None,
1880 callback_filter: None,
1881 callback_on_complete: None,
1882 callback_on_fail: None,
1883 callback_transform: None,
1884 progress: payload.progress,
1885 })
1886 }
1887
1888 fn into_retry_ready_row(
1889 self,
1890 queue: String,
1891 priority: i16,
1892 run_at: DateTime<Utc>,
1893 payload: serde_json::Value,
1894 ) -> ExistingReadyRow {
1895 ExistingReadyRow {
1896 job_id: self.job_id,
1897 kind: self.kind,
1898 queue,
1899 args: self.args,
1900 priority,
1901 attempt: 0,
1902 run_lease: self.run_lease,
1903 max_attempts: self.max_attempts,
1904 run_at,
1905 attempted_at: None,
1906 created_at: self.created_at,
1907 unique_key: self.unique_key,
1908 unique_states: self.unique_states,
1909 payload,
1910 }
1911 }
1912
1913 fn into_retry_deferred_row(
1914 self,
1915 queue: String,
1916 priority: i16,
1917 run_at: DateTime<Utc>,
1918 payload: serde_json::Value,
1919 ) -> DeferredJobRow {
1920 DeferredJobRow {
1921 job_id: self.job_id,
1922 kind: self.kind,
1923 queue,
1924 args: self.args,
1925 state: JobState::Scheduled,
1926 priority,
1927 attempt: 0,
1928 run_lease: self.run_lease,
1929 max_attempts: self.max_attempts,
1930 run_at,
1931 attempted_at: None,
1932 finalized_at: None,
1933 created_at: self.created_at,
1934 unique_key: self.unique_key,
1935 unique_states: self.unique_states,
1936 payload,
1937 }
1938 }
1939}
1940
1941#[derive(Debug, Clone)]
1942struct ExistingReadyRow {
1943 job_id: i64,
1944 kind: String,
1945 queue: String,
1946 args: serde_json::Value,
1947 priority: i16,
1948 attempt: i16,
1949 run_lease: i64,
1950 max_attempts: i16,
1951 run_at: DateTime<Utc>,
1952 attempted_at: Option<DateTime<Utc>>,
1953 created_at: DateTime<Utc>,
1954 unique_key: Option<Vec<u8>>,
1955 unique_states: Option<String>,
1956 payload: serde_json::Value,
1957}
1958
1959#[derive(Debug, Clone, sqlx::FromRow)]
1960struct DeletedLeaseRow {
1961 ready_slot: i32,
1962 ready_generation: i64,
1963 job_id: i64,
1964 queue: String,
1965 state: JobState,
1966 priority: i16,
1967 attempt: i16,
1968 run_lease: i64,
1969 max_attempts: i16,
1970 lane_seq: i64,
1971 enqueue_shard: i16,
1972 attempted_at: Option<DateTime<Utc>>,
1973}
1974
1975#[derive(Debug, Clone, sqlx::FromRow)]
1976struct ReadySnapshotRow {
1977 ready_slot: i32,
1978 ready_generation: i64,
1979 job_id: i64,
1980 kind: String,
1981 queue: String,
1982 args: serde_json::Value,
1983 lane_seq: i64,
1984 enqueue_shard: i16,
1985 run_at: DateTime<Utc>,
1986 created_at: DateTime<Utc>,
1987 unique_key: Option<Vec<u8>>,
1988 unique_states: Option<String>,
1989 payload: serde_json::Value,
1990}
1991
1992#[derive(Debug, Clone, sqlx::FromRow)]
1993struct AttemptStateRow {
1994 job_id: i64,
1995 run_lease: i64,
1996 progress: Option<serde_json::Value>,
1997 callback_filter: Option<String>,
1998 callback_on_complete: Option<String>,
1999 callback_on_fail: Option<String>,
2000 callback_transform: Option<String>,
2001 callback_result: Option<serde_json::Value>,
2002}
2003
2004#[derive(Debug, Clone)]
2005struct LeaseTransitionRow {
2006 ready_slot: i32,
2007 ready_generation: i64,
2008 job_id: i64,
2009 kind: String,
2010 queue: String,
2011 args: serde_json::Value,
2012 state: JobState,
2013 priority: i16,
2014 attempt: i16,
2015 run_lease: i64,
2016 max_attempts: i16,
2017 lane_seq: i64,
2018 enqueue_shard: i16,
2019 run_at: DateTime<Utc>,
2020 attempted_at: Option<DateTime<Utc>>,
2021 created_at: DateTime<Utc>,
2022 unique_key: Option<Vec<u8>>,
2023 unique_states: Option<String>,
2024 payload: serde_json::Value,
2025 progress: Option<serde_json::Value>,
2026}
2027
2028impl LeaseTransitionRow {
2029 fn into_done_row(
2030 self,
2031 state: JobState,
2032 finalized_at: DateTime<Utc>,
2033 payload: serde_json::Value,
2034 ) -> DoneJobRow {
2035 DoneJobRow {
2036 ready_slot: self.ready_slot,
2037 ready_generation: self.ready_generation,
2038 job_id: self.job_id,
2039 kind: self.kind,
2040 queue: self.queue,
2041 args: self.args,
2042 state,
2043 priority: self.priority,
2044 attempt: self.attempt,
2045 run_lease: self.run_lease,
2046 max_attempts: self.max_attempts,
2047 lane_seq: self.lane_seq,
2048 enqueue_shard: self.enqueue_shard,
2049 run_at: self.run_at,
2050 attempted_at: self.attempted_at,
2051 finalized_at,
2052 created_at: self.created_at,
2053 unique_key: self.unique_key,
2054 unique_states: self.unique_states,
2055 payload,
2056 }
2057 }
2058
2059 fn into_deferred_row(
2060 self,
2061 state: JobState,
2062 run_at: DateTime<Utc>,
2063 finalized_at: Option<DateTime<Utc>>,
2064 payload: serde_json::Value,
2065 ) -> DeferredJobRow {
2066 DeferredJobRow {
2067 job_id: self.job_id,
2068 kind: self.kind,
2069 queue: self.queue,
2070 args: self.args,
2071 state,
2072 priority: self.priority,
2073 attempt: self.attempt,
2074 run_lease: self.run_lease,
2075 max_attempts: self.max_attempts,
2076 run_at,
2077 attempted_at: self.attempted_at,
2078 finalized_at,
2079 created_at: self.created_at,
2080 unique_key: self.unique_key,
2081 unique_states: self.unique_states,
2082 payload,
2083 }
2084 }
2085
2086 fn into_ready_row(self, run_at: DateTime<Utc>, payload: serde_json::Value) -> ExistingReadyRow {
2087 ExistingReadyRow {
2088 job_id: self.job_id,
2089 kind: self.kind,
2090 queue: self.queue,
2091 args: self.args,
2092 priority: self.priority,
2093 attempt: self.attempt,
2094 run_lease: self.run_lease,
2095 max_attempts: self.max_attempts,
2096 run_at,
2097 attempted_at: self.attempted_at,
2098 created_at: self.created_at,
2099 unique_key: self.unique_key,
2100 unique_states: self.unique_states,
2101 payload,
2102 }
2103 }
2104
2105 fn into_dlq_row(
2106 self,
2107 finalized_at: DateTime<Utc>,
2108 payload: serde_json::Value,
2109 dlq_reason: String,
2110 dlq_at: DateTime<Utc>,
2111 ) -> DlqJobRow {
2112 DlqJobRow {
2113 job_id: self.job_id,
2114 kind: self.kind,
2115 queue: self.queue,
2116 args: self.args,
2117 state: JobState::Failed,
2118 priority: self.priority,
2119 attempt: self.attempt,
2120 run_lease: self.run_lease,
2121 max_attempts: self.max_attempts,
2122 run_at: self.run_at,
2123 attempted_at: self.attempted_at,
2124 finalized_at,
2125 created_at: self.created_at,
2126 unique_key: self.unique_key,
2127 unique_states: self.unique_states,
2128 payload,
2129 dlq_reason,
2130 dlq_at,
2131 original_run_lease: self.run_lease,
2132 }
2133 }
2134}
2135
2136#[derive(Debug, Clone, sqlx::FromRow)]
2137struct LeaseJobRow {
2138 job_id: i64,
2139 kind: String,
2140 queue: String,
2141 args: serde_json::Value,
2142 state: JobState,
2143 priority: i16,
2144 attempt: i16,
2145 run_lease: i64,
2146 max_attempts: i16,
2147 run_at: DateTime<Utc>,
2148 heartbeat_at: Option<DateTime<Utc>>,
2149 deadline_at: Option<DateTime<Utc>>,
2150 attempted_at: Option<DateTime<Utc>>,
2151 finalized_at: Option<DateTime<Utc>>,
2152 created_at: DateTime<Utc>,
2153 unique_key: Option<Vec<u8>>,
2154 callback_id: Option<Uuid>,
2155 callback_timeout_at: Option<DateTime<Utc>>,
2156 callback_filter: Option<String>,
2157 callback_on_complete: Option<String>,
2158 callback_on_fail: Option<String>,
2159 callback_transform: Option<String>,
2160 payload: serde_json::Value,
2161 progress: Option<serde_json::Value>,
2162 callback_result: Option<serde_json::Value>,
2163}
2164
2165impl LeaseJobRow {
2166 fn into_job_row(self) -> Result<JobRow, AwaError> {
2167 let payload = QueueStorage::materialize_runtime_payload(
2168 self.payload,
2169 self.progress,
2170 self.callback_result,
2171 )?;
2172 Ok(JobRow {
2173 id: self.job_id,
2174 kind: self.kind,
2175 queue: self.queue,
2176 args: self.args,
2177 state: self.state,
2178 priority: self.priority,
2179 attempt: self.attempt,
2180 run_lease: self.run_lease,
2181 max_attempts: self.max_attempts,
2182 run_at: self.run_at,
2183 heartbeat_at: self.heartbeat_at,
2184 deadline_at: self.deadline_at,
2185 attempted_at: self.attempted_at,
2186 finalized_at: self.finalized_at,
2187 created_at: self.created_at,
2188 errors: payload.errors_option(),
2189 metadata: payload.metadata,
2190 tags: payload.tags,
2191 unique_key: self.unique_key,
2192 unique_states: None,
2193 callback_id: self.callback_id,
2194 callback_timeout_at: self.callback_timeout_at,
2195 callback_filter: self.callback_filter,
2196 callback_on_complete: self.callback_on_complete,
2197 callback_on_fail: self.callback_on_fail,
2198 callback_transform: self.callback_transform,
2199 progress: payload.progress,
2200 })
2201 }
2202}
2203
2204#[derive(Debug, Clone, sqlx::FromRow)]
2205struct DeferredJobRow {
2206 job_id: i64,
2207 kind: String,
2208 queue: String,
2209 args: serde_json::Value,
2210 state: JobState,
2211 priority: i16,
2212 attempt: i16,
2213 run_lease: i64,
2214 max_attempts: i16,
2215 run_at: DateTime<Utc>,
2216 attempted_at: Option<DateTime<Utc>>,
2217 finalized_at: Option<DateTime<Utc>>,
2218 created_at: DateTime<Utc>,
2219 unique_key: Option<Vec<u8>>,
2220 unique_states: Option<String>,
2221 payload: serde_json::Value,
2222}
2223
2224impl DeferredJobRow {
2225 fn into_job_row(self) -> Result<JobRow, AwaError> {
2226 let payload = RuntimePayload::from_json(self.payload)?;
2227 Ok(JobRow {
2228 id: self.job_id,
2229 kind: self.kind,
2230 queue: self.queue,
2231 args: self.args,
2232 state: self.state,
2233 priority: self.priority,
2234 attempt: self.attempt,
2235 run_lease: self.run_lease,
2236 max_attempts: self.max_attempts,
2237 run_at: self.run_at,
2238 heartbeat_at: None,
2239 deadline_at: None,
2240 attempted_at: self.attempted_at,
2241 finalized_at: self.finalized_at,
2242 created_at: self.created_at,
2243 errors: payload.errors_option(),
2244 metadata: payload.metadata,
2245 tags: payload.tags,
2246 unique_key: self.unique_key,
2247 unique_states: None,
2248 callback_id: None,
2249 callback_timeout_at: None,
2250 callback_filter: None,
2251 callback_on_complete: None,
2252 callback_on_fail: None,
2253 callback_transform: None,
2254 progress: payload.progress,
2255 })
2256 }
2257}
2258
2259#[derive(Debug)]
2270pub struct QueueStorage {
2271 config: QueueStorageConfig,
2272 next_stripe_probe: AtomicUsize,
2273 shard_rotor: AtomicU16,
2277 enqueue_shards_cache: Mutex<HashMap<String, i16>>,
2282 ensured_lanes: Mutex<HashSet<(String, i16, i16)>>,
2290 prune_lock_timeout: Duration,
2291}
2292
2293impl QueueStorage {
2294 pub fn new(config: QueueStorageConfig) -> Result<Self, AwaError> {
2295 if config.queue_slot_count < 4 {
2296 return Err(AwaError::Validation(
2297 "queue storage requires at least 4 queue slots".into(),
2298 ));
2299 }
2300 if config.lease_slot_count < 2 {
2301 return Err(AwaError::Validation(
2302 "queue storage requires at least 2 lease slots".into(),
2303 ));
2304 }
2305 if config.claim_slot_count < 2 {
2306 return Err(AwaError::Validation(
2307 "queue storage requires at least 2 claim slots".into(),
2308 ));
2309 }
2310 if config.queue_stripe_count == 0 {
2311 return Err(AwaError::Validation(
2312 "queue storage requires at least 1 queue stripe".into(),
2313 ));
2314 }
2315 if config.schema == DEFAULT_SCHEMA
2316 && (config.queue_slot_count != DEFAULT_QUEUE_SLOT_COUNT
2317 || config.lease_slot_count != DEFAULT_LEASE_SLOT_COUNT
2318 || config.claim_slot_count != DEFAULT_CLAIM_SLOT_COUNT
2319 || !config.lease_claim_receipts)
2320 {
2321 return Err(AwaError::Validation(
2322 "the default `awa` queue-storage schema must use the default slot counts and \
2323 lease_claim_receipts=true"
2324 .into(),
2325 ));
2326 }
2327 validate_ident(&config.schema)?;
2328 Ok(Self {
2329 config,
2330 next_stripe_probe: AtomicUsize::new(0),
2331 shard_rotor: AtomicU16::new(0),
2332 enqueue_shards_cache: Mutex::new(HashMap::new()),
2333 ensured_lanes: Mutex::new(HashSet::new()),
2334 prune_lock_timeout: DEFAULT_PRUNE_LOCK_TIMEOUT,
2335 })
2336 }
2337
2338 #[doc(hidden)]
2339 pub fn with_prune_lock_timeout(mut self, timeout: Duration) -> Result<Self, AwaError> {
2340 if timeout.is_zero() {
2341 return Err(AwaError::Validation(
2342 "queue storage prune lock timeout must be greater than zero".into(),
2343 ));
2344 }
2345 self.prune_lock_timeout = timeout;
2346 Ok(self)
2347 }
2348
2349 pub fn from_existing_schema(schema: impl Into<String>) -> Result<Self, AwaError> {
2350 Self::new(QueueStorageConfig {
2351 schema: schema.into(),
2352 ..Default::default()
2353 })
2354 }
2355
2356 pub fn schema(&self) -> &str {
2357 &self.config.schema
2358 }
2359
2360 pub fn slot_count(&self) -> usize {
2361 self.queue_slot_count()
2362 }
2363
2364 pub fn queue_slot_count(&self) -> usize {
2365 self.config.queue_slot_count
2366 }
2367
2368 pub fn lease_slot_count(&self) -> usize {
2369 self.config.lease_slot_count
2370 }
2371
2372 pub fn claim_slot_count(&self) -> usize {
2373 self.config.claim_slot_count
2374 }
2375
2376 pub fn queue_stripe_count(&self) -> usize {
2377 self.config.queue_stripe_count
2378 }
2379
2380 pub fn lease_claim_receipts(&self) -> bool {
2381 self.config.lease_claim_receipts
2382 }
2383
2384 fn uses_queue_striping(&self) -> bool {
2385 self.queue_stripe_count() > 1
2386 }
2387
2388 fn is_physical_stripe_queue(&self, queue: &str) -> bool {
2389 self.uses_queue_striping()
2390 && queue
2391 .rsplit_once(QUEUE_STRIPE_DELIMITER)
2392 .is_some_and(|(_, suffix)| suffix.parse::<usize>().is_ok())
2393 }
2394
2395 fn physical_queue_for_stripe(&self, queue: &str, stripe: usize) -> String {
2396 format!("{queue}{QUEUE_STRIPE_DELIMITER}{stripe}")
2397 }
2398
2399 fn physical_queues_for_logical(&self, queue: &str) -> Vec<String> {
2400 if !self.uses_queue_striping() || self.is_physical_stripe_queue(queue) {
2401 return vec![queue.to_string()];
2402 }
2403 (0..self.queue_stripe_count())
2404 .map(|stripe| self.physical_queue_for_stripe(queue, stripe))
2405 .collect()
2406 }
2407
2408 fn stripe_probe_start(&self, stripe_count: usize) -> usize {
2409 if stripe_count <= 1 {
2410 return 0;
2411 }
2412 self.next_stripe_probe.fetch_add(1, Ordering::Relaxed) % stripe_count
2413 }
2414
2415 fn logical_queue_name<'a>(&self, queue: &'a str) -> &'a str {
2416 if !self.uses_queue_striping() {
2417 return queue;
2418 }
2419 queue
2420 .rsplit_once(QUEUE_STRIPE_DELIMITER)
2421 .and_then(|(prefix, suffix)| suffix.parse::<usize>().ok().map(|_| prefix))
2422 .unwrap_or(queue)
2423 }
2424
2425 fn queue_stripe_for_enqueue(
2426 &self,
2427 queue: &str,
2428 unique_key: &Option<Vec<u8>>,
2429 salt: i64,
2430 ) -> String {
2431 if !self.uses_queue_striping() || self.is_physical_stripe_queue(queue) {
2432 return queue.to_string();
2433 }
2434
2435 let stripe = if let Some(key) = unique_key {
2436 let mut hasher = DefaultHasher::new();
2437 key.hash(&mut hasher);
2438 (hasher.finish() as usize) % self.queue_stripe_count()
2439 } else {
2440 salt.rem_euclid(self.queue_stripe_count() as i64) as usize
2441 };
2442 self.physical_queue_for_stripe(queue, stripe)
2443 }
2444
2445 fn use_lease_claim_receipts_for_runtime(&self, _deadline_duration: Duration) -> bool {
2446 self.lease_claim_receipts()
2452 }
2453
2454 pub fn ready_child_relname(&self, slot: usize) -> String {
2455 format!("ready_entries_{slot}")
2456 }
2457
2458 pub fn done_child_relname(&self, slot: usize) -> String {
2459 format!("done_entries_{slot}")
2460 }
2461
2462 pub fn leases_relname(&self) -> &'static str {
2463 "leases"
2464 }
2465
2466 pub fn lease_claims_relname(&self) -> &'static str {
2467 "lease_claims"
2468 }
2469
2470 pub fn lease_claim_closures_relname(&self) -> &'static str {
2471 "lease_claim_closures"
2472 }
2473
2474 pub fn leases_child_relname(&self, slot: usize) -> String {
2475 format!("leases_{slot}")
2476 }
2477
2478 pub fn attempt_state_relname(&self) -> &'static str {
2479 "attempt_state"
2480 }
2481
2482 pub async fn active_schema(pool: &PgPool) -> Result<Option<String>, AwaError> {
2483 sqlx::query_scalar(
2484 "SELECT schema_name FROM awa.runtime_storage_backends WHERE backend = 'queue_storage'",
2485 )
2486 .fetch_optional(pool)
2487 .await
2488 .map_err(map_sqlx_error)
2489 }
2490
2491 pub async fn active_schema_in_tx(
2495 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
2496 ) -> Result<Option<String>, AwaError> {
2497 sqlx::query_scalar(
2498 "SELECT schema_name FROM awa.runtime_storage_backends WHERE backend = 'queue_storage'",
2499 )
2500 .fetch_optional(tx.as_mut())
2501 .await
2502 .map_err(map_sqlx_error)
2503 }
2504
2505 fn materialize_runtime_payload(
2506 payload: serde_json::Value,
2507 progress: Option<serde_json::Value>,
2508 callback_result: Option<serde_json::Value>,
2509 ) -> Result<RuntimePayload, AwaError> {
2510 let mut payload = RuntimePayload::from_json(payload)?;
2511 if let Some(progress) = progress {
2512 payload.set_progress(Some(progress));
2513 }
2514 if let Some(callback_result) = callback_result {
2515 payload.insert_callback_result(Some(callback_result));
2516 }
2517 Ok(payload)
2518 }
2519
2520 fn payload_with_attempt_state(
2521 payload: serde_json::Value,
2522 progress: Option<serde_json::Value>,
2523 ) -> Result<serde_json::Value, AwaError> {
2524 let mut payload = RuntimePayload::from_json(payload)?;
2525 if let Some(progress) = progress {
2526 payload.set_progress(Some(progress));
2527 }
2528 Ok(payload.into_json())
2529 }
2530
2531 fn payload_from_parts(
2532 metadata: serde_json::Value,
2533 tags: Vec<String>,
2534 errors: Option<Vec<serde_json::Value>>,
2535 progress: Option<serde_json::Value>,
2536 ) -> Result<serde_json::Value, AwaError> {
2537 Ok(RuntimePayload {
2538 metadata,
2539 tags,
2540 errors: errors.unwrap_or_default(),
2541 progress,
2542 }
2543 .into_json())
2544 }
2545
2546 async fn sync_unique_claim<'a>(
2547 &self,
2548 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
2549 job_id: i64,
2550 unique_key: &Option<Vec<u8>>,
2551 unique_states: Option<&str>,
2552 old_state: Option<JobState>,
2553 new_state: Option<JobState>,
2554 ) -> Result<(), AwaError> {
2555 let old_claim = old_state.is_some_and(|state| unique_state_claims(unique_states, state));
2556 let new_claim = new_state.is_some_and(|state| unique_state_claims(unique_states, state));
2557
2558 if old_claim && !new_claim {
2559 if let Some(key) = unique_key {
2560 sqlx::query(
2561 "DELETE FROM awa.job_unique_claims WHERE unique_key = $1 AND job_id = $2",
2562 )
2563 .bind(key)
2564 .bind(job_id)
2565 .execute(tx.as_mut())
2566 .await
2567 .map_err(map_sqlx_error)?;
2568 }
2569 }
2570
2571 if new_claim && !old_claim {
2572 if let Some(key) = unique_key {
2573 let result = sqlx::query(
2574 r#"
2575 INSERT INTO awa.job_unique_claims (unique_key, job_id)
2576 VALUES ($1, $2)
2577 ON CONFLICT (unique_key)
2578 DO UPDATE SET job_id = EXCLUDED.job_id
2579 WHERE awa.job_unique_claims.job_id = EXCLUDED.job_id
2580 "#,
2581 )
2582 .bind(key)
2583 .bind(job_id)
2584 .execute(tx.as_mut())
2585 .await
2586 .map_err(map_sqlx_error)?;
2587
2588 if result.rows_affected() == 0 {
2589 return Err(AwaError::UniqueConflict {
2590 constraint: Some("idx_awa_jobs_unique".to_string()),
2591 });
2592 }
2593 }
2594 }
2595
2596 Ok(())
2597 }
2598
2599 async fn sync_enqueue_unique_claims<'a>(
2600 &self,
2601 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
2602 claims: Vec<(Vec<u8>, i64)>,
2603 ) -> Result<(), AwaError> {
2604 if claims.is_empty() {
2605 return Ok(());
2606 }
2607
2608 let mut seen: HashSet<&[u8]> = HashSet::with_capacity(claims.len());
2609 for (key, _) in &claims {
2610 if !seen.insert(key.as_slice()) {
2611 return Err(AwaError::UniqueConflict {
2612 constraint: Some("idx_awa_jobs_unique".to_string()),
2613 });
2614 }
2615 }
2616
2617 let (keys, job_ids): (Vec<Vec<u8>>, Vec<i64>) = claims.into_iter().unzip();
2618 let (requested, applied): (i64, i64) = sqlx::query_as(
2619 r#"
2620 WITH input(unique_key, job_id) AS (
2621 SELECT * FROM unnest($1::bytea[], $2::bigint[])
2622 ),
2623 inserted AS (
2624 INSERT INTO awa.job_unique_claims (unique_key, job_id)
2625 SELECT unique_key, job_id FROM input
2626 ON CONFLICT (unique_key)
2627 DO UPDATE SET job_id = EXCLUDED.job_id
2628 WHERE awa.job_unique_claims.job_id = EXCLUDED.job_id
2629 RETURNING unique_key
2630 )
2631 SELECT
2632 (SELECT count(*)::bigint FROM input) AS requested,
2633 (SELECT count(*)::bigint FROM inserted) AS applied
2634 "#,
2635 )
2636 .bind(keys)
2637 .bind(job_ids)
2638 .fetch_one(tx.as_mut())
2639 .await
2640 .map_err(map_sqlx_error)?;
2641
2642 if applied != requested {
2643 return Err(AwaError::UniqueConflict {
2644 constraint: Some("idx_awa_jobs_unique".to_string()),
2645 });
2646 }
2647
2648 Ok(())
2649 }
2650
2651 async fn sync_ready_enqueue_unique_claims<'a>(
2655 &self,
2656 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
2657 rows: &[RuntimeReadyInsert],
2658 ) -> Result<(), AwaError> {
2659 let claims = rows
2660 .iter()
2661 .filter(|row| unique_state_claims(row.unique_states.as_deref(), JobState::Available))
2662 .filter_map(|row| row.unique_key.as_ref().map(|key| (key.clone(), row.job_id)))
2663 .collect();
2664 self.sync_enqueue_unique_claims(tx, claims).await
2665 }
2666
2667 async fn sync_deferred_enqueue_unique_claims<'a>(
2668 &self,
2669 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
2670 rows: &[DeferredJobRow],
2671 ) -> Result<(), AwaError> {
2672 let claims = rows
2673 .iter()
2674 .filter(|row| unique_state_claims(row.unique_states.as_deref(), row.state))
2675 .filter_map(|row| row.unique_key.as_ref().map(|key| (key.clone(), row.job_id)))
2676 .collect();
2677 self.sync_enqueue_unique_claims(tx, claims).await
2678 }
2679
2680 #[tracing::instrument(skip(self, pool), name = "queue_storage.prepare_schema")]
2698 pub async fn prepare_schema(&self, pool: &PgPool) -> Result<(), AwaError> {
2699 let schema = self.schema();
2700 let install_lock_name = format!("awa.queue_storage.install:{schema}");
2701 let mut install_tx = pool.begin().await.map_err(map_sqlx_error)?;
2702
2703 sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))")
2704 .bind(&install_lock_name)
2705 .execute(install_tx.as_mut())
2706 .await
2707 .map_err(map_sqlx_error)?;
2708
2709 let install_result = async {
2710 sqlx::query(&format!("CREATE SCHEMA IF NOT EXISTS {schema}"))
2731 .execute(install_tx.as_mut())
2732 .await
2733 .map_err(map_sqlx_error)?;
2734
2735 let open_receipt_claims_exists: bool = sqlx::query_scalar(
2742 r#"
2743 SELECT EXISTS (
2744 SELECT 1 FROM pg_class c
2745 JOIN pg_namespace n ON n.oid = c.relnamespace
2746 WHERE n.nspname = $1 AND c.relname = 'open_receipt_claims'
2747 )
2748 "#,
2749 )
2750 .bind(schema)
2751 .fetch_one(install_tx.as_mut())
2752 .await
2753 .map_err(map_sqlx_error)?;
2754 if open_receipt_claims_exists {
2755 let row_count: i64 = sqlx::query_scalar(&format!(
2756 "SELECT count(*)::bigint FROM {schema}.open_receipt_claims"
2757 ))
2758 .fetch_one(install_tx.as_mut())
2759 .await
2760 .map_err(map_sqlx_error)?;
2761 if row_count > 0 {
2762 return Err(AwaError::Validation(format!(
2763 "{schema}.open_receipt_claims has {row_count} rows but the runtime no \
2764 longer reads or writes this table. Run the ADR-023 reverse migration \
2765 (recreate from lease_claims minus durable closure evidence) to drain it, \
2766 then re-run prepare_schema."
2767 )));
2768 }
2769 sqlx::query(&format!(
2770 "DROP TABLE IF EXISTS {schema}.open_receipt_claims CASCADE"
2771 ))
2772 .execute(install_tx.as_mut())
2773 .await
2774 .map_err(map_sqlx_error)?;
2775 }
2776
2777 let lease_claims_relkind: Option<String> = sqlx::query_scalar(
2785 r#"
2786 SELECT c.relkind::text
2787 FROM pg_class c
2788 JOIN pg_namespace n ON n.oid = c.relnamespace
2789 WHERE n.nspname = $1 AND c.relname = 'lease_claims'
2790 "#,
2791 )
2792 .bind(schema)
2793 .fetch_optional(install_tx.as_mut())
2794 .await
2795 .map_err(map_sqlx_error)?;
2796
2797 let closures_relkind: Option<String> = sqlx::query_scalar(
2798 r#"
2799 SELECT c.relkind::text
2800 FROM pg_class c
2801 JOIN pg_namespace n ON n.oid = c.relnamespace
2802 WHERE n.nspname = $1 AND c.relname = 'lease_claim_closures'
2803 "#,
2804 )
2805 .bind(schema)
2806 .fetch_optional(install_tx.as_mut())
2807 .await
2808 .map_err(map_sqlx_error)?;
2809
2810 if lease_claims_relkind.as_deref() == Some("r") {
2811 sqlx::query(&format!(
2812 "ALTER TABLE {schema}.lease_claims RENAME TO lease_claims_legacy"
2813 ))
2814 .execute(install_tx.as_mut())
2815 .await
2816 .map_err(map_sqlx_error)?;
2817 }
2818 if closures_relkind.as_deref() == Some("r") {
2819 sqlx::query(&format!(
2820 "ALTER TABLE {schema}.lease_claim_closures RENAME TO lease_claim_closures_legacy"
2821 ))
2822 .execute(install_tx.as_mut())
2823 .await
2824 .map_err(map_sqlx_error)?;
2825 }
2826
2827 sqlx::query(&format!(
2834 "DROP TABLE IF EXISTS {schema}.queue_count_snapshots"
2835 ))
2836 .execute(install_tx.as_mut())
2837 .await
2838 .map_err(map_sqlx_error)?;
2839
2840 sqlx::query("SELECT awa.install_queue_storage_substrate($1, $2, $3, $4, $5)")
2847 .bind(schema)
2848 .bind(self.queue_slot_count() as i32)
2849 .bind(self.lease_slot_count() as i32)
2850 .bind(self.claim_slot_count() as i32)
2851 .bind(self.lease_claim_receipts())
2852 .execute(install_tx.as_mut())
2853 .await
2854 .map_err(map_sqlx_error)?;
2855
2856 sqlx::query("SELECT awa.apply_receipt_plane_fillfactor($1)")
2867 .bind(schema)
2868 .execute(install_tx.as_mut())
2869 .await
2870 .map_err(map_sqlx_error)?;
2871
2872 let lease_claims_legacy_exists: bool = sqlx::query_scalar(
2881 r#"
2882 SELECT EXISTS (
2883 SELECT 1 FROM pg_class c
2884 JOIN pg_namespace n ON n.oid = c.relnamespace
2885 WHERE n.nspname = $1 AND c.relname = 'lease_claims_legacy'
2886 )
2887 "#,
2888 )
2889 .bind(schema)
2890 .fetch_one(install_tx.as_mut())
2891 .await
2892 .map_err(map_sqlx_error)?;
2893
2894 let closures_legacy_exists: bool = sqlx::query_scalar(
2895 r#"
2896 SELECT EXISTS (
2897 SELECT 1 FROM pg_class c
2898 JOIN pg_namespace n ON n.oid = c.relnamespace
2899 WHERE n.nspname = $1 AND c.relname = 'lease_claim_closures_legacy'
2900 )
2901 "#,
2902 )
2903 .bind(schema)
2904 .fetch_one(install_tx.as_mut())
2905 .await
2906 .map_err(map_sqlx_error)?;
2907
2908 let legacy_claim_slot: Option<i32> =
2909 if lease_claims_legacy_exists || closures_legacy_exists {
2910 Some(
2911 sqlx::query_scalar(&format!(
2912 "SELECT current_slot FROM {schema}.claim_ring_state WHERE singleton"
2913 ))
2914 .fetch_one(install_tx.as_mut())
2915 .await
2916 .map_err(map_sqlx_error)?,
2917 )
2918 } else {
2919 None
2920 };
2921
2922 if lease_claims_legacy_exists {
2923 sqlx::query(&format!(
2924 "ALTER TABLE {schema}.lease_claims_legacy ADD COLUMN IF NOT EXISTS enqueue_shard SMALLINT NOT NULL DEFAULT 0"
2925 ))
2926 .execute(install_tx.as_mut())
2927 .await
2928 .map_err(map_sqlx_error)?;
2929 sqlx::query(&format!(
2930 "ALTER TABLE {schema}.lease_claims_legacy ADD COLUMN IF NOT EXISTS deadline_at TIMESTAMPTZ"
2931 ))
2932 .execute(install_tx.as_mut())
2933 .await
2934 .map_err(map_sqlx_error)?;
2935
2936 sqlx::query(&format!(
2937 r#"
2938 INSERT INTO {schema}.lease_claims (
2939 claim_slot, job_id, run_lease, ready_slot, ready_generation,
2940 queue, priority, attempt, max_attempts, lane_seq,
2941 enqueue_shard, claimed_at, materialized_at, deadline_at
2942 )
2943 SELECT
2944 $1,
2945 job_id, run_lease, ready_slot, ready_generation,
2946 queue, priority, attempt, max_attempts, lane_seq,
2947 enqueue_shard, claimed_at, materialized_at, deadline_at
2948 FROM {schema}.lease_claims_legacy
2949 ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING
2950 "#
2951 ))
2952 .bind(legacy_claim_slot.expect("legacy claim slot should be present"))
2953 .execute(install_tx.as_mut())
2954 .await
2955 .map_err(map_sqlx_error)?;
2956
2957 sqlx::query(&format!(
2958 "DROP TABLE {schema}.lease_claims_legacy"
2959 ))
2960 .execute(install_tx.as_mut())
2961 .await
2962 .map_err(map_sqlx_error)?;
2963 }
2964
2965 if closures_legacy_exists {
2966 sqlx::query(&format!(
2967 r#"
2968 INSERT INTO {schema}.lease_claim_closures (
2969 claim_slot, job_id, run_lease, outcome, closed_at
2970 )
2971 SELECT
2972 $1,
2973 job_id, run_lease, outcome, closed_at
2974 FROM {schema}.lease_claim_closures_legacy
2975 ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING
2976 "#
2977 ))
2978 .bind(legacy_claim_slot.expect("legacy claim slot should be present"))
2979 .execute(install_tx.as_mut())
2980 .await
2981 .map_err(map_sqlx_error)?;
2982
2983 sqlx::query(&format!(
2984 "DROP TABLE {schema}.lease_claim_closures_legacy"
2985 ))
2986 .execute(install_tx.as_mut())
2987 .await
2988 .map_err(map_sqlx_error)?;
2989 }
2990
2991 Ok(())
2992 }
2993 .await;
2994
2995 match install_result {
2996 Ok(()) => install_tx.commit().await.map_err(map_sqlx_error),
2997 Err(err) => {
2998 let _ = install_tx.rollback().await;
2999 Err(err)
3000 }
3001 }
3002 }
3003
3004 #[tracing::instrument(skip(self, pool), name = "queue_storage.activate_backend")]
3005 pub async fn activate_backend(&self, pool: &PgPool) -> Result<(), AwaError> {
3006 let schema = self.schema();
3007 let details = serde_json::json!({ "schema": schema });
3008
3009 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
3010
3011 sqlx::query(
3017 r#"
3018 INSERT INTO awa.runtime_storage_backends (backend, schema_name, updated_at)
3019 VALUES ('queue_storage', $1, now())
3020 ON CONFLICT (backend)
3021 DO UPDATE SET schema_name = EXCLUDED.schema_name, updated_at = EXCLUDED.updated_at
3022 "#,
3023 )
3024 .bind(schema)
3025 .execute(tx.as_mut())
3026 .await
3027 .map_err(map_sqlx_error)?;
3028
3029 let activation_result = sqlx::query(
3030 r#"
3031 UPDATE awa.storage_transition_state AS sts
3032 SET
3033 current_engine = 'queue_storage',
3034 prepared_engine = NULL,
3035 state = 'active',
3036 transition_epoch = CASE
3037 WHEN sts.current_engine = 'queue_storage'
3038 AND sts.prepared_engine IS NULL
3039 AND sts.state = 'active'
3040 AND sts.details = $1
3041 THEN sts.transition_epoch
3042 ELSE sts.transition_epoch + 1
3043 END,
3044 details = $1,
3045 entered_at = CASE
3046 WHEN sts.current_engine = 'queue_storage'
3047 AND sts.prepared_engine IS NULL
3048 AND sts.state = 'active'
3049 AND sts.details = $1
3050 THEN sts.entered_at
3051 ELSE now()
3052 END,
3053 updated_at = now(),
3054 finalized_at = CASE
3055 WHEN sts.current_engine = 'queue_storage'
3056 AND sts.prepared_engine IS NULL
3057 AND sts.state = 'active'
3058 AND sts.details = $1
3059 THEN COALESCE(sts.finalized_at, now())
3060 ELSE now()
3061 END
3062 WHERE sts.singleton
3063 "#,
3064 )
3065 .bind(details)
3066 .execute(tx.as_mut())
3067 .await
3068 .map_err(map_sqlx_error)?;
3069
3070 if activation_result.rows_affected() != 1 {
3071 return Err(AwaError::Validation(
3072 "queue storage activation requires the storage transition state row".into(),
3073 ));
3074 }
3075
3076 tx.commit().await.map_err(map_sqlx_error)?;
3077
3078 Ok(())
3079 }
3080
3081 #[tracing::instrument(skip(self, pool), name = "queue_storage.install")]
3082 pub async fn install(&self, pool: &PgPool) -> Result<(), AwaError> {
3083 self.prepare_schema(pool).await?;
3084 self.activate_backend(pool).await
3085 }
3086
3087 pub async fn reset(&self, pool: &PgPool) -> Result<(), AwaError> {
3088 let schema = self.schema();
3089 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
3090
3091 sqlx::query(&format!(
3099 "DROP TABLE IF EXISTS {schema}.lease_claims_legacy"
3100 ))
3101 .execute(tx.as_mut())
3102 .await
3103 .map_err(map_sqlx_error)?;
3104
3105 sqlx::query(&format!(
3106 "DROP TABLE IF EXISTS {schema}.lease_claim_closures_legacy"
3107 ))
3108 .execute(tx.as_mut())
3109 .await
3110 .map_err(map_sqlx_error)?;
3111
3112 sqlx::query(&format!(
3113 r#"
3114 TRUNCATE
3115 {schema}.ready_entries,
3116 {schema}.ready_claim_attempt_batches,
3117 {schema}.ready_tombstones,
3118 {schema}.ready_segments,
3119 {schema}.done_entries,
3120 {schema}.receipt_completion_batches,
3121 {schema}.receipt_completion_tombstones,
3122 {schema}.dlq_entries,
3123 {schema}.leases,
3124 {schema}.lease_claims,
3125 {schema}.lease_claim_batches,
3126 {schema}.lease_claim_closures,
3127 {schema}.lease_claim_closure_batches,
3128 {schema}.attempt_state,
3129 {schema}.deferred_jobs,
3130 {schema}.queue_lanes,
3131 {schema}.queue_terminal_rollups,
3132 {schema}.queue_terminal_live_counts,
3133 {schema}.queue_terminal_count_deltas,
3134 {schema}.queue_claimer_leases,
3135 {schema}.queue_claimer_state,
3136 {schema}.queue_ring_slots,
3137 {schema}.lease_ring_slots,
3138 {schema}.claim_ring_slots
3139 "#
3140 ))
3141 .execute(tx.as_mut())
3142 .await
3143 .map_err(map_sqlx_error)?;
3144
3145 sqlx::query(&format!(
3146 "ALTER SEQUENCE {schema}.job_id_seq RESTART WITH 1"
3147 ))
3148 .execute(tx.as_mut())
3149 .await
3150 .map_err(map_sqlx_error)?;
3151
3152 sqlx::query(&format!(
3153 r#"
3154 UPDATE {schema}.queue_ring_state
3155 SET current_slot = 0,
3156 generation = 0,
3157 slot_count = $1
3158 WHERE singleton = TRUE
3159 "#
3160 ))
3161 .bind(self.queue_slot_count() as i32)
3162 .execute(tx.as_mut())
3163 .await
3164 .map_err(map_sqlx_error)?;
3165
3166 sqlx::query(&format!(
3167 r#"
3168 UPDATE {schema}.lease_ring_state
3169 SET current_slot = 0,
3170 generation = 0,
3171 slot_count = $1
3172 WHERE singleton = TRUE
3173 "#
3174 ))
3175 .bind(self.lease_slot_count() as i32)
3176 .execute(tx.as_mut())
3177 .await
3178 .map_err(map_sqlx_error)?;
3179
3180 sqlx::query(&format!(
3181 r#"
3182 UPDATE {schema}.claim_ring_state
3183 SET current_slot = 0,
3184 generation = 0,
3185 slot_count = $1
3186 WHERE singleton = TRUE
3187 "#
3188 ))
3189 .bind(self.claim_slot_count() as i32)
3190 .execute(tx.as_mut())
3191 .await
3192 .map_err(map_sqlx_error)?;
3193
3194 for slot in 0..self.queue_slot_count() {
3195 sqlx::query(&format!(
3196 r#"
3197 INSERT INTO {schema}.queue_ring_slots (slot, generation)
3198 VALUES ($1, $2)
3199 "#
3200 ))
3201 .bind(slot as i32)
3202 .bind(if slot == 0 { 0_i64 } else { -1_i64 })
3203 .execute(tx.as_mut())
3204 .await
3205 .map_err(map_sqlx_error)?;
3206 }
3207
3208 for slot in 0..self.lease_slot_count() {
3209 sqlx::query(&format!(
3210 r#"
3211 INSERT INTO {schema}.lease_ring_slots (slot, generation)
3212 VALUES ($1, $2)
3213 "#
3214 ))
3215 .bind(slot as i32)
3216 .bind(if slot == 0 { 0_i64 } else { -1_i64 })
3217 .execute(tx.as_mut())
3218 .await
3219 .map_err(map_sqlx_error)?;
3220 }
3221
3222 for slot in 0..self.claim_slot_count() {
3223 sqlx::query(&format!(
3224 r#"
3225 INSERT INTO {schema}.claim_ring_slots (slot, generation)
3226 VALUES ($1, $2)
3227 "#
3228 ))
3229 .bind(slot as i32)
3230 .bind(if slot == 0 { 0_i64 } else { -1_i64 })
3231 .execute(tx.as_mut())
3232 .await
3233 .map_err(map_sqlx_error)?;
3234 }
3235
3236 tx.commit().await.map_err(map_sqlx_error)?;
3237
3238 self.clear_lane_cache();
3243 self.enqueue_shards_cache
3244 .lock()
3245 .expect("enqueue_shards_cache poisoned")
3246 .clear();
3247 Ok(())
3248 }
3249
3250 async fn ensure_lane<'a>(
3251 &self,
3252 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
3253 queue: &str,
3254 priority: i16,
3255 enqueue_shard: i16,
3256 ) -> Result<(), AwaError> {
3257 if self.lane_is_cached(queue, priority, enqueue_shard) {
3263 let schema = self.schema();
3264 let visible: bool = sqlx::query_scalar(&format!(
3265 r#"
3266 SELECT EXISTS (
3267 SELECT 1
3268 FROM {schema}.queue_enqueue_heads
3269 WHERE queue = $1
3270 AND priority = $2
3271 AND enqueue_shard = $3
3272 )
3273 "#
3274 ))
3275 .bind(queue)
3276 .bind(priority)
3277 .bind(enqueue_shard)
3278 .fetch_one(tx.as_mut())
3279 .await
3280 .map_err(map_sqlx_error)?;
3281
3282 if visible {
3283 return Ok(());
3284 }
3285
3286 self.invalidate_cached_lane(queue, priority, enqueue_shard);
3287 }
3288
3289 self.ensure_lane_inserts(tx, queue, priority, enqueue_shard)
3290 .await
3291 }
3292
3293 async fn ensure_lane_inserts<'a>(
3299 &self,
3300 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
3301 queue: &str,
3302 priority: i16,
3303 enqueue_shard: i16,
3304 ) -> Result<(), AwaError> {
3305 let schema = self.schema();
3306 let lane_lock_key = format!("{schema}:{queue}:{priority}:{enqueue_shard}");
3307 sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))")
3308 .bind(lane_lock_key)
3309 .execute(tx.as_mut())
3310 .await
3311 .map_err(map_sqlx_error)?;
3312
3313 sqlx::query(&format!(
3314 r#"
3315 INSERT INTO {schema}.queue_lanes (queue, priority)
3316 VALUES ($1, $2)
3317 ON CONFLICT (queue, priority) DO NOTHING
3318 "#
3319 ))
3320 .bind(queue)
3321 .bind(priority)
3322 .execute(tx.as_mut())
3323 .await
3324 .map_err(map_sqlx_error)?;
3325
3326 sqlx::query(&format!(
3327 r#"
3328 INSERT INTO {schema}.queue_enqueue_heads (queue, priority, enqueue_shard)
3329 VALUES ($1, $2, $3)
3330 ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING
3331 "#
3332 ))
3333 .bind(queue)
3334 .bind(priority)
3335 .bind(enqueue_shard)
3336 .execute(tx.as_mut())
3337 .await
3338 .map_err(map_sqlx_error)?;
3339
3340 sqlx::query(&format!(
3341 r#"
3342 INSERT INTO {schema}.queue_claim_heads (queue, priority, enqueue_shard)
3343 VALUES ($1, $2, $3)
3344 ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING
3345 "#
3346 ))
3347 .bind(queue)
3348 .bind(priority)
3349 .bind(enqueue_shard)
3350 .execute(tx.as_mut())
3351 .await
3352 .map_err(map_sqlx_error)?;
3353
3354 sqlx::query(&format!(
3355 r#"
3356 SELECT {schema}.ensure_lane_sequences($1, $2, $3)
3357 "#
3358 ))
3359 .bind(queue)
3360 .bind(priority)
3361 .bind(enqueue_shard)
3362 .execute(tx.as_mut())
3363 .await
3364 .map_err(map_sqlx_error)?;
3365
3366 self.mark_lane_ensured(queue, priority, enqueue_shard);
3367 Ok(())
3368 }
3369
3370 async fn shard_for_enqueue(
3384 &self,
3385 pool_executor: impl sqlx::PgExecutor<'_>,
3386 queue: &str,
3387 ordering_key: Option<&[u8]>,
3388 ) -> Result<i16, AwaError> {
3389 if let Some(cached) = self
3390 .enqueue_shards_cache
3391 .lock()
3392 .expect("enqueue_shards_cache poisoned")
3393 .get(queue)
3394 .copied()
3395 {
3396 return Ok(self.pick_shard(cached, ordering_key));
3397 }
3398
3399 let shards: i16 = sqlx::query_scalar(
3400 r#"
3401 SELECT COALESCE(MAX(enqueue_shards), 1)::smallint
3402 FROM awa.queue_meta
3403 WHERE queue = $1
3404 "#,
3405 )
3406 .bind(queue)
3407 .fetch_one(pool_executor)
3408 .await
3409 .map_err(map_sqlx_error)?;
3410
3411 let shards = shards.max(1);
3412 self.enqueue_shards_cache
3413 .lock()
3414 .expect("enqueue_shards_cache poisoned")
3415 .insert(queue.to_string(), shards);
3416
3417 Ok(self.pick_shard(shards, ordering_key))
3418 }
3419
3420 fn pick_shard(&self, shards: i16, ordering_key: Option<&[u8]>) -> i16 {
3430 if shards <= 1 {
3431 return 0;
3432 }
3433 match ordering_key {
3434 Some(key) => shard_for_ordering_key(key, shards),
3435 None => {
3436 let raw = self.shard_rotor.fetch_add(1, Ordering::Relaxed) as i32;
3437 raw.rem_euclid(shards as i32) as i16
3438 }
3439 }
3440 }
3441
3442 fn lane_is_cached(&self, queue: &str, priority: i16, enqueue_shard: i16) -> bool {
3443 let cache = self.ensured_lanes.lock().expect("ensured_lanes mutex");
3444 cache.contains(&(queue.to_string(), priority, enqueue_shard))
3445 }
3446
3447 fn mark_lane_ensured(&self, queue: &str, priority: i16, enqueue_shard: i16) {
3448 self.ensured_lanes
3449 .lock()
3450 .expect("ensured_lanes mutex")
3451 .insert((queue.to_string(), priority, enqueue_shard));
3452 }
3453
3454 fn invalidate_cached_lane(&self, queue: &str, priority: i16, enqueue_shard: i16) {
3455 self.ensured_lanes
3456 .lock()
3457 .expect("ensured_lanes mutex")
3458 .remove(&(queue.to_string(), priority, enqueue_shard));
3459 }
3460
3461 fn clear_lane_cache(&self) {
3462 self.ensured_lanes
3463 .lock()
3464 .expect("ensured_lanes mutex")
3465 .clear();
3466 }
3467
3468 async fn advance_enqueue_head<'a>(
3477 &self,
3478 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
3479 queue: &str,
3480 priority: i16,
3481 enqueue_shard: i16,
3482 count: i64,
3483 ) -> Result<i64, AwaError> {
3484 let schema = self.schema();
3485 let sql = format!(
3486 r#"
3487 SELECT {schema}.reserve_enqueue_seq($1, $2, $3, $4)
3488 FROM {schema}.queue_enqueue_heads
3489 WHERE queue = $1 AND priority = $2 AND enqueue_shard = $3
3490 "#
3491 );
3492
3493 let maybe_start: Option<i64> = sqlx::query_scalar(&sql)
3494 .bind(queue)
3495 .bind(priority)
3496 .bind(enqueue_shard)
3497 .bind(count)
3498 .fetch_optional(tx.as_mut())
3499 .await
3500 .map_err(map_sqlx_error)?;
3501
3502 if let Some(start) = maybe_start {
3503 return Ok(start);
3504 }
3505
3506 self.invalidate_cached_lane(queue, priority, enqueue_shard);
3514 self.ensure_lane_inserts(tx, queue, priority, enqueue_shard)
3515 .await?;
3516 let start: i64 = sqlx::query_scalar(&sql)
3517 .bind(queue)
3518 .bind(priority)
3519 .bind(enqueue_shard)
3520 .bind(count)
3521 .fetch_one(tx.as_mut())
3522 .await
3523 .map_err(map_sqlx_error)?;
3524 Ok(start)
3525 }
3526
3527 async fn current_queue_ring<'a>(
3528 &self,
3529 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
3530 ) -> Result<(i32, i64), AwaError> {
3531 let schema = self.schema();
3532 sqlx::query_as(&format!(
3533 r#"
3534 SELECT current_slot, generation
3535 FROM {schema}.queue_ring_state
3536 WHERE singleton = TRUE
3537 "#
3538 ))
3539 .fetch_one(tx.as_mut())
3540 .await
3541 .map_err(map_sqlx_error)
3542 }
3543
3544 async fn next_job_ids<'a>(
3545 &self,
3546 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
3547 count: usize,
3548 ) -> Result<Vec<i64>, AwaError> {
3549 if count == 0 {
3550 return Ok(Vec::new());
3551 }
3552
3553 let query = format!(
3554 "SELECT nextval('{}')::bigint FROM generate_series(1, $1::int)",
3555 self.job_id_sequence()
3556 );
3557
3558 sqlx::query_scalar(&query)
3559 .bind(count as i32)
3560 .fetch_all(tx.as_mut())
3561 .await
3562 .map_err(map_sqlx_error)
3563 }
3564
3565 async fn current_timestamp_tx<'a>(
3566 &self,
3567 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
3568 ) -> Result<DateTime<Utc>, AwaError> {
3569 sqlx::query_scalar("SELECT clock_timestamp()")
3570 .fetch_one(tx.as_mut())
3571 .await
3572 .map_err(map_sqlx_error)
3573 }
3574
3575 async fn claim_ready_rows_tx<'a>(
3576 &self,
3577 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
3578 queue: &str,
3579 max_batch: i64,
3580 deadline_duration: Duration,
3581 aging_interval: Duration,
3582 ) -> Result<Vec<ReadyJobLeaseRow>, AwaError> {
3583 let schema = self.schema();
3584 sqlx::query_as(&format!(
3585 r#"
3586 SELECT
3587 ready_slot,
3588 ready_generation,
3589 lane_seq,
3590 enqueue_shard,
3591 lease_slot,
3592 lease_generation,
3593 claim_slot,
3594 receipt_id,
3595 claim_batch_id,
3596 claim_batch_index,
3597 job_id,
3598 kind,
3599 queue,
3600 args,
3601 lane_priority,
3602 priority,
3603 attempt,
3604 run_lease,
3605 max_attempts,
3606 run_at,
3607 heartbeat_at,
3608 deadline_at,
3609 attempted_at,
3610 created_at,
3611 unique_key,
3612 unique_states,
3613 COALESCE(payload, '{{}}'::jsonb) AS payload
3614 FROM {schema}.claim_ready_runtime($1, $2, $3, $4)
3615 "#
3616 ))
3617 .bind(queue)
3618 .bind(max_batch)
3619 .bind(deadline_duration.as_secs_f64())
3620 .bind(aging_interval.as_secs_f64())
3621 .fetch_all(tx.as_mut())
3622 .await
3623 .map_err(map_sqlx_error)
3624 }
3625
3626 fn claim_cursor_advances(rows: &[ReadyJobLeaseRow]) -> Vec<ClaimCursorAdvance> {
3627 let mut next_by_lane: BTreeMap<ClaimCursorLaneKey, i64> = BTreeMap::new();
3628 for row in rows {
3629 let key = (row.queue.clone(), row.lane_priority, row.enqueue_shard);
3630 let next = row.lane_seq + 1;
3631 next_by_lane
3632 .entry(key)
3633 .and_modify(|current| *current = (*current).max(next))
3634 .or_insert(next);
3635 }
3636
3637 next_by_lane
3638 .into_iter()
3639 .map(
3640 |((queue, priority, enqueue_shard), next_seq)| ClaimCursorAdvance {
3641 queue,
3642 priority,
3643 enqueue_shard,
3644 next_seq,
3645 only_if_current: None,
3646 },
3647 )
3648 .collect()
3649 }
3650
3651 fn normalize_claim_cursor_advances(advances: &[ClaimCursorAdvance]) -> Vec<ClaimCursorAdvance> {
3652 let mut grouped: GroupedClaimCursorAdvances = BTreeMap::new();
3653
3654 for advance in advances {
3655 let key = (
3656 advance.queue.clone(),
3657 advance.priority,
3658 advance.enqueue_shard,
3659 );
3660 let (unconditional, conditional) = grouped.entry(key).or_default();
3661 if let Some(only_if_current) = advance.only_if_current {
3662 conditional
3663 .entry(only_if_current)
3664 .and_modify(|next| *next = (*next).max(advance.next_seq))
3665 .or_insert(advance.next_seq);
3666 } else {
3667 *unconditional = Some(
3668 unconditional
3669 .map(|next| next.max(advance.next_seq))
3670 .unwrap_or(advance.next_seq),
3671 );
3672 }
3673 }
3674
3675 let mut normalized = Vec::with_capacity(advances.len());
3676 for ((queue, priority, enqueue_shard), (unconditional, conditional)) in grouped {
3677 if let Some(next_seq) = unconditional {
3678 normalized.push(ClaimCursorAdvance {
3679 queue,
3680 priority,
3681 enqueue_shard,
3682 next_seq,
3683 only_if_current: None,
3684 });
3685 continue;
3686 }
3687
3688 for (only_if_current, next_seq) in conditional {
3689 normalized.push(ClaimCursorAdvance {
3690 queue: queue.clone(),
3691 priority,
3692 enqueue_shard,
3693 next_seq,
3694 only_if_current: Some(only_if_current),
3695 });
3696 }
3697 }
3698
3699 normalized
3700 }
3701
3702 async fn advance_claim_cursors(&self, pool: &PgPool, advances: &[ClaimCursorAdvance]) {
3703 let advances = Self::normalize_claim_cursor_advances(advances);
3704 for attempt in 1..=3 {
3708 match self.advance_claim_cursors_strict(pool, &advances).await {
3709 Ok(()) => return,
3710 Err(err) if attempt < 3 => {
3711 tracing::warn!(
3712 error = ?err,
3713 lanes = advances.len(),
3714 attempt,
3715 "failed to advance queue-storage claim cursors after committed state change; retrying"
3716 );
3717 tokio::time::sleep(Duration::from_millis(25 * attempt as u64)).await;
3718 }
3719 Err(err) => {
3720 tracing::warn!(
3721 error = ?err,
3722 lanes = advances.len(),
3723 attempts = attempt,
3724 "failed to advance queue-storage claim cursors after committed state change"
3725 );
3726 return;
3727 }
3728 }
3729 }
3730 }
3731
3732 async fn advance_claim_cursors_strict(
3733 &self,
3734 pool: &PgPool,
3735 advances: &[ClaimCursorAdvance],
3736 ) -> Result<(), AwaError> {
3737 if advances.is_empty() {
3738 return Ok(());
3739 }
3740
3741 let schema = self.schema();
3742 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
3743 for advance in advances {
3744 sqlx::query(&format!(
3745 r#"
3746 WITH head AS MATERIALIZED (
3747 SELECT seq_name
3748 FROM {schema}.queue_claim_heads
3749 WHERE queue = $1
3750 AND priority = $2
3751 AND enqueue_shard = $3
3752 FOR UPDATE
3753 )
3754 SELECT {schema}.set_sequence_next(seq_name, $4)
3755 FROM head
3756 WHERE $5::bigint IS NULL
3757 OR {schema}.sequence_next_value(seq_name) = $5
3758 "#
3759 ))
3760 .bind(&advance.queue)
3761 .bind(advance.priority)
3762 .bind(advance.enqueue_shard)
3763 .bind(advance.next_seq)
3764 .bind(advance.only_if_current)
3765 .execute(tx.as_mut())
3766 .await
3767 .map_err(map_sqlx_error)?;
3768 }
3769
3770 tx.commit().await.map_err(map_sqlx_error)?;
3771 Ok(())
3772 }
3773
3774 async fn execute_ready_inserts_tx<'a>(
3775 &self,
3776 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
3777 ring: (i32, i64),
3778 rows: &[RuntimeReadyInsert],
3779 ) -> Result<usize, AwaError> {
3780 if rows.is_empty() {
3781 return Ok(0);
3782 }
3783
3784 let schema = self.schema();
3785 let mut builder = QueryBuilder::<Postgres>::new(format!(
3786 "INSERT INTO {schema}.ready_entries (ready_slot, ready_generation, job_id, kind, queue, args, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, run_at, attempted_at, created_at, unique_key, unique_states, payload) "
3787 ));
3788 builder.push_values(rows.iter(), |mut b, row| {
3789 b.push_bind(ring.0)
3790 .push_bind(ring.1)
3791 .push_bind(row.job_id)
3792 .push_bind(&row.kind)
3793 .push_bind(&row.queue)
3794 .push_bind(&row.args)
3795 .push_bind(row.priority)
3796 .push_bind(row.attempt)
3797 .push_bind(row.run_lease)
3798 .push_bind(row.max_attempts)
3799 .push_bind(row.lane_seq)
3800 .push_bind(row.enqueue_shard)
3801 .push_bind(row.run_at)
3802 .push_bind(row.attempted_at)
3803 .push_bind(row.created_at)
3804 .push_bind(&row.unique_key)
3805 .push_bind(&row.unique_states)
3806 .push_bind(storage_payload(&row.payload));
3807 });
3808 builder
3809 .build()
3810 .execute(tx.as_mut())
3811 .await
3812 .map_err(map_sqlx_error)?;
3813
3814 Ok(rows.len())
3815 }
3816
3817 fn ready_segments_from_rows(rows: &[RuntimeReadyInsert]) -> Vec<ReadySegmentInsert> {
3818 if rows.is_empty() {
3819 return Vec::new();
3820 }
3821
3822 let mut segments = Vec::new();
3823 let mut start_index = 0usize;
3824
3825 for index in 1..=rows.len() {
3826 let split = if index == rows.len() {
3827 true
3828 } else {
3829 let previous = &rows[index - 1];
3830 let current = &rows[index];
3831 previous.queue != current.queue
3832 || previous.priority != current.priority
3833 || previous.enqueue_shard != current.enqueue_shard
3834 || previous.lane_seq + 1 != current.lane_seq
3835 || previous.run_at != current.run_at
3836 };
3837
3838 if split {
3839 let first = &rows[start_index];
3840 let last = &rows[index - 1];
3841 segments.push(ReadySegmentInsert {
3842 queue: first.queue.clone(),
3843 priority: first.priority,
3844 enqueue_shard: first.enqueue_shard,
3845 first_lane_seq: first.lane_seq,
3846 next_lane_seq: last.lane_seq + 1,
3847 first_run_at: first.run_at,
3848 });
3849 start_index = index;
3850 }
3851 }
3852
3853 segments
3854 }
3855
3856 async fn insert_ready_segments_tx<'a>(
3857 &self,
3858 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
3859 ring: (i32, i64),
3860 rows: &[RuntimeReadyInsert],
3861 ) -> Result<usize, AwaError> {
3862 let segments = Self::ready_segments_from_rows(rows);
3863 if segments.is_empty() {
3864 return Ok(0);
3865 }
3866
3867 let schema = self.schema();
3868 let mut builder = QueryBuilder::<Postgres>::new(format!(
3869 "INSERT INTO {schema}.ready_segments (ready_slot, ready_generation, queue, priority, enqueue_shard, first_lane_seq, next_lane_seq, first_run_at) "
3870 ));
3871 builder.push_values(segments.iter(), |mut b, segment| {
3872 b.push_bind(ring.0)
3873 .push_bind(ring.1)
3874 .push_bind(&segment.queue)
3875 .push_bind(segment.priority)
3876 .push_bind(segment.enqueue_shard)
3877 .push_bind(segment.first_lane_seq)
3878 .push_bind(segment.next_lane_seq)
3879 .push_bind(segment.first_run_at);
3880 });
3881 builder.push(
3882 " ON CONFLICT (ready_slot, ready_generation, queue, priority, enqueue_shard, first_lane_seq) DO NOTHING",
3883 );
3884 builder
3885 .build()
3886 .execute(tx.as_mut())
3887 .await
3888 .map_err(map_sqlx_error)?;
3889
3890 Ok(segments.len())
3891 }
3892
3893 async fn execute_ready_copy_tx<'a>(
3894 &self,
3895 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
3896 ring: (i32, i64),
3897 rows: &[RuntimeReadyInsert],
3898 ) -> Result<usize, AwaError> {
3899 if rows.is_empty() {
3900 return Ok(0);
3901 }
3902
3903 let schema = self.schema();
3904 let copy_sql = format!(
3905 "COPY {schema}.ready_entries (ready_slot, ready_generation, job_id, kind, queue, args, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, run_at, attempted_at, created_at, unique_key, unique_states, payload) FROM STDIN WITH (FORMAT csv, NULL '{COPY_NULL_SENTINEL}')"
3906 );
3907 let mut copy_in = tx
3908 .as_mut()
3909 .copy_in_raw(©_sql)
3910 .await
3911 .map_err(map_sqlx_error)?;
3912 let mut csv_buf = Vec::with_capacity(rows.len().min(1024) * 320);
3915 for row in rows {
3916 write_ready_copy_row(&mut csv_buf, ring.0, ring.1, row);
3917 if csv_buf.len() >= COPY_CHUNK_TARGET_BYTES {
3918 let chunk =
3919 std::mem::replace(&mut csv_buf, Vec::with_capacity(COPY_CHUNK_TARGET_BYTES));
3920 copy_in.send(chunk).await.map_err(map_sqlx_error)?;
3921 }
3922 }
3923 if !csv_buf.is_empty() {
3924 copy_in.send(csv_buf).await.map_err(map_sqlx_error)?;
3925 }
3926 copy_in.finish().await.map_err(map_sqlx_error)?;
3927
3928 Ok(rows.len())
3929 }
3930
3931 async fn insert_ready_rows_tx<'a>(
3932 &self,
3933 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
3934 rows: Vec<RuntimeReadyRow>,
3935 ) -> Result<usize, AwaError> {
3936 if rows.is_empty() {
3937 return Ok(0);
3938 }
3939
3940 let grouped = self.group_ready_rows_by_shard(tx, rows).await?;
3941 let total_rows: usize = grouped.values().map(Vec::len).sum();
3942 let job_ids = self.next_job_ids(tx, total_rows).await?;
3943 let mut job_id_iter = job_ids.into_iter();
3944
3945 let mut ready_rows = Vec::with_capacity(total_rows);
3946 let mut lane_ranges = Vec::with_capacity(grouped.len());
3947
3948 for ((queue, priority, enqueue_shard), lane_rows) in grouped {
3949 let range_start = ready_rows.len();
3950
3951 for row in lane_rows {
3952 let job_id = job_id_iter.next().ok_or_else(|| {
3953 AwaError::Validation("queue storage job id allocation underflow".to_string())
3954 })?;
3955 ready_rows.push(RuntimeReadyInsert {
3956 job_id,
3957 kind: row.kind,
3958 queue: row.queue,
3959 args: row.args,
3960 priority: row.priority,
3961 attempt: row.attempt,
3962 run_lease: row.run_lease,
3963 max_attempts: row.max_attempts,
3964 run_at: row.run_at,
3965 attempted_at: row.attempted_at,
3966 lane_seq: 0,
3967 enqueue_shard,
3968 created_at: row.created_at,
3969 unique_key: row.unique_key,
3970 unique_states: row.unique_states,
3971 payload: row.payload,
3972 });
3973 }
3974 lane_ranges.push((
3975 queue,
3976 priority,
3977 enqueue_shard,
3978 range_start,
3979 ready_rows.len(),
3980 ));
3981 }
3982
3983 self.sync_ready_enqueue_unique_claims(tx, &ready_rows)
3984 .await?;
3985 for (queue, priority, enqueue_shard, range_start, range_end) in lane_ranges {
3986 self.ensure_lane(tx, &queue, priority, enqueue_shard)
3987 .await?;
3988
3989 let count = (range_end - range_start) as i64;
3990 let start_seq = self
3991 .advance_enqueue_head(tx, &queue, priority, enqueue_shard, count)
3992 .await?;
3993
3994 for (offset, row) in ready_rows[range_start..range_end].iter_mut().enumerate() {
3995 row.lane_seq = start_seq + offset as i64;
3996 }
3997 }
3998 let ring = self.current_queue_ring(tx).await?;
3999 self.execute_ready_inserts_tx(tx, ring, &ready_rows).await?;
4000 self.insert_ready_segments_tx(tx, ring, &ready_rows).await?;
4001 Ok(total_rows)
4002 }
4003
4004 async fn group_ready_rows_by_shard<'a>(
4025 &self,
4026 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4027 rows: Vec<RuntimeReadyRow>,
4028 ) -> Result<BTreeMap<(String, i16, i16), Vec<RuntimeReadyRow>>, AwaError> {
4029 let mut by_queue_priority: BTreeMap<(String, i16), Vec<RuntimeReadyRow>> = BTreeMap::new();
4033 for row in rows {
4034 by_queue_priority
4035 .entry((row.queue.clone(), row.priority))
4036 .or_default()
4037 .push(row);
4038 }
4039
4040 let mut grouped: BTreeMap<(String, i16, i16), Vec<RuntimeReadyRow>> = BTreeMap::new();
4041 for ((queue, priority), bucket) in by_queue_priority {
4042 let mut rotor_rows: Vec<RuntimeReadyRow> = Vec::with_capacity(bucket.len());
4043 for row in bucket {
4044 if row.ordering_key.is_some() {
4045 let shard = self
4046 .shard_for_enqueue(tx.as_mut(), &queue, row.ordering_key.as_deref())
4047 .await?;
4048 grouped
4049 .entry((queue.clone(), priority, shard))
4050 .or_default()
4051 .push(row);
4052 } else {
4053 rotor_rows.push(row);
4054 }
4055 }
4056 if !rotor_rows.is_empty() {
4057 let shard = self.shard_for_enqueue(tx.as_mut(), &queue, None).await?;
4058 grouped
4059 .entry((queue.clone(), priority, shard))
4060 .or_default()
4061 .extend(rotor_rows);
4062 }
4063 }
4064 Ok(grouped)
4065 }
4066
4067 async fn insert_ready_rows_copy_tx<'a>(
4068 &self,
4069 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4070 rows: Vec<RuntimeReadyRow>,
4071 job_ids: Vec<i64>,
4072 ) -> Result<usize, AwaError> {
4073 if rows.is_empty() {
4074 return Ok(0);
4075 }
4076
4077 let grouped = self.group_ready_rows_by_shard(tx, rows).await?;
4078
4079 let total_rows: usize = grouped.values().map(Vec::len).sum();
4080 if job_ids.len() != total_rows {
4081 return Err(AwaError::Validation(
4082 "queue storage job id allocation count mismatch".to_string(),
4083 ));
4084 }
4085 let mut job_id_iter = job_ids.into_iter();
4086
4087 let mut ready_rows = Vec::with_capacity(total_rows);
4088 let mut lane_ranges = Vec::with_capacity(grouped.len());
4089
4090 for ((queue, priority, enqueue_shard), lane_rows) in grouped {
4091 let range_start = ready_rows.len();
4092
4093 for row in lane_rows {
4094 let job_id = job_id_iter.next().ok_or_else(|| {
4095 AwaError::Validation("queue storage job id allocation underflow".to_string())
4096 })?;
4097 ready_rows.push(RuntimeReadyInsert {
4098 job_id,
4099 kind: row.kind,
4100 queue: row.queue,
4101 args: row.args,
4102 priority: row.priority,
4103 attempt: row.attempt,
4104 run_lease: row.run_lease,
4105 max_attempts: row.max_attempts,
4106 run_at: row.run_at,
4107 attempted_at: row.attempted_at,
4108 lane_seq: 0,
4109 enqueue_shard,
4110 created_at: row.created_at,
4111 unique_key: row.unique_key,
4112 unique_states: row.unique_states,
4113 payload: row.payload,
4114 });
4115 }
4116 lane_ranges.push((
4117 queue,
4118 priority,
4119 enqueue_shard,
4120 range_start,
4121 ready_rows.len(),
4122 ));
4123 }
4124
4125 self.sync_ready_enqueue_unique_claims(tx, &ready_rows)
4126 .await?;
4127 for (queue, priority, enqueue_shard, range_start, range_end) in lane_ranges {
4128 self.ensure_lane(tx, &queue, priority, enqueue_shard)
4129 .await?;
4130
4131 let count = (range_end - range_start) as i64;
4132 let start_seq = self
4133 .advance_enqueue_head(tx, &queue, priority, enqueue_shard, count)
4134 .await?;
4135
4136 for (offset, row) in ready_rows[range_start..range_end].iter_mut().enumerate() {
4137 row.lane_seq = start_seq + offset as i64;
4138 }
4139 }
4140 let ring = self.current_queue_ring(tx).await?;
4141 self.execute_ready_copy_tx(tx, ring, &ready_rows).await?;
4142 self.insert_ready_segments_tx(tx, ring, &ready_rows).await?;
4143 Ok(total_rows)
4144 }
4145
4146 async fn insert_existing_ready_rows_tx<'a>(
4147 &self,
4148 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4149 rows: Vec<ExistingReadyRow>,
4150 old_state: Option<JobState>,
4151 ) -> Result<usize, AwaError> {
4152 if rows.is_empty() {
4153 return Ok(0);
4154 }
4155
4156 let mut grouped: BTreeMap<(String, i16), Vec<ExistingReadyRow>> = BTreeMap::new();
4157 for row in rows {
4158 grouped
4159 .entry((row.queue.clone(), row.priority))
4160 .or_default()
4161 .push(row);
4162 }
4163
4164 let total_rows: usize = grouped.values().map(Vec::len).sum();
4165 let mut ready_rows = Vec::with_capacity(total_rows);
4166
4167 for ((queue, priority), lane_rows) in grouped {
4168 let enqueue_shard = self.shard_for_enqueue(tx.as_mut(), &queue, None).await?;
4177 self.ensure_lane(tx, &queue, priority, enqueue_shard)
4178 .await?;
4179
4180 let count = lane_rows.len() as i64;
4181 let start_seq = self
4182 .advance_enqueue_head(tx, &queue, priority, enqueue_shard, count)
4183 .await?;
4184
4185 for (offset, row) in lane_rows.into_iter().enumerate() {
4186 self.sync_unique_claim(
4187 tx,
4188 row.job_id,
4189 &row.unique_key,
4190 row.unique_states.as_deref(),
4191 old_state,
4192 Some(JobState::Available),
4193 )
4194 .await?;
4195 ready_rows.push(RuntimeReadyInsert {
4196 job_id: row.job_id,
4197 kind: row.kind,
4198 queue: row.queue,
4199 args: row.args,
4200 priority: row.priority,
4201 attempt: row.attempt,
4202 run_lease: row.run_lease,
4203 max_attempts: row.max_attempts,
4204 run_at: row.run_at,
4205 attempted_at: row.attempted_at,
4206 lane_seq: start_seq + offset as i64,
4207 enqueue_shard,
4208 created_at: row.created_at,
4209 unique_key: row.unique_key,
4210 unique_states: row.unique_states,
4211 payload: row.payload,
4212 });
4213 }
4214 }
4215
4216 let ring = self.current_queue_ring(tx).await?;
4217 self.execute_ready_inserts_tx(tx, ring, &ready_rows).await?;
4218 self.insert_ready_segments_tx(tx, ring, &ready_rows).await?;
4219 Ok(total_rows)
4220 }
4221
4222 async fn insert_deferred_rows_tx<'a>(
4223 &self,
4224 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4225 rows: Vec<DeferredJobRow>,
4226 old_state: Option<JobState>,
4227 ) -> Result<usize, AwaError> {
4228 if rows.is_empty() {
4229 return Ok(0);
4230 }
4231
4232 if old_state.is_none() {
4233 self.sync_deferred_enqueue_unique_claims(tx, &rows).await?;
4234 } else {
4235 for row in &rows {
4236 self.sync_unique_claim(
4237 tx,
4238 row.job_id,
4239 &row.unique_key,
4240 row.unique_states.as_deref(),
4241 old_state,
4242 Some(row.state),
4243 )
4244 .await?;
4245 }
4246 }
4247
4248 let schema = self.schema();
4249 let mut builder = QueryBuilder::<Postgres>::new(format!(
4250 "INSERT INTO {schema}.deferred_jobs (job_id, kind, queue, args, state, priority, attempt, run_lease, max_attempts, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, payload) "
4251 ));
4252 builder.push_values(rows.iter(), |mut b, row| {
4253 b.push_bind(row.job_id)
4254 .push_bind(&row.kind)
4255 .push_bind(&row.queue)
4256 .push_bind(&row.args)
4257 .push_bind(row.state)
4258 .push_bind(row.priority)
4259 .push_bind(row.attempt)
4260 .push_bind(row.run_lease)
4261 .push_bind(row.max_attempts)
4262 .push_bind(row.run_at)
4263 .push_bind(row.attempted_at)
4264 .push_bind(row.finalized_at)
4265 .push_bind(row.created_at)
4266 .push_bind(&row.unique_key)
4267 .push_bind(&row.unique_states)
4268 .push_bind(storage_payload(&row.payload));
4269 });
4270 builder
4271 .build()
4272 .execute(tx.as_mut())
4273 .await
4274 .map_err(map_sqlx_error)?;
4275
4276 Ok(rows.len())
4277 }
4278
4279 async fn insert_deferred_rows_copy_tx<'a>(
4280 &self,
4281 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4282 rows: Vec<DeferredJobRow>,
4283 ) -> Result<usize, AwaError> {
4284 if rows.is_empty() {
4285 return Ok(0);
4286 }
4287
4288 self.sync_deferred_enqueue_unique_claims(tx, &rows).await?;
4289
4290 let schema = self.schema();
4291 let copy_sql = format!(
4292 "COPY {schema}.deferred_jobs (job_id, kind, queue, args, state, priority, attempt, run_lease, max_attempts, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, payload) FROM STDIN WITH (FORMAT csv, NULL '{COPY_NULL_SENTINEL}')"
4293 );
4294 let mut copy_in = tx
4295 .as_mut()
4296 .copy_in_raw(©_sql)
4297 .await
4298 .map_err(map_sqlx_error)?;
4299 let mut csv_buf = Vec::with_capacity(rows.len().min(1024) * 320);
4302 for row in &rows {
4303 write_deferred_copy_row(&mut csv_buf, row);
4304 if csv_buf.len() >= COPY_CHUNK_TARGET_BYTES {
4305 let chunk =
4306 std::mem::replace(&mut csv_buf, Vec::with_capacity(COPY_CHUNK_TARGET_BYTES));
4307 copy_in.send(chunk).await.map_err(map_sqlx_error)?;
4308 }
4309 }
4310 if !csv_buf.is_empty() {
4311 copy_in.send(csv_buf).await.map_err(map_sqlx_error)?;
4312 }
4313 copy_in.finish().await.map_err(map_sqlx_error)?;
4314
4315 Ok(rows.len())
4316 }
4317
4318 async fn insert_done_rows_tx<'a>(
4319 &self,
4320 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4321 rows: &[DoneJobRow],
4322 old_state: Option<JobState>,
4323 ) -> Result<usize, AwaError> {
4324 if rows.is_empty() {
4325 return Ok(0);
4326 }
4327
4328 for row in rows {
4329 self.sync_unique_claim(
4330 tx,
4331 row.job_id,
4332 &row.unique_key,
4333 row.unique_states.as_deref(),
4334 old_state,
4335 Some(row.state),
4336 )
4337 .await?;
4338 }
4339
4340 let schema = self.schema();
4341 let keep_ready_backing = matches!(
4342 old_state,
4343 Some(JobState::Running | JobState::WaitingExternal)
4344 );
4345 let ready_payloads = if rows
4346 .iter()
4347 .any(|row| !is_storage_payload_empty(&row.payload))
4348 {
4349 self.ready_payloads_for_done_rows_tx(tx, rows).await?
4350 } else {
4351 HashMap::new()
4352 };
4353 let mut ordered_rows: Vec<&DoneJobRow> = rows.iter().collect();
4354 ordered_rows.sort_unstable_by_key(|row| {
4355 (
4356 row.ready_slot,
4357 row.ready_generation,
4358 row.queue.as_str(),
4359 row.priority,
4360 row.enqueue_shard,
4361 row.lane_seq,
4362 row.job_id,
4363 )
4364 });
4365 let (ready_backed, synthetic): (Vec<_>, Vec<_>) = ordered_rows
4366 .into_iter()
4367 .partition(|row| keep_ready_backing && row.lane_seq >= 0);
4368
4369 if !ready_backed.is_empty() {
4370 let mut builder = QueryBuilder::<Postgres>::new(format!(
4371 "INSERT INTO {schema}.done_entries (ready_slot, ready_generation, job_id, kind, queue, state, priority, attempt, run_lease, lane_seq, enqueue_shard, attempted_at, finalized_at, payload) "
4372 ));
4373 builder.push_values(ready_backed, |mut b, row| {
4374 let ready_key = (
4375 row.ready_slot,
4376 row.ready_generation,
4377 row.queue.as_str(),
4378 row.priority,
4379 row.enqueue_shard,
4380 row.lane_seq,
4381 );
4382 let ready_payload = ready_payloads.get(&ready_key);
4383 b.push_bind(row.ready_slot)
4384 .push_bind(row.ready_generation)
4385 .push_bind(row.job_id)
4386 .push_bind(&row.kind)
4387 .push_bind(&row.queue)
4388 .push_bind(row.state)
4389 .push_bind(row.priority)
4390 .push_bind(row.attempt)
4391 .push_bind(row.run_lease)
4392 .push_bind(row.lane_seq)
4393 .push_bind(row.enqueue_shard)
4394 .push_bind(row.attempted_at)
4395 .push_bind(row.finalized_at)
4396 .push_bind(terminal_storage_payload(&row.payload, ready_payload));
4397 });
4398 builder
4399 .build()
4400 .execute(tx.as_mut())
4401 .await
4402 .map_err(map_sqlx_error)?;
4403 }
4404
4405 if !synthetic.is_empty() {
4406 let mut builder = QueryBuilder::<Postgres>::new(format!(
4407 "INSERT INTO {schema}.done_entries (ready_slot, ready_generation, job_id, kind, queue, args, state, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, payload) "
4408 ));
4409 builder.push_values(synthetic, |mut b, row| {
4410 let ready_key = (
4411 row.ready_slot,
4412 row.ready_generation,
4413 row.queue.as_str(),
4414 row.priority,
4415 row.enqueue_shard,
4416 row.lane_seq,
4417 );
4418 let ready_payload = ready_payloads.get(&ready_key);
4419 b.push_bind(row.ready_slot)
4420 .push_bind(row.ready_generation)
4421 .push_bind(row.job_id)
4422 .push_bind(&row.kind)
4423 .push_bind(&row.queue)
4424 .push_bind(&row.args)
4425 .push_bind(row.state)
4426 .push_bind(row.priority)
4427 .push_bind(row.attempt)
4428 .push_bind(row.run_lease)
4429 .push_bind(row.max_attempts)
4430 .push_bind(row.lane_seq)
4431 .push_bind(row.enqueue_shard)
4432 .push_bind(row.run_at)
4433 .push_bind(row.attempted_at)
4434 .push_bind(row.finalized_at)
4435 .push_bind(row.created_at)
4436 .push_bind(&row.unique_key)
4437 .push_bind(&row.unique_states)
4438 .push_bind(terminal_storage_payload(&row.payload, ready_payload));
4439 });
4440 builder
4441 .build()
4442 .execute(tx.as_mut())
4443 .await
4444 .map_err(map_sqlx_error)?;
4445 }
4446
4447 self.increment_live_terminal_counters_tx(tx, rows).await?;
4451
4452 Ok(rows.len())
4453 }
4454
4455 async fn increment_live_terminal_counters_tx<'a>(
4471 &self,
4472 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4473 rows: &[DoneJobRow],
4474 ) -> Result<(), AwaError> {
4475 if rows.is_empty() {
4476 return Ok(());
4477 }
4478 let mut by_group: BTreeMap<TerminalCounterKey, i64> = BTreeMap::new();
4479 for row in rows {
4480 let key = (
4481 row.ready_slot,
4482 row.ready_generation,
4483 row.queue.clone(),
4484 row.priority,
4485 row.enqueue_shard,
4486 terminal_counter_bucket(row.job_id),
4487 );
4488 *by_group.entry(key).or_insert(0) += 1;
4489 }
4490 self.append_terminal_count_deltas_tx(tx, by_group).await
4491 }
4492
4493 async fn decrement_live_terminal_counters_tx<'a>(
4501 &self,
4502 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4503 rows: &[TerminalCounterKey],
4504 ) -> Result<(), AwaError> {
4505 if rows.is_empty() {
4506 return Ok(());
4507 }
4508 let mut by_group: BTreeMap<TerminalCounterKey, i64> = BTreeMap::new();
4509 for key in rows {
4510 *by_group.entry(key.clone()).or_insert(0) -= 1;
4511 }
4512 self.append_terminal_count_deltas_tx(tx, by_group).await
4513 }
4514
4515 async fn append_terminal_count_deltas_tx<'a>(
4516 &self,
4517 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4518 by_group: BTreeMap<TerminalCounterKey, i64>,
4519 ) -> Result<(), AwaError> {
4520 if by_group.is_empty() {
4521 return Ok(());
4522 }
4523
4524 let mut ready_slots: Vec<i32> = Vec::with_capacity(by_group.len());
4525 let mut ready_generations: Vec<i64> = Vec::with_capacity(by_group.len());
4526 let mut queues: Vec<String> = Vec::with_capacity(by_group.len());
4527 let mut priorities: Vec<i16> = Vec::with_capacity(by_group.len());
4528 let mut enqueue_shards: Vec<i16> = Vec::with_capacity(by_group.len());
4529 let mut counter_buckets: Vec<i16> = Vec::with_capacity(by_group.len());
4530 let mut deltas: Vec<i64> = Vec::with_capacity(by_group.len());
4531 for ((slot, generation, queue, prio, shard, bucket), delta) in by_group {
4532 if delta == 0 {
4533 continue;
4534 }
4535 ready_slots.push(slot);
4536 ready_generations.push(generation);
4537 queues.push(queue);
4538 priorities.push(prio);
4539 enqueue_shards.push(shard);
4540 counter_buckets.push(bucket);
4541 deltas.push(delta);
4542 }
4543
4544 if deltas.is_empty() {
4545 return Ok(());
4546 }
4547
4548 let schema = self.schema();
4549 sqlx::query(&format!(
4550 r#"
4551 INSERT INTO {schema}.queue_terminal_count_deltas (
4552 ready_slot,
4553 ready_generation,
4554 queue,
4555 priority,
4556 enqueue_shard,
4557 counter_bucket,
4558 terminal_delta
4559 )
4560 SELECT
4561 ready_slot,
4562 ready_generation,
4563 queue,
4564 priority,
4565 enqueue_shard,
4566 counter_bucket,
4567 terminal_delta
4568 FROM unnest(
4569 $1::int[],
4570 $2::bigint[],
4571 $3::text[],
4572 $4::smallint[],
4573 $5::smallint[],
4574 $6::smallint[],
4575 $7::bigint[]
4576 ) AS d(
4577 ready_slot,
4578 ready_generation,
4579 queue,
4580 priority,
4581 enqueue_shard,
4582 counter_bucket,
4583 terminal_delta
4584 )
4585 ORDER BY
4586 ready_slot,
4587 ready_generation,
4588 queue,
4589 priority,
4590 enqueue_shard,
4591 counter_bucket
4592 "#
4593 ))
4594 .bind(&ready_slots)
4595 .bind(&ready_generations)
4596 .bind(&queues)
4597 .bind(&priorities)
4598 .bind(&enqueue_shards)
4599 .bind(&counter_buckets)
4600 .bind(&deltas)
4601 .execute(tx.as_mut())
4602 .await
4603 .map_err(map_sqlx_error)?;
4604 Ok(())
4605 }
4606
4607 fn done_rows_to_counter_keys(rows: &[DoneJobRow]) -> Vec<TerminalCounterKey> {
4612 rows.iter()
4613 .map(|row| {
4614 (
4615 row.ready_slot,
4616 row.ready_generation,
4617 row.queue.clone(),
4618 row.priority,
4619 row.enqueue_shard,
4620 terminal_counter_bucket(row.job_id),
4621 )
4622 })
4623 .collect()
4624 }
4625
4626 async fn ensure_terminal_removed_receipt_closures_tx<'a>(
4627 &self,
4628 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4629 rows: &[DoneJobRow],
4630 ) -> Result<(), AwaError> {
4631 if rows.is_empty() || !self.lease_claim_receipts() {
4632 return Ok(());
4633 }
4634
4635 let receipt_pairs: Vec<(i64, i64)> =
4636 rows.iter().map(|row| (row.job_id, row.run_lease)).collect();
4637 self.lock_receipt_attempts_tx(tx, &receipt_pairs).await?;
4638
4639 let schema = self.schema();
4640 let job_ids: Vec<i64> = rows.iter().map(|row| row.job_id).collect();
4641 let run_leases: Vec<i64> = rows.iter().map(|row| row.run_lease).collect();
4642 sqlx::query(&format!(
4643 r#"
4644 WITH refs(job_id, run_lease) AS (
4645 SELECT * FROM unnest($1::bigint[], $2::bigint[])
4646 ),
4647 inserted AS (
4648 INSERT INTO {schema}.lease_claim_closures (
4649 claim_slot,
4650 job_id,
4651 run_lease,
4652 outcome,
4653 closed_at
4654 )
4655 SELECT
4656 claims.claim_slot,
4657 claims.job_id,
4658 claims.run_lease,
4659 'terminal_removed',
4660 clock_timestamp()
4661 FROM {schema}.lease_claims AS claims
4662 JOIN refs
4663 ON refs.job_id = claims.job_id
4664 AND refs.run_lease = claims.run_lease
4665 ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING
4666 RETURNING claim_slot, job_id, run_lease, closed_at
4667 ),
4668 marked AS (
4669 UPDATE {schema}.lease_claims AS claims
4670 SET closed_at = COALESCE(claims.closed_at, inserted.closed_at)
4671 FROM inserted
4672 WHERE claims.claim_slot = inserted.claim_slot
4673 AND claims.job_id = inserted.job_id
4674 AND claims.run_lease = inserted.run_lease
4675 RETURNING claims.job_id
4676 )
4677 SELECT count(*) FROM marked
4678 "#
4679 ))
4680 .bind(&job_ids)
4681 .bind(&run_leases)
4682 .execute(tx.as_mut())
4683 .await
4684 .map_err(map_sqlx_error)?;
4685 Ok(())
4686 }
4687
4688 async fn ready_payloads_for_done_rows_tx<'a, 'r>(
4689 &self,
4690 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4691 rows: &'r [DoneJobRow],
4692 ) -> Result<HashMap<(i32, i64, &'r str, i16, i16, i64), serde_json::Value>, AwaError> {
4693 if rows.is_empty() {
4694 return Ok(HashMap::new());
4695 }
4696
4697 let schema = self.schema();
4698 let ready_slots: Vec<i32> = rows.iter().map(|row| row.ready_slot).collect();
4699 let ready_generations: Vec<i64> = rows.iter().map(|row| row.ready_generation).collect();
4700 let queues: Vec<&str> = rows.iter().map(|row| row.queue.as_str()).collect();
4701 let priorities: Vec<i16> = rows.iter().map(|row| row.priority).collect();
4702 let enqueue_shards: Vec<i16> = rows.iter().map(|row| row.enqueue_shard).collect();
4703 let lane_seqs: Vec<i64> = rows.iter().map(|row| row.lane_seq).collect();
4704
4705 let payload_rows: Vec<(i32, i64, String, i16, i16, i64, serde_json::Value)> =
4706 sqlx::query_as(&format!(
4707 r#"
4708 WITH refs(ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq) AS (
4709 SELECT * FROM unnest($1::int[], $2::bigint[], $3::text[], $4::smallint[], $5::smallint[], $6::bigint[])
4710 )
4711 SELECT
4712 ready.ready_slot,
4713 ready.ready_generation,
4714 ready.queue,
4715 ready.priority,
4716 ready.enqueue_shard,
4717 ready.lane_seq,
4718 COALESCE(ready.payload, '{{}}'::jsonb) AS payload
4719 FROM refs
4720 JOIN {schema}.ready_entries AS ready
4721 ON ready.ready_slot = refs.ready_slot
4722 AND ready.ready_generation = refs.ready_generation
4723 AND ready.queue = refs.queue
4724 AND ready.priority = refs.priority
4725 AND ready.enqueue_shard = refs.enqueue_shard
4726 AND ready.lane_seq = refs.lane_seq
4727 "#
4728 ))
4729 .bind(&ready_slots)
4730 .bind(&ready_generations)
4731 .bind(&queues)
4732 .bind(&priorities)
4733 .bind(&enqueue_shards)
4734 .bind(&lane_seqs)
4735 .fetch_all(tx.as_mut())
4736 .await
4737 .map_err(map_sqlx_error)?;
4738
4739 let mut payload_by_owned_key = HashMap::with_capacity(payload_rows.len());
4740 for (ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, payload) in
4741 payload_rows
4742 {
4743 payload_by_owned_key.insert(
4744 (
4745 ready_slot,
4746 ready_generation,
4747 queue,
4748 priority,
4749 enqueue_shard,
4750 lane_seq,
4751 ),
4752 payload,
4753 );
4754 }
4755
4756 let mut payloads = HashMap::with_capacity(payload_by_owned_key.len());
4757 for row in rows {
4758 if let Some(payload) = payload_by_owned_key.remove(&(
4759 row.ready_slot,
4760 row.ready_generation,
4761 row.queue.clone(),
4762 row.priority,
4763 row.enqueue_shard,
4764 row.lane_seq,
4765 )) {
4766 payloads.insert(
4767 (
4768 row.ready_slot,
4769 row.ready_generation,
4770 row.queue.as_str(),
4771 row.priority,
4772 row.enqueue_shard,
4773 row.lane_seq,
4774 ),
4775 payload,
4776 );
4777 }
4778 }
4779 Ok(payloads)
4780 }
4781
4782 async fn insert_dlq_rows_tx<'a>(
4783 &self,
4784 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4785 rows: &[DlqJobRow],
4786 old_state: Option<JobState>,
4787 ) -> Result<usize, AwaError> {
4788 if rows.is_empty() {
4789 return Ok(0);
4790 }
4791
4792 for row in rows {
4793 self.sync_unique_claim(
4794 tx,
4795 row.job_id,
4796 &row.unique_key,
4797 row.unique_states.as_deref(),
4798 old_state,
4799 Some(JobState::Failed),
4800 )
4801 .await?;
4802 }
4803
4804 let schema = self.schema();
4805 let mut builder = QueryBuilder::<Postgres>::new(format!(
4806 "INSERT INTO {schema}.dlq_entries (job_id, kind, queue, args, state, priority, attempt, run_lease, max_attempts, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, payload, dlq_reason, dlq_at, original_run_lease) "
4807 ));
4808 builder.push_values(rows.iter(), |mut b, row| {
4809 b.push_bind(row.job_id)
4810 .push_bind(&row.kind)
4811 .push_bind(&row.queue)
4812 .push_bind(&row.args)
4813 .push_bind(row.state)
4814 .push_bind(row.priority)
4815 .push_bind(row.attempt)
4816 .push_bind(row.run_lease)
4817 .push_bind(row.max_attempts)
4818 .push_bind(row.run_at)
4819 .push_bind(row.attempted_at)
4820 .push_bind(row.finalized_at)
4821 .push_bind(row.created_at)
4822 .push_bind(&row.unique_key)
4823 .push_bind(&row.unique_states)
4824 .push_bind(storage_payload(&row.payload))
4825 .push_bind(&row.dlq_reason)
4826 .push_bind(row.dlq_at)
4827 .push_bind(row.original_run_lease);
4828 });
4829 builder
4830 .build()
4831 .execute(tx.as_mut())
4832 .await
4833 .map_err(map_sqlx_error)?;
4834
4835 Ok(rows.len())
4836 }
4837
4838 async fn adjust_terminal_rollups_batch<'a, I>(
4839 &self,
4840 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4841 deltas: I,
4842 ) -> Result<(), AwaError>
4843 where
4844 I: IntoIterator<Item = (String, i16, i64, i64)>,
4845 {
4846 let mut grouped: BTreeMap<(String, i16), (i64, i64)> = BTreeMap::new();
4847 for (queue, priority, pruned_completed_delta, pruned_failed_delta) in deltas {
4848 if pruned_completed_delta == 0 && pruned_failed_delta == 0 {
4849 continue;
4850 }
4851 let entry = grouped.entry((queue, priority)).or_insert((0_i64, 0_i64));
4852 entry.0 += pruned_completed_delta;
4853 entry.1 += pruned_failed_delta;
4854 }
4855
4856 if grouped.is_empty() {
4857 return Ok(());
4858 }
4859
4860 let schema = self.schema();
4861 let mut queues = Vec::with_capacity(grouped.len());
4862 let mut priorities = Vec::with_capacity(grouped.len());
4863 let mut pruned_completed_deltas = Vec::with_capacity(grouped.len());
4864 let mut pruned_failed_deltas = Vec::with_capacity(grouped.len());
4865
4866 for ((queue, priority), (pruned_completed_delta, pruned_failed_delta)) in grouped {
4867 queues.push(queue);
4868 priorities.push(priority);
4869 pruned_completed_deltas.push(pruned_completed_delta);
4870 pruned_failed_deltas.push(pruned_failed_delta);
4871 }
4872
4873 sqlx::query(&format!(
4874 r#"
4875 WITH deltas(queue, priority, pruned_completed_delta, pruned_failed_delta) AS (
4876 SELECT *
4877 FROM unnest(
4878 $1::text[],
4879 $2::smallint[],
4880 $3::bigint[],
4881 $4::bigint[]
4882 )
4883 )
4884 INSERT INTO {schema}.queue_terminal_rollups AS rollups (
4885 queue,
4886 priority,
4887 pruned_completed_count,
4888 pruned_failed_count
4889 )
4890 SELECT
4891 deltas.queue,
4892 deltas.priority,
4893 deltas.pruned_completed_delta,
4894 deltas.pruned_failed_delta
4895 FROM deltas
4896 ON CONFLICT (queue, priority) DO UPDATE
4897 SET pruned_completed_count = GREATEST(
4898 0,
4899 rollups.pruned_completed_count + EXCLUDED.pruned_completed_count
4900 ),
4901 pruned_failed_count = GREATEST(
4902 0,
4903 rollups.pruned_failed_count + EXCLUDED.pruned_failed_count
4904 )
4905 "#
4906 ))
4907 .bind(&queues)
4908 .bind(&priorities)
4909 .bind(&pruned_completed_deltas)
4910 .bind(&pruned_failed_deltas)
4911 .execute(tx.as_mut())
4912 .await
4913 .map_err(map_sqlx_error)?;
4914 Ok(())
4915 }
4916
4917 async fn enqueue_runtime_rows(
4918 &self,
4919 pool: &PgPool,
4920 rows: Vec<RuntimeReadyRow>,
4921 ) -> Result<usize, AwaError> {
4922 if rows.is_empty() {
4923 return Ok(0);
4924 }
4925
4926 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
4927 let total_rows = self.insert_ready_rows_tx(&mut tx, rows.clone()).await?;
4928
4929 let queues_to_notify: Vec<String> = rows.iter().map(|row| row.queue.clone()).collect();
4930 self.notify_queues_tx(&mut tx, queues_to_notify).await?;
4931
4932 tx.commit().await.map_err(map_sqlx_error)?;
4933 Ok(total_rows)
4934 }
4935
4936 pub async fn enqueue_batch(
4937 &self,
4938 pool: &PgPool,
4939 queue: &str,
4940 priority: i16,
4941 count: i64,
4942 ) -> Result<i64, AwaError> {
4943 if count <= 0 {
4944 return Ok(0);
4945 }
4946
4947 let rows: Vec<_> = (0..count)
4948 .map(|seq| RuntimeReadyRow {
4949 kind: "bench_job".to_string(),
4950 queue: if self.uses_queue_striping() && !self.is_physical_stripe_queue(queue) {
4951 self.physical_queue_for_stripe(
4952 queue,
4953 seq.rem_euclid(self.queue_stripe_count() as i64) as usize,
4954 )
4955 } else {
4956 queue.to_string()
4957 },
4958 args: serde_json::json!({ "seq": seq }),
4959 priority,
4960 attempt: 0,
4961 run_lease: 0,
4962 max_attempts: 25,
4963 run_at: Utc::now(),
4964 attempted_at: None,
4965 created_at: Utc::now(),
4966 unique_key: None,
4967 unique_states: None,
4968 payload: RuntimePayload::default().into_json(),
4969 ordering_key: None,
4970 })
4971 .collect();
4972 self.enqueue_runtime_rows(pool, rows)
4973 .await
4974 .map(|count| count as i64)
4975 }
4976
4977 pub async fn enqueue_params_batch(
4978 &self,
4979 pool: &PgPool,
4980 jobs: &[InsertParams],
4981 ) -> Result<usize, AwaError> {
4982 if jobs.is_empty() {
4983 return Ok(0);
4984 }
4985
4986 let now = Utc::now();
4987 let mut ready_rows = Vec::new();
4988 let mut deferred_rows = Vec::new();
4989
4990 for (idx, job) in jobs.iter().enumerate() {
4991 let prepared = prepare_row_raw(job.kind.clone(), job.args.clone(), job.opts.clone())?;
4992 let payload = Self::payload_from_parts(prepared.metadata, prepared.tags, None, None)?;
4993 let queue =
4994 self.queue_stripe_for_enqueue(&prepared.queue, &prepared.unique_key, idx as i64);
4995
4996 let ready_row = RuntimeReadyRow {
4997 kind: prepared.kind,
4998 queue: queue.clone(),
4999 args: prepared.args,
5000 priority: prepared.priority,
5001 attempt: 0,
5002 run_lease: 0,
5003 max_attempts: prepared.max_attempts,
5004 run_at: prepared.run_at.unwrap_or(now),
5005 attempted_at: None,
5006 created_at: now,
5007 unique_key: prepared.unique_key,
5008 unique_states: prepared.unique_states,
5009 payload: payload.clone(),
5010 ordering_key: prepared.ordering_key,
5011 };
5012
5013 match prepared.state {
5014 JobState::Available => ready_rows.push(ready_row),
5015 JobState::Scheduled => deferred_rows.push(DeferredJobRow {
5016 job_id: 0,
5017 kind: ready_row.kind,
5018 queue,
5019 args: ready_row.args,
5020 state: JobState::Scheduled,
5021 priority: ready_row.priority,
5022 attempt: ready_row.attempt,
5023 run_lease: ready_row.run_lease,
5024 max_attempts: ready_row.max_attempts,
5025 run_at: ready_row.run_at,
5026 attempted_at: ready_row.attempted_at,
5027 finalized_at: None,
5028 created_at: ready_row.created_at,
5029 unique_key: ready_row.unique_key,
5030 unique_states: ready_row.unique_states,
5031 payload: payload.clone(),
5032 }),
5033 other => {
5034 return Err(AwaError::Validation(format!(
5035 "queue storage does not support initial state {other}"
5036 )));
5037 }
5038 }
5039 }
5040
5041 let queues_to_notify: Vec<String> =
5042 ready_rows.iter().map(|row| row.queue.clone()).collect();
5043
5044 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
5045 let mut total = 0usize;
5046 if !ready_rows.is_empty() {
5047 total += self
5048 .insert_ready_rows_tx(&mut tx, ready_rows.clone())
5049 .await?;
5050 }
5051 if !deferred_rows.is_empty() {
5052 let ids = self.next_job_ids(&mut tx, deferred_rows.len()).await?;
5053 let deferred_rows: Vec<_> = deferred_rows
5054 .into_iter()
5055 .zip(ids)
5056 .map(|(row, id)| DeferredJobRow { job_id: id, ..row })
5057 .collect();
5058 total += self
5059 .insert_deferred_rows_tx(&mut tx, deferred_rows, None)
5060 .await?;
5061 }
5062
5063 self.notify_queues_tx(&mut tx, queues_to_notify).await?;
5064
5065 tx.commit().await.map_err(map_sqlx_error)?;
5066 Ok(total)
5067 }
5068
5069 #[tracing::instrument(skip(self, pool, jobs), fields(job.count = jobs.len()), name = "queue_storage.enqueue_params_copy")]
5076 pub async fn enqueue_params_copy(
5077 &self,
5078 pool: &PgPool,
5079 jobs: &[InsertParams],
5080 ) -> Result<usize, AwaError> {
5081 if jobs.is_empty() {
5082 return Ok(0);
5083 }
5084
5085 let now = Utc::now();
5086 let mut ready_rows = Vec::new();
5087 let mut deferred_rows = Vec::new();
5088
5089 for (idx, job) in jobs.iter().enumerate() {
5090 let prepared = prepare_row_raw(job.kind.clone(), job.args.clone(), job.opts.clone())?;
5091 let payload = Self::payload_from_parts(prepared.metadata, prepared.tags, None, None)?;
5092 let queue =
5093 self.queue_stripe_for_enqueue(&prepared.queue, &prepared.unique_key, idx as i64);
5094
5095 let ready_row = RuntimeReadyRow {
5096 kind: prepared.kind,
5097 queue: queue.clone(),
5098 args: prepared.args,
5099 priority: prepared.priority,
5100 attempt: 0,
5101 run_lease: 0,
5102 max_attempts: prepared.max_attempts,
5103 run_at: prepared.run_at.unwrap_or(now),
5104 attempted_at: None,
5105 created_at: now,
5106 unique_key: prepared.unique_key,
5107 unique_states: prepared.unique_states,
5108 payload: payload.clone(),
5109 ordering_key: prepared.ordering_key,
5110 };
5111
5112 match prepared.state {
5113 JobState::Available => ready_rows.push(ready_row),
5114 JobState::Scheduled => deferred_rows.push(DeferredJobRow {
5115 job_id: 0,
5116 kind: ready_row.kind,
5117 queue,
5118 args: ready_row.args,
5119 state: JobState::Scheduled,
5120 priority: ready_row.priority,
5121 attempt: ready_row.attempt,
5122 run_lease: ready_row.run_lease,
5123 max_attempts: ready_row.max_attempts,
5124 run_at: ready_row.run_at,
5125 attempted_at: ready_row.attempted_at,
5126 finalized_at: None,
5127 created_at: ready_row.created_at,
5128 unique_key: ready_row.unique_key,
5129 unique_states: ready_row.unique_states,
5130 payload: payload.clone(),
5131 }),
5132 other => {
5133 return Err(AwaError::Validation(format!(
5134 "queue storage does not support initial state {other}"
5135 )));
5136 }
5137 }
5138 }
5139
5140 let queues_to_notify: Vec<String> =
5141 ready_rows.iter().map(|row| row.queue.clone()).collect();
5142
5143 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
5144 let mut total = 0usize;
5145 let job_ids = self
5146 .next_job_ids(&mut tx, ready_rows.len() + deferred_rows.len())
5147 .await?;
5148 let (ready_job_ids, deferred_job_ids) = job_ids.split_at(ready_rows.len());
5149 if !ready_rows.is_empty() {
5150 total += self
5151 .insert_ready_rows_copy_tx(&mut tx, ready_rows, ready_job_ids.to_vec())
5152 .await?;
5153 }
5154 if !deferred_rows.is_empty() {
5155 let deferred_rows: Vec<_> = deferred_rows
5156 .into_iter()
5157 .zip(deferred_job_ids.iter().copied())
5158 .map(|(row, id)| DeferredJobRow { job_id: id, ..row })
5159 .collect();
5160 total += self
5161 .insert_deferred_rows_copy_tx(&mut tx, deferred_rows)
5162 .await?;
5163 }
5164
5165 self.notify_queues_tx(&mut tx, queues_to_notify).await?;
5166
5167 tx.commit().await.map_err(map_sqlx_error)?;
5168 Ok(total)
5169 }
5170
5171 #[tracing::instrument(skip(self, pool), name = "queue_storage.claim_batch")]
5172 pub async fn claim_batch(
5173 &self,
5174 pool: &PgPool,
5175 queue: &str,
5176 max_batch: i64,
5177 ) -> Result<Vec<ClaimedEntry>, AwaError> {
5178 if max_batch <= 0 {
5179 return Ok(Vec::new());
5180 }
5181
5182 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
5183 let mut claimed_rows = Vec::new();
5184 let stripe_queues = self.physical_queues_for_logical(queue);
5185 let start = self.stripe_probe_start(stripe_queues.len());
5186 for offset in 0..stripe_queues.len() {
5187 if claimed_rows.len() >= max_batch as usize {
5188 break;
5189 }
5190 let stripe_queue = &stripe_queues[(start + offset) % stripe_queues.len()];
5191 let remaining = max_batch - claimed_rows.len() as i64;
5192 claimed_rows.extend(
5193 self.claim_ready_rows_tx(
5194 &mut tx,
5195 stripe_queue,
5196 remaining,
5197 Duration::ZERO,
5198 Duration::ZERO,
5199 )
5200 .await?,
5201 );
5202 }
5203 let claim_cursor_advances = Self::claim_cursor_advances(&claimed_rows);
5204 let claimed = claimed_rows
5205 .into_iter()
5206 .map(|row| row.claim_ref(self.lease_claim_receipts()))
5207 .collect();
5208
5209 tx.commit().await.map_err(map_sqlx_error)?;
5210 self.advance_claim_cursors(pool, &claim_cursor_advances)
5211 .await;
5212 Ok(claimed)
5213 }
5214
5215 #[tracing::instrument(skip(self, pool), fields(queue = %queue), name = "queue_storage.claim_runtime_batch")]
5216 pub async fn claim_runtime_batch(
5217 &self,
5218 pool: &PgPool,
5219 queue: &str,
5220 max_batch: i64,
5221 deadline_duration: Duration,
5222 ) -> Result<Vec<ClaimedRuntimeJob>, AwaError> {
5223 self.claim_runtime_batch_with_aging(
5224 pool,
5225 queue,
5226 max_batch,
5227 deadline_duration,
5228 Duration::ZERO,
5229 )
5230 .await
5231 }
5232
5233 #[tracing::instrument(skip(self, pool), fields(queue = %queue), name = "queue_storage.claim_runtime_batch_with_aging")]
5234 pub async fn claim_runtime_batch_with_aging(
5235 &self,
5236 pool: &PgPool,
5237 queue: &str,
5238 max_batch: i64,
5239 deadline_duration: Duration,
5240 aging_interval: Duration,
5241 ) -> Result<Vec<ClaimedRuntimeJob>, AwaError> {
5242 if max_batch <= 0 {
5243 return Ok(Vec::new());
5244 }
5245
5246 let stripe_queues = self.physical_queues_for_logical(queue);
5247 if stripe_queues.len() > 1 {
5248 let mut claimed = Vec::new();
5249 let start = self.stripe_probe_start(stripe_queues.len());
5250 for offset in 0..stripe_queues.len() {
5251 if claimed.len() >= max_batch as usize {
5252 break;
5253 }
5254 let stripe_queue = &stripe_queues[(start + offset) % stripe_queues.len()];
5255 let remaining = max_batch - claimed.len() as i64;
5256 match self
5257 .claim_runtime_batch_with_aging_physical(
5258 pool,
5259 stripe_queue,
5260 remaining,
5261 deadline_duration,
5262 aging_interval,
5263 )
5264 .await
5265 {
5266 Ok(stripe_claims) => claimed.extend(stripe_claims),
5267 Err(err) if claimed.is_empty() => return Err(err),
5268 Err(err) => {
5269 tracing::warn!(
5270 queue = %queue,
5271 stripe_queue = %stripe_queue,
5272 claimed = claimed.len(),
5273 error = ?err,
5274 "returning already-claimed runtime jobs after striped claim error"
5275 );
5276 break;
5277 }
5278 }
5279 }
5280 return Ok(claimed);
5281 }
5282
5283 self.claim_runtime_batch_with_aging_physical(
5284 pool,
5285 &stripe_queues[0],
5286 max_batch,
5287 deadline_duration,
5288 aging_interval,
5289 )
5290 .await
5291 }
5292
5293 async fn claim_runtime_batch_with_aging_physical(
5294 &self,
5295 pool: &PgPool,
5296 queue: &str,
5297 max_batch: i64,
5298 deadline_duration: Duration,
5299 aging_interval: Duration,
5300 ) -> Result<Vec<ClaimedRuntimeJob>, AwaError> {
5301 if max_batch <= 0 {
5302 return Ok(Vec::new());
5303 }
5304
5305 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
5306 let mut claimed = Vec::new();
5307 claimed.extend(
5308 self.claim_ready_rows_tx(&mut tx, queue, max_batch, deadline_duration, aging_interval)
5309 .await?,
5310 );
5311
5312 for row in &claimed {
5313 self.sync_unique_claim(
5314 &mut tx,
5315 row.job_id,
5316 &row.unique_key,
5317 row.unique_states.as_deref(),
5318 Some(JobState::Available),
5319 Some(JobState::Running),
5320 )
5321 .await?;
5322 }
5323
5324 let use_lease_claim_receipts = self.use_lease_claim_receipts_for_runtime(deadline_duration);
5325 if !use_lease_claim_receipts && deadline_duration.is_zero() {
5326 let converted = claimed
5331 .iter()
5332 .cloned()
5333 .map(|row| row.into_claimed_runtime_job(use_lease_claim_receipts))
5334 .collect::<Result<Vec<_>, _>>()?;
5335 let claim_cursor_advances = Self::claim_cursor_advances(&claimed);
5336 tx.commit().await.map_err(map_sqlx_error)?;
5337 self.advance_claim_cursors(pool, &claim_cursor_advances)
5338 .await;
5339 return Ok(converted);
5340 }
5341
5342 let claim_cursor_advances = Self::claim_cursor_advances(&claimed);
5343
5344 tx.commit().await.map_err(map_sqlx_error)?;
5347 self.advance_claim_cursors(pool, &claim_cursor_advances)
5348 .await;
5349
5350 claimed
5351 .into_iter()
5352 .map(|row| row.into_claimed_runtime_job(use_lease_claim_receipts))
5353 .collect()
5354 }
5355
5356 #[tracing::instrument(skip(self, pool), fields(queue = %queue, instance_id = %instance_id), name = "queue_storage.acquire_queue_claimer")]
5357 pub async fn acquire_queue_claimer(
5358 &self,
5359 pool: &PgPool,
5360 queue: &str,
5361 instance_id: Uuid,
5362 max_claimers: i16,
5363 lease_ttl: Duration,
5364 idle_threshold: Duration,
5365 ) -> Result<Option<QueueClaimerLease>, AwaError> {
5366 Ok(self
5367 .acquire_queue_claimer_row(
5368 pool,
5369 queue,
5370 instance_id,
5371 max_claimers,
5372 lease_ttl,
5373 idle_threshold,
5374 )
5375 .await?
5376 .map(QueueClaimerLeaseRow::lease))
5377 }
5378
5379 async fn acquire_queue_claimer_row(
5380 &self,
5381 pool: &PgPool,
5382 queue: &str,
5383 instance_id: Uuid,
5384 max_claimers: i16,
5385 lease_ttl: Duration,
5386 idle_threshold: Duration,
5387 ) -> Result<Option<QueueClaimerLeaseRow>, AwaError> {
5388 if max_claimers <= 0 {
5389 return Ok(None);
5390 }
5391
5392 let schema = self.schema();
5393 let now = Utc::now();
5394 let expires_at = now
5395 + TimeDelta::from_std(lease_ttl)
5396 .map_err(|err| AwaError::Validation(format!("invalid claimer lease ttl: {err}")))?;
5397 let idle_cutoff = now
5398 - TimeDelta::from_std(idle_threshold).map_err(|err| {
5399 AwaError::Validation(format!("invalid claimer idle threshold: {err}"))
5400 })?;
5401 let probe_start = if max_claimers > 1 {
5402 ((instance_id.as_u128() ^ (now.timestamp_millis() as u128)) % (max_claimers as u128))
5403 as i16
5404 } else {
5405 0
5406 };
5407
5408 if let Some(owned) = sqlx::query_as::<_, QueueClaimerLeaseRow>(&format!(
5409 r#"
5410 SELECT claimer_slot, lease_epoch, last_claimed_at, expires_at
5411 FROM {schema}.queue_claimer_leases
5412 WHERE queue = $1
5413 AND owner_instance_id = $2
5414 AND expires_at > $3
5415 ORDER BY claimer_slot
5416 LIMIT 1
5417 "#
5418 ))
5419 .bind(queue)
5420 .bind(instance_id)
5421 .bind(now)
5422 .fetch_optional(pool)
5423 .await
5424 .map_err(map_sqlx_error)?
5425 {
5426 return Ok(Some(owned));
5427 }
5428
5429 for offset in 0..max_claimers {
5430 let slot = (probe_start + offset) % max_claimers;
5431 if let Some(updated) = sqlx::query_as::<_, QueueClaimerLeaseRow>(&format!(
5432 r#"
5433 UPDATE {schema}.queue_claimer_leases
5434 SET owner_instance_id = $3,
5435 lease_epoch = CASE
5436 WHEN owner_instance_id = $3 THEN lease_epoch
5437 ELSE lease_epoch + 1
5438 END,
5439 leased_at = $4,
5440 last_claimed_at = $4,
5441 expires_at = $5
5442 WHERE queue = $1
5443 AND claimer_slot = $2
5444 AND (
5445 owner_instance_id = $3
5446 OR expires_at <= $4
5447 OR last_claimed_at <= $6
5448 )
5449 RETURNING claimer_slot, lease_epoch, last_claimed_at, expires_at
5450 "#
5451 ))
5452 .bind(queue)
5453 .bind(slot)
5454 .bind(instance_id)
5455 .bind(now)
5456 .bind(expires_at)
5457 .bind(idle_cutoff)
5458 .fetch_optional(pool)
5459 .await
5460 .map_err(map_sqlx_error)?
5461 {
5462 return Ok(Some(updated));
5463 }
5464
5465 if let Some(inserted) = sqlx::query_as::<_, QueueClaimerLeaseRow>(&format!(
5466 r#"
5467 INSERT INTO {schema}.queue_claimer_leases (
5468 queue,
5469 claimer_slot,
5470 owner_instance_id,
5471 lease_epoch,
5472 leased_at,
5473 last_claimed_at,
5474 expires_at
5475 )
5476 VALUES ($1, $2, $3, 0, $4, $4, $5)
5477 ON CONFLICT (queue, claimer_slot) DO NOTHING
5478 RETURNING claimer_slot, lease_epoch, last_claimed_at, expires_at
5479 "#
5480 ))
5481 .bind(queue)
5482 .bind(slot)
5483 .bind(instance_id)
5484 .bind(now)
5485 .bind(expires_at)
5486 .fetch_optional(pool)
5487 .await
5488 .map_err(map_sqlx_error)?
5489 {
5490 return Ok(Some(inserted));
5491 }
5492 }
5493
5494 Ok(None)
5495 }
5496
5497 #[tracing::instrument(skip(self, pool), fields(queue = %queue, instance_id = %instance_id, claimer_slot = lease.claimer_slot), name = "queue_storage.mark_queue_claimer_active")]
5498 pub async fn mark_queue_claimer_active(
5499 &self,
5500 pool: &PgPool,
5501 queue: &str,
5502 instance_id: Uuid,
5503 lease: QueueClaimerLease,
5504 lease_ttl: Duration,
5505 ) -> Result<bool, AwaError> {
5506 let schema = self.schema();
5507 let now = Utc::now();
5508 let expires_at = now
5509 + TimeDelta::from_std(lease_ttl)
5510 .map_err(|err| AwaError::Validation(format!("invalid claimer lease ttl: {err}")))?;
5511
5512 let result = sqlx::query(&format!(
5513 r#"
5514 UPDATE {schema}.queue_claimer_leases
5515 SET last_claimed_at = $5,
5516 expires_at = $6
5517 WHERE queue = $1
5518 AND claimer_slot = $2
5519 AND owner_instance_id = $3
5520 AND lease_epoch = $4
5521 "#
5522 ))
5523 .bind(queue)
5524 .bind(lease.claimer_slot)
5525 .bind(instance_id)
5526 .bind(lease.lease_epoch)
5527 .bind(now)
5528 .bind(expires_at)
5529 .execute(pool)
5530 .await
5531 .map_err(map_sqlx_error)?;
5532
5533 Ok(result.rows_affected() == 1)
5534 }
5535
5536 fn desired_queue_claimer_target(
5537 &self,
5538 current_target: Option<i16>,
5539 signal: &AvailableSignal,
5540 max_claimers: i16,
5541 ) -> i16 {
5542 let available = signal.available.max(0) as u64;
5550 let backlog = available;
5551 let current = current_target.unwrap_or(1).clamp(1, max_claimers.max(1));
5552 let max_four = 4.min(max_claimers.max(1));
5553 let max_two = 2.min(max_claimers.max(1));
5554
5555 match current {
5556 4.. => {
5557 if available >= 32 || backlog >= 16 {
5558 max_four
5559 } else if available >= 8 || backlog >= 4 {
5560 max_two
5561 } else {
5562 1
5563 }
5564 }
5565 2..=3 => {
5566 if available >= 128 || backlog >= 64 {
5567 max_four
5568 } else if available >= 4 || backlog >= 2 {
5569 max_two
5570 } else {
5571 1
5572 }
5573 }
5574 _ => {
5575 if available >= 64 || backlog >= 32 {
5576 max_four
5577 } else if available >= 8 || backlog >= 4 {
5578 max_two
5579 } else {
5580 1
5581 }
5582 }
5583 }
5584 }
5585
5586 async fn queue_claimer_target(
5587 &self,
5588 pool: &PgPool,
5589 queue: &str,
5590 max_claimers: i16,
5591 control_interval: Duration,
5592 ) -> Result<i16, AwaError> {
5593 let schema = self.schema();
5594 let now = Utc::now();
5595 let stale_cutoff = now
5596 - TimeDelta::from_std(control_interval).map_err(|err| {
5597 AwaError::Validation(format!("invalid claimer control interval: {err}"))
5598 })?;
5599
5600 if let Some(target) = sqlx::query_scalar::<_, i16>(&format!(
5601 r#"
5602 SELECT target_claimers
5603 FROM {schema}.queue_claimer_state
5604 WHERE queue = $1
5605 AND updated_at > $2
5606 "#
5607 ))
5608 .bind(queue)
5609 .bind(stale_cutoff)
5610 .fetch_optional(pool)
5611 .await
5612 .map_err(map_sqlx_error)?
5613 {
5614 return Ok(target.clamp(1, max_claimers.max(1)));
5615 }
5616
5617 let current_target = sqlx::query_scalar::<_, i16>(&format!(
5618 r#"
5619 SELECT target_claimers
5620 FROM {schema}.queue_claimer_state
5621 WHERE queue = $1
5622 "#
5623 ))
5624 .bind(queue)
5625 .fetch_optional(pool)
5626 .await
5627 .map_err(map_sqlx_error)?;
5628
5629 let signal = self.queue_claimer_signal(pool, queue).await?;
5630 let desired = self.desired_queue_claimer_target(current_target, &signal, max_claimers);
5631
5632 if let Some(updated) = sqlx::query_scalar::<_, i16>(&format!(
5633 r#"
5634 INSERT INTO {schema}.queue_claimer_state (queue, target_claimers, updated_at)
5635 VALUES ($1, $2, $3)
5636 ON CONFLICT (queue) DO UPDATE
5637 SET target_claimers = EXCLUDED.target_claimers,
5638 updated_at = EXCLUDED.updated_at
5639 WHERE {schema}.queue_claimer_state.updated_at <= $4
5640 RETURNING target_claimers
5641 "#
5642 ))
5643 .bind(queue)
5644 .bind(desired)
5645 .bind(now)
5646 .bind(stale_cutoff)
5647 .fetch_optional(pool)
5648 .await
5649 .map_err(map_sqlx_error)?
5650 {
5651 return Ok(updated.clamp(1, max_claimers.max(1)));
5652 }
5653
5654 Ok(current_target
5655 .unwrap_or(desired)
5656 .clamp(1, max_claimers.max(1)))
5657 }
5658
5659 async fn queue_claimer_signal(
5673 &self,
5674 pool: &PgPool,
5675 queue: &str,
5676 ) -> Result<AvailableSignal, AwaError> {
5677 let schema = self.schema();
5678 let queues = self.physical_queues_for_logical(queue);
5679 let available: i64 = sqlx::query_scalar(&format!(
5680 r#"
5681 SELECT COALESCE(
5682 sum(GREATEST(
5683 {schema}.sequence_next_value(qe.seq_name)
5684 - {schema}.sequence_next_value(qc.seq_name),
5685 0
5686 )),
5687 0
5688 )::bigint
5689 FROM {schema}.queue_enqueue_heads AS qe
5690 JOIN {schema}.queue_claim_heads AS qc
5691 ON qc.queue = qe.queue
5692 AND qc.priority = qe.priority
5693 AND qc.enqueue_shard = qe.enqueue_shard
5694 WHERE qe.queue = ANY($1)
5695 "#
5696 ))
5697 .bind(&queues)
5698 .fetch_one(pool)
5699 .await
5700 .map_err(map_sqlx_error)?;
5701
5702 Ok(AvailableSignal { available })
5703 }
5704
5705 #[allow(clippy::too_many_arguments)]
5706 #[tracing::instrument(skip(self, pool), fields(queue = %queue, instance_id = %instance_id), name = "queue_storage.claim_runtime_batch_with_aging_for_instance")]
5707 pub async fn claim_runtime_batch_with_aging_for_instance(
5708 &self,
5709 pool: &PgPool,
5710 queue: &str,
5711 max_batch: i64,
5712 deadline_duration: Duration,
5713 aging_interval: Duration,
5714 instance_id: Uuid,
5715 max_claimers: i16,
5716 lease_ttl: Duration,
5717 idle_threshold: Duration,
5718 ) -> Result<Vec<ClaimedRuntimeJob>, AwaError> {
5719 let target_claimers = self
5720 .queue_claimer_target(pool, queue, max_claimers, Duration::from_millis(500))
5721 .await?;
5722
5723 let Some(lease) = self
5724 .acquire_queue_claimer_row(
5725 pool,
5726 queue,
5727 instance_id,
5728 target_claimers,
5729 lease_ttl,
5730 idle_threshold,
5731 )
5732 .await?
5733 else {
5734 return Ok(Vec::new());
5735 };
5736
5737 let claimed = self
5738 .claim_runtime_batch_with_aging(
5739 pool,
5740 queue,
5741 max_batch,
5742 deadline_duration,
5743 aging_interval,
5744 )
5745 .await?;
5746
5747 if !claimed.is_empty() && lease.needs_refresh(Utc::now(), lease_ttl, idle_threshold) {
5748 let _ = self
5749 .mark_queue_claimer_active(pool, queue, instance_id, lease.lease(), lease_ttl)
5750 .await?;
5751 }
5752
5753 Ok(claimed)
5754 }
5755
5756 #[tracing::instrument(skip(self, pool), fields(queue = %queue), name = "queue_storage.claim_job_batch")]
5757 pub async fn claim_job_batch(
5758 &self,
5759 pool: &PgPool,
5760 queue: &str,
5761 max_batch: i64,
5762 deadline_duration: Duration,
5763 ) -> Result<Vec<JobRow>, AwaError> {
5764 self.claim_runtime_batch(pool, queue, max_batch, deadline_duration)
5765 .await
5766 .map(|claimed| claimed.into_iter().map(|row| row.job).collect())
5767 }
5768
5769 #[tracing::instrument(skip(self, pool, claimed), name = "queue_storage.complete_batch")]
5770 pub async fn complete_batch(
5771 &self,
5772 pool: &PgPool,
5773 claimed: &[ClaimedEntry],
5774 ) -> Result<usize, AwaError> {
5775 self.complete_claimed_batch(pool, claimed)
5776 .await
5777 .map(|updated| updated.len())
5778 }
5779
5780 #[tracing::instrument(
5781 skip(self, pool, claimed),
5782 name = "queue_storage.complete_claimed_batch"
5783 )]
5784 pub async fn complete_claimed_batch(
5785 &self,
5786 pool: &PgPool,
5787 claimed: &[ClaimedEntry],
5788 ) -> Result<Vec<(i64, i64)>, AwaError> {
5789 if claimed.is_empty() {
5790 return Ok(Vec::new());
5791 }
5792
5793 let schema = self.schema();
5794 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
5795
5796 let lease_slots: Vec<i32> = claimed.iter().map(|entry| entry.lease_slot).collect();
5797 let queues: Vec<String> = claimed.iter().map(|entry| entry.queue.clone()).collect();
5798 let priorities: Vec<i16> = claimed.iter().map(|entry| entry.priority).collect();
5799 let enqueue_shards: Vec<i16> = claimed.iter().map(|entry| entry.enqueue_shard).collect();
5800 let lane_seqs: Vec<i64> = claimed.iter().map(|entry| entry.lane_seq).collect();
5801
5802 let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
5803 r#"
5804 WITH completed(lease_slot, queue, priority, enqueue_shard, lane_seq) AS (
5805 SELECT * FROM unnest($1::int[], $2::text[], $3::smallint[], $4::smallint[], $5::bigint[])
5806 )
5807 DELETE FROM {schema}.leases AS leases
5808 USING completed
5809 WHERE leases.lease_slot = completed.lease_slot
5810 AND leases.queue = completed.queue
5811 AND leases.priority = completed.priority
5812 AND leases.enqueue_shard = completed.enqueue_shard
5813 AND leases.lane_seq = completed.lane_seq
5814 RETURNING
5815 leases.ready_slot,
5816 leases.ready_generation,
5817 leases.job_id,
5818 leases.queue,
5819 leases.state,
5820 leases.priority,
5821 leases.attempt,
5822 leases.run_lease,
5823 leases.max_attempts,
5824 leases.lane_seq,
5825 leases.enqueue_shard,
5826 leases.heartbeat_at,
5827 leases.deadline_at,
5828 leases.attempted_at,
5829 leases.callback_id,
5830 leases.callback_timeout_at
5831 "#
5832 ))
5833 .bind(&lease_slots)
5834 .bind(&queues)
5835 .bind(&priorities)
5836 .bind(&enqueue_shards)
5837 .bind(&lane_seqs)
5838 .fetch_all(tx.as_mut())
5839 .await
5840 .map_err(map_sqlx_error)?;
5841
5842 if deleted.is_empty() {
5843 tx.commit().await.map_err(map_sqlx_error)?;
5844 return Ok(Vec::new());
5845 }
5846
5847 let completed_pairs: Vec<(i64, i64)> = deleted
5848 .iter()
5849 .map(|row| (row.job_id, row.run_lease))
5850 .collect();
5851 self.close_receipt_pairs_tx(&mut tx, &completed_pairs, "completed")
5852 .await?;
5853
5854 let moved = self.hydrate_deleted_leases_tx(&mut tx, deleted).await?;
5855
5856 let finalized_at = Utc::now();
5857 let mut done_rows = Vec::with_capacity(moved.len());
5858 for entry in moved.iter().cloned() {
5859 let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state(
5860 entry.payload.clone(),
5861 entry.progress.clone(),
5862 )?)?;
5863 payload.set_progress(None);
5864 done_rows.push(entry.into_done_row(
5865 JobState::Completed,
5866 finalized_at,
5867 payload.into_json(),
5868 ));
5869 }
5870
5871 self.insert_done_rows_tx(&mut tx, &done_rows, Some(JobState::Running))
5872 .await?;
5873 tx.commit().await.map_err(map_sqlx_error)?;
5874 Ok(moved
5875 .into_iter()
5876 .map(|entry| (entry.job_id, entry.run_lease))
5877 .collect())
5878 }
5879
5880 fn receipt_fast_complete_candidate(entry: &ClaimedRuntimeJob) -> bool {
5881 entry.claim.lease_claim_receipt
5882 && entry.claim.receipt_id.is_some()
5883 && entry.job.unique_key.is_none()
5884 && is_compact_receipt_completion_metadata(&entry.job.metadata)
5885 && entry.job.tags.is_empty()
5886 && entry.job.errors.as_ref().is_none_or(Vec::is_empty)
5887 }
5888
5889 async fn complete_receipt_runtime_batch_fast(
5890 &self,
5891 pool: &PgPool,
5892 claimed: &[ClaimedRuntimeJob],
5893 ) -> Result<Vec<(i64, i64)>, AwaError> {
5894 if claimed.is_empty() {
5895 return Ok(Vec::new());
5896 }
5897
5898 let schema = self.schema();
5899 let finalized_at = Utc::now();
5900 let job_ids: Vec<i64> = claimed.iter().map(|entry| entry.job.id).collect();
5901 let run_leases: Vec<i64> = claimed.iter().map(|entry| entry.job.run_lease).collect();
5902 let mut by_partition: BTreeMap<usize, Vec<&ClaimedRuntimeJob>> = BTreeMap::new();
5903 for entry in claimed {
5904 let claim_slot_index =
5905 ring_slot_index(entry.claim.claim_slot, self.claim_slot_count(), "claim")?;
5906 by_partition
5907 .entry(claim_slot_index)
5908 .or_default()
5909 .push(entry);
5910 }
5911
5912 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
5913
5914 let receipt_pairs: Vec<(i64, i64)> = job_ids
5920 .iter()
5921 .zip(run_leases.iter())
5922 .map(|(job_id, run_lease)| (*job_id, *run_lease))
5923 .collect();
5924 if let Err(err) = self.lock_receipt_attempts_tx(&mut tx, &receipt_pairs).await {
5925 let _ = tx.rollback().await;
5926 return Err(err);
5927 }
5928
5929 let mut rows = Vec::with_capacity(claimed.len());
5930 for (claim_slot_index, group) in by_partition {
5931 let claim_rel = claim_child_name(schema, claim_slot_index);
5932 let claim_batch_rel = claim_batch_child_name(schema, claim_slot_index);
5933 let closure_rel = closure_child_name(schema, claim_slot_index);
5934 let closure_batch_rel = claim_closure_batch_child_name(schema, claim_slot_index);
5935 let receipt_batch_rel = format!("{schema}.receipt_completion_batches");
5936 let closed_evidence =
5937 receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims");
5938
5939 let claim_slots: Vec<i32> = group.iter().map(|entry| entry.claim.claim_slot).collect();
5940 let ready_slots: Vec<i32> = group.iter().map(|entry| entry.claim.ready_slot).collect();
5941 let ready_generations: Vec<i64> = group
5942 .iter()
5943 .map(|entry| entry.claim.ready_generation)
5944 .collect();
5945 let job_ids: Vec<i64> = group.iter().map(|entry| entry.job.id).collect();
5946 let queues: Vec<String> = group.iter().map(|entry| entry.job.queue.clone()).collect();
5947 let priorities: Vec<i16> = group.iter().map(|entry| entry.claim.priority).collect();
5948 let attempts: Vec<i16> = group.iter().map(|entry| entry.job.attempt).collect();
5949 let run_leases: Vec<i64> = group.iter().map(|entry| entry.job.run_lease).collect();
5950 let receipt_ids: Vec<i64> = group
5951 .iter()
5952 .map(|entry| {
5953 entry
5954 .claim
5955 .receipt_id
5956 .expect("receipt fast completion requires receipt_id")
5957 })
5958 .collect();
5959 let claim_batch_ids: Vec<Option<i64>> = group
5960 .iter()
5961 .map(|entry| entry.claim.claim_batch_id)
5962 .collect();
5963 let claim_batch_indices: Vec<Option<i32>> = group
5964 .iter()
5965 .map(|entry| entry.claim.claim_batch_index)
5966 .collect();
5967 let lane_seqs: Vec<i64> = group.iter().map(|entry| entry.claim.lane_seq).collect();
5968 let enqueue_shards: Vec<i16> = group
5969 .iter()
5970 .map(|entry| entry.claim.enqueue_shard)
5971 .collect();
5972 let attempted_ats: Vec<Option<DateTime<Utc>>> =
5973 group.iter().map(|entry| entry.job.attempted_at).collect();
5974 let finalized_ats: Vec<DateTime<Utc>> = vec![finalized_at; group.len()];
5975
5976 let completed: Vec<(i64, i64)> = match sqlx::query_as(&format!(
5977 r#"
5978 WITH completed(
5979 claim_slot,
5980 ready_slot,
5981 ready_generation,
5982 job_id,
5983 queue,
5984 priority,
5985 attempt,
5986 run_lease,
5987 receipt_id,
5988 claim_batch_id,
5989 claim_batch_index,
5990 lane_seq,
5991 enqueue_shard,
5992 attempted_at,
5993 finalized_at
5994 ) AS (
5995 SELECT *
5996 FROM unnest(
5997 $1::int[],
5998 $2::int[],
5999 $3::bigint[],
6000 $4::bigint[],
6001 $5::text[],
6002 $6::smallint[],
6003 $7::smallint[],
6004 $8::bigint[],
6005 $9::bigint[],
6006 $10::bigint[],
6007 $11::int[],
6008 $12::bigint[],
6009 $13::smallint[],
6010 $14::timestamptz[],
6011 $15::timestamptz[]
6012 )
6013 ),
6014 row_claim_refs AS (
6015 SELECT claims.claim_slot, claims.job_id, claims.run_lease, claims.receipt_id
6016 FROM {claim_rel} AS claims
6017 JOIN completed
6018 ON completed.claim_slot = claims.claim_slot
6019 AND completed.job_id = claims.job_id
6020 AND completed.run_lease = claims.run_lease
6021 AND completed.receipt_id = claims.receipt_id
6022 AND completed.claim_batch_id IS NULL
6023 WHERE NOT {closed_evidence}
6024 AND NOT EXISTS (
6025 SELECT 1
6026 FROM {schema}.leases AS lease
6027 WHERE lease.job_id = claims.job_id
6028 AND lease.run_lease = claims.run_lease
6029 )
6030 ),
6031 batch_claim_refs AS (
6032 SELECT
6033 claim_batches.claim_slot,
6034 items.job_id,
6035 items.run_lease,
6036 items.receipt_id
6037 FROM completed
6038 JOIN {claim_batch_rel} AS claim_batches
6039 ON claim_batches.claim_slot = completed.claim_slot
6040 AND claim_batches.batch_id = completed.claim_batch_id
6041 CROSS JOIN LATERAL (
6042 SELECT
6043 claim_batches.job_ids[completed.claim_batch_index] AS job_id,
6044 claim_batches.run_leases[completed.claim_batch_index] AS run_lease,
6045 claim_batches.receipt_ids[completed.claim_batch_index] AS receipt_id
6046 ) AS items
6047 WHERE completed.claim_batch_id IS NOT NULL
6048 AND completed.claim_batch_index IS NOT NULL
6049 AND completed.claim_batch_index BETWEEN 1 AND claim_batches.claimed_count
6050 AND completed.job_id = items.job_id
6051 AND completed.run_lease = items.run_lease
6052 AND completed.receipt_id = items.receipt_id
6053 AND NOT EXISTS (
6054 SELECT 1
6055 FROM {closure_rel} AS closures
6056 WHERE closures.claim_slot = claim_batches.claim_slot
6057 AND closures.job_id = items.job_id
6058 AND closures.run_lease = items.run_lease
6059 )
6060 -- Compact successful completion treats receipt closure
6061 -- evidence as the same-attempt disposition fence. The
6062 -- non-success paths write explicit closure rows before
6063 -- moving the job to done/deferred/DLQ, so probing those
6064 -- terminal ledgers here only adds hot-path work.
6065 AND NOT EXISTS (
6066 SELECT 1
6067 FROM {closure_batch_rel} AS closure_batches
6068 WHERE closure_batches.receipt_ranges @> items.receipt_id
6069 )
6070 AND NOT EXISTS (
6071 SELECT 1
6072 FROM {schema}.leases AS lease
6073 WHERE lease.job_id = items.job_id
6074 AND lease.run_lease = items.run_lease
6075 )
6076 ),
6077 claim_refs AS (
6078 SELECT claim_slot, job_id, run_lease, receipt_id FROM row_claim_refs
6079 UNION ALL
6080 SELECT claim_slot, job_id, run_lease, receipt_id FROM batch_claim_refs
6081 ),
6082 deleted_attempts AS (
6083 DELETE FROM {schema}.attempt_state AS attempt
6084 USING claim_refs
6085 WHERE attempt.job_id = claim_refs.job_id
6086 AND attempt.run_lease = claim_refs.run_lease
6087 RETURNING attempt.job_id
6088 ),
6089 completed_rows AS (
6090 SELECT completed.*
6091 FROM completed
6092 JOIN claim_refs
6093 ON claim_refs.claim_slot = completed.claim_slot
6094 AND claim_refs.job_id = completed.job_id
6095 AND claim_refs.run_lease = completed.run_lease
6096 AND claim_refs.receipt_id = completed.receipt_id
6097 ),
6098 claim_closure_batches AS (
6099 INSERT INTO {closure_batch_rel} (
6100 claim_slot,
6101 ready_slot,
6102 ready_generation,
6103 outcome,
6104 closed_count,
6105 receipt_ids,
6106 receipt_ranges,
6107 closed_at
6108 )
6109 SELECT
6110 completed.claim_slot,
6111 completed.ready_slot,
6112 completed.ready_generation,
6113 'completed',
6114 count(*)::int AS closed_count,
6115 array_agg(completed.receipt_id ORDER BY completed.lane_seq, completed.job_id),
6116 range_agg(int8range(completed.receipt_id, completed.receipt_id + 1, '[)') ORDER BY completed.receipt_id),
6117 max(completed.finalized_at)
6118 FROM completed_rows AS completed
6119 GROUP BY
6120 completed.claim_slot,
6121 completed.ready_slot,
6122 completed.ready_generation
6123 RETURNING claim_slot, ready_slot, ready_generation, receipt_ids
6124 ),
6125 terminal AS (
6126 INSERT INTO {receipt_batch_rel} (
6127 ready_slot,
6128 ready_generation,
6129 claim_slot,
6130 queue,
6131 priority,
6132 enqueue_shard,
6133 completed_count,
6134 job_ids,
6135 run_leases,
6136 lane_seqs,
6137 attempts,
6138 attempted_ats,
6139 finalized_at
6140 )
6141 SELECT
6142 completed.ready_slot,
6143 completed.ready_generation,
6144 completed.claim_slot,
6145 completed.queue,
6146 completed.priority,
6147 completed.enqueue_shard,
6148 count(*)::int AS completed_count,
6149 array_agg(completed.job_id ORDER BY completed.lane_seq, completed.job_id),
6150 array_agg(completed.run_lease ORDER BY completed.lane_seq, completed.job_id),
6151 array_agg(completed.lane_seq ORDER BY completed.lane_seq, completed.job_id),
6152 array_agg(completed.attempt ORDER BY completed.lane_seq, completed.job_id),
6153 array_agg(completed.attempted_at ORDER BY completed.lane_seq, completed.job_id),
6154 max(completed.finalized_at)
6155 FROM completed_rows AS completed
6156 GROUP BY
6157 completed.ready_slot,
6158 completed.ready_generation,
6159 completed.claim_slot,
6160 completed.queue,
6161 completed.priority,
6162 completed.enqueue_shard
6163 RETURNING
6164 ready_slot,
6165 ready_generation,
6166 claim_slot,
6167 queue,
6168 priority,
6169 enqueue_shard,
6170 job_ids,
6171 run_leases
6172 )
6173 SELECT completed_rows.job_id, completed_rows.run_lease
6174 FROM completed_rows
6175 CROSS JOIN (SELECT count(*) FROM claim_closure_batches) AS closure_batch_write
6176 CROSS JOIN (SELECT count(*) FROM terminal) AS terminal_write
6177 "#
6178 ))
6179 .bind(&claim_slots)
6180 .bind(&ready_slots)
6181 .bind(&ready_generations)
6182 .bind(&job_ids)
6183 .bind(&queues)
6184 .bind(&priorities)
6185 .bind(&attempts)
6186 .bind(&run_leases)
6187 .bind(&receipt_ids)
6188 .bind(&claim_batch_ids)
6189 .bind(&claim_batch_indices)
6190 .bind(&lane_seqs)
6191 .bind(&enqueue_shards)
6192 .bind(&attempted_ats)
6193 .bind(&finalized_ats)
6194 .fetch_all(tx.as_mut())
6195 .await
6196 {
6197 Ok(rows) => rows,
6198 Err(err) => {
6199 let _ = tx.rollback().await;
6200 return Err(map_sqlx_error(err));
6201 }
6202 };
6203
6204 rows.extend(completed);
6205 }
6206
6207 tx.commit().await.map_err(map_sqlx_error)?;
6208 Ok(rows)
6209 }
6210
6211 #[tracing::instrument(
6212 skip(self, pool, claimed),
6213 name = "queue_storage.complete_runtime_batch"
6214 )]
6215 pub async fn complete_runtime_batch(
6216 &self,
6217 pool: &PgPool,
6218 claimed: &[ClaimedRuntimeJob],
6219 ) -> Result<Vec<(i64, i64)>, AwaError> {
6220 if claimed.is_empty() {
6221 return Ok(Vec::new());
6222 }
6223
6224 if self.lease_claim_receipts() && claimed.iter().all(Self::receipt_fast_complete_candidate)
6225 {
6226 let mut updated = self
6227 .complete_receipt_runtime_batch_fast(pool, claimed)
6228 .await?;
6229 if updated.len() == claimed.len() {
6230 return Ok(updated);
6231 }
6232
6233 let updated_pairs: BTreeSet<(i64, i64)> = updated.iter().copied().collect();
6234 let missed: Vec<_> = claimed
6235 .iter()
6236 .filter(|entry| !updated_pairs.contains(&(entry.job.id, entry.job.run_lease)))
6237 .cloned()
6238 .collect();
6239 if !missed.is_empty() {
6240 updated.extend(self.complete_runtime_batch_slow(pool, &missed).await?);
6241 }
6242 return Ok(updated);
6243 }
6244
6245 self.complete_runtime_batch_slow(pool, claimed).await
6246 }
6247
6248 async fn complete_runtime_batch_slow(
6249 &self,
6250 pool: &PgPool,
6251 claimed: &[ClaimedRuntimeJob],
6252 ) -> Result<Vec<(i64, i64)>, AwaError> {
6253 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
6254 let result = self
6255 .complete_runtime_batch_slow_in_tx(&mut tx, claimed)
6256 .await?;
6257 tx.commit().await.map_err(map_sqlx_error)?;
6258 Ok(result)
6259 }
6260
6261 pub async fn complete_runtime_batch_slow_in_tx(
6267 &self,
6268 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
6269 claimed: &[ClaimedRuntimeJob],
6270 ) -> Result<Vec<(i64, i64)>, AwaError> {
6271 if claimed.is_empty() {
6272 return Ok(Vec::new());
6273 }
6274
6275 let schema = self.schema();
6276
6277 let claimed_map: BTreeMap<(i64, i64), ClaimedRuntimeJob> = claimed
6278 .iter()
6279 .cloned()
6280 .map(|entry| ((entry.job.id, entry.job.run_lease), entry))
6281 .collect();
6282
6283 if self.lease_claim_receipts() {
6284 let (mut receipt_claimed, mut materialized_claimed): (Vec<_>, Vec<_>) = claimed
6285 .iter()
6286 .cloned()
6287 .partition(|entry| entry.claim.lease_claim_receipt);
6288 let mut updated_all = Vec::new();
6289
6290 if !receipt_claimed.is_empty() {
6291 let receipt_pairs: Vec<(i64, i64)> = receipt_claimed
6292 .iter()
6293 .map(|entry| (entry.job.id, entry.job.run_lease))
6294 .collect();
6295 self.lock_receipt_attempts_tx(tx, &receipt_pairs).await?;
6296
6297 let receipt_claim_slots: Vec<i32> = receipt_claimed
6301 .iter()
6302 .map(|entry| entry.claim.claim_slot)
6303 .collect();
6304 let receipt_job_ids: Vec<i64> =
6305 receipt_claimed.iter().map(|entry| entry.job.id).collect();
6306 let receipt_run_leases: Vec<i64> = receipt_claimed
6307 .iter()
6308 .map(|entry| entry.job.run_lease)
6309 .collect();
6310 let receipt_receipt_ids: Vec<i64> = receipt_claimed
6311 .iter()
6312 .map(|entry| {
6313 entry
6314 .claim
6315 .receipt_id
6316 .expect("receipt-backed slow completion requires receipt_id")
6317 })
6318 .collect();
6319 let closure_rel = format!("{schema}.lease_claim_closures");
6320 let closure_batch_rel = format!("{schema}.lease_claim_closure_batches");
6321 let closed_evidence =
6322 receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims");
6323 let updated: Vec<(i64, i64)> = sqlx::query_as(&format!(
6324 r#"
6325 WITH completed(claim_slot, job_id, run_lease, receipt_id) AS (
6326 SELECT * FROM unnest($1::int[], $2::bigint[], $3::bigint[], $4::bigint[])
6327 ),
6328 locked_row_claims AS (
6329 SELECT claims.claim_slot, claims.job_id, claims.run_lease
6330 FROM {schema}.lease_claims AS claims
6331 JOIN completed
6332 ON completed.claim_slot = claims.claim_slot
6333 AND completed.job_id = claims.job_id
6334 AND completed.run_lease = claims.run_lease
6335 AND completed.receipt_id = claims.receipt_id
6336 WHERE NOT {closed_evidence}
6337 AND NOT EXISTS (
6338 SELECT 1
6339 FROM {schema}.leases AS lease
6340 WHERE lease.job_id = claims.job_id
6341 AND lease.run_lease = claims.run_lease
6342 )
6343 FOR UPDATE OF claims
6344 ),
6345 locked_batch_claims AS (
6346 SELECT
6347 claim_batches.claim_slot,
6348 claim_batches.ready_slot,
6349 claim_batches.ready_generation,
6350 items.job_id,
6351 items.run_lease,
6352 items.receipt_id
6353 FROM {schema}.lease_claim_batches AS claim_batches
6354 CROSS JOIN LATERAL unnest(
6355 claim_batches.job_ids,
6356 claim_batches.run_leases,
6357 claim_batches.receipt_ids
6358 ) AS items(job_id, run_lease, receipt_id)
6359 JOIN completed
6360 ON completed.claim_slot = claim_batches.claim_slot
6361 AND completed.job_id = items.job_id
6362 AND completed.run_lease = items.run_lease
6363 AND completed.receipt_id = items.receipt_id
6364 WHERE NOT EXISTS (
6365 SELECT 1
6366 FROM {schema}.lease_claim_closures AS closures
6367 WHERE closures.claim_slot = claim_batches.claim_slot
6368 AND closures.job_id = items.job_id
6369 AND closures.run_lease = items.run_lease
6370 )
6371 AND NOT EXISTS (
6372 SELECT 1
6373 FROM {schema}.lease_claim_closure_batches AS closure_batches
6374 WHERE closure_batches.claim_slot = claim_batches.claim_slot
6375 AND closure_batches.receipt_ranges @> items.receipt_id
6376 )
6377 AND NOT EXISTS (
6378 SELECT 1
6379 FROM {schema}.leases AS lease
6380 WHERE lease.job_id = items.job_id
6381 AND lease.run_lease = items.run_lease
6382 )
6383 AND NOT EXISTS (
6384 SELECT 1
6385 FROM {schema}.done_entries AS done
6386 WHERE done.job_id = items.job_id
6387 AND done.run_lease = items.run_lease
6388 )
6389 AND NOT EXISTS (
6390 SELECT 1
6391 FROM {schema}.deferred_jobs AS deferred
6392 WHERE deferred.job_id = items.job_id
6393 AND deferred.run_lease = items.run_lease
6394 )
6395 AND NOT EXISTS (
6396 SELECT 1
6397 FROM {schema}.dlq_entries AS dlq
6398 WHERE dlq.job_id = items.job_id
6399 AND dlq.run_lease = items.run_lease
6400 )
6401 FOR UPDATE OF claim_batches
6402 ),
6403 locked_claims AS (
6404 SELECT claim_slot, job_id, run_lease FROM locked_row_claims
6405 UNION ALL
6406 SELECT claim_slot, job_id, run_lease FROM locked_batch_claims
6407 ),
6408 deleted_attempts AS (
6409 DELETE FROM {schema}.attempt_state AS attempt
6410 USING locked_claims
6411 WHERE attempt.job_id = locked_claims.job_id
6412 AND attempt.run_lease = locked_claims.run_lease
6413 RETURNING attempt.job_id
6414 ),
6415 -- Row-sourced (deadline lease) claims close into the
6416 -- explicit ledger so the prune gates balance them
6417 -- against the lease_claims row via the closure JOIN.
6418 closed_claims AS (
6419 INSERT INTO {schema}.lease_claim_closures (
6420 claim_slot,
6421 job_id,
6422 run_lease,
6423 outcome,
6424 closed_at
6425 )
6426 SELECT
6427 locked_row_claims.claim_slot,
6428 locked_row_claims.job_id,
6429 locked_row_claims.run_lease,
6430 'completed',
6431 clock_timestamp()
6432 FROM locked_row_claims
6433 ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING
6434 RETURNING claim_slot, job_id, run_lease, closed_at
6435 ),
6436 -- Compact batch-sourced claims have no lease_claims row
6437 -- to JOIN, so close them into the batch ledger that the
6438 -- queue prune gate counts via compact_count.
6439 closed_batches AS (
6440 INSERT INTO {schema}.lease_claim_closure_batches (
6441 claim_slot,
6442 ready_slot,
6443 ready_generation,
6444 outcome,
6445 closed_count,
6446 receipt_ids,
6447 receipt_ranges,
6448 closed_at
6449 )
6450 SELECT
6451 locked_batch_claims.claim_slot,
6452 locked_batch_claims.ready_slot,
6453 locked_batch_claims.ready_generation,
6454 'completed',
6455 count(*)::int,
6456 array_agg(locked_batch_claims.receipt_id ORDER BY locked_batch_claims.receipt_id),
6457 range_agg(int8range(locked_batch_claims.receipt_id, locked_batch_claims.receipt_id + 1, '[)') ORDER BY locked_batch_claims.receipt_id),
6458 clock_timestamp()
6459 FROM locked_batch_claims
6460 GROUP BY
6461 locked_batch_claims.claim_slot,
6462 locked_batch_claims.ready_slot,
6463 locked_batch_claims.ready_generation
6464 RETURNING claim_slot
6465 ),
6466 marked_claims AS (
6467 UPDATE {schema}.lease_claims AS claims
6468 SET closed_at = COALESCE(claims.closed_at, closed_claims.closed_at)
6469 FROM closed_claims
6470 WHERE claims.claim_slot = closed_claims.claim_slot
6471 AND claims.job_id = closed_claims.job_id
6472 AND claims.run_lease = closed_claims.run_lease
6473 RETURNING claims.job_id
6474 ),
6475 -- Force the batch insert to run and report the compact
6476 -- claims it closed so the caller finalizes them too.
6477 closed_batch_pairs AS (
6478 SELECT locked_batch_claims.job_id, locked_batch_claims.run_lease
6479 FROM locked_batch_claims
6480 WHERE EXISTS (SELECT 1 FROM closed_batches)
6481 )
6482 SELECT job_id, run_lease
6483 FROM closed_claims
6484 UNION ALL
6485 SELECT job_id, run_lease
6486 FROM closed_batch_pairs
6487 "#
6488 ))
6489 .bind(&receipt_claim_slots)
6490 .bind(&receipt_job_ids)
6491 .bind(&receipt_run_leases)
6492 .bind(&receipt_receipt_ids)
6493 .fetch_all(tx.as_mut())
6494 .await
6495 .map_err(map_sqlx_error)?;
6496
6497 if !updated.is_empty() {
6498 let finalized_at = Utc::now();
6499 let mut done_rows = Vec::with_capacity(updated.len());
6500 for (job_id, run_lease) in &updated {
6501 if let Some(runtime_job) = claimed_map.get(&(*job_id, *run_lease)).cloned()
6502 {
6503 done_rows.push(runtime_job.into_done_row(finalized_at)?);
6504 }
6505 }
6506
6507 self.insert_done_rows_tx(tx, &done_rows, Some(JobState::Running))
6508 .await?;
6509 updated_all.extend(updated);
6510 }
6511
6512 let updated_pairs: BTreeSet<(i64, i64)> = updated_all.iter().copied().collect();
6513 let mut escalated_receipts = Vec::new();
6514 for entry in receipt_claimed.drain(..) {
6515 if !updated_pairs.contains(&(entry.job.id, entry.job.run_lease)) {
6516 escalated_receipts.push(entry);
6517 }
6518 }
6519 materialized_claimed.extend(escalated_receipts);
6520 }
6521
6522 if !materialized_claimed.is_empty() {
6523 let ready_slots: Vec<i32> = materialized_claimed
6524 .iter()
6525 .map(|entry| entry.claim.ready_slot)
6526 .collect();
6527 let ready_generations: Vec<i64> = materialized_claimed
6528 .iter()
6529 .map(|entry| entry.claim.ready_generation)
6530 .collect();
6531 let job_ids: Vec<i64> = materialized_claimed
6532 .iter()
6533 .map(|entry| entry.job.id)
6534 .collect();
6535 let queues: Vec<String> = materialized_claimed
6536 .iter()
6537 .map(|entry| entry.claim.queue.clone())
6538 .collect();
6539 let priorities: Vec<i16> = materialized_claimed
6540 .iter()
6541 .map(|entry| entry.claim.priority)
6542 .collect();
6543 let enqueue_shards: Vec<i16> = materialized_claimed
6544 .iter()
6545 .map(|entry| entry.claim.enqueue_shard)
6546 .collect();
6547 let lane_seqs: Vec<i64> = materialized_claimed
6548 .iter()
6549 .map(|entry| entry.claim.lane_seq)
6550 .collect();
6551 let run_leases: Vec<i64> = materialized_claimed
6552 .iter()
6553 .map(|entry| entry.job.run_lease)
6554 .collect();
6555
6556 let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
6562 r#"
6563 WITH completed(ready_slot, ready_generation, job_id, queue, priority, enqueue_shard, lane_seq, run_lease) AS (
6564 SELECT * FROM unnest($1::int[], $2::bigint[], $3::bigint[], $4::text[], $5::smallint[], $6::smallint[], $7::bigint[], $8::bigint[])
6565 ),
6566 deleted AS (
6567 DELETE FROM {schema}.leases AS leases
6568 USING completed
6569 WHERE leases.ready_slot = completed.ready_slot
6570 AND leases.ready_generation = completed.ready_generation
6571 AND leases.job_id = completed.job_id
6572 AND leases.queue = completed.queue
6573 AND leases.priority = completed.priority
6574 AND leases.enqueue_shard = completed.enqueue_shard
6575 AND leases.lane_seq = completed.lane_seq
6576 AND leases.run_lease = completed.run_lease
6577 RETURNING
6578 leases.ready_slot,
6579 leases.ready_generation,
6580 leases.job_id,
6581 leases.queue,
6582 leases.state,
6583 leases.priority,
6584 leases.attempt,
6585 leases.run_lease,
6586 leases.max_attempts,
6587 leases.lane_seq,
6588 leases.enqueue_shard,
6589 leases.heartbeat_at,
6590 leases.deadline_at,
6591 leases.attempted_at,
6592 leases.callback_id,
6593 leases.callback_timeout_at
6594 ),
6595 del_attempts AS (
6596 DELETE FROM {schema}.attempt_state AS attempt
6597 USING deleted
6598 WHERE attempt.job_id = deleted.job_id
6599 AND attempt.run_lease = deleted.run_lease
6600 RETURNING attempt.job_id
6601 )
6602 SELECT
6603 ready_slot,
6604 ready_generation,
6605 job_id,
6606 queue,
6607 state,
6608 priority,
6609 attempt,
6610 run_lease,
6611 max_attempts,
6612 lane_seq,
6613 enqueue_shard,
6614 heartbeat_at,
6615 deadline_at,
6616 attempted_at,
6617 callback_id,
6618 callback_timeout_at
6619 FROM deleted
6620 "#
6621 ))
6622 .bind(&ready_slots)
6623 .bind(&ready_generations)
6624 .bind(&job_ids)
6625 .bind(&queues)
6626 .bind(&priorities)
6627 .bind(&enqueue_shards)
6628 .bind(&lane_seqs)
6629 .bind(&run_leases)
6630 .fetch_all(tx.as_mut())
6631 .await
6632 .map_err(map_sqlx_error)?;
6633
6634 if !deleted.is_empty() {
6635 let completed_pairs: Vec<(i64, i64)> = deleted
6636 .iter()
6637 .map(|row| (row.job_id, row.run_lease))
6638 .collect();
6639 self.close_receipt_pairs_tx(tx, &completed_pairs, "completed")
6640 .await?;
6641
6642 let finalized_at = Utc::now();
6643 let mut done_rows = Vec::with_capacity(deleted.len());
6644 for deleted_row in deleted {
6645 if let Some(runtime_job) = claimed_map
6646 .get(&(deleted_row.job_id, deleted_row.run_lease))
6647 .cloned()
6648 {
6649 done_rows.push(runtime_job.into_done_row(finalized_at)?);
6650 updated_all.push((deleted_row.job_id, deleted_row.run_lease));
6651 }
6652 }
6653
6654 self.insert_done_rows_tx(tx, &done_rows, Some(JobState::Running))
6655 .await?;
6656 }
6657 }
6658
6659 return Ok(updated_all);
6660 }
6661
6662 let lease_slots: Vec<i32> = claimed.iter().map(|entry| entry.claim.lease_slot).collect();
6663 let queues: Vec<String> = claimed
6664 .iter()
6665 .map(|entry| entry.claim.queue.clone())
6666 .collect();
6667 let priorities: Vec<i16> = claimed.iter().map(|entry| entry.claim.priority).collect();
6668 let enqueue_shards: Vec<i16> = claimed
6669 .iter()
6670 .map(|entry| entry.claim.enqueue_shard)
6671 .collect();
6672 let lane_seqs: Vec<i64> = claimed.iter().map(|entry| entry.claim.lane_seq).collect();
6673 let run_leases: Vec<i64> = claimed.iter().map(|entry| entry.job.run_lease).collect();
6674
6675 let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
6682 r#"
6683 WITH completed(lease_slot, queue, priority, enqueue_shard, lane_seq, run_lease) AS (
6684 SELECT * FROM unnest($1::int[], $2::text[], $3::smallint[], $4::smallint[], $5::bigint[], $6::bigint[])
6685 ),
6686 deleted AS (
6687 DELETE FROM {schema}.leases AS leases
6688 USING completed
6689 WHERE leases.lease_slot = completed.lease_slot
6690 AND leases.queue = completed.queue
6691 AND leases.priority = completed.priority
6692 AND leases.enqueue_shard = completed.enqueue_shard
6693 AND leases.lane_seq = completed.lane_seq
6694 AND leases.run_lease = completed.run_lease
6695 RETURNING
6696 leases.ready_slot,
6697 leases.ready_generation,
6698 leases.job_id,
6699 leases.queue,
6700 leases.state,
6701 leases.priority,
6702 leases.attempt,
6703 leases.run_lease,
6704 leases.max_attempts,
6705 leases.lane_seq,
6706 leases.enqueue_shard,
6707 leases.heartbeat_at,
6708 leases.deadline_at,
6709 leases.attempted_at,
6710 leases.callback_id,
6711 leases.callback_timeout_at
6712 ),
6713 del_attempts AS (
6714 DELETE FROM {schema}.attempt_state AS attempt
6715 USING deleted
6716 WHERE attempt.job_id = deleted.job_id
6717 AND attempt.run_lease = deleted.run_lease
6718 RETURNING attempt.job_id
6719 )
6720 SELECT
6721 ready_slot,
6722 ready_generation,
6723 job_id,
6724 queue,
6725 state,
6726 priority,
6727 attempt,
6728 run_lease,
6729 max_attempts,
6730 lane_seq,
6731 enqueue_shard,
6732 heartbeat_at,
6733 deadline_at,
6734 attempted_at,
6735 callback_id,
6736 callback_timeout_at
6737 FROM deleted
6738 "#
6739 ))
6740 .bind(&lease_slots)
6741 .bind(&queues)
6742 .bind(&priorities)
6743 .bind(&enqueue_shards)
6744 .bind(&lane_seqs)
6745 .bind(&run_leases)
6746 .fetch_all(tx.as_mut())
6747 .await
6748 .map_err(map_sqlx_error)?;
6749
6750 if deleted.is_empty() {
6751 return Ok(Vec::new());
6752 }
6753
6754 let completed_pairs: Vec<(i64, i64)> = deleted
6755 .iter()
6756 .map(|row| (row.job_id, row.run_lease))
6757 .collect();
6758 self.close_receipt_pairs_tx(tx, &completed_pairs, "completed")
6759 .await?;
6760
6761 let finalized_at = Utc::now();
6762 let mut done_rows = Vec::with_capacity(deleted.len());
6763 let mut updated = Vec::with_capacity(deleted.len());
6764 for deleted_row in deleted {
6765 if let Some(runtime_job) = claimed_map
6766 .get(&(deleted_row.job_id, deleted_row.run_lease))
6767 .cloned()
6768 {
6769 done_rows.push(runtime_job.into_done_row(finalized_at)?);
6770 updated.push((deleted_row.job_id, deleted_row.run_lease));
6771 }
6772 }
6773
6774 self.insert_done_rows_tx(tx, &done_rows, Some(JobState::Running))
6775 .await?;
6776 Ok(updated)
6777 }
6778
6779 #[tracing::instrument(
6780 skip(self, pool, completions),
6781 name = "queue_storage.complete_job_batch_by_id"
6782 )]
6783 pub async fn complete_job_batch_by_id(
6784 &self,
6785 pool: &PgPool,
6786 completions: &[(i64, i64)],
6787 ) -> Result<Vec<(i64, i64)>, AwaError> {
6788 if completions.is_empty() {
6789 return Ok(Vec::new());
6790 }
6791 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
6792 let result = self
6793 .complete_job_batch_by_id_in_tx(&mut tx, completions)
6794 .await?;
6795 tx.commit().await.map_err(map_sqlx_error)?;
6796 Ok(result)
6797 }
6798
6799 pub async fn complete_job_batch_by_id_in_tx(
6804 &self,
6805 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
6806 completions: &[(i64, i64)],
6807 ) -> Result<Vec<(i64, i64)>, AwaError> {
6808 if completions.is_empty() {
6809 return Ok(Vec::new());
6810 }
6811
6812 let schema = self.schema();
6813
6814 let job_ids: Vec<i64> = completions.iter().map(|(job_id, _)| *job_id).collect();
6815 let run_leases: Vec<i64> = completions
6816 .iter()
6817 .map(|(_, run_lease)| *run_lease)
6818 .collect();
6819
6820 let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
6821 r#"
6822 WITH completed(job_id, run_lease) AS (
6823 SELECT * FROM unnest($1::bigint[], $2::bigint[])
6824 )
6825 DELETE FROM {schema}.leases AS leases
6826 USING completed
6827 WHERE leases.job_id = completed.job_id
6828 AND leases.run_lease = completed.run_lease
6829 RETURNING
6830 leases.ready_slot,
6831 leases.ready_generation,
6832 leases.job_id,
6833 leases.queue,
6834 leases.state,
6835 leases.priority,
6836 leases.attempt,
6837 leases.run_lease,
6838 leases.max_attempts,
6839 leases.lane_seq,
6840 leases.enqueue_shard,
6841 leases.heartbeat_at,
6842 leases.deadline_at,
6843 leases.attempted_at,
6844 leases.callback_id,
6845 leases.callback_timeout_at
6846 "#
6847 ))
6848 .bind(&job_ids)
6849 .bind(&run_leases)
6850 .fetch_all(tx.as_mut())
6851 .await
6852 .map_err(map_sqlx_error)?;
6853
6854 if deleted.is_empty() {
6855 return Ok(Vec::new());
6856 }
6857
6858 let completed_pairs: Vec<(i64, i64)> = deleted
6859 .iter()
6860 .map(|row| (row.job_id, row.run_lease))
6861 .collect();
6862 self.close_receipt_pairs_tx(tx, &completed_pairs, "completed")
6863 .await?;
6864
6865 let moved = self.hydrate_deleted_leases_tx(tx, deleted).await?;
6866
6867 let finalized_at = Utc::now();
6868 let mut done_rows = Vec::with_capacity(moved.len());
6869 for entry in moved.iter().cloned() {
6870 let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state(
6871 entry.payload.clone(),
6872 entry.progress.clone(),
6873 )?)?;
6874 payload.set_progress(None);
6875 done_rows.push(entry.into_done_row(
6876 JobState::Completed,
6877 finalized_at,
6878 payload.into_json(),
6879 ));
6880 }
6881
6882 self.insert_done_rows_tx(tx, &done_rows, Some(JobState::Running))
6883 .await?;
6884 Ok(moved
6885 .into_iter()
6886 .map(|entry| (entry.job_id, entry.run_lease))
6887 .collect())
6888 }
6889
6890 async fn queue_counts_exact(
6906 &self,
6907 pool: &PgPool,
6908 queue: &str,
6909 ) -> Result<QueueCounts, AwaError> {
6910 let schema = self.schema();
6911 let queues = self.physical_queues_for_logical(queue);
6912 let closure_rel = format!("{schema}.lease_claim_closures");
6913 let closure_batch_rel = format!("{schema}.lease_claim_closure_batches");
6914 let closed_evidence =
6915 receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims");
6916 let counter_trusted = self.terminal_counter_trusted(pool).await?;
6917 let live_terminal_cte = if counter_trusted {
6922 format!(
6923 "live_terminal AS (
6924 SELECT GREATEST(
6925 0,
6926 COALESCE((
6927 SELECT SUM(live_terminal_count)
6928 FROM {schema}.queue_terminal_live_counts
6929 WHERE queue = ANY($1)
6930 ), 0)
6931 +
6932 COALESCE((
6933 SELECT SUM(terminal_delta)
6934 FROM {schema}.queue_terminal_count_deltas
6935 WHERE queue = ANY($1)
6936 ), 0)
6937 +
6938 COALESCE((
6939 SELECT SUM(completed_count)
6940 FROM {schema}.receipt_completion_batches
6941 WHERE queue = ANY($1)
6942 ), 0)
6943 -
6944 COALESCE((
6945 SELECT count(*)::bigint
6946 FROM {schema}.receipt_completion_tombstones
6947 WHERE queue = ANY($1)
6948 ), 0)
6949 )::bigint AS terminal
6950 )"
6951 )
6952 } else {
6953 format!(
6954 "live_terminal AS (
6955 SELECT count(*)::bigint AS terminal
6956 FROM {schema}.terminal_jobs
6957 WHERE queue = ANY($1)
6958 )"
6959 )
6960 };
6961 let row: (i64, i64, i64, i64) = sqlx::query_as(&format!(
6962 r#"
6963 WITH lane_counts AS (
6964 -- Exact count: a ready row is available iff its
6965 -- lane_seq has not yet been passed by the lane's
6966 -- claim sequence cursor. Each shard within a (queue,
6967 -- priority) lane carries its own sequence, so the
6968 -- join matches on shard too — otherwise a ready row
6969 -- in shard A could be incorrectly compared against
6970 -- shard B's claim cursor.
6971 SELECT COALESCE(count(*)::bigint, 0) AS available
6972 FROM {schema}.ready_entries AS ready
6973 JOIN {schema}.queue_claim_heads AS claims
6974 ON claims.queue = ready.queue
6975 AND claims.priority = ready.priority
6976 AND claims.enqueue_shard = ready.enqueue_shard
6977 WHERE ready.queue = ANY($1)
6978 AND ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name)
6979 AND NOT EXISTS (
6980 SELECT 1 FROM {schema}.ready_tombstones AS tomb
6981 WHERE tomb.queue = ready.queue
6982 AND tomb.priority = ready.priority
6983 AND tomb.enqueue_shard = ready.enqueue_shard
6984 AND tomb.lane_seq = ready.lane_seq
6985 AND tomb.ready_slot = ready.ready_slot
6986 AND tomb.ready_generation = ready.ready_generation
6987 )
6988 ),
6989 pruned_terminal AS (
6990 -- The GREATEST legacy dedupe applies to the completed
6991 -- column only: queue_lanes never carried a failed
6992 -- column, so failed counts come from the rollups alone.
6993 SELECT
6994 COALESCE(
6995 sum(
6996 GREATEST(
6997 COALESCE(lanes.pruned_completed_count, 0),
6998 COALESCE(rollups.pruned_completed_count, 0)
6999 )
7000 ),
7001 0
7002 )::bigint AS completed,
7003 COALESCE(
7004 sum(COALESCE(rollups.pruned_failed_count, 0)),
7005 0
7006 )::bigint AS failed
7007 FROM (
7008 SELECT queue, priority, pruned_completed_count
7009 FROM {schema}.queue_lanes
7010 WHERE queue = ANY($1)
7011 ) AS lanes
7012 FULL OUTER JOIN (
7013 SELECT queue, priority, pruned_completed_count, pruned_failed_count
7014 FROM {schema}.queue_terminal_rollups
7015 WHERE queue = ANY($1)
7016 ) AS rollups
7017 USING (queue, priority)
7018 ),
7019 live_running AS (
7020 SELECT (
7021 COALESCE((
7022 SELECT count(*)::bigint
7023 FROM {schema}.leases
7024 WHERE queue = ANY($1)
7025 AND state = 'running'
7026 ), 0)
7027 +
7028 -- Derive the receipt-backed running count from
7029 -- lease_claims anti-joined with every durable
7030 -- closure evidence shape.
7031 COALESCE((
7032 SELECT count(*)::bigint
7033 FROM {schema}.lease_claims AS claims
7034 WHERE claims.queue = ANY($1)
7035 AND NOT {closed_evidence}
7036 AND NOT EXISTS (
7037 SELECT 1
7038 FROM {schema}.leases AS lease
7039 WHERE lease.job_id = claims.job_id
7040 AND lease.run_lease = claims.run_lease
7041 )
7042 ), 0)
7043 +
7044 -- Zero-deadline compact receipt claims live in
7045 -- lease_claim_batches until they complete or a cold
7046 -- path materializes them. Expand only for exact
7047 -- admin-grade counts.
7048 COALESCE((
7049 SELECT count(*)::bigint
7050 FROM {schema}.lease_claim_batches AS batches
7051 CROSS JOIN LATERAL unnest(
7052 batches.job_ids,
7053 batches.run_leases,
7054 batches.receipt_ids
7055 ) AS items(job_id, run_lease, receipt_id)
7056 WHERE batches.queue = ANY($1)
7057 AND NOT EXISTS (
7058 SELECT 1
7059 FROM {schema}.lease_claim_closures AS closures
7060 WHERE closures.claim_slot = batches.claim_slot
7061 AND closures.job_id = items.job_id
7062 AND closures.run_lease = items.run_lease
7063 )
7064 AND NOT EXISTS (
7065 SELECT 1
7066 FROM {schema}.lease_claim_closure_batches AS closure_batches
7067 WHERE closure_batches.claim_slot = batches.claim_slot
7068 AND closure_batches.receipt_ranges @> items.receipt_id
7069 )
7070 AND NOT EXISTS (
7071 SELECT 1
7072 FROM {schema}.leases AS lease
7073 WHERE lease.job_id = items.job_id
7074 AND lease.run_lease = items.run_lease
7075 )
7076 AND NOT EXISTS (
7077 SELECT 1 FROM {schema}.done_entries AS done
7078 WHERE done.job_id = items.job_id
7079 AND done.run_lease = items.run_lease
7080 )
7081 AND NOT EXISTS (
7082 SELECT 1 FROM {schema}.deferred_jobs AS deferred
7083 WHERE deferred.job_id = items.job_id
7084 AND deferred.run_lease = items.run_lease
7085 )
7086 AND NOT EXISTS (
7087 SELECT 1 FROM {schema}.dlq_entries AS dlq
7088 WHERE dlq.job_id = items.job_id
7089 AND dlq.run_lease = items.run_lease
7090 )
7091 ), 0)
7092 )::bigint AS running
7093 ),
7094 {live_terminal_cte}
7095 SELECT
7096 lane_counts.available,
7097 live_running.running,
7098 pruned_terminal.completed + pruned_terminal.failed
7099 + live_terminal.terminal AS terminal,
7100 pruned_terminal.failed AS pruned_failed
7101 FROM lane_counts
7102 CROSS JOIN pruned_terminal
7103 CROSS JOIN live_running
7104 CROSS JOIN live_terminal
7105 "#
7106 ))
7107 .bind(&queues)
7108 .fetch_one(pool)
7109 .await
7110 .map_err(map_sqlx_error)?;
7111
7112 let (available, running, terminal, pruned_failed) = row;
7113 Ok(QueueCounts {
7114 available,
7115 running,
7116 terminal,
7117 pruned_failed,
7118 })
7119 }
7120
7121 #[tracing::instrument(skip(self, pool), fields(queue = %queue), name = "queue_storage.queue_counts")]
7122 pub async fn queue_counts(&self, pool: &PgPool, queue: &str) -> Result<QueueCounts, AwaError> {
7123 self.queue_counts_exact(pool, queue).await
7124 }
7125
7126 #[tracing::instrument(skip(self, pool), fields(queue = %queue), name = "queue_storage.queue_counts_fast")]
7163 pub async fn queue_counts_fast(
7164 &self,
7165 pool: &PgPool,
7166 queue: &str,
7167 ) -> Result<QueueCounts, AwaError> {
7168 let schema = self.schema();
7169 let queues = self.physical_queues_for_logical(queue);
7170 let available = self.queue_claimer_signal(pool, queue).await?.available;
7173 let running: i64 = sqlx::query_scalar(&format!(
7178 r#"
7179 SELECT COALESCE(count(*)::bigint, 0)
7180 FROM {schema}.leases
7181 WHERE queue = ANY($1)
7182 AND state = 'running'
7183 "#
7184 ))
7185 .bind(&queues)
7186 .fetch_one(pool)
7187 .await
7188 .map_err(map_sqlx_error)?;
7189 let (pruned_completed, pruned_failed): (i64, i64) = sqlx::query_as(&format!(
7194 r#"
7195 SELECT
7196 COALESCE(sum(GREATEST(
7197 COALESCE(lanes.pruned_completed_count, 0),
7198 COALESCE(rollups.pruned_completed_count, 0)
7199 )), 0)::bigint,
7200 COALESCE(sum(COALESCE(rollups.pruned_failed_count, 0)), 0)::bigint
7201 FROM (
7202 SELECT queue, priority, pruned_completed_count
7203 FROM {schema}.queue_lanes
7204 WHERE queue = ANY($1)
7205 ) AS lanes
7206 FULL OUTER JOIN (
7207 SELECT queue, priority, pruned_completed_count, pruned_failed_count
7208 FROM {schema}.queue_terminal_rollups
7209 WHERE queue = ANY($1)
7210 ) AS rollups
7211 USING (queue, priority)
7212 "#
7213 ))
7214 .bind(&queues)
7215 .fetch_one(pool)
7216 .await
7217 .map_err(map_sqlx_error)?;
7218 Ok(QueueCounts {
7219 available,
7220 running,
7221 terminal: pruned_completed + pruned_failed,
7222 pruned_failed,
7223 })
7224 }
7225
7226 pub async fn pruned_failed_count_for_queue(
7232 &self,
7233 pool: &PgPool,
7234 queue: &str,
7235 ) -> Result<u64, AwaError> {
7236 let schema = self.schema();
7237 let queues = self.physical_queues_for_logical(queue);
7238 let pruned_failed: i64 = sqlx::query_scalar(&format!(
7239 r#"
7240 SELECT COALESCE(sum(pruned_failed_count), 0)::bigint
7241 FROM {schema}.queue_terminal_rollups
7242 WHERE queue = ANY($1)
7243 "#
7244 ))
7245 .bind(&queues)
7246 .fetch_one(pool)
7247 .await
7248 .map_err(map_sqlx_error)?;
7249 Ok(pruned_failed.max(0) as u64)
7250 }
7251
7252 async fn retry_job_tx<'a>(
7253 &self,
7254 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
7255 job_id: i64,
7256 ) -> Result<Option<JobRow>, AwaError> {
7257 let schema = self.schema();
7258 let deleted_waiting: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
7259 r#"
7260 DELETE FROM {schema}.leases
7261 WHERE job_id = $1
7262 AND state = 'waiting_external'
7263 RETURNING
7264 ready_slot,
7265 ready_generation,
7266 job_id,
7267 queue,
7268 state,
7269 priority,
7270 attempt,
7271 run_lease,
7272 max_attempts,
7273 lane_seq,
7274 enqueue_shard,
7275 heartbeat_at,
7276 deadline_at,
7277 attempted_at,
7278 callback_id,
7279 callback_timeout_at
7280 "#
7281 ))
7282 .bind(job_id)
7283 .fetch_all(tx.as_mut())
7284 .await
7285 .map_err(map_sqlx_error)?;
7286
7287 if !deleted_waiting.is_empty() {
7288 let waiting = self
7289 .hydrate_deleted_leases_tx(tx, deleted_waiting)
7290 .await?
7291 .into_iter()
7292 .next()
7293 .expect("deleted waiting lease");
7294 let ready_payload = Self::payload_with_attempt_state(
7295 waiting.payload.clone(),
7296 waiting.progress.clone(),
7297 )?;
7298 let ready_row = ExistingReadyRow {
7299 attempt: 0,
7300 run_at: Utc::now(),
7301 attempted_at: None,
7302 ..waiting.clone().into_ready_row(Utc::now(), ready_payload)
7303 };
7304 self.insert_existing_ready_rows_tx(tx, vec![ready_row.clone()], Some(waiting.state))
7305 .await?;
7306 self.notify_queues_tx(tx, std::iter::once(waiting.queue.clone()))
7307 .await?;
7308 return Ok(Some(
7309 ReadyJobRow {
7310 job_id: ready_row.job_id,
7311 kind: ready_row.kind,
7312 queue: ready_row.queue,
7313 args: ready_row.args,
7314 priority: ready_row.priority,
7315 attempt: ready_row.attempt,
7316 run_lease: ready_row.run_lease,
7317 max_attempts: ready_row.max_attempts,
7318 run_at: ready_row.run_at,
7319 attempted_at: ready_row.attempted_at,
7320 created_at: ready_row.created_at,
7321 unique_key: ready_row.unique_key,
7322 payload: ready_row.payload,
7323 }
7324 .into_job_row()?,
7325 ));
7326 }
7327
7328 let done_projection = done_row_projection("done", "ready");
7329 let ready_join = done_ready_join(schema, "done", "ready");
7330 let terminal: Option<DoneJobRow> = sqlx::query_as(&format!(
7331 r#"
7332 WITH deleted AS (
7333 DELETE FROM {schema}.done_entries
7334 WHERE (job_id, finalized_at) IN (
7335 SELECT job_id, finalized_at
7336 FROM {schema}.done_entries
7337 WHERE job_id = $1
7338 AND state IN ('failed', 'cancelled')
7339 ORDER BY finalized_at DESC
7340 LIMIT 1
7341 FOR UPDATE SKIP LOCKED
7342 )
7343 RETURNING *
7344 )
7345 SELECT {done_projection}
7346 FROM deleted AS done
7347 {ready_join}
7348 "#
7349 ))
7350 .bind(job_id)
7351 .fetch_optional(tx.as_mut())
7352 .await
7353 .map_err(map_sqlx_error)?;
7354
7355 if let Some(terminal) = terminal {
7356 self.ensure_terminal_removed_receipt_closures_tx(tx, std::slice::from_ref(&terminal))
7357 .await?;
7358 self.decrement_live_terminal_counters_tx(
7361 tx,
7362 &Self::done_rows_to_counter_keys(std::slice::from_ref(&terminal)),
7363 )
7364 .await?;
7365 let ready_row = ExistingReadyRow {
7366 job_id: terminal.job_id,
7367 kind: terminal.kind,
7368 queue: terminal.queue.clone(),
7369 args: terminal.args,
7370 priority: terminal.priority,
7371 attempt: 0,
7372 run_lease: terminal.run_lease,
7373 max_attempts: terminal.max_attempts,
7374 run_at: Utc::now(),
7375 attempted_at: None,
7376 created_at: terminal.created_at,
7377 unique_key: terminal.unique_key,
7378 unique_states: terminal.unique_states,
7379 payload: terminal.payload,
7380 };
7381 self.insert_existing_ready_rows_tx(tx, vec![ready_row.clone()], Some(terminal.state))
7382 .await?;
7383 self.notify_queues_tx(tx, std::iter::once(terminal.queue.clone()))
7384 .await?;
7385 return Ok(Some(
7386 ReadyJobRow {
7387 job_id: ready_row.job_id,
7388 kind: ready_row.kind,
7389 queue: ready_row.queue,
7390 args: ready_row.args,
7391 priority: ready_row.priority,
7392 attempt: ready_row.attempt,
7393 run_lease: ready_row.run_lease,
7394 max_attempts: ready_row.max_attempts,
7395 run_at: ready_row.run_at,
7396 attempted_at: ready_row.attempted_at,
7397 created_at: ready_row.created_at,
7398 unique_key: ready_row.unique_key,
7399 payload: ready_row.payload,
7400 }
7401 .into_job_row()?,
7402 ));
7403 }
7404
7405 Ok(None)
7406 }
7407
7408 pub async fn retry_job(&self, pool: &PgPool, job_id: i64) -> Result<Option<JobRow>, AwaError> {
7409 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
7410 let row = self.retry_job_tx(&mut tx, job_id).await?;
7411 tx.commit().await.map_err(map_sqlx_error)?;
7412 Ok(row)
7413 }
7414
7415 pub async fn retry_jobs_by_ids(
7423 &self,
7424 pool: &PgPool,
7425 ids: &[i64],
7426 ) -> Result<(Vec<JobRow>, u64), AwaError> {
7427 let unique_ids: BTreeSet<i64> = ids.iter().copied().collect();
7428 if unique_ids.is_empty() {
7429 return Ok((Vec::new(), 0));
7430 }
7431
7432 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
7433 let mut rows = Vec::with_capacity(unique_ids.len());
7434 for job_id in &unique_ids {
7435 if let Some(row) = self.retry_job_tx(&mut tx, *job_id).await? {
7436 rows.push(row);
7437 }
7438 }
7439 tx.commit().await.map_err(map_sqlx_error)?;
7440 Ok((rows, unique_ids.len() as u64))
7441 }
7442
7443 async fn close_receipt_tx<'a>(
7459 &self,
7460 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
7461 job_id: i64,
7462 run_lease: i64,
7463 outcome: &str,
7464 ) -> Result<(), AwaError> {
7465 self.close_receipt_pairs_tx(tx, &[(job_id, run_lease)], outcome)
7466 .await
7467 }
7468
7469 async fn close_receipt_pairs_tx<'a>(
7470 &self,
7471 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
7472 pairs: &[(i64, i64)],
7473 outcome: &str,
7474 ) -> Result<(), AwaError> {
7475 if pairs.is_empty() || !self.lease_claim_receipts() {
7476 return Ok(());
7477 }
7478
7479 let unique_pairs: BTreeSet<(i64, i64)> = pairs.iter().copied().collect();
7480 let unique_jobs: Vec<(i64, i64)> = unique_pairs.iter().copied().collect();
7481 self.lock_receipt_attempts_tx(tx, &unique_jobs).await?;
7482
7483 let job_ids: Vec<i64> = unique_pairs.iter().map(|(job_id, _)| *job_id).collect();
7484 let run_leases: Vec<i64> = unique_pairs
7485 .iter()
7486 .map(|(_, run_lease)| *run_lease)
7487 .collect();
7488 let schema = self.schema();
7489 let closure_rel = format!("{schema}.lease_claim_closures");
7490 let closure_batch_rel = format!("{schema}.lease_claim_closure_batches");
7491 let closed_evidence =
7492 receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims");
7493
7494 sqlx::query(&format!(
7495 r#"
7496 WITH refs(job_id, run_lease) AS (
7497 SELECT * FROM unnest($1::bigint[], $2::bigint[])
7498 ),
7499 locked_claims AS (
7500 SELECT claims.claim_slot, claims.job_id, claims.run_lease
7501 FROM {schema}.lease_claims AS claims
7502 JOIN refs
7503 ON refs.job_id = claims.job_id
7504 AND refs.run_lease = claims.run_lease
7505 WHERE NOT {closed_evidence}
7506 FOR UPDATE OF claims
7507 ),
7508 locked_batch_claims AS (
7509 SELECT
7510 claim_batches.claim_slot,
7511 claim_batches.ready_slot,
7512 claim_batches.ready_generation,
7513 items.job_id,
7514 items.run_lease,
7515 items.receipt_id
7516 FROM {schema}.lease_claim_batches AS claim_batches
7517 CROSS JOIN LATERAL unnest(
7518 claim_batches.job_ids,
7519 claim_batches.run_leases,
7520 claim_batches.receipt_ids
7521 ) AS items(job_id, run_lease, receipt_id)
7522 JOIN refs
7523 ON refs.job_id = items.job_id
7524 AND refs.run_lease = items.run_lease
7525 WHERE NOT EXISTS (
7526 SELECT 1
7527 FROM {schema}.lease_claim_closures AS closures
7528 WHERE closures.claim_slot = claim_batches.claim_slot
7529 AND closures.job_id = items.job_id
7530 AND closures.run_lease = items.run_lease
7531 )
7532 AND NOT EXISTS (
7533 SELECT 1
7534 FROM {schema}.lease_claim_closure_batches AS closure_batches
7535 WHERE closure_batches.claim_slot = claim_batches.claim_slot
7536 AND closure_batches.receipt_ranges @> items.receipt_id
7537 )
7538 AND NOT EXISTS (
7539 SELECT 1
7540 FROM {schema}.done_entries AS done
7541 WHERE done.job_id = items.job_id
7542 AND done.run_lease = items.run_lease
7543 )
7544 AND NOT EXISTS (
7545 SELECT 1
7546 FROM {schema}.deferred_jobs AS deferred
7547 WHERE deferred.job_id = items.job_id
7548 AND deferred.run_lease = items.run_lease
7549 )
7550 AND NOT EXISTS (
7551 SELECT 1
7552 FROM {schema}.dlq_entries AS dlq
7553 WHERE dlq.job_id = items.job_id
7554 AND dlq.run_lease = items.run_lease
7555 )
7556 FOR UPDATE OF claim_batches
7557 ),
7558 -- Row-sourced (deadline lease) claims keep their explicit
7559 -- closure: the queue/claim prune gates balance those against
7560 -- the lease_claims row via the closure JOIN. Compact
7561 -- batch-sourced claims have no lease_claims row to JOIN, so
7562 -- they are closed into lease_claim_closure_batches below.
7563 closing_claims AS (
7564 SELECT
7565 locked_claims.claim_slot,
7566 locked_claims.job_id,
7567 locked_claims.run_lease,
7568 $3 AS outcome,
7569 clock_timestamp() AS closed_at
7570 FROM locked_claims
7571 ),
7572 inserted AS (
7573 INSERT INTO {schema}.lease_claim_closures (
7574 claim_slot,
7575 job_id,
7576 run_lease,
7577 outcome,
7578 closed_at
7579 )
7580 SELECT
7581 claim_slot,
7582 job_id,
7583 run_lease,
7584 outcome,
7585 closed_at
7586 FROM closing_claims
7587 ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING
7588 RETURNING claim_slot, job_id, run_lease, closed_at
7589 ),
7590 inserted_batches AS (
7591 INSERT INTO {schema}.lease_claim_closure_batches (
7592 claim_slot,
7593 ready_slot,
7594 ready_generation,
7595 outcome,
7596 closed_count,
7597 receipt_ids,
7598 receipt_ranges,
7599 closed_at
7600 )
7601 SELECT
7602 locked_batch_claims.claim_slot,
7603 locked_batch_claims.ready_slot,
7604 locked_batch_claims.ready_generation,
7605 $3,
7606 count(*)::int,
7607 array_agg(locked_batch_claims.receipt_id ORDER BY locked_batch_claims.receipt_id),
7608 range_agg(int8range(locked_batch_claims.receipt_id, locked_batch_claims.receipt_id + 1, '[)') ORDER BY locked_batch_claims.receipt_id),
7609 clock_timestamp()
7610 FROM locked_batch_claims
7611 GROUP BY
7612 locked_batch_claims.claim_slot,
7613 locked_batch_claims.ready_slot,
7614 locked_batch_claims.ready_generation
7615 RETURNING claim_slot
7616 ),
7617 marked AS (
7618 UPDATE {schema}.lease_claims AS claims
7619 SET closed_at = COALESCE(claims.closed_at, inserted.closed_at)
7620 FROM inserted
7621 WHERE claims.claim_slot = inserted.claim_slot
7622 AND claims.job_id = inserted.job_id
7623 AND claims.run_lease = inserted.run_lease
7624 RETURNING claims.job_id
7625 )
7626 SELECT
7627 (SELECT count(*) FROM marked)
7628 + (SELECT count(*) FROM inserted_batches)
7629 "#
7630 ))
7631 .bind(&job_ids)
7632 .bind(&run_leases)
7633 .bind(outcome)
7634 .execute(tx.as_mut())
7635 .await
7636 .map_err(map_sqlx_error)?;
7637
7638 Ok(())
7639 }
7640
7641 async fn notify_cancellation_tx<'a>(
7647 &self,
7648 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
7649 job_id: i64,
7650 run_lease: i64,
7651 ) -> Result<(), AwaError> {
7652 let payload = serde_json::json!({ "job_id": job_id, "run_lease": run_lease }).to_string();
7653 sqlx::query("SELECT pg_notify('awa:cancel', $1)")
7654 .bind(payload)
7655 .execute(tx.as_mut())
7656 .await
7657 .map_err(map_sqlx_error)?;
7658 Ok(())
7659 }
7660
7661 async fn cancel_job_tx<'a>(
7662 &self,
7663 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
7664 job_id: i64,
7665 ) -> Result<Option<CancelJobTxResult>, AwaError> {
7666 let schema = self.schema();
7667 let ready: Option<ReadyTransitionRow> = sqlx::query_as(&format!(
7668 r#"
7669 WITH target AS (
7670 SELECT ready.*
7671 FROM {schema}.ready_entries AS ready
7672 JOIN {schema}.queue_claim_heads AS claims
7673 ON claims.queue = ready.queue
7674 AND claims.priority = ready.priority
7675 AND claims.enqueue_shard = ready.enqueue_shard
7676 WHERE ready.job_id = $1
7677 AND ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name)
7678 AND NOT EXISTS (
7679 SELECT 1 FROM {schema}.ready_tombstones AS tomb
7680 WHERE tomb.queue = ready.queue
7681 AND tomb.priority = ready.priority
7682 AND tomb.enqueue_shard = ready.enqueue_shard
7683 AND tomb.lane_seq = ready.lane_seq
7684 AND tomb.ready_slot = ready.ready_slot
7685 AND tomb.ready_generation = ready.ready_generation
7686 )
7687 ORDER BY ready.lane_seq DESC
7688 LIMIT 1
7689 FOR UPDATE OF ready SKIP LOCKED
7690 ),
7691 tombstone AS (
7692 INSERT INTO {schema}.ready_tombstones (
7693 ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, job_id
7694 )
7695 SELECT ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, job_id
7696 FROM target
7697 ON CONFLICT DO NOTHING
7698 )
7699 SELECT
7700 ready_slot,
7701 ready_generation,
7702 job_id,
7703 kind,
7704 queue,
7705 args,
7706 priority,
7707 attempt,
7708 run_lease,
7709 max_attempts,
7710 lane_seq,
7711 enqueue_shard,
7712 run_at,
7713 attempted_at,
7714 created_at,
7715 unique_key,
7716 unique_states,
7717 COALESCE(payload, '{{}}'::jsonb) AS payload
7718 FROM target
7719 "#
7720 ))
7721 .bind(job_id)
7722 .fetch_optional(tx.as_mut())
7723 .await
7724 .map_err(map_sqlx_error)?;
7725
7726 if let Some(ready) = ready {
7727 let done =
7728 ready
7729 .clone()
7730 .into_done_row(JobState::Cancelled, Utc::now(), ready.payload.clone());
7731 self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(JobState::Available))
7732 .await?;
7733 let claim_cursor_advance = ClaimCursorAdvance {
7742 queue: ready.queue.clone(),
7743 priority: ready.priority,
7744 enqueue_shard: ready.enqueue_shard,
7745 next_seq: ready.lane_seq + 1,
7746 only_if_current: Some(ready.lane_seq),
7747 };
7748 return Ok(Some(CancelJobTxResult {
7749 row: done.into_job_row()?,
7750 claim_cursor_advance: Some(claim_cursor_advance),
7751 }));
7752 }
7753
7754 let deleted_lease: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
7755 r#"
7756 DELETE FROM {schema}.leases
7757 WHERE job_id = $1
7758 AND state IN ('running', 'waiting_external')
7759 RETURNING
7760 ready_slot,
7761 ready_generation,
7762 job_id,
7763 queue,
7764 state,
7765 priority,
7766 attempt,
7767 run_lease,
7768 max_attempts,
7769 lane_seq,
7770 enqueue_shard,
7771 heartbeat_at,
7772 deadline_at,
7773 attempted_at,
7774 callback_id,
7775 callback_timeout_at
7776 "#
7777 ))
7778 .bind(job_id)
7779 .fetch_all(tx.as_mut())
7780 .await
7781 .map_err(map_sqlx_error)?;
7782
7783 if !deleted_lease.is_empty() {
7784 let lease = self
7785 .hydrate_deleted_leases_tx(tx, deleted_lease)
7786 .await?
7787 .into_iter()
7788 .next()
7789 .expect("deleted running lease");
7790 let done_payload =
7791 Self::payload_with_attempt_state(lease.payload.clone(), lease.progress.clone())?;
7792 let done = lease
7793 .clone()
7794 .into_done_row(JobState::Cancelled, Utc::now(), done_payload);
7795 self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(lease.state))
7796 .await?;
7797 self.close_receipt_tx(tx, lease.job_id, lease.run_lease, "cancelled")
7801 .await?;
7802 self.notify_cancellation_tx(tx, lease.job_id, lease.run_lease)
7804 .await?;
7805 return Ok(Some(CancelJobTxResult {
7806 row: done.into_job_row()?,
7807 claim_cursor_advance: None,
7808 }));
7809 }
7810
7811 if self.lease_claim_receipts() {
7817 let closure_rel = format!("{schema}.lease_claim_closures");
7818 let closure_batch_rel = format!("{schema}.lease_claim_closure_batches");
7819 let closed_evidence =
7820 receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims");
7821 type ReceiptCancelRow = (
7822 i32,
7823 i64,
7824 i32,
7825 i64,
7826 String,
7827 i16,
7828 i16,
7829 i16,
7830 i64,
7831 DateTime<Utc>,
7832 i64,
7833 bool,
7834 );
7835 let receipt: Option<ReceiptCancelRow> = sqlx::query_as(&format!(
7836 r#"
7837 WITH row_receipt AS (
7838 SELECT
7839 claims.claim_slot,
7840 claims.run_lease,
7841 claims.ready_slot,
7842 claims.ready_generation,
7843 claims.queue,
7844 claims.priority,
7845 claims.attempt,
7846 claims.max_attempts,
7847 claims.lane_seq,
7848 claims.claimed_at,
7849 claims.receipt_id,
7850 false AS compact_batch
7851 FROM {schema}.lease_claims AS claims
7852 WHERE claims.job_id = $1
7853 AND NOT {closed_evidence}
7854 ORDER BY claims.run_lease DESC
7855 LIMIT 1
7856 FOR UPDATE OF claims SKIP LOCKED
7857 ),
7858 batch_receipt AS (
7859 SELECT
7860 claim_batches.claim_slot,
7861 items.run_lease,
7862 claim_batches.ready_slot,
7863 claim_batches.ready_generation,
7864 claim_batches.queue,
7865 claim_batches.priority,
7866 items.attempt,
7867 items.max_attempts,
7868 items.lane_seq,
7869 claim_batches.claimed_at,
7870 items.receipt_id,
7871 true AS compact_batch
7872 FROM {schema}.lease_claim_batches AS claim_batches
7873 CROSS JOIN LATERAL unnest(
7874 claim_batches.job_ids,
7875 claim_batches.run_leases,
7876 claim_batches.receipt_ids,
7877 claim_batches.lane_seqs,
7878 claim_batches.attempts,
7879 claim_batches.max_attempts
7880 ) AS items(job_id, run_lease, receipt_id, lane_seq, attempt, max_attempts)
7881 WHERE items.job_id = $1
7882 AND NOT EXISTS (
7883 SELECT 1
7884 FROM {schema}.lease_claim_closures AS closures
7885 WHERE closures.claim_slot = claim_batches.claim_slot
7886 AND closures.job_id = items.job_id
7887 AND closures.run_lease = items.run_lease
7888 )
7889 AND NOT EXISTS (
7890 SELECT 1
7891 FROM {schema}.lease_claim_closure_batches AS closure_batches
7892 WHERE closure_batches.claim_slot = claim_batches.claim_slot
7893 AND closure_batches.receipt_ranges @> items.receipt_id
7894 )
7895 AND NOT EXISTS (
7896 SELECT 1
7897 FROM {schema}.leases AS lease
7898 WHERE lease.job_id = items.job_id
7899 AND lease.run_lease = items.run_lease
7900 )
7901 AND NOT EXISTS (
7902 SELECT 1 FROM {schema}.done_entries AS done
7903 WHERE done.job_id = items.job_id
7904 AND done.run_lease = items.run_lease
7905 )
7906 AND NOT EXISTS (
7907 SELECT 1 FROM {schema}.deferred_jobs AS deferred
7908 WHERE deferred.job_id = items.job_id
7909 AND deferred.run_lease = items.run_lease
7910 )
7911 AND NOT EXISTS (
7912 SELECT 1 FROM {schema}.dlq_entries AS dlq
7913 WHERE dlq.job_id = items.job_id
7914 AND dlq.run_lease = items.run_lease
7915 )
7916 ORDER BY items.run_lease DESC
7917 LIMIT 1
7918 FOR UPDATE OF claim_batches SKIP LOCKED
7919 )
7920 SELECT *
7921 FROM (
7922 SELECT * FROM row_receipt
7923 UNION ALL
7924 SELECT * FROM batch_receipt
7925 ) AS receipt
7926 ORDER BY run_lease DESC
7927 LIMIT 1
7928 "#
7929 ))
7930 .bind(job_id)
7931 .fetch_optional(tx.as_mut())
7932 .await
7933 .map_err(map_sqlx_error)?;
7934
7935 if let Some((
7936 claim_slot,
7937 run_lease,
7938 ready_slot,
7939 ready_generation,
7940 queue,
7941 priority,
7942 attempt,
7943 max_attempts,
7944 lane_seq,
7945 claimed_at,
7946 receipt_id,
7947 compact_batch,
7948 )) = receipt
7949 {
7950 let ready_match: Option<ReadyTransitionRow> = sqlx::query_as(&format!(
7953 r#"
7954 SELECT
7955 ready_slot,
7956 ready_generation,
7957 job_id,
7958 kind,
7959 queue,
7960 args,
7961 priority,
7962 attempt,
7963 run_lease,
7964 max_attempts,
7965 lane_seq,
7966 enqueue_shard,
7967 run_at,
7968 attempted_at,
7969 created_at,
7970 unique_key,
7971 unique_states,
7972 COALESCE(payload, '{{}}'::jsonb) AS payload
7973 FROM {schema}.ready_entries
7974 WHERE job_id = $1
7975 AND ready_slot = $2
7976 AND ready_generation = $3
7977 AND queue = $4
7978 AND lane_seq = $5
7979 "#
7980 ))
7981 .bind(job_id)
7982 .bind(ready_slot)
7983 .bind(ready_generation)
7984 .bind(&queue)
7985 .bind(lane_seq)
7986 .fetch_optional(tx.as_mut())
7987 .await
7988 .map_err(map_sqlx_error)?;
7989
7990 let Some(ready) = ready_match else {
7991 return Ok(None);
7995 };
7996
7997 let done = DoneJobRow {
7998 ready_slot,
7999 ready_generation,
8000 job_id,
8001 kind: ready.kind,
8002 queue: queue.clone(),
8003 args: ready.args,
8004 state: JobState::Cancelled,
8005 priority,
8006 attempt,
8007 run_lease,
8008 max_attempts,
8009 lane_seq,
8010 enqueue_shard: ready.enqueue_shard,
8011 run_at: ready.run_at,
8012 attempted_at: Some(claimed_at),
8013 finalized_at: Utc::now(),
8014 created_at: ready.created_at,
8015 unique_key: ready.unique_key,
8016 unique_states: ready.unique_states,
8017 payload: ready.payload,
8018 };
8019 self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(JobState::Running))
8020 .await?;
8021 if compact_batch {
8028 sqlx::query(&format!(
8029 r#"
8030 INSERT INTO {schema}.lease_claim_closure_batches (
8031 claim_slot,
8032 ready_slot,
8033 ready_generation,
8034 outcome,
8035 closed_count,
8036 receipt_ids,
8037 receipt_ranges,
8038 closed_at
8039 )
8040 VALUES (
8041 $1,
8042 $2,
8043 $3,
8044 'cancelled',
8045 1,
8046 ARRAY[$4::bigint],
8047 int8multirange(int8range($4::bigint, $4::bigint + 1, '[)')),
8048 clock_timestamp()
8049 )
8050 "#
8051 ))
8052 .bind(claim_slot)
8053 .bind(ready_slot)
8054 .bind(ready_generation)
8055 .bind(receipt_id)
8056 .execute(tx.as_mut())
8057 .await
8058 .map_err(map_sqlx_error)?;
8059 } else {
8060 sqlx::query(&format!(
8061 r#"
8062 WITH inserted AS (
8063 INSERT INTO {schema}.lease_claim_closures (claim_slot, job_id, run_lease, outcome, closed_at)
8064 VALUES ($1, $2, $3, 'cancelled', clock_timestamp())
8065 ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING
8066 RETURNING claim_slot, job_id, run_lease, closed_at
8067 ),
8068 marked AS (
8069 UPDATE {schema}.lease_claims AS claims
8070 SET closed_at = COALESCE(claims.closed_at, inserted.closed_at)
8071 FROM inserted
8072 WHERE claims.claim_slot = inserted.claim_slot
8073 AND claims.job_id = inserted.job_id
8074 AND claims.run_lease = inserted.run_lease
8075 RETURNING claims.job_id
8076 )
8077 SELECT count(*) FROM marked
8078 "#
8079 ))
8080 .bind(claim_slot)
8081 .bind(job_id)
8082 .bind(run_lease)
8083 .execute(tx.as_mut())
8084 .await
8085 .map_err(map_sqlx_error)?;
8086 }
8087 sqlx::query(&format!(
8098 "DELETE FROM {schema}.leases WHERE job_id = $1 AND run_lease = $2"
8099 ))
8100 .bind(job_id)
8101 .bind(run_lease)
8102 .execute(tx.as_mut())
8103 .await
8104 .map_err(map_sqlx_error)?;
8105 self.notify_cancellation_tx(tx, job_id, run_lease).await?;
8106 return Ok(Some(CancelJobTxResult {
8107 row: done.into_job_row()?,
8108 claim_cursor_advance: None,
8109 }));
8110 }
8111 }
8112
8113 let deferred: Option<DeferredJobRow> = sqlx::query_as(&format!(
8114 r#"
8115 DELETE FROM {schema}.deferred_jobs
8116 WHERE job_id = $1
8117 AND state IN ('scheduled', 'retryable')
8118 RETURNING
8119 job_id,
8120 kind,
8121 queue,
8122 args,
8123 state,
8124 priority,
8125 attempt,
8126 run_lease,
8127 max_attempts,
8128 run_at,
8129 attempted_at,
8130 finalized_at,
8131 created_at,
8132 unique_key,
8133 unique_states,
8134 COALESCE(payload, '{{}}'::jsonb) AS payload
8135 "#
8136 ))
8137 .bind(job_id)
8138 .fetch_optional(tx.as_mut())
8139 .await
8140 .map_err(map_sqlx_error)?;
8141
8142 if let Some(deferred) = deferred {
8143 let (ready_slot, ready_generation) = self.current_queue_ring(tx).await?;
8144 self.ensure_lane(tx, &deferred.queue, deferred.priority, 0)
8152 .await?;
8153 let done = DoneJobRow {
8154 ready_slot,
8155 ready_generation,
8156 job_id: deferred.job_id,
8157 kind: deferred.kind,
8158 queue: deferred.queue.clone(),
8159 args: deferred.args,
8160 state: JobState::Cancelled,
8161 priority: deferred.priority,
8162 attempt: deferred.attempt,
8163 run_lease: deferred.run_lease,
8164 max_attempts: deferred.max_attempts,
8165 lane_seq: -deferred.job_id,
8166 enqueue_shard: 0,
8167 run_at: deferred.run_at,
8168 attempted_at: deferred.attempted_at,
8169 finalized_at: Utc::now(),
8170 created_at: deferred.created_at,
8171 unique_key: deferred.unique_key,
8172 unique_states: deferred.unique_states,
8173 payload: deferred.payload,
8174 };
8175 self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(deferred.state))
8176 .await?;
8177 return Ok(Some(CancelJobTxResult {
8178 row: done.into_job_row()?,
8179 claim_cursor_advance: None,
8180 }));
8181 }
8182
8183 Ok(None)
8184 }
8185
8186 pub async fn cancel_job(&self, pool: &PgPool, job_id: i64) -> Result<Option<JobRow>, AwaError> {
8187 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
8188 let result = self.cancel_job_tx(&mut tx, job_id).await?;
8189 tx.commit().await.map_err(map_sqlx_error)?;
8190 if let Some(result) = result {
8191 if let Some(advance) = result.claim_cursor_advance.as_ref() {
8192 self.advance_claim_cursors(pool, std::slice::from_ref(advance))
8193 .await;
8194 }
8195 Ok(Some(result.row))
8196 } else {
8197 Ok(None)
8198 }
8199 }
8200
8201 pub(crate) async fn cancel_job_in_tx<'a>(
8207 &self,
8208 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
8209 job_id: i64,
8210 ) -> Result<Option<JobRow>, AwaError> {
8211 Ok(self
8212 .cancel_job_tx(tx, job_id)
8213 .await?
8214 .map(|result| result.row))
8215 }
8216
8217 pub async fn cancel_jobs_by_ids(
8218 &self,
8219 pool: &PgPool,
8220 ids: &[i64],
8221 ) -> Result<Vec<JobRow>, AwaError> {
8222 if ids.is_empty() {
8223 return Ok(Vec::new());
8224 }
8225
8226 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
8227 let mut rows = Vec::with_capacity(ids.len());
8228 let mut claim_cursor_advances = Vec::new();
8229 for job_id in ids {
8230 if let Some(result) = self.cancel_job_tx(&mut tx, *job_id).await? {
8231 if let Some(advance) = result.claim_cursor_advance {
8232 claim_cursor_advances.push(advance);
8233 }
8234 rows.push(result.row);
8235 }
8236 }
8237 tx.commit().await.map_err(map_sqlx_error)?;
8238 self.advance_claim_cursors(pool, &claim_cursor_advances)
8239 .await;
8240 Ok(rows)
8241 }
8242
8243 pub async fn set_priority(
8244 &self,
8245 pool: &PgPool,
8246 job_id: i64,
8247 priority: i16,
8248 ) -> Result<bool, AwaError> {
8249 if !(1..=4).contains(&priority) {
8250 return Err(AwaError::Validation(
8251 "priority must be between 1 and 4".to_string(),
8252 ));
8253 }
8254
8255 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
8256 let result = self.set_priority_tx(&mut tx, job_id, priority).await?;
8257 tx.commit().await.map_err(map_sqlx_error)?;
8258 Ok(result)
8259 }
8260
8261 pub async fn set_priority_tx<'a>(
8262 &self,
8263 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
8264 job_id: i64,
8265 priority: i16,
8266 ) -> Result<bool, AwaError> {
8267 if !(1..=4).contains(&priority) {
8268 return Err(AwaError::Validation(
8269 "priority must be between 1 and 4".to_string(),
8270 ));
8271 }
8272
8273 if self
8274 .update_deferred_batch_fields_tx(tx, job_id, None, Some(priority))
8275 .await?
8276 {
8277 return Ok(true);
8278 }
8279
8280 let result = self
8281 .move_ready_batch_fields_tx(tx, job_id, None, Some(priority))
8282 .await?;
8283 Ok(result.moved)
8284 }
8285
8286 pub async fn move_queue(
8287 &self,
8288 pool: &PgPool,
8289 job_id: i64,
8290 queue: &str,
8291 priority: Option<i16>,
8292 ) -> Result<bool, AwaError> {
8293 if queue.is_empty() || queue.len() > 200 {
8294 return Err(AwaError::Validation(
8295 "destination queue must be 1..=200 characters".to_string(),
8296 ));
8297 }
8298 if let Some(priority) = priority {
8299 if !(1..=4).contains(&priority) {
8300 return Err(AwaError::Validation(
8301 "priority must be between 1 and 4".to_string(),
8302 ));
8303 }
8304 }
8305
8306 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
8307 let result = self.move_queue_tx(&mut tx, job_id, queue, priority).await?;
8308 tx.commit().await.map_err(map_sqlx_error)?;
8309 Ok(result)
8310 }
8311
8312 pub async fn move_queue_tx<'a>(
8313 &self,
8314 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
8315 job_id: i64,
8316 queue: &str,
8317 priority: Option<i16>,
8318 ) -> Result<bool, AwaError> {
8319 if queue.is_empty() || queue.len() > 200 {
8320 return Err(AwaError::Validation(
8321 "destination queue must be 1..=200 characters".to_string(),
8322 ));
8323 }
8324 if let Some(priority) = priority {
8325 if !(1..=4).contains(&priority) {
8326 return Err(AwaError::Validation(
8327 "priority must be between 1 and 4".to_string(),
8328 ));
8329 }
8330 }
8331
8332 if self
8333 .update_deferred_batch_fields_tx(tx, job_id, Some(queue), priority)
8334 .await?
8335 {
8336 return Ok(true);
8337 }
8338
8339 let result = self
8340 .move_ready_batch_fields_tx(tx, job_id, Some(queue), priority)
8341 .await?;
8342 Ok(result.moved)
8343 }
8344
8345 async fn update_deferred_batch_fields_tx<'a>(
8346 &self,
8347 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
8348 job_id: i64,
8349 queue: Option<&str>,
8350 priority: Option<i16>,
8351 ) -> Result<bool, AwaError> {
8352 let schema = self.schema();
8353 let row: Option<DeferredJobRow> = sqlx::query_as(&format!(
8354 r#"
8355 SELECT
8356 job_id,
8357 kind,
8358 queue,
8359 args,
8360 state,
8361 priority,
8362 attempt,
8363 run_lease,
8364 max_attempts,
8365 run_at,
8366 attempted_at,
8367 finalized_at,
8368 created_at,
8369 unique_key,
8370 unique_states,
8371 COALESCE(payload, '{{}}'::jsonb) AS payload
8372 FROM {schema}.deferred_jobs
8373 WHERE job_id = $1
8374 AND state = 'scheduled'
8375 FOR UPDATE SKIP LOCKED
8376 "#
8377 ))
8378 .bind(job_id)
8379 .fetch_optional(tx.as_mut())
8380 .await
8381 .map_err(map_sqlx_error)?;
8382
8383 let Some(row) = row else {
8384 return Ok(false);
8385 };
8386
8387 let old_queue = row.queue.clone();
8388 let old_priority = row.priority;
8389 let requested_queue = queue.unwrap_or(&old_queue);
8390 let old_logical_queue = self.logical_queue_name(&old_queue).to_string();
8391 let new_queue = if queue.is_some()
8392 && requested_queue != old_queue
8393 && requested_queue != old_logical_queue
8394 {
8395 self.queue_stripe_for_enqueue(requested_queue, &row.unique_key, row.job_id)
8396 } else {
8397 old_queue.clone()
8398 };
8399 let new_priority = priority.unwrap_or(old_priority);
8400 if new_queue == old_queue && new_priority == old_priority {
8401 return Ok(false);
8402 }
8403 let mut payload = RuntimePayload::from_json(row.payload)?;
8404 let metadata = payload.metadata.as_object_mut().ok_or_else(|| {
8405 AwaError::Validation("queue storage payload metadata must be a JSON object".to_string())
8406 })?;
8407 if queue.is_some() {
8408 metadata
8409 .entry("_awa_original_queue".to_string())
8410 .or_insert_with(|| serde_json::Value::from(old_logical_queue));
8411 }
8412 if priority.is_some() {
8413 metadata
8414 .entry("_awa_original_priority".to_string())
8415 .or_insert_with(|| serde_json::Value::from(i64::from(old_priority)));
8416 }
8417
8418 sqlx::query(&format!(
8419 r#"
8420 UPDATE {schema}.deferred_jobs
8421 SET queue = $2,
8422 priority = $3,
8423 payload = $4
8424 WHERE job_id = $1
8425 "#
8426 ))
8427 .bind(job_id)
8428 .bind(new_queue)
8429 .bind(new_priority)
8430 .bind(storage_payload(&payload.into_json()))
8431 .execute(tx.as_mut())
8432 .await
8433 .map_err(map_sqlx_error)?;
8434 Ok(true)
8435 }
8436
8437 async fn move_ready_batch_fields_tx<'a>(
8438 &self,
8439 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
8440 job_id: i64,
8441 queue: Option<&str>,
8442 priority: Option<i16>,
8443 ) -> Result<ReadyBatchMoveResult, AwaError> {
8444 let schema = self.schema();
8445 let ready: Option<ReadyTransitionRow> = sqlx::query_as(&format!(
8446 r#"
8447 WITH target AS (
8448 SELECT ready.*
8449 FROM {schema}.ready_entries AS ready
8450 JOIN {schema}.queue_claim_heads AS claims
8451 ON claims.queue = ready.queue
8452 AND claims.priority = ready.priority
8453 AND claims.enqueue_shard = ready.enqueue_shard
8454 WHERE ready.job_id = $1
8455 AND ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name)
8456 AND NOT EXISTS (
8457 SELECT 1 FROM {schema}.ready_tombstones AS tomb
8458 WHERE tomb.queue = ready.queue
8459 AND tomb.priority = ready.priority
8460 AND tomb.enqueue_shard = ready.enqueue_shard
8461 AND tomb.lane_seq = ready.lane_seq
8462 AND tomb.ready_slot = ready.ready_slot
8463 AND tomb.ready_generation = ready.ready_generation
8464 )
8465 ORDER BY ready.lane_seq DESC
8466 LIMIT 1
8467 FOR UPDATE OF ready SKIP LOCKED
8468 )
8469 SELECT
8470 ready_slot,
8471 ready_generation,
8472 job_id,
8473 kind,
8474 queue,
8475 args,
8476 priority,
8477 attempt,
8478 run_lease,
8479 max_attempts,
8480 lane_seq,
8481 enqueue_shard,
8482 run_at,
8483 attempted_at,
8484 created_at,
8485 unique_key,
8486 unique_states,
8487 COALESCE(payload, '{{}}'::jsonb) AS payload
8488 FROM target
8489 "#
8490 ))
8491 .bind(job_id)
8492 .fetch_optional(tx.as_mut())
8493 .await
8494 .map_err(map_sqlx_error)?;
8495
8496 let Some(ready) = ready else {
8497 return Ok(ReadyBatchMoveResult { moved: false });
8498 };
8499
8500 let old_queue = ready.queue.clone();
8501 let old_priority = ready.priority;
8502 let requested_queue = queue.unwrap_or(&old_queue);
8503 let old_logical_queue = self.logical_queue_name(&old_queue).to_string();
8504 let new_queue = if queue.is_some()
8505 && requested_queue != old_queue
8506 && requested_queue != old_logical_queue
8507 {
8508 self.queue_stripe_for_enqueue(requested_queue, &ready.unique_key, ready.job_id)
8509 } else {
8510 old_queue.clone()
8511 };
8512 let new_priority = priority.unwrap_or(old_priority);
8513 if new_queue == old_queue && new_priority == old_priority {
8514 return Ok(ReadyBatchMoveResult { moved: false });
8515 }
8516 sqlx::query(&format!(
8517 r#"
8518 INSERT INTO {schema}.ready_tombstones (
8519 ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, job_id
8520 )
8521 VALUES ($1, $2, $3, $4, $5, $6, $7)
8522 ON CONFLICT DO NOTHING
8523 "#
8524 ))
8525 .bind(ready.ready_slot)
8526 .bind(ready.ready_generation)
8527 .bind(&ready.queue)
8528 .bind(ready.priority)
8529 .bind(ready.enqueue_shard)
8530 .bind(ready.lane_seq)
8531 .bind(ready.job_id)
8532 .execute(tx.as_mut())
8533 .await
8534 .map_err(map_sqlx_error)?;
8535
8536 let mut payload = RuntimePayload::from_json(ready.payload.clone())?;
8537 let metadata = payload.metadata.as_object_mut().ok_or_else(|| {
8538 AwaError::Validation("queue storage payload metadata must be a JSON object".to_string())
8539 })?;
8540 if queue.is_some() {
8541 metadata
8542 .entry("_awa_original_queue".to_string())
8543 .or_insert_with(|| serde_json::Value::from(old_logical_queue));
8544 }
8545 if priority.is_some() {
8546 metadata
8547 .entry("_awa_original_priority".to_string())
8548 .or_insert_with(|| serde_json::Value::from(i64::from(old_priority)));
8549 }
8550
8551 let notify_queue = new_queue.clone();
8552 let ready_row = ready.into_existing_ready_row(new_queue, new_priority, payload.into_json());
8553 self.insert_existing_ready_rows_tx(tx, vec![ready_row], Some(JobState::Available))
8554 .await?;
8555 self.notify_queues_tx(tx, std::iter::once(notify_queue))
8556 .await?;
8557
8558 Ok(ReadyBatchMoveResult { moved: true })
8559 }
8560
8561 pub async fn age_waiting_priorities(
8562 &self,
8563 pool: &PgPool,
8564 aging_interval: Duration,
8565 limit: i64,
8566 ) -> Result<Vec<i64>, AwaError> {
8567 if limit <= 0 {
8568 return Ok(Vec::new());
8569 }
8570
8571 let cutoff = Utc::now()
8572 - TimeDelta::from_std(aging_interval)
8573 .map_err(|err| AwaError::Validation(format!("invalid aging interval: {err}")))?;
8574 let schema = self.schema();
8575 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
8576
8577 let moved: Vec<ReadyTransitionRow> = sqlx::query_as(&format!(
8578 r#"
8579 WITH target AS (
8580 SELECT ready.*
8581 FROM {schema}.ready_entries AS ready
8582 JOIN {schema}.queue_claim_heads AS claims
8583 ON claims.queue = ready.queue
8584 AND claims.priority = ready.priority
8585 AND claims.enqueue_shard = ready.enqueue_shard
8586 WHERE ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name)
8587 AND ready.priority > 1
8588 AND ready.run_at <= $1
8589 AND NOT EXISTS (
8590 SELECT 1 FROM {schema}.ready_tombstones AS tomb
8591 WHERE tomb.queue = ready.queue
8592 AND tomb.priority = ready.priority
8593 AND tomb.enqueue_shard = ready.enqueue_shard
8594 AND tomb.lane_seq = ready.lane_seq
8595 AND tomb.ready_slot = ready.ready_slot
8596 AND tomb.ready_generation = ready.ready_generation
8597 )
8598 ORDER BY ready.run_at ASC, ready.lane_seq ASC
8599 LIMIT $2
8600 FOR UPDATE OF ready SKIP LOCKED
8601 ),
8602 tombstones AS (
8603 INSERT INTO {schema}.ready_tombstones (
8604 ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, job_id
8605 )
8606 SELECT ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, job_id
8607 FROM target
8608 ON CONFLICT DO NOTHING
8609 )
8610 SELECT
8611 ready_slot,
8612 ready_generation,
8613 job_id,
8614 kind,
8615 queue,
8616 args,
8617 priority,
8618 attempt,
8619 run_lease,
8620 max_attempts,
8621 lane_seq,
8622 enqueue_shard,
8623 run_at,
8624 attempted_at,
8625 created_at,
8626 unique_key,
8627 unique_states,
8628 COALESCE(payload, '{{}}'::jsonb) AS payload
8629 FROM target
8630 "#
8631 ))
8632 .bind(cutoff)
8633 .bind(limit)
8634 .fetch_all(tx.as_mut())
8635 .await
8636 .map_err(map_sqlx_error)?;
8637
8638 if moved.is_empty() {
8639 tx.commit().await.map_err(map_sqlx_error)?;
8640 return Ok(Vec::new());
8641 }
8642
8643 let mut ids = Vec::with_capacity(moved.len());
8644 let mut queues = BTreeSet::new();
8645 let mut ready_rows = Vec::with_capacity(moved.len());
8646
8647 for row in moved {
8648 ids.push(row.job_id);
8649 queues.insert(row.queue.clone());
8650
8651 let mut payload = RuntimePayload::from_json(row.payload)?;
8652 let metadata = payload.metadata.as_object_mut().ok_or_else(|| {
8653 AwaError::Validation(
8654 "queue storage payload metadata must be a JSON object".to_string(),
8655 )
8656 })?;
8657 metadata
8658 .entry("_awa_original_priority".to_string())
8659 .or_insert_with(|| serde_json::Value::from(i64::from(row.priority)));
8660
8661 ready_rows.push(ExistingReadyRow {
8662 job_id: row.job_id,
8663 kind: row.kind,
8664 queue: row.queue,
8665 args: row.args,
8666 priority: row.priority - 1,
8667 attempt: row.attempt,
8668 run_lease: row.run_lease,
8669 max_attempts: row.max_attempts,
8670 run_at: row.run_at,
8671 attempted_at: row.attempted_at,
8672 created_at: row.created_at,
8673 unique_key: row.unique_key,
8674 unique_states: row.unique_states,
8675 payload: payload.into_json(),
8676 });
8677 }
8678
8679 self.insert_existing_ready_rows_tx(&mut tx, ready_rows, Some(JobState::Available))
8684 .await?;
8685 self.notify_queues_tx(&mut tx, queues).await?;
8686 tx.commit().await.map_err(map_sqlx_error)?;
8687 Ok(ids)
8688 }
8689
8690 fn with_progress(
8691 payload: serde_json::Value,
8692 progress: Option<serde_json::Value>,
8693 ) -> Result<serde_json::Value, AwaError> {
8694 let mut payload = RuntimePayload::from_json(payload)?;
8695 payload.set_progress(progress);
8696 Ok(payload.into_json())
8697 }
8698
8699 async fn take_callback_result(
8700 &self,
8701 pool: &PgPool,
8702 job_id: i64,
8703 run_lease: i64,
8704 ) -> Result<serde_json::Value, AwaError> {
8705 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
8706 let mut row: Option<AttemptStateRow> = sqlx::query_as(&format!(
8707 r#"
8708 SELECT
8709 job_id,
8710 run_lease,
8711 progress,
8712 callback_filter,
8713 callback_on_complete,
8714 callback_on_fail,
8715 callback_transform,
8716 callback_result
8717 FROM {}
8718 WHERE job_id = $1
8719 AND run_lease = $2
8720 FOR UPDATE
8721 "#,
8722 self.attempt_state_table()
8723 ))
8724 .bind(job_id)
8725 .bind(run_lease)
8726 .fetch_optional(tx.as_mut())
8727 .await
8728 .map_err(map_sqlx_error)?;
8729
8730 let Some(mut row) = row.take() else {
8731 tx.commit().await.map_err(map_sqlx_error)?;
8732 return Ok(serde_json::Value::Null);
8733 };
8734
8735 let result = row
8736 .callback_result
8737 .take()
8738 .unwrap_or(serde_json::Value::Null);
8739
8740 if row.progress.is_none()
8741 && row.callback_filter.is_none()
8742 && row.callback_on_complete.is_none()
8743 && row.callback_on_fail.is_none()
8744 && row.callback_transform.is_none()
8745 {
8746 sqlx::query(&format!(
8747 "DELETE FROM {} WHERE job_id = $1 AND run_lease = $2",
8748 self.attempt_state_table()
8749 ))
8750 .bind(job_id)
8751 .bind(run_lease)
8752 .execute(tx.as_mut())
8753 .await
8754 .map_err(map_sqlx_error)?;
8755 } else {
8756 sqlx::query(&format!(
8757 "UPDATE {} SET callback_result = NULL, updated_at = clock_timestamp() WHERE job_id = $1 AND run_lease = $2",
8758 self.attempt_state_table()
8759 ))
8760 .bind(job_id)
8761 .bind(run_lease)
8762 .execute(tx.as_mut())
8763 .await
8764 .map_err(map_sqlx_error)?;
8765 }
8766
8767 tx.commit().await.map_err(map_sqlx_error)?;
8768 Ok(result)
8769 }
8770
8771 async fn backoff_at_tx<'a>(
8772 &self,
8773 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
8774 attempt: i16,
8775 max_attempts: i16,
8776 ) -> Result<DateTime<Utc>, AwaError> {
8777 sqlx::query_scalar("SELECT clock_timestamp() + awa.backoff_duration($1, $2)")
8778 .bind(attempt)
8779 .bind(max_attempts)
8780 .fetch_one(tx.as_mut())
8781 .await
8782 .map_err(map_sqlx_error)
8783 }
8784
8785 async fn notify_queues_tx<'a>(
8786 &self,
8787 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
8788 queues: impl IntoIterator<Item = String>,
8789 ) -> Result<(), AwaError> {
8790 let channels: Vec<String> = queues
8794 .into_iter()
8795 .map(|queue| format!("awa:{}", self.logical_queue_name(&queue)))
8796 .collect::<BTreeSet<String>>()
8797 .into_iter()
8798 .collect();
8799 if channels.is_empty() {
8800 return Ok(());
8801 }
8802 sqlx::query("SELECT pg_notify(channel, '') FROM unnest($1::text[]) AS channel")
8803 .bind(&channels)
8804 .execute(tx.as_mut())
8805 .await
8806 .map_err(map_sqlx_error)?;
8807 Ok(())
8808 }
8809
8810 async fn lock_receipt_attempts_tx<'a>(
8811 &self,
8812 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
8813 jobs: &[(i64, i64)],
8814 ) -> Result<(), AwaError> {
8815 if jobs.is_empty() {
8816 return Ok(());
8817 }
8818
8819 let unique_jobs: BTreeSet<(i64, i64)> = jobs.iter().copied().collect();
8820 let job_ids: Vec<i64> = unique_jobs.iter().map(|(job_id, _)| *job_id).collect();
8821 let run_leases: Vec<i64> = unique_jobs
8822 .iter()
8823 .map(|(_, run_lease)| *run_lease)
8824 .collect();
8825
8826 sqlx::query(
8827 r#"
8828 WITH locks AS MATERIALIZED (
8829 SELECT DISTINCT pg_catalog.hashtextextended(
8830 format('awa.receipt.complete:%s:%s', job_id, run_lease),
8831 0
8832 ) AS lock_key
8833 FROM unnest($1::bigint[], $2::bigint[]) AS input(job_id, run_lease)
8834 ORDER BY lock_key
8835 )
8836 SELECT pg_catalog.pg_advisory_xact_lock(lock_key)
8837 FROM locks
8838 "#,
8839 )
8840 .bind(&job_ids)
8841 .bind(&run_leases)
8842 .execute(tx.as_mut())
8843 .await
8844 .map_err(map_sqlx_error)?;
8845
8846 Ok(())
8847 }
8848
8849 async fn ensure_running_leases_from_receipts_tx<'a>(
8850 &self,
8851 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
8852 jobs: &[(i64, i64)],
8853 ) -> Result<usize, AwaError> {
8854 if jobs.is_empty() {
8855 return Ok(0);
8856 }
8857
8858 self.lock_receipt_attempts_tx(tx, jobs).await?;
8859
8860 let schema = self.schema();
8861 let closure_rel = format!("{schema}.lease_claim_closures");
8862 let closure_batch_rel = format!("{schema}.lease_claim_closure_batches");
8863 let closed_evidence =
8864 receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims");
8865 let job_ids: Vec<i64> = jobs.iter().map(|(job_id, _)| *job_id).collect();
8866 let run_leases: Vec<i64> = jobs.iter().map(|(_, run_lease)| *run_lease).collect();
8867 let inserted: i64 = sqlx::query_scalar(&format!(
8868 r#"
8869 WITH inflight(job_id, run_lease) AS (
8870 SELECT * FROM unnest($1::bigint[], $2::bigint[])
8871 ),
8872 lease_ring AS (
8873 SELECT current_slot AS lease_slot, generation AS lease_generation
8874 FROM {schema}.lease_ring_state
8875 WHERE singleton = TRUE
8876 ),
8877 row_claim_refs AS (
8878 -- Source claim metadata directly from the partitioned
8879 -- lease_claims table anti-joined against every durable
8880 -- closure evidence shape.
8881 SELECT
8882 claims.claim_slot,
8883 claims.job_id,
8884 claims.run_lease,
8885 claims.ready_slot,
8886 claims.ready_generation,
8887 claims.queue,
8888 claims.priority,
8889 claims.attempt,
8890 claims.max_attempts,
8891 claims.lane_seq,
8892 claims.enqueue_shard,
8893 claims.claimed_at,
8894 claims.deadline_at
8895 FROM {schema}.lease_claims AS claims
8896 JOIN inflight
8897 ON inflight.job_id = claims.job_id
8898 AND inflight.run_lease = claims.run_lease
8899 WHERE NOT {closed_evidence}
8900 FOR UPDATE OF claims
8901 ),
8902 batch_claim_refs AS (
8903 SELECT
8904 claim_batches.claim_slot,
8905 items.job_id,
8906 items.run_lease,
8907 claim_batches.ready_slot,
8908 claim_batches.ready_generation,
8909 claim_batches.queue,
8910 claim_batches.priority,
8911 items.attempt,
8912 items.max_attempts,
8913 items.lane_seq,
8914 claim_batches.enqueue_shard,
8915 claim_batches.claimed_at,
8916 claim_batches.deadline_at
8917 FROM {schema}.lease_claim_batches AS claim_batches
8918 CROSS JOIN LATERAL unnest(
8919 claim_batches.job_ids,
8920 claim_batches.run_leases,
8921 claim_batches.receipt_ids,
8922 claim_batches.lane_seqs,
8923 claim_batches.attempts,
8924 claim_batches.max_attempts
8925 ) AS items(job_id, run_lease, receipt_id, lane_seq, attempt, max_attempts)
8926 JOIN inflight
8927 ON inflight.job_id = items.job_id
8928 AND inflight.run_lease = items.run_lease
8929 WHERE NOT EXISTS (
8930 SELECT 1
8931 FROM {schema}.lease_claim_closures AS closures
8932 WHERE closures.claim_slot = claim_batches.claim_slot
8933 AND closures.job_id = items.job_id
8934 AND closures.run_lease = items.run_lease
8935 )
8936 AND NOT EXISTS (
8937 SELECT 1
8938 FROM {schema}.lease_claim_closure_batches AS closure_batches
8939 WHERE closure_batches.claim_slot = claim_batches.claim_slot
8940 AND closure_batches.receipt_ranges @> items.receipt_id
8941 )
8942 AND NOT EXISTS (
8943 SELECT 1
8944 FROM {schema}.done_entries AS done
8945 WHERE done.job_id = items.job_id
8946 AND done.run_lease = items.run_lease
8947 )
8948 AND NOT EXISTS (
8949 SELECT 1
8950 FROM {schema}.deferred_jobs AS deferred
8951 WHERE deferred.job_id = items.job_id
8952 AND deferred.run_lease = items.run_lease
8953 )
8954 AND NOT EXISTS (
8955 SELECT 1
8956 FROM {schema}.dlq_entries AS dlq
8957 WHERE dlq.job_id = items.job_id
8958 AND dlq.run_lease = items.run_lease
8959 )
8960 FOR UPDATE OF claim_batches
8961 ),
8962 claim_refs AS (
8963 SELECT * FROM row_claim_refs
8964 UNION ALL
8965 SELECT * FROM batch_claim_refs
8966 ),
8967 already_live AS (
8968 SELECT claim_refs.job_id, claim_refs.run_lease
8969 FROM claim_refs
8970 WHERE EXISTS (
8971 SELECT 1
8972 FROM {schema}.leases AS lease
8973 WHERE lease.job_id = claim_refs.job_id
8974 AND lease.run_lease = claim_refs.run_lease
8975 )
8976 ),
8977 inserted AS (
8978 INSERT INTO {schema}.leases (
8979 lease_slot,
8980 lease_generation,
8981 ready_slot,
8982 ready_generation,
8983 job_id,
8984 queue,
8985 state,
8986 priority,
8987 attempt,
8988 run_lease,
8989 max_attempts,
8990 lane_seq,
8991 enqueue_shard,
8992 heartbeat_at,
8993 deadline_at,
8994 attempted_at
8995 )
8996 SELECT
8997 lease_ring.lease_slot,
8998 lease_ring.lease_generation,
8999 claim_refs.ready_slot,
9000 claim_refs.ready_generation,
9001 claim_refs.job_id,
9002 claim_refs.queue,
9003 'running'::awa.job_state,
9004 claim_refs.priority,
9005 claim_refs.attempt,
9006 claim_refs.run_lease,
9007 claim_refs.max_attempts,
9008 claim_refs.lane_seq,
9009 claim_refs.enqueue_shard,
9010 clock_timestamp(),
9011 -- Preserve the per-claim deadline so the lease-side
9012 -- deadline rescue path picks up materialized claims
9013 -- without an extra hop. NULL when receipts mode is
9014 -- on with `deadline_duration = 0` (the short-job
9015 -- shape that needs no deadline at all).
9016 claim_refs.deadline_at,
9017 claim_refs.claimed_at
9018 FROM claim_refs
9019 CROSS JOIN lease_ring
9020 WHERE NOT EXISTS (
9021 SELECT 1
9022 FROM {schema}.leases AS lease
9023 WHERE lease.job_id = claim_refs.job_id
9024 AND lease.run_lease = claim_refs.run_lease
9025 )
9026 RETURNING job_id, run_lease
9027 ),
9028 marked AS (
9029 UPDATE {schema}.lease_claims AS claims
9030 SET materialized_at = clock_timestamp()
9031 FROM (
9032 SELECT job_id, run_lease FROM inserted
9033 UNION
9034 SELECT job_id, run_lease FROM already_live
9035 ) AS moved
9036 WHERE claims.job_id = moved.job_id
9037 AND claims.run_lease = moved.run_lease
9038 RETURNING claims.job_id
9039 )
9040 SELECT count(*)::bigint
9041 FROM (
9042 SELECT job_id, run_lease FROM inserted
9043 UNION
9044 SELECT job_id, run_lease FROM already_live
9045 ) AS moved
9046 "#
9047 ))
9048 .bind(&job_ids)
9049 .bind(&run_leases)
9050 .fetch_one(tx.as_mut())
9051 .await
9052 .map_err(map_sqlx_error)?;
9053 Ok(inserted as usize)
9054 }
9055
9056 async fn ensure_mutable_running_attempt_tx<'a>(
9057 &self,
9058 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
9059 job_id: i64,
9060 run_lease: i64,
9061 ) -> Result<(), AwaError> {
9062 if self.lease_claim_receipts() {
9063 self.ensure_running_leases_from_receipts_tx(tx, &[(job_id, run_lease)])
9064 .await?;
9065 }
9066 Ok(())
9067 }
9068
9069 async fn upsert_attempt_state_from_receipts_tx<'a>(
9070 &self,
9071 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
9072 jobs: &[(i64, i64)],
9073 ) -> Result<usize, AwaError> {
9074 if jobs.is_empty() {
9075 return Ok(0);
9076 }
9077
9078 self.lock_receipt_attempts_tx(tx, jobs).await?;
9079
9080 let schema = self.schema();
9081 let closure_rel = format!("{schema}.lease_claim_closures");
9082 let closure_batch_rel = format!("{schema}.lease_claim_closure_batches");
9083 let closed_evidence =
9084 receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims");
9085 let job_ids: Vec<i64> = jobs.iter().map(|(job_id, _)| *job_id).collect();
9086 let run_leases: Vec<i64> = jobs.iter().map(|(_, run_lease)| *run_lease).collect();
9087 let updated: i64 = sqlx::query_scalar(&format!(
9088 r#"
9089 WITH inflight(job_id, run_lease) AS (
9090 SELECT * FROM unnest($1::bigint[], $2::bigint[])
9091 ),
9092 row_claim_refs AS (
9093 -- Source open-claim identity from lease_claims
9094 -- anti-joined against every durable closure evidence
9095 -- shape.
9096 SELECT claims.job_id, claims.run_lease
9097 FROM {schema}.lease_claims AS claims
9098 JOIN inflight
9099 ON inflight.job_id = claims.job_id
9100 AND inflight.run_lease = claims.run_lease
9101 WHERE NOT {closed_evidence}
9102 FOR UPDATE OF claims
9103 ),
9104 batch_claim_refs AS (
9105 SELECT items.job_id, items.run_lease
9106 FROM {schema}.lease_claim_batches AS claim_batches
9107 CROSS JOIN LATERAL unnest(
9108 claim_batches.job_ids,
9109 claim_batches.run_leases,
9110 claim_batches.receipt_ids
9111 ) AS items(job_id, run_lease, receipt_id)
9112 JOIN inflight
9113 ON inflight.job_id = items.job_id
9114 AND inflight.run_lease = items.run_lease
9115 WHERE NOT EXISTS (
9116 SELECT 1
9117 FROM {schema}.lease_claim_closures AS closures
9118 WHERE closures.claim_slot = claim_batches.claim_slot
9119 AND closures.job_id = items.job_id
9120 AND closures.run_lease = items.run_lease
9121 )
9122 AND NOT EXISTS (
9123 SELECT 1
9124 FROM {schema}.lease_claim_closure_batches AS closure_batches
9125 WHERE closure_batches.claim_slot = claim_batches.claim_slot
9126 AND closure_batches.receipt_ranges @> items.receipt_id
9127 )
9128 AND NOT EXISTS (
9129 SELECT 1 FROM {schema}.done_entries AS done
9130 WHERE done.job_id = items.job_id
9131 AND done.run_lease = items.run_lease
9132 )
9133 AND NOT EXISTS (
9134 SELECT 1 FROM {schema}.deferred_jobs AS deferred
9135 WHERE deferred.job_id = items.job_id
9136 AND deferred.run_lease = items.run_lease
9137 )
9138 AND NOT EXISTS (
9139 SELECT 1 FROM {schema}.dlq_entries AS dlq
9140 WHERE dlq.job_id = items.job_id
9141 AND dlq.run_lease = items.run_lease
9142 )
9143 FOR UPDATE OF claim_batches
9144 ),
9145 claim_refs AS (
9146 SELECT * FROM row_claim_refs
9147 UNION ALL
9148 SELECT * FROM batch_claim_refs
9149 ),
9150 upserted AS (
9151 INSERT INTO {schema}.attempt_state (job_id, run_lease, heartbeat_at, updated_at)
9152 SELECT claim_refs.job_id, claim_refs.run_lease, clock_timestamp(), clock_timestamp()
9153 FROM claim_refs
9154 ON CONFLICT (job_id, run_lease)
9155 DO UPDATE SET
9156 heartbeat_at = clock_timestamp(),
9157 updated_at = clock_timestamp()
9158 RETURNING job_id, run_lease
9159 ),
9160 marked AS (
9161 UPDATE {schema}.lease_claims AS claims
9162 SET materialized_at = COALESCE(claims.materialized_at, clock_timestamp())
9163 FROM row_claim_refs
9164 WHERE claims.job_id = row_claim_refs.job_id
9165 AND claims.run_lease = row_claim_refs.run_lease
9166 RETURNING claims.job_id
9167 )
9168 SELECT count(*)::bigint FROM upserted
9169 "#
9170 ))
9171 .bind(&job_ids)
9172 .bind(&run_leases)
9173 .fetch_one(tx.as_mut())
9174 .await
9175 .map_err(map_sqlx_error)?;
9176 Ok(updated as usize)
9177 }
9178
9179 async fn upsert_attempt_state_progress_from_receipts_tx<'a>(
9180 &self,
9181 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
9182 jobs: &[(i64, i64, serde_json::Value)],
9183 ) -> Result<usize, AwaError> {
9184 if jobs.is_empty() {
9185 return Ok(0);
9186 }
9187
9188 let receipt_jobs: Vec<(i64, i64)> = jobs
9189 .iter()
9190 .map(|(job_id, run_lease, _)| (*job_id, *run_lease))
9191 .collect();
9192 self.lock_receipt_attempts_tx(tx, &receipt_jobs).await?;
9193
9194 let schema = self.schema();
9195 let closure_rel = format!("{schema}.lease_claim_closures");
9196 let closure_batch_rel = format!("{schema}.lease_claim_closure_batches");
9197 let closed_evidence =
9198 receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims");
9199 let job_ids: Vec<i64> = jobs.iter().map(|(job_id, _, _)| *job_id).collect();
9200 let run_leases: Vec<i64> = jobs.iter().map(|(_, run_lease, _)| *run_lease).collect();
9201 let progress: Vec<serde_json::Value> = jobs
9202 .iter()
9203 .map(|(_, _, progress)| progress.clone())
9204 .collect();
9205 let updated: i64 = sqlx::query_scalar(&format!(
9206 r#"
9207 WITH inflight(job_id, run_lease, progress) AS (
9208 SELECT * FROM unnest($1::bigint[], $2::bigint[], $3::jsonb[])
9209 ),
9210 row_claim_refs AS (
9211 -- Same anti-join pattern as the heartbeat-only path
9212 -- above.
9213 SELECT claims.job_id, claims.run_lease, inflight.progress
9214 FROM {schema}.lease_claims AS claims
9215 JOIN inflight
9216 ON inflight.job_id = claims.job_id
9217 AND inflight.run_lease = claims.run_lease
9218 WHERE NOT {closed_evidence}
9219 FOR UPDATE OF claims
9220 ),
9221 batch_claim_refs AS (
9222 SELECT items.job_id, items.run_lease, inflight.progress
9223 FROM {schema}.lease_claim_batches AS claim_batches
9224 CROSS JOIN LATERAL unnest(
9225 claim_batches.job_ids,
9226 claim_batches.run_leases,
9227 claim_batches.receipt_ids
9228 ) AS items(job_id, run_lease, receipt_id)
9229 JOIN inflight
9230 ON inflight.job_id = items.job_id
9231 AND inflight.run_lease = items.run_lease
9232 WHERE NOT EXISTS (
9233 SELECT 1
9234 FROM {schema}.lease_claim_closures AS closures
9235 WHERE closures.claim_slot = claim_batches.claim_slot
9236 AND closures.job_id = items.job_id
9237 AND closures.run_lease = items.run_lease
9238 )
9239 AND NOT EXISTS (
9240 SELECT 1
9241 FROM {schema}.lease_claim_closure_batches AS closure_batches
9242 WHERE closure_batches.claim_slot = claim_batches.claim_slot
9243 AND closure_batches.receipt_ranges @> items.receipt_id
9244 )
9245 AND NOT EXISTS (
9246 SELECT 1 FROM {schema}.done_entries AS done
9247 WHERE done.job_id = items.job_id
9248 AND done.run_lease = items.run_lease
9249 )
9250 AND NOT EXISTS (
9251 SELECT 1 FROM {schema}.deferred_jobs AS deferred
9252 WHERE deferred.job_id = items.job_id
9253 AND deferred.run_lease = items.run_lease
9254 )
9255 AND NOT EXISTS (
9256 SELECT 1 FROM {schema}.dlq_entries AS dlq
9257 WHERE dlq.job_id = items.job_id
9258 AND dlq.run_lease = items.run_lease
9259 )
9260 FOR UPDATE OF claim_batches
9261 ),
9262 claim_refs AS (
9263 SELECT * FROM row_claim_refs
9264 UNION ALL
9265 SELECT * FROM batch_claim_refs
9266 ),
9267 upserted AS (
9268 INSERT INTO {schema}.attempt_state (
9269 job_id,
9270 run_lease,
9271 heartbeat_at,
9272 progress,
9273 updated_at
9274 )
9275 SELECT
9276 claim_refs.job_id,
9277 claim_refs.run_lease,
9278 clock_timestamp(),
9279 claim_refs.progress,
9280 clock_timestamp()
9281 FROM claim_refs
9282 ON CONFLICT (job_id, run_lease)
9283 DO UPDATE SET
9284 heartbeat_at = clock_timestamp(),
9285 progress = EXCLUDED.progress,
9286 updated_at = clock_timestamp()
9287 RETURNING job_id, run_lease
9288 ),
9289 marked AS (
9290 UPDATE {schema}.lease_claims AS claims
9291 SET materialized_at = COALESCE(claims.materialized_at, clock_timestamp())
9292 FROM row_claim_refs
9293 WHERE claims.job_id = row_claim_refs.job_id
9294 AND claims.run_lease = row_claim_refs.run_lease
9295 RETURNING claims.job_id
9296 )
9297 SELECT count(*)::bigint FROM upserted
9298 "#
9299 ))
9300 .bind(&job_ids)
9301 .bind(&run_leases)
9302 .bind(&progress)
9303 .fetch_one(tx.as_mut())
9304 .await
9305 .map_err(map_sqlx_error)?;
9306 Ok(updated as usize)
9307 }
9308
9309 async fn hydrate_deleted_leases_tx<'a>(
9310 &self,
9311 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
9312 deleted: Vec<DeletedLeaseRow>,
9313 ) -> Result<Vec<LeaseTransitionRow>, AwaError> {
9314 if deleted.is_empty() {
9315 return Ok(Vec::new());
9316 }
9317
9318 let schema = self.schema();
9319 let ready_slots: Vec<i32> = deleted.iter().map(|row| row.ready_slot).collect();
9320 let ready_generations: Vec<i64> = deleted.iter().map(|row| row.ready_generation).collect();
9321 let queues: Vec<String> = deleted.iter().map(|row| row.queue.clone()).collect();
9322 let enqueue_shards: Vec<i16> = deleted.iter().map(|row| row.enqueue_shard).collect();
9323 let lane_seqs: Vec<i64> = deleted.iter().map(|row| row.lane_seq).collect();
9324 let job_ids: Vec<i64> = deleted.iter().map(|row| row.job_id).collect();
9325 let run_leases: Vec<i64> = deleted.iter().map(|row| row.run_lease).collect();
9326
9327 let ready_rows: Vec<ReadySnapshotRow> = sqlx::query_as(&format!(
9328 r#"
9329 WITH refs(ready_slot, ready_generation, queue, enqueue_shard, lane_seq, job_id) AS (
9330 SELECT * FROM unnest($1::int[], $2::bigint[], $3::text[], $4::smallint[], $5::bigint[], $6::bigint[])
9331 )
9332 SELECT
9333 ready.ready_slot,
9334 ready.ready_generation,
9335 ready.job_id,
9336 ready.kind,
9337 ready.queue,
9338 ready.args,
9339 ready.lane_seq,
9340 ready.enqueue_shard,
9341 ready.run_at,
9342 ready.created_at,
9343 ready.unique_key,
9344 ready.unique_states,
9345 COALESCE(ready.payload, '{{}}'::jsonb) AS payload
9346 FROM refs
9347 JOIN {schema}.ready_entries AS ready
9348 ON ready.ready_slot = refs.ready_slot
9349 AND ready.ready_generation = refs.ready_generation
9350 AND ready.queue = refs.queue
9351 AND ready.enqueue_shard = refs.enqueue_shard
9352 AND ready.lane_seq = refs.lane_seq
9353 AND ready.job_id = refs.job_id
9354 "#
9355 ))
9356 .bind(&ready_slots)
9357 .bind(&ready_generations)
9358 .bind(&queues)
9359 .bind(&enqueue_shards)
9360 .bind(&lane_seqs)
9361 .bind(&job_ids)
9362 .fetch_all(tx.as_mut())
9363 .await
9364 .map_err(map_sqlx_error)?;
9365
9366 let attempt_rows: Vec<AttemptStateRow> = sqlx::query_as(&format!(
9367 r#"
9368 WITH refs(job_id, run_lease) AS (
9369 SELECT * FROM unnest($1::bigint[], $2::bigint[])
9370 )
9371 DELETE FROM {schema}.attempt_state AS attempt
9372 USING refs
9373 WHERE attempt.job_id = refs.job_id
9374 AND attempt.run_lease = refs.run_lease
9375 RETURNING
9376 attempt.job_id,
9377 attempt.run_lease,
9378 attempt.progress,
9379 attempt.callback_filter,
9380 attempt.callback_on_complete,
9381 attempt.callback_on_fail,
9382 attempt.callback_transform,
9383 attempt.callback_result
9384 "#
9385 ))
9386 .bind(&job_ids)
9387 .bind(&run_leases)
9388 .fetch_all(tx.as_mut())
9389 .await
9390 .map_err(map_sqlx_error)?;
9391
9392 sqlx::query(&format!(
9407 r#"
9408 WITH refs(job_id, run_lease) AS (
9409 SELECT * FROM unnest($1::bigint[], $2::bigint[])
9410 ),
9411 row_claims AS (
9412 SELECT claims.claim_slot, claims.job_id, claims.run_lease,
9413 'rescue', clock_timestamp()
9414 FROM {schema}.lease_claims AS claims
9415 JOIN refs
9416 ON refs.job_id = claims.job_id
9417 AND refs.run_lease = claims.run_lease
9418 ),
9419 -- Compact batch-sourced claims (including those materialized
9420 -- into `leases` without a lease_claims row) have no row to
9421 -- JOIN, so they close into the batch ledger the queue prune
9422 -- gate counts via compact_count. The double-close guards
9423 -- below keep a re-hydrated rescued receipt from writing a
9424 -- second batch closure.
9425 batch_claims AS (
9426 SELECT
9427 claim_batches.claim_slot,
9428 claim_batches.ready_slot,
9429 claim_batches.ready_generation,
9430 items.job_id,
9431 items.run_lease,
9432 items.receipt_id
9433 FROM {schema}.lease_claim_batches AS claim_batches
9434 CROSS JOIN LATERAL unnest(
9435 claim_batches.job_ids,
9436 claim_batches.run_leases,
9437 claim_batches.receipt_ids
9438 ) AS items(job_id, run_lease, receipt_id)
9439 JOIN refs
9440 ON refs.job_id = items.job_id
9441 AND refs.run_lease = items.run_lease
9442 WHERE NOT EXISTS (
9443 SELECT 1
9444 FROM {schema}.lease_claim_closures AS closures
9445 WHERE closures.claim_slot = claim_batches.claim_slot
9446 AND closures.job_id = items.job_id
9447 AND closures.run_lease = items.run_lease
9448 )
9449 AND NOT EXISTS (
9450 SELECT 1
9451 FROM {schema}.lease_claim_closure_batches AS closure_batches
9452 WHERE closure_batches.claim_slot = claim_batches.claim_slot
9453 AND closure_batches.receipt_ranges @> items.receipt_id
9454 )
9455 ),
9456 inserted AS (
9457 INSERT INTO {schema}.lease_claim_closures
9458 (claim_slot, job_id, run_lease, outcome, closed_at)
9459 SELECT * FROM row_claims
9460 ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING
9461 RETURNING claim_slot, job_id, run_lease, closed_at
9462 ),
9463 inserted_batches AS (
9464 INSERT INTO {schema}.lease_claim_closure_batches (
9465 claim_slot,
9466 ready_slot,
9467 ready_generation,
9468 outcome,
9469 closed_count,
9470 receipt_ids,
9471 receipt_ranges,
9472 closed_at
9473 )
9474 SELECT
9475 batch_claims.claim_slot,
9476 batch_claims.ready_slot,
9477 batch_claims.ready_generation,
9478 'rescue',
9479 count(*)::int,
9480 array_agg(batch_claims.receipt_id ORDER BY batch_claims.receipt_id),
9481 range_agg(int8range(batch_claims.receipt_id, batch_claims.receipt_id + 1, '[)') ORDER BY batch_claims.receipt_id),
9482 clock_timestamp()
9483 FROM batch_claims
9484 GROUP BY
9485 batch_claims.claim_slot,
9486 batch_claims.ready_slot,
9487 batch_claims.ready_generation
9488 RETURNING claim_slot
9489 ),
9490 marked AS (
9491 UPDATE {schema}.lease_claims AS claims
9492 SET closed_at = COALESCE(claims.closed_at, inserted.closed_at)
9493 FROM inserted
9494 WHERE claims.claim_slot = inserted.claim_slot
9495 AND claims.job_id = inserted.job_id
9496 AND claims.run_lease = inserted.run_lease
9497 RETURNING claims.job_id
9498 )
9499 SELECT
9500 (SELECT count(*) FROM marked)
9501 + (SELECT count(*) FROM inserted_batches)
9502 "#
9503 ))
9504 .bind(&job_ids)
9505 .bind(&run_leases)
9506 .execute(tx.as_mut())
9507 .await
9508 .map_err(map_sqlx_error)?;
9509
9510 let ready_map: BTreeMap<(i32, i64, String, i16, i64, i64), ReadySnapshotRow> = ready_rows
9511 .into_iter()
9512 .map(|row| {
9513 (
9514 (
9515 row.ready_slot,
9516 row.ready_generation,
9517 row.queue.clone(),
9518 row.enqueue_shard,
9519 row.lane_seq,
9520 row.job_id,
9521 ),
9522 row,
9523 )
9524 })
9525 .collect();
9526
9527 let attempt_map: BTreeMap<(i64, i64), AttemptStateRow> = attempt_rows
9528 .into_iter()
9529 .map(|row| ((row.job_id, row.run_lease), row))
9530 .collect();
9531
9532 let mut hydrated = Vec::with_capacity(deleted.len());
9533 for deleted_row in deleted {
9534 let ready = ready_map
9535 .get(&(
9536 deleted_row.ready_slot,
9537 deleted_row.ready_generation,
9538 deleted_row.queue.clone(),
9539 deleted_row.enqueue_shard,
9540 deleted_row.lane_seq,
9541 deleted_row.job_id,
9542 ))
9543 .ok_or_else(|| {
9544 AwaError::Validation(format!(
9545 "queue storage ready row missing for deleted lease job {} run_lease {}",
9546 deleted_row.job_id, deleted_row.run_lease
9547 ))
9548 })?;
9549 let attempt = attempt_map.get(&(deleted_row.job_id, deleted_row.run_lease));
9550
9551 hydrated.push(LeaseTransitionRow {
9552 ready_slot: deleted_row.ready_slot,
9553 ready_generation: deleted_row.ready_generation,
9554 job_id: deleted_row.job_id,
9555 kind: ready.kind.clone(),
9556 queue: ready.queue.clone(),
9557 args: ready.args.clone(),
9558 state: deleted_row.state,
9559 priority: deleted_row.priority,
9560 attempt: deleted_row.attempt,
9561 run_lease: deleted_row.run_lease,
9562 max_attempts: deleted_row.max_attempts,
9563 lane_seq: deleted_row.lane_seq,
9564 enqueue_shard: deleted_row.enqueue_shard,
9565 run_at: ready.run_at,
9566 attempted_at: deleted_row.attempted_at,
9567 created_at: ready.created_at,
9568 unique_key: ready.unique_key.clone(),
9569 unique_states: ready.unique_states.clone(),
9570 payload: ready.payload.clone(),
9571 progress: attempt.and_then(|row| row.progress.clone()),
9572 });
9573 }
9574
9575 Ok(hydrated)
9576 }
9577
9578 async fn close_open_receipt_claim_tx<'a>(
9579 &self,
9580 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
9581 job_id: i64,
9582 run_lease: i64,
9583 outcome: &str,
9584 ) -> Result<Option<LeaseTransitionRow>, AwaError> {
9585 if !self.lease_claim_receipts() {
9586 return Ok(None);
9587 }
9588
9589 self.lock_receipt_attempts_tx(tx, &[(job_id, run_lease)])
9590 .await?;
9591
9592 let schema = self.schema();
9593 let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
9594 r#"
9595 WITH row_target AS (
9596 -- Target is the open claim identified from the
9597 -- partitioned lease_claims table anti-joined against
9598 -- durable closure evidence.
9599 SELECT
9600 claims.claim_slot,
9601 NULL::bigint AS batch_id,
9602 claims.receipt_id,
9603 claims.ready_slot,
9604 claims.ready_generation,
9605 claims.job_id,
9606 claims.queue,
9607 'running'::awa.job_state AS state,
9608 claims.priority,
9609 claims.attempt,
9610 claims.run_lease,
9611 claims.max_attempts,
9612 claims.lane_seq,
9613 claims.enqueue_shard,
9614 claims.claimed_at AS attempted_at
9615 FROM {schema}.lease_claims AS claims
9616 WHERE claims.job_id = $1
9617 AND claims.run_lease = $2
9618 AND claims.closed_at IS NULL
9619 AND NOT EXISTS (
9620 SELECT 1 FROM {schema}.lease_claim_closures AS closures
9621 WHERE closures.claim_slot = claims.claim_slot
9622 AND closures.job_id = claims.job_id
9623 AND closures.run_lease = claims.run_lease
9624 )
9625 AND NOT EXISTS (
9626 SELECT 1
9627 FROM {schema}.lease_claim_closure_batches AS closure_batches
9628 WHERE closure_batches.receipt_ranges @> claims.receipt_id
9629 )
9630 AND NOT EXISTS (
9631 SELECT 1 FROM {schema}.done_entries AS done
9632 WHERE done.job_id = claims.job_id
9633 AND done.run_lease = claims.run_lease
9634 )
9635 AND NOT EXISTS (
9636 SELECT 1 FROM {schema}.deferred_jobs AS deferred
9637 WHERE deferred.job_id = claims.job_id
9638 AND deferred.run_lease = claims.run_lease
9639 )
9640 AND NOT EXISTS (
9641 SELECT 1 FROM {schema}.dlq_entries AS dlq
9642 WHERE dlq.job_id = claims.job_id
9643 AND dlq.run_lease = claims.run_lease
9644 )
9645 AND NOT EXISTS (
9646 SELECT 1 FROM {schema}.leases AS lease
9647 WHERE lease.job_id = claims.job_id
9648 AND lease.run_lease = claims.run_lease
9649 )
9650 FOR UPDATE OF claims
9651 ),
9652 batch_target AS (
9653 SELECT
9654 claim_batches.claim_slot,
9655 claim_batches.batch_id,
9656 items.receipt_id,
9657 claim_batches.ready_slot,
9658 claim_batches.ready_generation,
9659 items.job_id,
9660 claim_batches.queue,
9661 'running'::awa.job_state AS state,
9662 claim_batches.priority,
9663 items.attempt,
9664 items.run_lease,
9665 items.max_attempts,
9666 items.lane_seq,
9667 claim_batches.enqueue_shard,
9668 claim_batches.claimed_at AS attempted_at
9669 FROM {schema}.lease_claim_batches AS claim_batches
9670 CROSS JOIN LATERAL unnest(
9671 claim_batches.job_ids,
9672 claim_batches.run_leases,
9673 claim_batches.receipt_ids,
9674 claim_batches.lane_seqs,
9675 claim_batches.attempts,
9676 claim_batches.max_attempts
9677 ) AS items(job_id, run_lease, receipt_id, lane_seq, attempt, max_attempts)
9678 WHERE items.job_id = $1
9679 AND items.run_lease = $2
9680 AND NOT EXISTS (
9681 SELECT 1 FROM {schema}.lease_claim_closures AS closures
9682 WHERE closures.claim_slot = claim_batches.claim_slot
9683 AND closures.job_id = items.job_id
9684 AND closures.run_lease = items.run_lease
9685 )
9686 AND NOT EXISTS (
9687 SELECT 1
9688 FROM {schema}.lease_claim_closure_batches AS closure_batches
9689 WHERE closure_batches.claim_slot = claim_batches.claim_slot
9690 AND closure_batches.receipt_ranges @> items.receipt_id
9691 )
9692 AND NOT EXISTS (
9693 SELECT 1 FROM {schema}.done_entries AS done
9694 WHERE done.job_id = items.job_id
9695 AND done.run_lease = items.run_lease
9696 )
9697 AND NOT EXISTS (
9698 SELECT 1 FROM {schema}.deferred_jobs AS deferred
9699 WHERE deferred.job_id = items.job_id
9700 AND deferred.run_lease = items.run_lease
9701 )
9702 AND NOT EXISTS (
9703 SELECT 1 FROM {schema}.dlq_entries AS dlq
9704 WHERE dlq.job_id = items.job_id
9705 AND dlq.run_lease = items.run_lease
9706 )
9707 AND NOT EXISTS (
9708 SELECT 1 FROM {schema}.leases AS lease
9709 WHERE lease.job_id = items.job_id
9710 AND lease.run_lease = items.run_lease
9711 )
9712 FOR UPDATE OF claim_batches
9713 ),
9714 target AS (
9715 SELECT * FROM row_target
9716 UNION ALL
9717 SELECT * FROM batch_target
9718 LIMIT 1
9719 ),
9720 -- A row target (batch_id IS NULL) is a deadline lease claim:
9721 -- close it explicitly so the prune gates balance it against
9722 -- the lease_claims row. A batch target has no lease_claims
9723 -- row to JOIN, so it closes into the batch ledger below.
9724 inserted AS (
9725 INSERT INTO {schema}.lease_claim_closures (claim_slot, job_id, run_lease, outcome, closed_at)
9726 SELECT target.claim_slot, target.job_id, target.run_lease, $3, clock_timestamp()
9727 FROM target
9728 WHERE target.batch_id IS NULL
9729 ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING
9730 RETURNING claim_slot, job_id, run_lease, closed_at
9731 ),
9732 inserted_batches AS (
9733 INSERT INTO {schema}.lease_claim_closure_batches (
9734 claim_slot,
9735 ready_slot,
9736 ready_generation,
9737 outcome,
9738 closed_count,
9739 receipt_ids,
9740 receipt_ranges,
9741 closed_at
9742 )
9743 SELECT
9744 target.claim_slot,
9745 target.ready_slot,
9746 target.ready_generation,
9747 $3,
9748 count(*)::int,
9749 array_agg(target.receipt_id ORDER BY target.receipt_id),
9750 range_agg(int8range(target.receipt_id, target.receipt_id + 1, '[)') ORDER BY target.receipt_id),
9751 clock_timestamp()
9752 FROM target
9753 WHERE target.batch_id IS NOT NULL
9754 GROUP BY target.claim_slot, target.ready_slot, target.ready_generation
9755 RETURNING claim_slot
9756 ),
9757 marked AS (
9758 UPDATE {schema}.lease_claims AS claims
9759 SET closed_at = COALESCE(claims.closed_at, inserted.closed_at)
9760 FROM inserted
9761 WHERE claims.claim_slot = inserted.claim_slot
9762 AND claims.job_id = inserted.job_id
9763 AND claims.run_lease = inserted.run_lease
9764 RETURNING claims.job_id
9765 ),
9766 -- The batch ledger has no job_id/run_lease column to RETURN,
9767 -- so derive the closed batch target from `target` gated on
9768 -- the batch insert having fired (target is LIMIT 1).
9769 closed_target AS (
9770 SELECT claim_slot, job_id, run_lease FROM inserted
9771 UNION ALL
9772 SELECT target.claim_slot, target.job_id, target.run_lease
9773 FROM target
9774 WHERE target.batch_id IS NOT NULL
9775 AND EXISTS (SELECT 1 FROM inserted_batches)
9776 )
9777 SELECT
9778 target.ready_slot,
9779 target.ready_generation,
9780 target.job_id,
9781 target.queue,
9782 target.state,
9783 target.priority,
9784 target.attempt,
9785 target.run_lease,
9786 target.max_attempts,
9787 target.lane_seq,
9788 target.enqueue_shard,
9789 NULL::timestamptz AS heartbeat_at,
9790 NULL::timestamptz AS deadline_at,
9791 target.attempted_at,
9792 NULL::uuid AS callback_id,
9793 NULL::timestamptz AS callback_timeout_at
9794 FROM target
9795 JOIN closed_target
9796 ON closed_target.claim_slot = target.claim_slot
9797 AND closed_target.job_id = target.job_id
9798 AND closed_target.run_lease = target.run_lease
9799 "#
9800 ))
9801 .bind(job_id)
9802 .bind(run_lease)
9803 .bind(outcome)
9804 .fetch_all(tx.as_mut())
9805 .await
9806 .map_err(map_sqlx_error)?;
9807
9808 if deleted.is_empty() {
9809 return Ok(None);
9810 }
9811
9812 let moved = self.hydrate_deleted_leases_tx(tx, deleted).await?;
9813 Ok(moved.into_iter().next())
9814 }
9815
9816 async fn take_running_attempt_tx<'a>(
9817 &self,
9818 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
9819 job_id: i64,
9820 run_lease: i64,
9821 receipt_outcome: &str,
9822 ) -> Result<Option<LeaseTransitionRow>, AwaError> {
9823 if let Some(moved) = self
9824 .close_open_receipt_claim_tx(tx, job_id, run_lease, receipt_outcome)
9825 .await?
9826 {
9827 return Ok(Some(moved));
9828 }
9829
9830 let schema = self.schema();
9831 let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
9832 r#"
9833 DELETE FROM {schema}.leases
9834 WHERE job_id = $1
9835 AND run_lease = $2
9836 AND state = 'running'
9837 RETURNING
9838 ready_slot,
9839 ready_generation,
9840 job_id,
9841 queue,
9842 state,
9843 priority,
9844 attempt,
9845 run_lease,
9846 max_attempts,
9847 lane_seq,
9848 enqueue_shard,
9849 heartbeat_at,
9850 deadline_at,
9851 attempted_at,
9852 callback_id,
9853 callback_timeout_at
9854 "#
9855 ))
9856 .bind(job_id)
9857 .bind(run_lease)
9858 .fetch_all(tx.as_mut())
9859 .await
9860 .map_err(map_sqlx_error)?;
9861
9862 if deleted.is_empty() {
9863 return Ok(None);
9864 }
9865
9866 let moved = self.hydrate_deleted_leases_tx(tx, deleted).await?;
9867 Ok(moved.into_iter().next())
9868 }
9869
9870 async fn rescue_stale_receipt_claims_tx<'a>(
9871 &self,
9872 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
9873 cutoff: DateTime<Utc>,
9874 ) -> Result<Vec<DeletedLeaseRow>, AwaError> {
9875 let mut rescued = Vec::new();
9876 let mut remaining = RECEIPT_RESCUE_BATCH_LIMIT;
9877 let preferred_slot = self.oldest_initialized_claim_slot_tx(tx).await?;
9878
9879 if let Some(slot) = preferred_slot {
9880 let mut slot_rescued = self
9881 .rescue_stale_receipt_claims_for_slot_tx(tx, slot, cutoff, remaining)
9882 .await?;
9883 remaining = remaining.saturating_sub(slot_rescued.len() as i64);
9884 rescued.append(&mut slot_rescued);
9885 if remaining == 0 {
9886 return Ok(rescued);
9887 }
9888 }
9889
9890 for slot in 0..self.claim_slot_count() {
9891 let slot = slot as i32;
9892 if Some(slot) == preferred_slot {
9893 continue;
9894 }
9895
9896 let mut slot_rescued = self
9897 .rescue_stale_receipt_claims_for_slot_tx(tx, slot, cutoff, remaining)
9898 .await?;
9899 remaining = remaining.saturating_sub(slot_rescued.len() as i64);
9900 rescued.append(&mut slot_rescued);
9901
9902 if remaining == 0 {
9903 break;
9904 }
9905 }
9906
9907 Ok(rescued)
9908 }
9909
9910 async fn rescue_stale_receipt_claims_for_slot_tx<'a>(
9911 &self,
9912 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
9913 slot: i32,
9914 cutoff: DateTime<Utc>,
9915 rescue_limit: i64,
9916 ) -> Result<Vec<DeletedLeaseRow>, AwaError> {
9917 let schema = self.schema();
9918 let claim_child = claim_child_name(schema, slot as usize);
9919 let claim_batch_child = claim_batch_child_name(schema, slot as usize);
9920 let closure_child = closure_child_name(schema, slot as usize);
9921 let closure_batch_child = claim_closure_batch_child_name(schema, slot as usize);
9922 let rescued: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
9923 r#"
9924 WITH cursor_row AS (
9925 SELECT
9926 rescue_cursor_claimed_at,
9927 rescue_cursor_job_id,
9928 rescue_cursor_run_lease
9929 FROM {schema}.claim_ring_slots
9930 WHERE slot = $1
9931 FOR UPDATE
9932 ),
9933 claim_source AS MATERIALIZED (
9934 SELECT
9935 claims.claim_slot,
9936 NULL::bigint AS batch_id,
9937 claims.ready_slot,
9938 claims.ready_generation,
9939 claims.job_id,
9940 claims.queue,
9941 claims.priority,
9942 claims.attempt,
9943 claims.run_lease,
9944 claims.max_attempts,
9945 claims.lane_seq,
9946 claims.enqueue_shard,
9947 claims.receipt_id,
9948 claims.claimed_at,
9949 claims.closed_at,
9950 false AS compact_batch
9951 FROM {claim_child} AS claims
9952 WHERE claims.claim_slot = $1
9953 UNION ALL
9954 SELECT
9955 claim_batches.claim_slot,
9956 claim_batches.batch_id,
9957 claim_batches.ready_slot,
9958 claim_batches.ready_generation,
9959 items.job_id,
9960 claim_batches.queue,
9961 claim_batches.priority,
9962 items.attempt,
9963 items.run_lease,
9964 items.max_attempts,
9965 items.lane_seq,
9966 claim_batches.enqueue_shard,
9967 items.receipt_id,
9968 claim_batches.claimed_at,
9969 NULL::timestamptz AS closed_at,
9970 true AS compact_batch
9971 FROM {claim_batch_child} AS claim_batches
9972 CROSS JOIN LATERAL unnest(
9973 claim_batches.job_ids,
9974 claim_batches.run_leases,
9975 claim_batches.receipt_ids,
9976 claim_batches.lane_seqs,
9977 claim_batches.attempts,
9978 claim_batches.max_attempts
9979 ) AS items(job_id, run_lease, receipt_id, lane_seq, attempt, max_attempts)
9980 WHERE claim_batches.claim_slot = $1
9981 ),
9982 after_cursor AS MATERIALIZED (
9983 SELECT
9984 claim_source.claim_slot,
9985 claim_source.job_id,
9986 claim_source.run_lease,
9987 row_number() OVER (
9988 ORDER BY claim_source.claimed_at, claim_source.job_id, claim_source.run_lease
9989 ) AS rn
9990 FROM claim_source
9991 CROSS JOIN cursor_row
9992 WHERE claim_source.claimed_at < $2
9993 AND (claim_source.claimed_at, claim_source.job_id, claim_source.run_lease)
9994 > (
9995 cursor_row.rescue_cursor_claimed_at,
9996 cursor_row.rescue_cursor_job_id,
9997 cursor_row.rescue_cursor_run_lease
9998 )
9999 ORDER BY claim_source.claimed_at, claim_source.job_id, claim_source.run_lease
10000 LIMIT $3
10001 ),
10002 after_count AS (
10003 SELECT count(*)::bigint AS count FROM after_cursor
10004 ),
10005 before_cursor AS MATERIALIZED (
10006 SELECT
10007 claim_source.claim_slot,
10008 claim_source.job_id,
10009 claim_source.run_lease,
10010 after_count.count + row_number() OVER (
10011 ORDER BY claim_source.claimed_at, claim_source.job_id, claim_source.run_lease
10012 ) AS rn
10013 FROM claim_source
10014 CROSS JOIN cursor_row
10015 CROSS JOIN after_count
10016 WHERE after_count.count < $3
10017 AND claim_source.claimed_at < $2
10018 AND (claim_source.claimed_at, claim_source.job_id, claim_source.run_lease)
10019 <= (
10020 cursor_row.rescue_cursor_claimed_at,
10021 cursor_row.rescue_cursor_job_id,
10022 cursor_row.rescue_cursor_run_lease
10023 )
10024 ORDER BY claim_source.claimed_at, claim_source.job_id, claim_source.run_lease
10025 LIMIT (SELECT GREATEST($3 - count, 0) FROM after_count)
10026 ),
10027 candidate_keys AS MATERIALIZED (
10028 SELECT claim_slot, job_id, run_lease, rn FROM after_cursor
10029 UNION ALL
10030 SELECT claim_slot, job_id, run_lease, rn FROM before_cursor
10031 ),
10032 candidates AS MATERIALIZED (
10033 SELECT
10034 claim_source.claim_slot,
10035 claim_source.batch_id,
10036 claim_source.ready_slot,
10037 claim_source.ready_generation,
10038 claim_source.job_id,
10039 claim_source.queue,
10040 claim_source.priority,
10041 claim_source.attempt,
10042 claim_source.run_lease,
10043 claim_source.max_attempts,
10044 claim_source.lane_seq,
10045 claim_source.enqueue_shard,
10046 claim_source.receipt_id,
10047 claim_source.claimed_at,
10048 claim_source.closed_at,
10049 claim_source.compact_batch,
10050 COALESCE(attempt.heartbeat_at, claim_source.claimed_at) < $2 AS is_stale,
10051 (
10052 claim_source.closed_at IS NOT NULL
10053 OR EXISTS (
10054 SELECT 1 FROM {closure_child} AS closures
10055 WHERE closures.claim_slot = claim_source.claim_slot
10056 AND closures.job_id = claim_source.job_id
10057 AND closures.run_lease = claim_source.run_lease
10058 )
10059 OR EXISTS (
10060 SELECT 1
10061 FROM {closure_batch_child} AS closure_batches
10062 WHERE closure_batches.receipt_ranges @> claim_source.receipt_id
10063 )
10064 OR EXISTS (
10065 SELECT 1 FROM {schema}.done_entries AS done
10066 WHERE done.job_id = claim_source.job_id
10067 AND done.run_lease = claim_source.run_lease
10068 )
10069 OR EXISTS (
10070 SELECT 1 FROM {schema}.deferred_jobs AS deferred
10071 WHERE deferred.job_id = claim_source.job_id
10072 AND deferred.run_lease = claim_source.run_lease
10073 )
10074 OR EXISTS (
10075 SELECT 1 FROM {schema}.dlq_entries AS dlq
10076 WHERE dlq.job_id = claim_source.job_id
10077 AND dlq.run_lease = claim_source.run_lease
10078 )
10079 ) AS is_closed,
10080 EXISTS (
10081 SELECT 1 FROM {schema}.leases AS lease
10082 WHERE lease.job_id = claim_source.job_id
10083 AND lease.run_lease = claim_source.run_lease
10084 ) AS is_lease_managed,
10085 candidate_keys.rn
10086 FROM candidate_keys
10087 JOIN claim_source
10088 ON claim_source.claim_slot = candidate_keys.claim_slot
10089 AND claim_source.job_id = candidate_keys.job_id
10090 AND claim_source.run_lease = candidate_keys.run_lease
10091 LEFT JOIN {schema}.attempt_state AS attempt
10092 ON attempt.job_id = claim_source.job_id
10093 AND attempt.run_lease = claim_source.run_lease
10094 ),
10095 stale_candidates AS (
10096 SELECT candidates.*
10097 FROM candidates
10098 WHERE NOT is_closed
10099 AND NOT is_lease_managed
10100 AND is_stale
10101 ORDER BY rn
10102 LIMIT $4
10103 ),
10104 stale_row_locked AS (
10105 SELECT stale_candidates.*
10106 FROM stale_candidates
10107 JOIN {claim_child} AS claims
10108 ON claims.claim_slot = stale_candidates.claim_slot
10109 AND claims.job_id = stale_candidates.job_id
10110 AND claims.run_lease = stale_candidates.run_lease
10111 WHERE NOT stale_candidates.compact_batch
10112 AND claims.closed_at IS NULL
10113 AND NOT EXISTS (
10114 SELECT 1 FROM {closure_child} AS closures
10115 WHERE closures.claim_slot = claims.claim_slot
10116 AND closures.job_id = claims.job_id
10117 AND closures.run_lease = claims.run_lease
10118 )
10119 AND NOT EXISTS (
10120 SELECT 1
10121 FROM {closure_batch_child} AS closure_batches
10122 WHERE closure_batches.receipt_ranges @> claims.receipt_id
10123 )
10124 AND NOT EXISTS (
10125 SELECT 1 FROM {schema}.done_entries AS done
10126 WHERE done.job_id = claims.job_id
10127 AND done.run_lease = claims.run_lease
10128 )
10129 AND NOT EXISTS (
10130 SELECT 1 FROM {schema}.deferred_jobs AS deferred
10131 WHERE deferred.job_id = claims.job_id
10132 AND deferred.run_lease = claims.run_lease
10133 )
10134 AND NOT EXISTS (
10135 SELECT 1 FROM {schema}.dlq_entries AS dlq
10136 WHERE dlq.job_id = claims.job_id
10137 AND dlq.run_lease = claims.run_lease
10138 )
10139 AND pg_catalog.pg_try_advisory_xact_lock(
10140 pg_catalog.hashtextextended(
10141 format('awa.receipt.complete:%s:%s', claims.job_id, claims.run_lease),
10142 0
10143 )
10144 )
10145 -- A claim that already materialized into `leases` is
10146 -- on the lease-side heartbeat-rescue path (see
10147 -- `rescue_stale_heartbeats`). Rescuing it again here
10148 -- would write a second closure for an attempt the
10149 -- runtime is still tracking via its lease row.
10150 AND NOT EXISTS (
10151 SELECT 1 FROM {schema}.leases AS lease
10152 WHERE lease.job_id = claims.job_id
10153 AND lease.run_lease = claims.run_lease
10154 )
10155 FOR UPDATE OF claims SKIP LOCKED
10156 ),
10157 stale_batch_locked AS (
10158 SELECT stale_candidates.*
10159 FROM stale_candidates
10160 JOIN {claim_batch_child} AS claim_batches
10161 ON claim_batches.claim_slot = stale_candidates.claim_slot
10162 AND claim_batches.batch_id = stale_candidates.batch_id
10163 WHERE stale_candidates.compact_batch
10164 AND NOT EXISTS (
10165 SELECT 1 FROM {closure_child} AS closures
10166 WHERE closures.claim_slot = stale_candidates.claim_slot
10167 AND closures.job_id = stale_candidates.job_id
10168 AND closures.run_lease = stale_candidates.run_lease
10169 )
10170 AND NOT EXISTS (
10171 SELECT 1
10172 FROM {closure_batch_child} AS closure_batches
10173 WHERE closure_batches.receipt_ranges @> stale_candidates.receipt_id
10174 )
10175 AND NOT EXISTS (
10176 SELECT 1 FROM {schema}.done_entries AS done
10177 WHERE done.job_id = stale_candidates.job_id
10178 AND done.run_lease = stale_candidates.run_lease
10179 )
10180 AND NOT EXISTS (
10181 SELECT 1 FROM {schema}.deferred_jobs AS deferred
10182 WHERE deferred.job_id = stale_candidates.job_id
10183 AND deferred.run_lease = stale_candidates.run_lease
10184 )
10185 AND NOT EXISTS (
10186 SELECT 1 FROM {schema}.dlq_entries AS dlq
10187 WHERE dlq.job_id = stale_candidates.job_id
10188 AND dlq.run_lease = stale_candidates.run_lease
10189 )
10190 AND pg_catalog.pg_try_advisory_xact_lock(
10191 pg_catalog.hashtextextended(
10192 format('awa.receipt.complete:%s:%s', stale_candidates.job_id, stale_candidates.run_lease),
10193 0
10194 )
10195 )
10196 AND NOT EXISTS (
10197 SELECT 1 FROM {schema}.leases AS lease
10198 WHERE lease.job_id = stale_candidates.job_id
10199 AND lease.run_lease = stale_candidates.run_lease
10200 )
10201 FOR UPDATE OF claim_batches SKIP LOCKED
10202 ),
10203 stale_locked AS (
10204 SELECT * FROM stale_row_locked
10205 UNION ALL
10206 SELECT * FROM stale_batch_locked
10207 ),
10208 -- Row-sourced (deadline lease) claims close explicitly so the
10209 -- prune gates balance them against the lease_claims row.
10210 -- Compact batch-sourced claims have no lease_claims row to
10211 -- JOIN, so they close into the batch ledger the queue prune
10212 -- gate counts via compact_count.
10213 inserted AS (
10214 INSERT INTO {schema}.lease_claim_closures (claim_slot, job_id, run_lease, outcome, closed_at)
10215 SELECT stale_locked.claim_slot, stale_locked.job_id, stale_locked.run_lease, 'rescued', clock_timestamp()
10216 FROM stale_locked
10217 WHERE NOT stale_locked.compact_batch
10218 ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING
10219 RETURNING claim_slot, job_id, run_lease, closed_at
10220 ),
10221 inserted_batches AS (
10222 INSERT INTO {schema}.lease_claim_closure_batches (
10223 claim_slot,
10224 ready_slot,
10225 ready_generation,
10226 outcome,
10227 closed_count,
10228 receipt_ids,
10229 receipt_ranges,
10230 closed_at
10231 )
10232 SELECT
10233 stale_locked.claim_slot,
10234 stale_locked.ready_slot,
10235 stale_locked.ready_generation,
10236 'rescued',
10237 count(*)::int,
10238 array_agg(stale_locked.receipt_id ORDER BY stale_locked.receipt_id),
10239 range_agg(int8range(stale_locked.receipt_id, stale_locked.receipt_id + 1, '[)') ORDER BY stale_locked.receipt_id),
10240 clock_timestamp()
10241 FROM stale_locked
10242 WHERE stale_locked.compact_batch
10243 GROUP BY
10244 stale_locked.claim_slot,
10245 stale_locked.ready_slot,
10246 stale_locked.ready_generation
10247 RETURNING claim_slot
10248 ),
10249 closed_locked AS (
10250 SELECT claim_slot, job_id, run_lease FROM inserted
10251 UNION ALL
10252 SELECT stale_locked.claim_slot, stale_locked.job_id, stale_locked.run_lease
10253 FROM stale_locked
10254 WHERE stale_locked.compact_batch
10255 AND EXISTS (SELECT 1 FROM inserted_batches)
10256 ),
10257 marked AS (
10258 UPDATE {schema}.lease_claims AS claims
10259 SET closed_at = COALESCE(claims.closed_at, inserted.closed_at)
10260 FROM inserted
10261 WHERE claims.claim_slot = inserted.claim_slot
10262 AND claims.job_id = inserted.job_id
10263 AND claims.run_lease = inserted.run_lease
10264 RETURNING claims.job_id
10265 ),
10266 annotated AS (
10267 SELECT
10268 candidates.*,
10269 (
10270 candidates.is_closed
10271 OR candidates.is_lease_managed
10272 OR NOT candidates.is_stale
10273 OR EXISTS (
10274 SELECT 1 FROM closed_locked
10275 WHERE closed_locked.claim_slot = candidates.claim_slot
10276 AND closed_locked.job_id = candidates.job_id
10277 AND closed_locked.run_lease = candidates.run_lease
10278 )
10279 ) AS advanceable
10280 FROM candidates
10281 ),
10282 bounded AS (
10283 SELECT
10284 annotated.*,
10285 min(CASE WHEN NOT annotated.advanceable THEN annotated.rn END) OVER () AS first_blocked_rn
10286 FROM annotated
10287 ),
10288 advance_target AS (
10289 SELECT claimed_at, job_id, run_lease
10290 FROM bounded
10291 WHERE first_blocked_rn IS NULL OR rn < first_blocked_rn
10292 ORDER BY rn DESC
10293 LIMIT 1
10294 ),
10295 advance_cursor AS (
10296 UPDATE {schema}.claim_ring_slots AS slots
10297 SET rescue_cursor_claimed_at = advance_target.claimed_at,
10298 rescue_cursor_job_id = advance_target.job_id,
10299 rescue_cursor_run_lease = advance_target.run_lease
10300 FROM advance_target
10301 WHERE slots.slot = $1
10302 RETURNING slots.slot
10303 ),
10304 cursor_advance AS (
10305 SELECT count(*) FROM advance_cursor
10306 )
10307 SELECT
10308 stale_locked.ready_slot,
10309 stale_locked.ready_generation,
10310 stale_locked.job_id,
10311 stale_locked.queue,
10312 'running'::awa.job_state AS state,
10313 stale_locked.priority,
10314 stale_locked.attempt,
10315 stale_locked.run_lease,
10316 stale_locked.max_attempts,
10317 stale_locked.lane_seq,
10318 stale_locked.enqueue_shard,
10319 stale_locked.claimed_at AS attempted_at
10320 FROM stale_locked
10321 JOIN closed_locked
10322 ON closed_locked.claim_slot = stale_locked.claim_slot
10323 AND closed_locked.job_id = stale_locked.job_id
10324 AND closed_locked.run_lease = stale_locked.run_lease
10325 CROSS JOIN cursor_advance
10326 "#
10327 ))
10328 .bind(slot)
10329 .bind(cutoff)
10330 .bind(RECEIPT_RESCUE_CURSOR_SCAN_LIMIT)
10331 .bind(rescue_limit)
10332 .fetch_all(tx.as_mut())
10333 .await
10334 .map_err(map_sqlx_error)?;
10335 Ok(rescued)
10336 }
10337
10338 async fn rescue_expired_receipt_deadlines_tx<'a>(
10356 &self,
10357 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
10358 ) -> Result<Vec<DeletedLeaseRow>, AwaError> {
10359 let mut rescued = Vec::new();
10360 let mut remaining = RECEIPT_RESCUE_BATCH_LIMIT;
10361 let preferred_slot = self.oldest_initialized_claim_slot_tx(tx).await?;
10362
10363 if let Some(slot) = preferred_slot {
10364 let mut slot_rescued = self
10365 .rescue_expired_receipt_deadlines_for_slot_tx(tx, slot, remaining)
10366 .await?;
10367 remaining = remaining.saturating_sub(slot_rescued.len() as i64);
10368 rescued.append(&mut slot_rescued);
10369 if remaining == 0 {
10370 return Ok(rescued);
10371 }
10372 }
10373
10374 for slot in 0..self.claim_slot_count() {
10375 let slot = slot as i32;
10376 if Some(slot) == preferred_slot {
10377 continue;
10378 }
10379
10380 let mut slot_rescued = self
10381 .rescue_expired_receipt_deadlines_for_slot_tx(tx, slot, remaining)
10382 .await?;
10383 remaining = remaining.saturating_sub(slot_rescued.len() as i64);
10384 rescued.append(&mut slot_rescued);
10385
10386 if remaining == 0 {
10387 break;
10388 }
10389 }
10390
10391 Ok(rescued)
10392 }
10393
10394 async fn oldest_initialized_claim_slot_tx<'a>(
10395 &self,
10396 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
10397 ) -> Result<Option<i32>, AwaError> {
10398 let schema = self.schema();
10399 let preferred_slot = sqlx::query_as::<_, (i32, i64, i32)>(&format!(
10400 r#"
10401 SELECT current_slot, generation, slot_count
10402 FROM {schema}.claim_ring_state
10403 WHERE singleton = TRUE
10404 "#
10405 ))
10406 .fetch_optional(tx.as_mut())
10407 .await
10408 .map_err(map_sqlx_error)?
10409 .and_then(|(current_slot, generation, slot_count)| {
10410 oldest_initialized_ring_slot(current_slot, generation, slot_count)
10411 .map(|(slot, _generation)| slot)
10412 .filter(|slot| *slot >= 0 && (*slot as usize) < self.claim_slot_count())
10413 });
10414
10415 Ok(preferred_slot)
10416 }
10417
10418 async fn rescue_expired_receipt_deadlines_for_slot_tx<'a>(
10419 &self,
10420 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
10421 slot: i32,
10422 rescue_limit: i64,
10423 ) -> Result<Vec<DeletedLeaseRow>, AwaError> {
10424 if rescue_limit <= 0 {
10425 return Ok(Vec::new());
10426 }
10427
10428 let schema = self.schema();
10429 let claim_child = claim_child_name(schema, slot as usize);
10430 let closure_child = closure_child_name(schema, slot as usize);
10431 let closure_batch_child = claim_closure_batch_child_name(schema, slot as usize);
10432 let closed_evidence =
10433 receipt_closed_evidence_sql(schema, &closure_child, &closure_batch_child, "claims");
10434 let rescued: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
10435 r#"
10436 WITH cursor_row AS (
10437 SELECT
10438 deadline_cursor_deadline_at,
10439 deadline_cursor_job_id,
10440 deadline_cursor_run_lease
10441 FROM {schema}.claim_ring_slots
10442 WHERE slot = $1
10443 FOR UPDATE
10444 ),
10445 after_cursor AS MATERIALIZED (
10446 SELECT
10447 claims.claim_slot,
10448 claims.job_id,
10449 claims.run_lease,
10450 row_number() OVER (
10451 ORDER BY claims.deadline_at, claims.job_id, claims.run_lease
10452 ) AS rn
10453 FROM {claim_child} AS claims
10454 CROSS JOIN cursor_row
10455 WHERE claims.claim_slot = $1
10456 AND claims.deadline_at IS NOT NULL
10457 AND claims.deadline_at < clock_timestamp()
10458 AND (claims.deadline_at, claims.job_id, claims.run_lease)
10459 > (
10460 cursor_row.deadline_cursor_deadline_at,
10461 cursor_row.deadline_cursor_job_id,
10462 cursor_row.deadline_cursor_run_lease
10463 )
10464 ORDER BY claims.deadline_at, claims.job_id, claims.run_lease
10465 LIMIT $2
10466 ),
10467 after_count AS (
10468 SELECT count(*)::bigint AS count FROM after_cursor
10469 ),
10470 before_cursor AS MATERIALIZED (
10471 SELECT
10472 claims.claim_slot,
10473 claims.job_id,
10474 claims.run_lease,
10475 after_count.count + row_number() OVER (
10476 ORDER BY claims.deadline_at, claims.job_id, claims.run_lease
10477 ) AS rn
10478 FROM {claim_child} AS claims
10479 CROSS JOIN cursor_row
10480 CROSS JOIN after_count
10481 WHERE after_count.count < $2
10482 AND claims.claim_slot = $1
10483 AND claims.deadline_at IS NOT NULL
10484 AND claims.deadline_at < clock_timestamp()
10485 AND (claims.deadline_at, claims.job_id, claims.run_lease)
10486 <= (
10487 cursor_row.deadline_cursor_deadline_at,
10488 cursor_row.deadline_cursor_job_id,
10489 cursor_row.deadline_cursor_run_lease
10490 )
10491 ORDER BY claims.deadline_at, claims.job_id, claims.run_lease
10492 LIMIT (SELECT GREATEST($2 - count, 0) FROM after_count)
10493 ),
10494 candidate_keys AS MATERIALIZED (
10495 SELECT claim_slot, job_id, run_lease, rn FROM after_cursor
10496 UNION ALL
10497 SELECT claim_slot, job_id, run_lease, rn FROM before_cursor
10498 ),
10499 candidates AS MATERIALIZED (
10500 SELECT
10501 claims.claim_slot,
10502 claims.ready_slot,
10503 claims.ready_generation,
10504 claims.job_id,
10505 claims.queue,
10506 'running'::awa.job_state AS state,
10507 claims.priority,
10508 claims.attempt,
10509 claims.run_lease,
10510 claims.max_attempts,
10511 claims.lane_seq,
10512 claims.enqueue_shard,
10513 claims.claimed_at,
10514 claims.deadline_at,
10515 claims.deadline_at < clock_timestamp() AS is_expired,
10516 {closed_evidence} AS is_closed,
10517 EXISTS (
10518 SELECT 1 FROM {schema}.leases AS lease
10519 WHERE lease.job_id = claims.job_id
10520 AND lease.run_lease = claims.run_lease
10521 ) AS is_lease_managed,
10522 candidate_keys.rn
10523 FROM candidate_keys
10524 JOIN {claim_child} AS claims
10525 ON claims.claim_slot = candidate_keys.claim_slot
10526 AND claims.job_id = candidate_keys.job_id
10527 AND claims.run_lease = candidate_keys.run_lease
10528 ),
10529 expired_candidates AS (
10530 SELECT candidates.*
10531 FROM candidates
10532 WHERE NOT is_closed
10533 AND NOT is_lease_managed
10534 AND is_expired
10535 ORDER BY rn
10536 LIMIT $3
10537 ),
10538 expired_locked AS (
10539 SELECT expired_candidates.*
10540 FROM expired_candidates
10541 JOIN {claim_child} AS claims
10542 ON claims.claim_slot = expired_candidates.claim_slot
10543 AND claims.job_id = expired_candidates.job_id
10544 AND claims.run_lease = expired_candidates.run_lease
10545 WHERE NOT {closed_evidence}
10546 AND claims.deadline_at < clock_timestamp()
10547 AND pg_catalog.pg_try_advisory_xact_lock(
10548 pg_catalog.hashtextextended(
10549 format('awa.receipt.complete:%s:%s', claims.job_id, claims.run_lease),
10550 0
10551 )
10552 )
10553 AND NOT EXISTS (
10554 SELECT 1 FROM {schema}.leases AS lease
10555 WHERE lease.job_id = claims.job_id
10556 AND lease.run_lease = claims.run_lease
10557 )
10558 FOR UPDATE OF claims SKIP LOCKED
10559 ),
10560 inserted AS (
10561 INSERT INTO {schema}.lease_claim_closures (claim_slot, job_id, run_lease, outcome, closed_at)
10562 SELECT
10563 expired_locked.claim_slot,
10564 expired_locked.job_id,
10565 expired_locked.run_lease,
10566 'deadline_expired',
10567 clock_timestamp()
10568 FROM expired_locked
10569 ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING
10570 RETURNING claim_slot, job_id, run_lease, closed_at
10571 ),
10572 marked AS (
10573 UPDATE {schema}.lease_claims AS claims
10574 SET closed_at = COALESCE(claims.closed_at, inserted.closed_at)
10575 FROM inserted
10576 WHERE claims.claim_slot = inserted.claim_slot
10577 AND claims.job_id = inserted.job_id
10578 AND claims.run_lease = inserted.run_lease
10579 RETURNING claims.job_id
10580 ),
10581 annotated AS (
10582 SELECT
10583 candidates.*,
10584 (
10585 candidates.is_closed
10586 OR candidates.is_lease_managed
10587 OR EXISTS (
10588 SELECT 1 FROM inserted
10589 WHERE inserted.claim_slot = candidates.claim_slot
10590 AND inserted.job_id = candidates.job_id
10591 AND inserted.run_lease = candidates.run_lease
10592 )
10593 ) AS advanceable
10594 FROM candidates
10595 ),
10596 bounded AS (
10597 SELECT
10598 annotated.*,
10599 min(CASE WHEN NOT annotated.advanceable THEN annotated.rn END) OVER () AS first_blocked_rn
10600 FROM annotated
10601 ),
10602 advance_target AS (
10603 SELECT deadline_at, job_id, run_lease
10604 FROM bounded
10605 WHERE first_blocked_rn IS NULL OR rn < first_blocked_rn
10606 ORDER BY rn DESC
10607 LIMIT 1
10608 ),
10609 advance_cursor AS (
10610 UPDATE {schema}.claim_ring_slots AS slots
10611 SET deadline_cursor_deadline_at = advance_target.deadline_at,
10612 deadline_cursor_job_id = advance_target.job_id,
10613 deadline_cursor_run_lease = advance_target.run_lease
10614 FROM advance_target
10615 WHERE slots.slot = $1
10616 RETURNING slots.slot
10617 ),
10618 cursor_advance AS (
10619 SELECT count(*) FROM advance_cursor
10620 )
10621 SELECT
10622 expired_locked.ready_slot,
10623 expired_locked.ready_generation,
10624 expired_locked.job_id,
10625 expired_locked.queue,
10626 expired_locked.state,
10627 expired_locked.priority,
10628 expired_locked.attempt,
10629 expired_locked.run_lease,
10630 expired_locked.max_attempts,
10631 expired_locked.lane_seq,
10632 expired_locked.enqueue_shard,
10633 expired_locked.claimed_at AS attempted_at
10634 FROM expired_locked
10635 JOIN inserted
10636 ON inserted.claim_slot = expired_locked.claim_slot
10637 AND inserted.job_id = expired_locked.job_id
10638 AND inserted.run_lease = expired_locked.run_lease
10639 CROSS JOIN cursor_advance
10640 "#
10641 ))
10642 .bind(slot)
10643 .bind(RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT)
10644 .bind(rescue_limit)
10645 .fetch_all(tx.as_mut())
10646 .await
10647 .map_err(map_sqlx_error)?;
10648 Ok(rescued)
10649 }
10650
10651 pub async fn load_job(&self, pool: &PgPool, job_id: i64) -> Result<Option<JobRow>, AwaError> {
10652 let schema = self.schema();
10653 let closure_rel = format!("{schema}.lease_claim_closures");
10654 let closure_batch_rel = format!("{schema}.lease_claim_closure_batches");
10655 let closed_evidence =
10656 receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims");
10657 let mut candidates = Vec::new();
10658
10659 let ready_rows: Vec<ReadyJobRow> = sqlx::query_as(&format!(
10660 r#"
10661 SELECT
10662 job_id,
10663 kind,
10664 queue,
10665 args,
10666 priority,
10667 attempt,
10668 run_lease,
10669 max_attempts,
10670 run_at,
10671 attempted_at,
10672 created_at,
10673 unique_key,
10674 unique_states,
10675 COALESCE(payload, '{{}}'::jsonb) AS payload
10676 FROM {schema}.ready_entries
10677 WHERE job_id = $1
10678 ORDER BY run_lease DESC, attempted_at DESC NULLS LAST, run_at DESC
10679 "#,
10680 ))
10681 .bind(job_id)
10682 .fetch_all(pool)
10683 .await
10684 .map_err(map_sqlx_error)?;
10685 for row in ready_rows {
10686 candidates.push(row.into_job_row()?);
10687 }
10688
10689 let deferred_rows: Vec<DeferredJobRow> = sqlx::query_as(&format!(
10690 r#"
10691 SELECT
10692 job_id,
10693 kind,
10694 queue,
10695 args,
10696 state,
10697 priority,
10698 attempt,
10699 run_lease,
10700 max_attempts,
10701 run_at,
10702 attempted_at,
10703 finalized_at,
10704 created_at,
10705 unique_key,
10706 unique_states,
10707 COALESCE(payload, '{{}}'::jsonb) AS payload
10708 FROM {schema}.deferred_jobs
10709 WHERE job_id = $1
10710 "#,
10711 ))
10712 .bind(job_id)
10713 .fetch_all(pool)
10714 .await
10715 .map_err(map_sqlx_error)?;
10716 for row in deferred_rows {
10717 candidates.push(row.into_job_row()?);
10718 }
10719
10720 let lease_rows: Vec<LeaseJobRow> = sqlx::query_as(&format!(
10721 r#"
10722 SELECT
10723 lease.ready_slot,
10724 lease.ready_generation,
10725 lease.job_id,
10726 ready.kind,
10727 ready.queue,
10728 ready.args,
10729 lease.state,
10730 lease.priority,
10731 lease.attempt,
10732 lease.run_lease,
10733 lease.max_attempts,
10734 lease.lane_seq,
10735 ready.run_at,
10736 COALESCE(attempt.heartbeat_at, lease.heartbeat_at) AS heartbeat_at,
10737 lease.deadline_at,
10738 lease.attempted_at,
10739 NULL::timestamptz AS finalized_at,
10740 ready.created_at,
10741 ready.unique_key,
10742 ready.unique_states,
10743 lease.callback_id,
10744 lease.callback_timeout_at,
10745 attempt.callback_filter,
10746 attempt.callback_on_complete,
10747 attempt.callback_on_fail,
10748 attempt.callback_transform,
10749 COALESCE(ready.payload, '{{}}'::jsonb) AS payload,
10750 attempt.progress,
10751 attempt.callback_result
10752 FROM {schema}.leases AS lease
10753 JOIN {schema}.ready_entries AS ready
10754 ON ready.ready_slot = lease.ready_slot
10755 AND ready.ready_generation = lease.ready_generation
10756 AND ready.queue = lease.queue
10757 AND ready.priority = lease.priority
10758 AND ready.enqueue_shard = lease.enqueue_shard
10759 AND ready.lane_seq = lease.lane_seq
10760 LEFT JOIN {schema}.attempt_state AS attempt
10761 ON attempt.job_id = lease.job_id
10762 AND attempt.run_lease = lease.run_lease
10763 WHERE lease.job_id = $1
10764 ORDER BY lease.run_lease DESC
10765 "#,
10766 ))
10767 .bind(job_id)
10768 .fetch_all(pool)
10769 .await
10770 .map_err(map_sqlx_error)?;
10771 for row in lease_rows {
10772 candidates.push(row.into_job_row()?);
10773 }
10774
10775 let lease_claim_rows: Vec<LeaseJobRow> = sqlx::query_as(&format!(
10778 r#"
10779 SELECT
10780 claims.ready_slot,
10781 claims.ready_generation,
10782 claims.job_id,
10783 ready.kind,
10784 ready.queue,
10785 ready.args,
10786 'running'::awa.job_state AS state,
10787 claims.priority,
10788 claims.attempt,
10789 claims.run_lease,
10790 claims.max_attempts,
10791 claims.lane_seq,
10792 ready.run_at,
10793 attempt.heartbeat_at,
10794 claims.deadline_at,
10795 claims.claimed_at AS attempted_at,
10796 NULL::timestamptz AS finalized_at,
10797 ready.created_at,
10798 ready.unique_key,
10799 ready.unique_states,
10800 NULL::uuid AS callback_id,
10801 NULL::timestamptz AS callback_timeout_at,
10802 attempt.callback_filter,
10803 attempt.callback_on_complete,
10804 attempt.callback_on_fail,
10805 attempt.callback_transform,
10806 COALESCE(ready.payload, '{{}}'::jsonb) AS payload,
10807 attempt.progress,
10808 attempt.callback_result
10809 FROM {schema}.lease_claims AS claims
10810 JOIN {schema}.ready_entries AS ready
10811 ON ready.ready_slot = claims.ready_slot
10812 AND ready.ready_generation = claims.ready_generation
10813 AND ready.queue = claims.queue
10814 AND ready.priority = claims.priority
10815 AND ready.enqueue_shard = claims.enqueue_shard
10816 AND ready.lane_seq = claims.lane_seq
10817 LEFT JOIN {schema}.attempt_state AS attempt
10818 ON attempt.job_id = claims.job_id
10819 AND attempt.run_lease = claims.run_lease
10820 WHERE claims.job_id = $1
10821 AND NOT {closed_evidence}
10822 -- Exclude claims that have already been materialized into
10823 -- leases — the lease-backed branch above already reports
10824 -- those.
10825 AND NOT EXISTS (
10826 SELECT 1 FROM {schema}.leases AS lease
10827 WHERE lease.job_id = claims.job_id
10828 AND lease.run_lease = claims.run_lease
10829 )
10830 ORDER BY claims.run_lease DESC
10831 "#,
10832 ))
10833 .bind(job_id)
10834 .fetch_all(pool)
10835 .await
10836 .map_err(map_sqlx_error)?;
10837 for row in lease_claim_rows {
10838 candidates.push(row.into_job_row()?);
10839 }
10840
10841 let lease_claim_batch_rows: Vec<LeaseJobRow> = sqlx::query_as(&format!(
10846 r#"
10847 SELECT
10848 claim_batches.ready_slot,
10849 claim_batches.ready_generation,
10850 items.job_id,
10851 ready.kind,
10852 ready.queue,
10853 ready.args,
10854 'running'::awa.job_state AS state,
10855 claim_batches.priority,
10856 items.attempt,
10857 items.run_lease,
10858 items.max_attempts,
10859 items.lane_seq,
10860 ready.run_at,
10861 attempt.heartbeat_at,
10862 claim_batches.deadline_at,
10863 claim_batches.claimed_at AS attempted_at,
10864 NULL::timestamptz AS finalized_at,
10865 ready.created_at,
10866 ready.unique_key,
10867 ready.unique_states,
10868 NULL::uuid AS callback_id,
10869 NULL::timestamptz AS callback_timeout_at,
10870 attempt.callback_filter,
10871 attempt.callback_on_complete,
10872 attempt.callback_on_fail,
10873 attempt.callback_transform,
10874 COALESCE(ready.payload, '{{}}'::jsonb) AS payload,
10875 attempt.progress,
10876 attempt.callback_result
10877 FROM {schema}.lease_claim_batches AS claim_batches
10878 CROSS JOIN LATERAL unnest(
10879 claim_batches.job_ids,
10880 claim_batches.run_leases,
10881 claim_batches.receipt_ids,
10882 claim_batches.lane_seqs,
10883 claim_batches.attempts,
10884 claim_batches.max_attempts
10885 ) AS items(job_id, run_lease, receipt_id, lane_seq, attempt, max_attempts)
10886 JOIN {schema}.ready_entries AS ready
10887 ON ready.ready_slot = claim_batches.ready_slot
10888 AND ready.ready_generation = claim_batches.ready_generation
10889 AND ready.queue = claim_batches.queue
10890 AND ready.enqueue_shard = claim_batches.enqueue_shard
10891 AND ready.lane_seq = items.lane_seq
10892 AND ready.job_id = items.job_id
10893 LEFT JOIN {schema}.attempt_state AS attempt
10894 ON attempt.job_id = items.job_id
10895 AND attempt.run_lease = items.run_lease
10896 WHERE items.job_id = $1
10897 AND NOT EXISTS (
10898 SELECT 1
10899 FROM {schema}.lease_claim_closures AS closures
10900 WHERE closures.claim_slot = claim_batches.claim_slot
10901 AND closures.job_id = items.job_id
10902 AND closures.run_lease = items.run_lease
10903 )
10904 AND NOT EXISTS (
10905 SELECT 1
10906 FROM {schema}.lease_claim_closure_batches AS closure_batches
10907 WHERE closure_batches.claim_slot = claim_batches.claim_slot
10908 AND closure_batches.receipt_ranges @> items.receipt_id
10909 )
10910 AND NOT EXISTS (
10911 SELECT 1
10912 FROM {schema}.leases AS lease
10913 WHERE lease.job_id = items.job_id
10914 AND lease.run_lease = items.run_lease
10915 )
10916 AND NOT EXISTS (
10917 SELECT 1 FROM {schema}.done_entries AS done
10918 WHERE done.job_id = items.job_id
10919 AND done.run_lease = items.run_lease
10920 )
10921 AND NOT EXISTS (
10922 SELECT 1 FROM {schema}.deferred_jobs AS deferred
10923 WHERE deferred.job_id = items.job_id
10924 AND deferred.run_lease = items.run_lease
10925 )
10926 AND NOT EXISTS (
10927 SELECT 1 FROM {schema}.dlq_entries AS dlq
10928 WHERE dlq.job_id = items.job_id
10929 AND dlq.run_lease = items.run_lease
10930 )
10931 ORDER BY items.run_lease DESC
10932 "#,
10933 ))
10934 .bind(job_id)
10935 .fetch_all(pool)
10936 .await
10937 .map_err(map_sqlx_error)?;
10938 for row in lease_claim_batch_rows {
10939 candidates.push(row.into_job_row()?);
10940 }
10941
10942 let done_rows: Vec<DoneJobRow> = sqlx::query_as(&format!(
10943 r#"
10944 SELECT
10945 ready_slot,
10946 ready_generation,
10947 job_id,
10948 kind,
10949 queue,
10950 args,
10951 state,
10952 priority,
10953 attempt,
10954 run_lease,
10955 max_attempts,
10956 lane_seq,
10957 enqueue_shard,
10958 run_at,
10959 attempted_at,
10960 finalized_at,
10961 created_at,
10962 unique_key,
10963 unique_states,
10964 payload
10965 FROM {schema}.terminal_jobs AS done
10966 WHERE done.job_id = $1
10967 ORDER BY done.run_lease DESC, done.finalized_at DESC
10968 "#,
10969 ))
10970 .bind(job_id)
10971 .fetch_all(pool)
10972 .await
10973 .map_err(map_sqlx_error)?;
10974 for row in done_rows {
10975 candidates.push(row.into_job_row()?);
10976 }
10977
10978 let dlq_rows: Vec<DlqJobRow> = sqlx::query_as(&format!(
10979 r#"
10980 SELECT
10981 job_id,
10982 kind,
10983 queue,
10984 args,
10985 state,
10986 priority,
10987 attempt,
10988 run_lease,
10989 max_attempts,
10990 run_at,
10991 attempted_at,
10992 finalized_at,
10993 created_at,
10994 unique_key,
10995 unique_states,
10996 COALESCE(payload, '{{}}'::jsonb) AS payload,
10997 dlq_reason,
10998 dlq_at,
10999 original_run_lease
11000 FROM {schema}.dlq_entries
11001 WHERE job_id = $1
11002 ORDER BY dlq_at DESC
11003 "#,
11004 ))
11005 .bind(job_id)
11006 .fetch_all(pool)
11007 .await
11008 .map_err(map_sqlx_error)?;
11009 for row in dlq_rows {
11010 candidates.push(row.into_job_row()?);
11011 }
11012
11013 Ok(candidates.into_iter().max_by_key(|job| {
11014 (
11015 job.run_lease,
11016 transition_timestamp(job),
11017 state_rank(job.state),
11018 )
11019 }))
11020 }
11021
11022 pub async fn register_callback(
11023 &self,
11024 pool: &PgPool,
11025 job_id: i64,
11026 run_lease: i64,
11027 timeout: Duration,
11028 ) -> Result<Uuid, AwaError> {
11029 let callback_id = Uuid::new_v4();
11030 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
11031 self.ensure_mutable_running_attempt_tx(&mut tx, job_id, run_lease)
11032 .await?;
11033 let updated = sqlx::query(&format!(
11034 r#"
11035 UPDATE {}
11036 SET callback_id = $2,
11037 callback_timeout_at = clock_timestamp() + make_interval(secs => $3)
11038 WHERE job_id = $1
11039 AND state = 'running'
11040 AND run_lease = $4
11041 "#,
11042 self.leases_table()
11043 ))
11044 .bind(job_id)
11045 .bind(callback_id)
11046 .bind(timeout.as_secs_f64())
11047 .bind(run_lease)
11048 .execute(tx.as_mut())
11049 .await
11050 .map_err(map_sqlx_error)?;
11051
11052 if updated.rows_affected() == 0 {
11053 tx.rollback().await.map_err(map_sqlx_error)?;
11054 return Err(AwaError::Validation("job is not in running state".into()));
11055 }
11056
11057 sqlx::query(&format!(
11058 r#"
11059 UPDATE {}
11060 SET callback_filter = NULL,
11061 callback_on_complete = NULL,
11062 callback_on_fail = NULL,
11063 callback_transform = NULL,
11064 updated_at = clock_timestamp()
11065 WHERE job_id = $1
11066 AND run_lease = $2
11067 "#,
11068 self.attempt_state_table()
11069 ))
11070 .bind(job_id)
11071 .bind(run_lease)
11072 .execute(tx.as_mut())
11073 .await
11074 .map_err(map_sqlx_error)?;
11075
11076 sqlx::query(&format!(
11077 r#"
11078 DELETE FROM {}
11079 WHERE job_id = $1
11080 AND run_lease = $2
11081 AND progress IS NULL
11082 AND callback_result IS NULL
11083 AND callback_filter IS NULL
11084 AND callback_on_complete IS NULL
11085 AND callback_on_fail IS NULL
11086 AND callback_transform IS NULL
11087 "#,
11088 self.attempt_state_table()
11089 ))
11090 .bind(job_id)
11091 .bind(run_lease)
11092 .execute(tx.as_mut())
11093 .await
11094 .map_err(map_sqlx_error)?;
11095
11096 tx.commit().await.map_err(map_sqlx_error)?;
11097 Ok(callback_id)
11098 }
11099
11100 pub async fn register_callback_with_config(
11101 &self,
11102 pool: &PgPool,
11103 job_id: i64,
11104 run_lease: i64,
11105 timeout: Duration,
11106 config: &CallbackConfig,
11107 ) -> Result<Uuid, AwaError> {
11108 if config.is_empty() {
11109 return self
11110 .register_callback(pool, job_id, run_lease, timeout)
11111 .await;
11112 }
11113
11114 #[cfg(feature = "cel")]
11115 {
11116 for (name, expr) in [
11117 ("filter", &config.filter),
11118 ("on_complete", &config.on_complete),
11119 ("on_fail", &config.on_fail),
11120 ("transform", &config.transform),
11121 ] {
11122 if let Some(src) = expr {
11123 let program = cel::Program::compile(src).map_err(|e| {
11124 AwaError::Validation(format!("invalid CEL expression for {name}: {e}"))
11125 })?;
11126 let references = program.references();
11127 let bad_vars: Vec<String> = references
11128 .variables()
11129 .into_iter()
11130 .filter(|v| *v != "payload")
11131 .map(str::to_string)
11132 .collect();
11133 if !bad_vars.is_empty() {
11134 return Err(AwaError::Validation(format!(
11135 "CEL expression for {name} references undeclared variable(s): {}; only 'payload' is available",
11136 bad_vars.join(", ")
11137 )));
11138 }
11139 }
11140 }
11141 }
11142
11143 #[cfg(not(feature = "cel"))]
11144 {
11145 if !config.is_empty() {
11146 return Err(AwaError::Validation(
11147 "CEL expressions require the 'cel' feature".into(),
11148 ));
11149 }
11150 }
11151
11152 let callback_id = Uuid::new_v4();
11153 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
11154 self.ensure_mutable_running_attempt_tx(&mut tx, job_id, run_lease)
11155 .await?;
11156 let updated = sqlx::query(&format!(
11157 r#"
11158 UPDATE {}
11159 SET callback_id = $2,
11160 callback_timeout_at = clock_timestamp() + make_interval(secs => $3)
11161 WHERE job_id = $1
11162 AND state = 'running'
11163 AND run_lease = $4
11164 "#,
11165 self.leases_table()
11166 ))
11167 .bind(job_id)
11168 .bind(callback_id)
11169 .bind(timeout.as_secs_f64())
11170 .bind(run_lease)
11171 .execute(tx.as_mut())
11172 .await
11173 .map_err(map_sqlx_error)?;
11174
11175 if updated.rows_affected() == 0 {
11176 tx.rollback().await.map_err(map_sqlx_error)?;
11177 return Err(AwaError::Validation("job is not in running state".into()));
11178 }
11179
11180 sqlx::query(&format!(
11181 r#"
11182 INSERT INTO {} (
11183 job_id,
11184 run_lease,
11185 callback_filter,
11186 callback_on_complete,
11187 callback_on_fail,
11188 callback_transform,
11189 updated_at
11190 )
11191 VALUES ($1, $2, $3, $4, $5, $6, clock_timestamp())
11192 ON CONFLICT (job_id, run_lease)
11193 DO UPDATE SET
11194 callback_filter = EXCLUDED.callback_filter,
11195 callback_on_complete = EXCLUDED.callback_on_complete,
11196 callback_on_fail = EXCLUDED.callback_on_fail,
11197 callback_transform = EXCLUDED.callback_transform,
11198 updated_at = clock_timestamp()
11199 "#,
11200 self.attempt_state_table()
11201 ))
11202 .bind(job_id)
11203 .bind(run_lease)
11204 .bind(&config.filter)
11205 .bind(&config.on_complete)
11206 .bind(&config.on_fail)
11207 .bind(&config.transform)
11208 .execute(tx.as_mut())
11209 .await
11210 .map_err(map_sqlx_error)?;
11211
11212 tx.commit().await.map_err(map_sqlx_error)?;
11213 Ok(callback_id)
11214 }
11215
11216 pub async fn cancel_callback(
11217 &self,
11218 pool: &PgPool,
11219 job_id: i64,
11220 run_lease: i64,
11221 ) -> Result<bool, AwaError> {
11222 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
11223 let result = sqlx::query(&format!(
11224 r#"
11225 UPDATE {}
11226 SET callback_id = NULL,
11227 callback_timeout_at = NULL
11228 WHERE job_id = $1
11229 AND callback_id IS NOT NULL
11230 AND state = 'running'
11231 AND run_lease = $2
11232 "#,
11233 self.leases_table()
11234 ))
11235 .bind(job_id)
11236 .bind(run_lease)
11237 .execute(tx.as_mut())
11238 .await
11239 .map_err(map_sqlx_error)?;
11240 if result.rows_affected() == 0 {
11241 tx.rollback().await.map_err(map_sqlx_error)?;
11242 return Ok(false);
11243 }
11244
11245 sqlx::query(&format!(
11246 r#"
11247 UPDATE {}
11248 SET callback_filter = NULL,
11249 callback_on_complete = NULL,
11250 callback_on_fail = NULL,
11251 callback_transform = NULL,
11252 updated_at = clock_timestamp()
11253 WHERE job_id = $1
11254 AND run_lease = $2
11255 "#,
11256 self.attempt_state_table()
11257 ))
11258 .bind(job_id)
11259 .bind(run_lease)
11260 .execute(tx.as_mut())
11261 .await
11262 .map_err(map_sqlx_error)?;
11263
11264 sqlx::query(&format!(
11265 r#"
11266 DELETE FROM {}
11267 WHERE job_id = $1
11268 AND run_lease = $2
11269 AND progress IS NULL
11270 AND callback_result IS NULL
11271 AND callback_filter IS NULL
11272 AND callback_on_complete IS NULL
11273 AND callback_on_fail IS NULL
11274 AND callback_transform IS NULL
11275 "#,
11276 self.attempt_state_table()
11277 ))
11278 .bind(job_id)
11279 .bind(run_lease)
11280 .execute(tx.as_mut())
11281 .await
11282 .map_err(map_sqlx_error)?;
11283
11284 tx.commit().await.map_err(map_sqlx_error)?;
11285 Ok(true)
11286 }
11287
11288 pub async fn load_active_lease_in_tx(
11295 &self,
11296 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
11297 job_id: i64,
11298 run_lease: i64,
11299 ) -> Result<Option<JobRow>, AwaError> {
11300 let schema = self.schema();
11301 let row: Option<LeaseJobRow> = sqlx::query_as(&format!(
11302 r#"
11303 SELECT
11304 lease.ready_slot,
11305 lease.ready_generation,
11306 lease.job_id,
11307 ready.kind,
11308 ready.queue,
11309 ready.args,
11310 lease.state,
11311 lease.priority,
11312 lease.attempt,
11313 lease.run_lease,
11314 lease.max_attempts,
11315 lease.lane_seq,
11316 ready.run_at,
11317 COALESCE(attempt.heartbeat_at, lease.heartbeat_at) AS heartbeat_at,
11318 lease.deadline_at,
11319 lease.attempted_at,
11320 NULL::timestamptz AS finalized_at,
11321 ready.created_at,
11322 ready.unique_key,
11323 ready.unique_states,
11324 lease.callback_id,
11325 lease.callback_timeout_at,
11326 attempt.callback_filter,
11327 attempt.callback_on_complete,
11328 attempt.callback_on_fail,
11329 attempt.callback_transform,
11330 COALESCE(ready.payload, '{{}}'::jsonb) AS payload,
11331 attempt.progress,
11332 attempt.callback_result
11333 FROM {schema}.leases AS lease
11334 JOIN {schema}.ready_entries AS ready
11335 ON ready.ready_slot = lease.ready_slot
11336 AND ready.ready_generation = lease.ready_generation
11337 AND ready.queue = lease.queue
11338 AND ready.priority = lease.priority
11339 AND ready.enqueue_shard = lease.enqueue_shard
11340 AND ready.lane_seq = lease.lane_seq
11341 LEFT JOIN {schema}.attempt_state AS attempt
11342 ON attempt.job_id = lease.job_id
11343 AND attempt.run_lease = lease.run_lease
11344 WHERE lease.job_id = $1
11345 AND lease.run_lease = $2
11346 "#,
11347 ))
11348 .bind(job_id)
11349 .bind(run_lease)
11350 .fetch_optional(tx.as_mut())
11351 .await
11352 .map_err(map_sqlx_error)?;
11353 row.map(LeaseJobRow::into_job_row).transpose()
11354 }
11355
11356 pub async fn enter_callback_wait(
11357 &self,
11358 pool: &PgPool,
11359 job_id: i64,
11360 run_lease: i64,
11361 callback_id: Uuid,
11362 ) -> Result<bool, AwaError> {
11363 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
11364 let entered = self
11365 .enter_callback_wait_in_tx(&mut tx, job_id, run_lease, callback_id)
11366 .await?;
11367 tx.commit().await.map_err(map_sqlx_error)?;
11368 Ok(entered)
11369 }
11370
11371 pub async fn enter_callback_wait_in_tx(
11374 &self,
11375 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
11376 job_id: i64,
11377 run_lease: i64,
11378 callback_id: Uuid,
11379 ) -> Result<bool, AwaError> {
11380 let result = sqlx::query(&format!(
11381 r#"
11382 UPDATE {}
11383 SET state = 'waiting_external',
11384 heartbeat_at = NULL,
11385 deadline_at = NULL
11386 WHERE job_id = $1
11387 AND state = 'running'
11388 AND run_lease = $2
11389 AND callback_id = $3
11390 "#,
11391 self.leases_table()
11392 ))
11393 .bind(job_id)
11394 .bind(run_lease)
11395 .bind(callback_id)
11396 .execute(tx.as_mut())
11397 .await
11398 .map_err(map_sqlx_error)?;
11399 Ok(result.rows_affected() > 0)
11400 }
11401
11402 pub async fn check_callback_state(
11403 &self,
11404 pool: &PgPool,
11405 job_id: i64,
11406 callback_id: Uuid,
11407 ) -> Result<CallbackPollResult, AwaError> {
11408 let row: Option<(JobState, Option<Uuid>, i64, Option<serde_json::Value>)> =
11409 sqlx::query_as(&format!(
11410 r#"
11411 SELECT
11412 lease.state,
11413 lease.callback_id,
11414 lease.run_lease,
11415 attempt.callback_result
11416 FROM {} AS lease
11417 LEFT JOIN {} AS attempt
11418 ON attempt.job_id = lease.job_id
11419 AND attempt.run_lease = lease.run_lease
11420 WHERE lease.job_id = $1
11421 ORDER BY lease.run_lease DESC
11422 LIMIT 1
11423 "#,
11424 self.leases_table(),
11425 self.attempt_state_table()
11426 ))
11427 .bind(job_id)
11428 .fetch_optional(pool)
11429 .await
11430 .map_err(map_sqlx_error)?;
11431
11432 match row {
11433 Some((JobState::Running, None, run_lease, Some(_))) => {
11434 let result = self.take_callback_result(pool, job_id, run_lease).await?;
11435 Ok(CallbackPollResult::Resolved(result))
11436 }
11437 Some((state, Some(current_callback_id), _, _))
11438 if current_callback_id != callback_id =>
11439 {
11440 Ok(CallbackPollResult::Stale {
11441 token: callback_id,
11442 current: current_callback_id,
11443 state,
11444 })
11445 }
11446 Some((JobState::WaitingExternal, Some(current), _, _)) if current == callback_id => {
11447 Ok(CallbackPollResult::Pending)
11448 }
11449 Some((state, _, _, _)) => Ok(CallbackPollResult::UnexpectedState {
11450 token: callback_id,
11451 state,
11452 }),
11453 None => {
11454 if let Some(job) = self.load_job(pool, job_id).await? {
11455 Ok(CallbackPollResult::UnexpectedState {
11456 token: callback_id,
11457 state: job.state,
11458 })
11459 } else {
11460 Ok(CallbackPollResult::NotFound)
11461 }
11462 }
11463 }
11464 }
11465
11466 pub async fn callback_job(
11467 &self,
11468 pool: &PgPool,
11469 callback_id: Uuid,
11470 run_lease: Option<i64>,
11471 ) -> Result<Option<JobRow>, AwaError> {
11472 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
11473 let result = self
11474 .callback_job_in_tx(&mut tx, callback_id, run_lease, false)
11475 .await?;
11476 tx.commit().await.map_err(map_sqlx_error)?;
11477 Ok(result)
11478 }
11479
11480 pub async fn callback_job_in_tx(
11486 &self,
11487 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
11488 callback_id: Uuid,
11489 run_lease: Option<i64>,
11490 for_update: bool,
11491 ) -> Result<Option<JobRow>, AwaError> {
11492 let lock_clause = if for_update {
11493 "FOR UPDATE OF lease"
11494 } else {
11495 ""
11496 };
11497 let row: Option<LeaseJobRow> = sqlx::query_as(&format!(
11498 r#"
11499 SELECT
11500 lease.ready_slot,
11501 lease.ready_generation,
11502 lease.job_id,
11503 ready.kind,
11504 ready.queue,
11505 ready.args,
11506 lease.state,
11507 lease.priority,
11508 lease.attempt,
11509 lease.run_lease,
11510 lease.max_attempts,
11511 lease.lane_seq,
11512 ready.run_at,
11513 COALESCE(attempt.heartbeat_at, lease.heartbeat_at) AS heartbeat_at,
11514 lease.deadline_at,
11515 lease.attempted_at,
11516 NULL::timestamptz AS finalized_at,
11517 ready.created_at,
11518 ready.unique_key,
11519 ready.unique_states,
11520 lease.callback_id,
11521 lease.callback_timeout_at,
11522 attempt.callback_filter,
11523 attempt.callback_on_complete,
11524 attempt.callback_on_fail,
11525 attempt.callback_transform,
11526 COALESCE(ready.payload, '{{}}'::jsonb) AS payload,
11527 attempt.progress,
11528 attempt.callback_result
11529 FROM {} AS lease
11530 JOIN {schema}.ready_entries AS ready
11531 ON ready.ready_slot = lease.ready_slot
11532 AND ready.ready_generation = lease.ready_generation
11533 AND ready.queue = lease.queue
11534 AND ready.priority = lease.priority
11535 AND ready.enqueue_shard = lease.enqueue_shard
11536 AND ready.lane_seq = lease.lane_seq
11537 LEFT JOIN {schema}.attempt_state AS attempt
11538 ON attempt.job_id = lease.job_id
11539 AND attempt.run_lease = lease.run_lease
11540 WHERE lease.callback_id = $1
11541 AND lease.state IN ('waiting_external', 'running')
11542 AND ($2::bigint IS NULL OR lease.run_lease = $2)
11543 ORDER BY lease.run_lease DESC
11544 LIMIT 1
11545 {lock_clause}
11546 "#,
11547 self.leases_table(),
11548 schema = self.schema(),
11549 ))
11550 .bind(callback_id)
11551 .bind(run_lease)
11552 .fetch_optional(tx.as_mut())
11553 .await
11554 .map_err(map_sqlx_error)?;
11555
11556 row.map(LeaseJobRow::into_job_row).transpose()
11557 }
11558
11559 #[tracing::instrument(skip(self, pool, payload), name = "queue_storage.complete_external")]
11560 pub async fn complete_external(
11561 &self,
11562 pool: &PgPool,
11563 callback_id: Uuid,
11564 payload: Option<serde_json::Value>,
11565 run_lease: Option<i64>,
11566 resume: bool,
11567 ) -> Result<JobRow, AwaError> {
11568 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
11569 let result = self
11570 .complete_external_in_tx(&mut tx, callback_id, payload, run_lease, resume)
11571 .await?;
11572 tx.commit().await.map_err(map_sqlx_error)?;
11573 Ok(result)
11574 }
11575
11576 pub async fn complete_external_in_tx(
11583 &self,
11584 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
11585 callback_id: Uuid,
11586 payload: Option<serde_json::Value>,
11587 run_lease: Option<i64>,
11588 resume: bool,
11589 ) -> Result<JobRow, AwaError> {
11590 if resume {
11591 let resumed: Option<(i64, i64)> = sqlx::query_as(&format!(
11592 r#"
11593 UPDATE {}
11594 SET state = 'running',
11595 callback_id = NULL,
11596 callback_timeout_at = NULL,
11597 heartbeat_at = clock_timestamp()
11598 WHERE callback_id = $1
11599 AND state IN ('waiting_external', 'running')
11600 AND ($2::bigint IS NULL OR run_lease = $2)
11601 RETURNING job_id, run_lease
11602 "#,
11603 self.leases_table()
11604 ))
11605 .bind(callback_id)
11606 .bind(run_lease)
11607 .fetch_optional(tx.as_mut())
11608 .await
11609 .map_err(map_sqlx_error)?;
11610
11611 let Some((job_id, resumed_run_lease)) = resumed else {
11612 return Err(AwaError::CallbackNotFound {
11613 callback_id: callback_id.to_string(),
11614 });
11615 };
11616
11617 sqlx::query(&format!(
11618 r#"
11619 INSERT INTO {} (
11620 job_id,
11621 run_lease,
11622 callback_filter,
11623 callback_on_complete,
11624 callback_on_fail,
11625 callback_transform,
11626 callback_result,
11627 updated_at
11628 )
11629 VALUES ($1, $2, NULL, NULL, NULL, NULL, $3, clock_timestamp())
11630 ON CONFLICT (job_id, run_lease)
11631 DO UPDATE SET
11632 callback_filter = NULL,
11633 callback_on_complete = NULL,
11634 callback_on_fail = NULL,
11635 callback_transform = NULL,
11636 callback_result = EXCLUDED.callback_result,
11637 updated_at = clock_timestamp()
11638 "#,
11639 self.attempt_state_table()
11640 ))
11641 .bind(job_id)
11642 .bind(resumed_run_lease)
11643 .bind(payload.unwrap_or(serde_json::Value::Null))
11644 .execute(tx.as_mut())
11645 .await
11646 .map_err(map_sqlx_error)?;
11647
11648 return self
11649 .load_active_lease_in_tx(tx, job_id, resumed_run_lease)
11650 .await?
11651 .ok_or(AwaError::CallbackNotFound {
11652 callback_id: callback_id.to_string(),
11653 });
11654 }
11655
11656 let schema = self.schema();
11657 let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
11658 r#"
11659 DELETE FROM {schema}.leases
11660 WHERE callback_id = $1
11661 AND state IN ('waiting_external', 'running')
11662 AND ($2::bigint IS NULL OR run_lease = $2)
11663 RETURNING
11664 ready_slot,
11665 ready_generation,
11666 job_id,
11667 queue,
11668 state,
11669 priority,
11670 attempt,
11671 run_lease,
11672 max_attempts,
11673 lane_seq,
11674 enqueue_shard,
11675 heartbeat_at,
11676 deadline_at,
11677 attempted_at,
11678 callback_id,
11679 callback_timeout_at
11680 "#
11681 ))
11682 .bind(callback_id)
11683 .bind(run_lease)
11684 .fetch_all(tx.as_mut())
11685 .await
11686 .map_err(map_sqlx_error)?;
11687
11688 if deleted.is_empty() {
11689 return Err(AwaError::CallbackNotFound {
11690 callback_id: callback_id.to_string(),
11691 });
11692 }
11693
11694 let completed_pairs: Vec<(i64, i64)> = deleted
11695 .iter()
11696 .map(|row| (row.job_id, row.run_lease))
11697 .collect();
11698 self.close_receipt_pairs_tx(tx, &completed_pairs, "completed")
11699 .await?;
11700
11701 let moved = self.hydrate_deleted_leases_tx(tx, deleted).await?;
11702 let moved = moved.into_iter().next().expect("deleted callback lease");
11703
11704 let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state(
11705 moved.payload.clone(),
11706 moved.progress.clone(),
11707 )?)?;
11708 payload.set_progress(None);
11709 let done_row =
11710 moved
11711 .clone()
11712 .into_done_row(JobState::Completed, Utc::now(), payload.into_json());
11713 self.insert_done_rows_tx(tx, std::slice::from_ref(&done_row), Some(moved.state))
11714 .await?;
11715 done_row.into_job_row()
11716 }
11717
11718 pub async fn fail_external(
11719 &self,
11720 pool: &PgPool,
11721 callback_id: Uuid,
11722 error: &str,
11723 run_lease: Option<i64>,
11724 ) -> Result<JobRow, AwaError> {
11725 self.fail_external_with_error_entry(
11726 pool,
11727 callback_id,
11728 serde_json::json!({ "error": error }),
11729 run_lease,
11730 )
11731 .await
11732 }
11733
11734 pub async fn fail_external_with_error_entry(
11735 &self,
11736 pool: &PgPool,
11737 callback_id: Uuid,
11738 error_entry: serde_json::Value,
11739 run_lease: Option<i64>,
11740 ) -> Result<JobRow, AwaError> {
11741 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
11742 let result = self
11743 .fail_external_with_error_entry_in_tx(&mut tx, callback_id, error_entry, run_lease)
11744 .await?;
11745 tx.commit().await.map_err(map_sqlx_error)?;
11746 Ok(result)
11747 }
11748
11749 pub async fn fail_external_with_error_entry_in_tx(
11752 &self,
11753 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
11754 callback_id: Uuid,
11755 error_entry: serde_json::Value,
11756 run_lease: Option<i64>,
11757 ) -> Result<JobRow, AwaError> {
11758 let schema = self.schema();
11759 let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
11760 r#"
11761 DELETE FROM {schema}.leases
11762 WHERE callback_id = $1
11763 AND state IN ('waiting_external', 'running')
11764 AND ($2::bigint IS NULL OR run_lease = $2)
11765 RETURNING
11766 ready_slot,
11767 ready_generation,
11768 job_id,
11769 queue,
11770 state,
11771 priority,
11772 attempt,
11773 run_lease,
11774 max_attempts,
11775 lane_seq,
11776 enqueue_shard,
11777 heartbeat_at,
11778 deadline_at,
11779 attempted_at,
11780 callback_id,
11781 callback_timeout_at
11782 "#
11783 ))
11784 .bind(callback_id)
11785 .bind(run_lease)
11786 .fetch_all(tx.as_mut())
11787 .await
11788 .map_err(map_sqlx_error)?;
11789
11790 if deleted.is_empty() {
11791 return Err(AwaError::CallbackNotFound {
11792 callback_id: callback_id.to_string(),
11793 });
11794 }
11795
11796 let failed_pairs: Vec<(i64, i64)> = deleted
11797 .iter()
11798 .map(|row| (row.job_id, row.run_lease))
11799 .collect();
11800 self.close_receipt_pairs_tx(tx, &failed_pairs, "failed")
11801 .await?;
11802
11803 let moved = self.hydrate_deleted_leases_tx(tx, deleted).await?;
11804 let moved = moved.into_iter().next().expect("deleted callback lease");
11805
11806 let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state(
11807 moved.payload.clone(),
11808 moved.progress.clone(),
11809 )?)?;
11810 let mut error_entry = match error_entry {
11811 serde_json::Value::Object(map) => serde_json::Value::Object(map),
11812 other => serde_json::json!({ "error": other }),
11813 };
11814 let error_obj = error_entry
11815 .as_object_mut()
11816 .ok_or_else(|| AwaError::Validation("callback error entry must be an object".into()))?;
11817 error_obj
11818 .entry("attempt".to_string())
11819 .or_insert_with(|| serde_json::Value::from(i64::from(moved.attempt)));
11820 error_obj
11821 .entry("at".to_string())
11822 .or_insert_with(|| serde_json::Value::String(Utc::now().to_rfc3339()));
11823 error_obj
11824 .entry("terminal".to_string())
11825 .or_insert(serde_json::Value::Bool(true));
11826 payload.push_error(error_entry);
11827 let done_row =
11828 moved
11829 .clone()
11830 .into_done_row(JobState::Failed, Utc::now(), payload.into_json());
11831 self.insert_done_rows_tx(tx, std::slice::from_ref(&done_row), Some(moved.state))
11832 .await?;
11833 done_row.into_job_row()
11834 }
11835
11836 pub async fn retry_external(
11837 &self,
11838 pool: &PgPool,
11839 callback_id: Uuid,
11840 run_lease: Option<i64>,
11841 ) -> Result<JobRow, AwaError> {
11842 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
11843 let result = self
11844 .retry_external_in_tx(&mut tx, callback_id, run_lease)
11845 .await?;
11846 tx.commit().await.map_err(map_sqlx_error)?;
11847 Ok(result)
11848 }
11849
11850 pub async fn retry_external_in_tx(
11852 &self,
11853 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
11854 callback_id: Uuid,
11855 run_lease: Option<i64>,
11856 ) -> Result<JobRow, AwaError> {
11857 let schema = self.schema();
11858 let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
11859 r#"
11860 DELETE FROM {schema}.leases
11861 WHERE callback_id = $1
11862 AND state = 'waiting_external'
11863 AND ($2::bigint IS NULL OR run_lease = $2)
11864 RETURNING
11865 ready_slot,
11866 ready_generation,
11867 job_id,
11868 queue,
11869 state,
11870 priority,
11871 attempt,
11872 run_lease,
11873 max_attempts,
11874 lane_seq,
11875 enqueue_shard,
11876 heartbeat_at,
11877 deadline_at,
11878 attempted_at,
11879 callback_id,
11880 callback_timeout_at
11881 "#
11882 ))
11883 .bind(callback_id)
11884 .bind(run_lease)
11885 .fetch_all(tx.as_mut())
11886 .await
11887 .map_err(map_sqlx_error)?;
11888
11889 if deleted.is_empty() {
11890 return Err(AwaError::CallbackNotFound {
11891 callback_id: callback_id.to_string(),
11892 });
11893 }
11894
11895 let retryable_pairs: Vec<(i64, i64)> = deleted
11896 .iter()
11897 .map(|row| (row.job_id, row.run_lease))
11898 .collect();
11899 self.close_receipt_pairs_tx(tx, &retryable_pairs, "retryable")
11900 .await?;
11901
11902 let moved = self.hydrate_deleted_leases_tx(tx, deleted).await?;
11903 let moved = moved.into_iter().next().expect("deleted callback lease");
11904
11905 let ready_payload =
11906 Self::payload_with_attempt_state(moved.payload.clone(), moved.progress.clone())?;
11907
11908 let ready_row = ExistingReadyRow {
11909 attempt: 0,
11910 run_at: Utc::now(),
11911 ..moved.clone().into_ready_row(Utc::now(), ready_payload)
11912 };
11913 self.insert_existing_ready_rows_tx(tx, vec![ready_row.clone()], Some(moved.state))
11914 .await?;
11915 self.notify_queues_tx(tx, std::iter::once(moved.queue.clone()))
11916 .await?;
11917 ReadyJobRow {
11918 job_id: ready_row.job_id,
11919 kind: ready_row.kind,
11920 queue: ready_row.queue,
11921 args: ready_row.args,
11922 priority: ready_row.priority,
11923 attempt: ready_row.attempt,
11924 run_lease: ready_row.run_lease,
11925 max_attempts: ready_row.max_attempts,
11926 run_at: ready_row.run_at,
11927 attempted_at: ready_row.attempted_at,
11928 created_at: ready_row.created_at,
11929 unique_key: ready_row.unique_key,
11930 payload: ready_row.payload,
11931 }
11932 .into_job_row()
11933 }
11934
11935 pub async fn heartbeat_callback(
11936 &self,
11937 pool: &PgPool,
11938 callback_id: Uuid,
11939 timeout: Duration,
11940 ) -> Result<JobRow, AwaError> {
11941 let updated: Option<(i64, i64)> = sqlx::query_as(&format!(
11942 r#"
11943 UPDATE {}
11944 SET callback_timeout_at = clock_timestamp() + make_interval(secs => $2)
11945 WHERE callback_id = $1
11946 AND state = 'waiting_external'
11947 RETURNING job_id, run_lease
11948 "#,
11949 self.leases_table()
11950 ))
11951 .bind(callback_id)
11952 .bind(timeout.as_secs_f64())
11953 .fetch_optional(pool)
11954 .await
11955 .map_err(map_sqlx_error)?;
11956
11957 let Some((job_id, _run_lease)) = updated else {
11958 return Err(AwaError::CallbackNotFound {
11959 callback_id: callback_id.to_string(),
11960 });
11961 };
11962
11963 self.load_job(pool, job_id)
11964 .await?
11965 .ok_or(AwaError::CallbackNotFound {
11966 callback_id: callback_id.to_string(),
11967 })
11968 }
11969
11970 pub async fn flush_progress(
11971 &self,
11972 pool: &PgPool,
11973 job_id: i64,
11974 run_lease: i64,
11975 progress: serde_json::Value,
11976 ) -> Result<(), AwaError> {
11977 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
11978 if self.lease_claim_receipts() {
11979 self.upsert_attempt_state_progress_from_receipts_tx(
11980 &mut tx,
11981 &[(job_id, run_lease, progress.clone())],
11982 )
11983 .await?;
11984 }
11985 sqlx::query(&format!(
11986 r#"
11987 INSERT INTO {} (job_id, run_lease, progress, updated_at)
11988 SELECT lease.job_id, lease.run_lease, $3, clock_timestamp()
11989 FROM {} AS lease
11990 WHERE lease.job_id = $1
11991 AND lease.run_lease = $2
11992 AND lease.state IN ('running', 'waiting_external')
11993 ON CONFLICT (job_id, run_lease)
11994 DO UPDATE SET
11995 progress = EXCLUDED.progress,
11996 updated_at = clock_timestamp()
11997 "#,
11998 self.attempt_state_table(),
11999 self.leases_table()
12000 ))
12001 .bind(job_id)
12002 .bind(run_lease)
12003 .bind(progress)
12004 .execute(tx.as_mut())
12005 .await
12006 .map_err(map_sqlx_error)?;
12007 tx.commit().await.map_err(map_sqlx_error)?;
12008 Ok(())
12009 }
12010
12011 pub async fn heartbeat_batch(
12012 &self,
12013 pool: &PgPool,
12014 jobs: &[(i64, i64)],
12015 ) -> Result<usize, AwaError> {
12016 if jobs.is_empty() {
12017 return Ok(0);
12018 }
12019
12020 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12021 let mut updated = 0_usize;
12022 if self.lease_claim_receipts() {
12023 updated += self
12032 .upsert_attempt_state_from_receipts_tx(&mut tx, jobs)
12033 .await?;
12034 } else {
12035 let job_ids: Vec<i64> = jobs.iter().map(|(job_id, _)| *job_id).collect();
12040 let run_leases: Vec<i64> = jobs.iter().map(|(_, run_lease)| *run_lease).collect();
12041 let result = sqlx::query(&format!(
12042 r#"
12043 WITH inflight AS (
12044 SELECT * FROM unnest($1::bigint[], $2::bigint[]) AS v(job_id, run_lease)
12045 )
12046 UPDATE {table}
12047 SET heartbeat_at = clock_timestamp()
12048 FROM inflight
12049 WHERE {table}.job_id = inflight.job_id
12050 AND {table}.run_lease = inflight.run_lease
12051 AND {table}.state = 'running'
12052 "#,
12053 table = self.leases_table(),
12054 ))
12055 .bind(&job_ids)
12056 .bind(&run_leases)
12057 .execute(tx.as_mut())
12058 .await
12059 .map_err(map_sqlx_error)?;
12060 updated += result.rows_affected() as usize;
12061 }
12062 tx.commit().await.map_err(map_sqlx_error)?;
12063 Ok(updated)
12064 }
12065
12066 pub async fn heartbeat_progress_batch(
12067 &self,
12068 pool: &PgPool,
12069 jobs: &[(i64, i64, serde_json::Value)],
12070 ) -> Result<usize, AwaError> {
12071 if jobs.is_empty() {
12072 return Ok(0);
12073 }
12074
12075 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12076 let updated = if self.lease_claim_receipts() {
12077 self.upsert_attempt_state_progress_from_receipts_tx(&mut tx, jobs)
12087 .await?
12088 } else {
12089 let schema = self.schema();
12093 let job_ids: Vec<i64> = jobs.iter().map(|(job_id, _, _)| *job_id).collect();
12094 let run_leases: Vec<i64> = jobs.iter().map(|(_, run_lease, _)| *run_lease).collect();
12095 let progress: Vec<serde_json::Value> =
12096 jobs.iter().map(|(_, _, value)| value.clone()).collect();
12097 let lease_updated: i64 = sqlx::query_scalar(&format!(
12098 r#"
12099 WITH inflight AS (
12100 SELECT * FROM unnest($1::bigint[], $2::bigint[], $3::jsonb[]) AS v(job_id, run_lease, progress)
12101 ),
12102 updated AS (
12103 UPDATE {table} AS lease
12104 SET heartbeat_at = clock_timestamp()
12105 FROM inflight
12106 WHERE lease.job_id = inflight.job_id
12107 AND lease.run_lease = inflight.run_lease
12108 AND lease.state = 'running'
12109 RETURNING lease.job_id, lease.run_lease, inflight.progress
12110 ),
12111 upsert_attempt AS (
12112 INSERT INTO {schema}.attempt_state (job_id, run_lease, progress, updated_at)
12113 SELECT job_id, run_lease, progress, clock_timestamp()
12114 FROM updated
12115 ON CONFLICT (job_id, run_lease)
12116 DO UPDATE SET
12117 progress = EXCLUDED.progress,
12118 updated_at = clock_timestamp()
12119 )
12120 SELECT count(*)::bigint FROM updated
12121 "#,
12122 table = self.leases_table(),
12123 ))
12124 .bind(&job_ids)
12125 .bind(&run_leases)
12126 .bind(&progress)
12127 .fetch_one(tx.as_mut())
12128 .await
12129 .map_err(map_sqlx_error)?;
12130 lease_updated as usize
12131 };
12132 tx.commit().await.map_err(map_sqlx_error)?;
12133 Ok(updated)
12134 }
12135
12136 pub async fn retry_after(
12137 &self,
12138 pool: &PgPool,
12139 job_id: i64,
12140 run_lease: i64,
12141 retry_after: Duration,
12142 progress: Option<serde_json::Value>,
12143 ) -> Result<Option<JobRow>, AwaError> {
12144 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12145 let result = self
12146 .retry_after_in_tx(&mut tx, job_id, run_lease, retry_after, progress)
12147 .await?;
12148 tx.commit().await.map_err(map_sqlx_error)?;
12149 Ok(result)
12150 }
12151
12152 pub async fn retry_after_in_tx(
12156 &self,
12157 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12158 job_id: i64,
12159 run_lease: i64,
12160 retry_after: Duration,
12161 progress: Option<serde_json::Value>,
12162 ) -> Result<Option<JobRow>, AwaError> {
12163 let Some(moved) = self
12164 .take_running_attempt_tx(tx, job_id, run_lease, "retryable")
12165 .await?
12166 else {
12167 return Ok(None);
12168 };
12169 let now = self.current_timestamp_tx(tx).await?;
12170
12171 let payload =
12172 Self::with_progress(moved.payload.clone(), progress.or(moved.progress.clone()))?;
12173 let deferred = moved.clone().into_deferred_row(
12174 JobState::Retryable,
12175 now + TimeDelta::from_std(retry_after).map_err(|err| {
12176 AwaError::Validation(format!("invalid retry_after duration: {err}"))
12177 })?,
12178 Some(now),
12179 payload,
12180 );
12181 self.insert_deferred_rows_tx(tx, vec![deferred.clone()], Some(moved.state))
12182 .await?;
12183 Ok(Some(deferred.into_job_row()?))
12184 }
12185
12186 pub async fn snooze(
12187 &self,
12188 pool: &PgPool,
12189 job_id: i64,
12190 run_lease: i64,
12191 snooze_for: Duration,
12192 progress: Option<serde_json::Value>,
12193 ) -> Result<Option<JobRow>, AwaError> {
12194 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12195 let Some(moved) = self
12196 .take_running_attempt_tx(&mut tx, job_id, run_lease, "scheduled")
12197 .await?
12198 else {
12199 tx.commit().await.map_err(map_sqlx_error)?;
12200 return Ok(None);
12201 };
12202 let now = self.current_timestamp_tx(&mut tx).await?;
12203
12204 let payload =
12205 Self::with_progress(moved.payload.clone(), progress.or(moved.progress.clone()))?;
12206 let mut deferred = moved.clone().into_deferred_row(
12207 JobState::Scheduled,
12208 now + TimeDelta::from_std(snooze_for)
12209 .map_err(|err| AwaError::Validation(format!("invalid snooze duration: {err}")))?,
12210 None,
12211 payload,
12212 );
12213 deferred.attempt = deferred.attempt.saturating_sub(1);
12214 self.insert_deferred_rows_tx(&mut tx, vec![deferred.clone()], Some(moved.state))
12215 .await?;
12216 tx.commit().await.map_err(map_sqlx_error)?;
12217 Ok(Some(deferred.into_job_row()?))
12218 }
12219
12220 async fn insert_rescued_deferred_or_cancel_tx(
12227 &self,
12228 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12229 row: LeaseTransitionRow,
12230 deferred: DeferredJobRow,
12231 duplicate_error: &str,
12232 ) -> Result<JobRow, AwaError> {
12233 let old_state = row.state;
12234 {
12235 let mut savepoint = tx.begin().await.map_err(map_sqlx_error)?;
12236 match self
12237 .insert_deferred_rows_tx(&mut savepoint, vec![deferred.clone()], Some(old_state))
12238 .await
12239 {
12240 Ok(_) => {
12241 savepoint.commit().await.map_err(map_sqlx_error)?;
12242 return deferred.into_job_row();
12243 }
12244 Err(AwaError::UniqueConflict { .. }) => {
12245 savepoint.rollback().await.map_err(map_sqlx_error)?;
12246 }
12247 Err(err) => return Err(err),
12248 }
12249 }
12250 self.cancel_rescued_duplicate_tx(tx, row, deferred.payload, duplicate_error)
12251 .await
12252 }
12253
12254 async fn insert_rescued_done_or_cancel_tx(
12258 &self,
12259 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12260 row: LeaseTransitionRow,
12261 done: DoneJobRow,
12262 duplicate_error: &str,
12263 ) -> Result<JobRow, AwaError> {
12264 let old_state = row.state;
12265 {
12266 let mut savepoint = tx.begin().await.map_err(map_sqlx_error)?;
12267 match self
12268 .insert_done_rows_tx(&mut savepoint, std::slice::from_ref(&done), Some(old_state))
12269 .await
12270 {
12271 Ok(_) => {
12272 savepoint.commit().await.map_err(map_sqlx_error)?;
12273 return done.into_job_row();
12274 }
12275 Err(AwaError::UniqueConflict { .. }) => {
12276 savepoint.rollback().await.map_err(map_sqlx_error)?;
12277 }
12278 Err(err) => return Err(err),
12279 }
12280 }
12281 self.cancel_rescued_duplicate_tx(tx, row, done.payload, duplicate_error)
12282 .await
12283 }
12284
12285 async fn cancel_rescued_duplicate_tx(
12293 &self,
12294 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12295 row: LeaseTransitionRow,
12296 base_payload: serde_json::Value,
12297 duplicate_error: &str,
12298 ) -> Result<JobRow, AwaError> {
12299 let old_state = row.state;
12300 let mut payload = RuntimePayload::from_json(base_payload)?;
12301 payload.push_error(lifecycle_error(duplicate_error, row.attempt, true));
12302 let done = row.into_done_row(JobState::Cancelled, Utc::now(), payload.into_json());
12303
12304 {
12305 let mut savepoint = tx.begin().await.map_err(map_sqlx_error)?;
12306 match self
12307 .insert_done_rows_tx(&mut savepoint, std::slice::from_ref(&done), Some(old_state))
12308 .await
12309 {
12310 Ok(_) => {
12311 savepoint.commit().await.map_err(map_sqlx_error)?;
12312 return done.into_job_row();
12313 }
12314 Err(AwaError::UniqueConflict { .. }) => {
12315 savepoint.rollback().await.map_err(map_sqlx_error)?;
12316 }
12317 Err(err) => return Err(err),
12318 }
12319 }
12320
12321 let mut stripped = done;
12322 stripped.unique_key = None;
12323 stripped.unique_states = None;
12324 self.insert_done_rows_tx(tx, std::slice::from_ref(&stripped), Some(old_state))
12325 .await?;
12326 stripped.into_job_row()
12327 }
12328
12329 pub async fn cancel_running(
12330 &self,
12331 pool: &PgPool,
12332 job_id: i64,
12333 run_lease: i64,
12334 reason: &str,
12335 progress: Option<serde_json::Value>,
12336 ) -> Result<Option<JobRow>, AwaError> {
12337 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12338 let result = self
12339 .cancel_running_in_tx(&mut tx, job_id, run_lease, reason, progress)
12340 .await?;
12341 tx.commit().await.map_err(map_sqlx_error)?;
12342 Ok(result)
12343 }
12344
12345 pub async fn cancel_running_in_tx(
12347 &self,
12348 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12349 job_id: i64,
12350 run_lease: i64,
12351 reason: &str,
12352 progress: Option<serde_json::Value>,
12353 ) -> Result<Option<JobRow>, AwaError> {
12354 let Some(moved) = self
12355 .take_running_attempt_tx(tx, job_id, run_lease, "cancelled")
12356 .await?
12357 else {
12358 return Ok(None);
12359 };
12360
12361 let mut payload = RuntimePayload::from_json(Self::with_progress(
12362 moved.payload.clone(),
12363 progress.or(moved.progress.clone()),
12364 )?)?;
12365 payload.push_error(lifecycle_error(
12366 format!("cancelled: {reason}"),
12367 moved.attempt,
12368 false,
12369 ));
12370 let done =
12371 moved
12372 .clone()
12373 .into_done_row(JobState::Cancelled, Utc::now(), payload.into_json());
12374 self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(moved.state))
12375 .await?;
12376 Ok(Some(done.into_job_row()?))
12377 }
12378
12379 pub async fn fail_terminal(
12380 &self,
12381 pool: &PgPool,
12382 job_id: i64,
12383 run_lease: i64,
12384 error: &str,
12385 progress: Option<serde_json::Value>,
12386 ) -> Result<Option<JobRow>, AwaError> {
12387 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12388 let result = self
12389 .fail_terminal_in_tx(&mut tx, job_id, run_lease, error, progress)
12390 .await?;
12391 tx.commit().await.map_err(map_sqlx_error)?;
12392 Ok(result)
12393 }
12394
12395 pub async fn fail_terminal_in_tx(
12397 &self,
12398 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12399 job_id: i64,
12400 run_lease: i64,
12401 error: &str,
12402 progress: Option<serde_json::Value>,
12403 ) -> Result<Option<JobRow>, AwaError> {
12404 let Some(moved) = self
12405 .take_running_attempt_tx(tx, job_id, run_lease, "failed")
12406 .await?
12407 else {
12408 return Ok(None);
12409 };
12410
12411 let mut payload = RuntimePayload::from_json(Self::with_progress(
12412 moved.payload.clone(),
12413 progress.or(moved.progress.clone()),
12414 )?)?;
12415 payload.push_error(lifecycle_error(error, moved.attempt, true));
12416 let done = moved
12417 .clone()
12418 .into_done_row(JobState::Failed, Utc::now(), payload.into_json());
12419 self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(moved.state))
12420 .await?;
12421 Ok(Some(done.into_job_row()?))
12422 }
12423
12424 pub async fn fail_to_dlq(
12425 &self,
12426 pool: &PgPool,
12427 job_id: i64,
12428 run_lease: i64,
12429 dlq_reason: &str,
12430 error: &str,
12431 progress: Option<serde_json::Value>,
12432 ) -> Result<Option<JobRow>, AwaError> {
12433 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12434 let result = self
12435 .fail_to_dlq_in_tx(&mut tx, job_id, run_lease, dlq_reason, error, progress)
12436 .await?;
12437 tx.commit().await.map_err(map_sqlx_error)?;
12438 Ok(result)
12439 }
12440
12441 pub async fn fail_to_dlq_in_tx(
12443 &self,
12444 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12445 job_id: i64,
12446 run_lease: i64,
12447 dlq_reason: &str,
12448 error: &str,
12449 progress: Option<serde_json::Value>,
12450 ) -> Result<Option<JobRow>, AwaError> {
12451 let Some(moved) = self
12452 .take_running_attempt_tx(tx, job_id, run_lease, "dlq")
12453 .await?
12454 else {
12455 return Ok(None);
12456 };
12457
12458 let finalized_at = Utc::now();
12459 let dlq_at = finalized_at;
12460 let mut payload = RuntimePayload::from_json(Self::with_progress(
12461 moved.payload.clone(),
12462 progress.or(moved.progress.clone()),
12463 )?)?;
12464 payload.push_error(lifecycle_error(error, moved.attempt, true));
12465 let dlq_row = moved.clone().into_dlq_row(
12466 finalized_at,
12467 payload.into_json(),
12468 dlq_reason.to_string(),
12469 dlq_at,
12470 );
12471 self.insert_dlq_rows_tx(tx, std::slice::from_ref(&dlq_row), Some(moved.state))
12472 .await?;
12473 Ok(Some(dlq_row.into_job_row()?))
12474 }
12475
12476 pub async fn move_failed_to_dlq(
12477 &self,
12478 pool: &PgPool,
12479 job_id: i64,
12480 dlq_reason: &str,
12481 ) -> Result<Option<JobRow>, AwaError> {
12482 let schema = self.schema();
12483 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12484 let done_projection = done_row_projection("done", "ready");
12485 let ready_join = done_ready_join(schema, "done", "ready");
12486 let moved: Option<DoneJobRow> = sqlx::query_as(&format!(
12487 r#"
12488 WITH deleted AS (
12489 DELETE FROM {schema}.done_entries
12490 WHERE (job_id, finalized_at) IN (
12491 SELECT job_id, finalized_at
12492 FROM {schema}.done_entries
12493 WHERE job_id = $1
12494 AND state = 'failed'
12495 ORDER BY finalized_at DESC
12496 LIMIT 1
12497 FOR UPDATE SKIP LOCKED
12498 )
12499 RETURNING *
12500 )
12501 SELECT {done_projection}
12502 FROM deleted AS done
12503 {ready_join}
12504 "#
12505 ))
12506 .bind(job_id)
12507 .fetch_optional(tx.as_mut())
12508 .await
12509 .map_err(map_sqlx_error)?;
12510
12511 let Some(moved) = moved else {
12512 tx.commit().await.map_err(map_sqlx_error)?;
12513 return Ok(None);
12514 };
12515
12516 self.ensure_terminal_removed_receipt_closures_tx(&mut tx, std::slice::from_ref(&moved))
12517 .await?;
12518 self.decrement_live_terminal_counters_tx(
12521 &mut tx,
12522 &Self::done_rows_to_counter_keys(std::slice::from_ref(&moved)),
12523 )
12524 .await?;
12525 let dlq_row = moved
12526 .clone()
12527 .into_dlq_row(dlq_reason.to_string(), Utc::now());
12528 self.insert_dlq_rows_tx(&mut tx, std::slice::from_ref(&dlq_row), Some(moved.state))
12529 .await?;
12530 tx.commit().await.map_err(map_sqlx_error)?;
12531 Ok(Some(dlq_row.into_job_row()?))
12532 }
12533
12534 #[tracing::instrument(
12535 skip(self, pool, dlq_reason),
12536 fields(kind = ?kind, queue = ?queue),
12537 name = "queue_storage.bulk_move_failed_to_dlq"
12538 )]
12539 pub async fn bulk_move_failed_to_dlq(
12540 &self,
12541 pool: &PgPool,
12542 kind: Option<&str>,
12543 queue: Option<&str>,
12544 dlq_reason: &str,
12545 ) -> Result<u64, AwaError> {
12546 let schema = self.schema();
12547 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12548 let done_projection = done_row_projection("done", "ready");
12549 let ready_join = done_ready_join(schema, "done", "ready");
12550 let moved: Vec<DoneJobRow> = sqlx::query_as(&format!(
12551 r#"
12552 WITH deleted AS (
12553 DELETE FROM {schema}.done_entries
12554 WHERE state = 'failed'
12555 AND ($1::text IS NULL OR kind = $1)
12556 AND ($2::text IS NULL OR queue = $2)
12557 RETURNING *
12558 )
12559 SELECT {done_projection}
12560 FROM deleted AS done
12561 {ready_join}
12562 "#
12563 ))
12564 .bind(kind)
12565 .bind(queue)
12566 .fetch_all(tx.as_mut())
12567 .await
12568 .map_err(map_sqlx_error)?;
12569
12570 if moved.is_empty() {
12571 tx.commit().await.map_err(map_sqlx_error)?;
12572 return Ok(0);
12573 }
12574
12575 self.ensure_terminal_removed_receipt_closures_tx(&mut tx, &moved)
12576 .await?;
12577 self.decrement_live_terminal_counters_tx(&mut tx, &Self::done_rows_to_counter_keys(&moved))
12580 .await?;
12581 let dlq_at = Utc::now();
12582 let rows: Vec<DlqJobRow> = moved
12583 .into_iter()
12584 .map(|row| row.into_dlq_row(dlq_reason.to_string(), dlq_at))
12585 .collect();
12586 self.insert_dlq_rows_tx(&mut tx, &rows, Some(JobState::Failed))
12587 .await?;
12588 tx.commit().await.map_err(map_sqlx_error)?;
12589 Ok(rows.len() as u64)
12590 }
12591
12592 pub async fn retry_from_dlq(
12593 &self,
12594 pool: &PgPool,
12595 job_id: i64,
12596 opts: &RetryFromDlqOpts,
12597 ) -> Result<Option<JobRow>, AwaError> {
12598 let schema = self.schema();
12599 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12600 let moved: Option<DlqJobRow> = sqlx::query_as(&format!(
12601 r#"
12602 DELETE FROM {schema}.dlq_entries
12603 WHERE job_id = $1
12604 RETURNING
12605 job_id,
12606 kind,
12607 queue,
12608 args,
12609 state,
12610 priority,
12611 attempt,
12612 run_lease,
12613 max_attempts,
12614 run_at,
12615 attempted_at,
12616 finalized_at,
12617 created_at,
12618 unique_key,
12619 unique_states,
12620 COALESCE(payload, '{{}}'::jsonb) AS payload,
12621 dlq_reason,
12622 dlq_at,
12623 original_run_lease
12624 "#
12625 ))
12626 .bind(job_id)
12627 .fetch_optional(tx.as_mut())
12628 .await
12629 .map_err(map_sqlx_error)?;
12630
12631 let Some(moved) = moved else {
12632 tx.commit().await.map_err(map_sqlx_error)?;
12633 return Ok(None);
12634 };
12635
12636 let queue = opts.queue.clone().unwrap_or_else(|| moved.queue.clone());
12637 let priority = opts.priority.unwrap_or(moved.priority);
12638 let mut payload = RuntimePayload::from_json(moved.payload.clone())?;
12639 payload.set_progress(None);
12640 let payload = payload.into_json();
12641
12642 if let Some(run_at) = opts.run_at.filter(|run_at| *run_at > Utc::now()) {
12643 let deferred = moved.into_retry_deferred_row(queue, priority, run_at, payload);
12644 self.insert_deferred_rows_tx(&mut tx, vec![deferred.clone()], Some(JobState::Failed))
12645 .await?;
12646 tx.commit().await.map_err(map_sqlx_error)?;
12647 return Ok(Some(deferred.into_job_row()?));
12648 }
12649
12650 let ready = moved.into_retry_ready_row(queue.clone(), priority, Utc::now(), payload);
12651 self.insert_existing_ready_rows_tx(&mut tx, vec![ready.clone()], Some(JobState::Failed))
12652 .await?;
12653 self.notify_queues_tx(&mut tx, std::iter::once(queue))
12654 .await?;
12655 tx.commit().await.map_err(map_sqlx_error)?;
12656 Ok(Some(
12657 ReadyJobRow {
12658 job_id: ready.job_id,
12659 kind: ready.kind,
12660 queue: ready.queue,
12661 args: ready.args,
12662 priority: ready.priority,
12663 attempt: ready.attempt,
12664 run_lease: ready.run_lease,
12665 max_attempts: ready.max_attempts,
12666 run_at: ready.run_at,
12667 attempted_at: ready.attempted_at,
12668 created_at: ready.created_at,
12669 unique_key: ready.unique_key,
12670 payload: ready.payload,
12671 }
12672 .into_job_row()?,
12673 ))
12674 }
12675
12676 #[tracing::instrument(
12677 skip(self, pool, filter),
12678 fields(kind = ?filter.kind, queue = ?filter.queue, tag = ?filter.tag),
12679 name = "queue_storage.bulk_retry_from_dlq"
12680 )]
12681 pub async fn bulk_retry_from_dlq(
12682 &self,
12683 pool: &PgPool,
12684 filter: &ListDlqFilter,
12685 ) -> Result<u64, AwaError> {
12686 let schema = self.schema();
12687 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12688 let moved: Vec<DlqJobRow> = sqlx::query_as(&format!(
12689 r#"
12690 DELETE FROM {schema}.dlq_entries
12691 WHERE ($1::text IS NULL OR kind = $1)
12692 AND ($2::text IS NULL OR queue = $2)
12693 AND ($3::text IS NULL OR payload -> 'tags' ? $3)
12694 AND (
12695 ($4::bigint IS NULL AND $5::timestamptz IS NULL)
12696 OR ($4::bigint IS NOT NULL AND $5::timestamptz IS NULL AND job_id < $4)
12697 OR ($4::bigint IS NULL AND $5::timestamptz IS NOT NULL AND dlq_at < $5)
12698 OR (
12699 $4::bigint IS NOT NULL
12700 AND $5::timestamptz IS NOT NULL
12701 AND (dlq_at, job_id) < ($5, $4)
12702 )
12703 )
12704 RETURNING
12705 job_id,
12706 kind,
12707 queue,
12708 args,
12709 state,
12710 priority,
12711 attempt,
12712 run_lease,
12713 max_attempts,
12714 run_at,
12715 attempted_at,
12716 finalized_at,
12717 created_at,
12718 unique_key,
12719 unique_states,
12720 COALESCE(payload, '{{}}'::jsonb) AS payload,
12721 dlq_reason,
12722 dlq_at,
12723 original_run_lease
12724 "#
12725 ))
12726 .bind(&filter.kind)
12727 .bind(&filter.queue)
12728 .bind(&filter.tag)
12729 .bind(filter.before_id)
12730 .bind(filter.before_dlq_at)
12731 .fetch_all(tx.as_mut())
12732 .await
12733 .map_err(map_sqlx_error)?;
12734
12735 if moved.is_empty() {
12736 tx.commit().await.map_err(map_sqlx_error)?;
12737 return Ok(0);
12738 }
12739
12740 let run_at = Utc::now();
12741 let mut queues = BTreeSet::new();
12742 let mut ready_rows = Vec::with_capacity(moved.len());
12743 for moved_row in moved {
12744 let queue = moved_row.queue.clone();
12745 let priority = moved_row.priority;
12746 queues.insert(queue.clone());
12747 let mut payload = RuntimePayload::from_json(moved_row.payload.clone())?;
12748 payload.set_progress(None);
12749 ready_rows.push(moved_row.into_retry_ready_row(
12750 queue,
12751 priority,
12752 run_at,
12753 payload.into_json(),
12754 ));
12755 }
12756
12757 let revived = ready_rows.len() as u64;
12758 self.insert_existing_ready_rows_tx(&mut tx, ready_rows, Some(JobState::Failed))
12759 .await?;
12760 self.notify_queues_tx(&mut tx, queues).await?;
12761 tx.commit().await.map_err(map_sqlx_error)?;
12762 Ok(revived)
12763 }
12764
12765 pub async fn discard_failed_by_kind(&self, pool: &PgPool, kind: &str) -> Result<u64, AwaError> {
12766 let schema = self.schema();
12767 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12768
12769 let done_projection = done_row_projection("done", "ready");
12770 let ready_join = done_ready_join(schema, "done", "ready");
12771 let deleted_done: Vec<DoneJobRow> = sqlx::query_as(&format!(
12772 r#"
12773 WITH deleted AS (
12774 DELETE FROM {schema}.done_entries
12775 WHERE kind = $1
12776 AND state = 'failed'
12777 RETURNING *
12778 )
12779 SELECT {done_projection}
12780 FROM deleted AS done
12781 {ready_join}
12782 "#
12783 ))
12784 .bind(kind)
12785 .fetch_all(tx.as_mut())
12786 .await
12787 .map_err(map_sqlx_error)?;
12788
12789 let deleted_dlq: Vec<DlqJobRow> = sqlx::query_as(&format!(
12790 r#"
12791 DELETE FROM {schema}.dlq_entries
12792 WHERE kind = $1
12793 RETURNING
12794 job_id,
12795 kind,
12796 queue,
12797 args,
12798 state,
12799 priority,
12800 attempt,
12801 run_lease,
12802 max_attempts,
12803 run_at,
12804 attempted_at,
12805 finalized_at,
12806 created_at,
12807 unique_key,
12808 unique_states,
12809 COALESCE(payload, '{{}}'::jsonb) AS payload,
12810 dlq_reason,
12811 dlq_at,
12812 original_run_lease
12813 "#
12814 ))
12815 .bind(kind)
12816 .fetch_all(tx.as_mut())
12817 .await
12818 .map_err(map_sqlx_error)?;
12819
12820 self.ensure_terminal_removed_receipt_closures_tx(&mut tx, &deleted_done)
12825 .await?;
12826 self.decrement_live_terminal_counters_tx(
12827 &mut tx,
12828 &Self::done_rows_to_counter_keys(&deleted_done),
12829 )
12830 .await?;
12831
12832 for row in &deleted_done {
12833 self.sync_unique_claim(
12834 &mut tx,
12835 row.job_id,
12836 &row.unique_key,
12837 row.unique_states.as_deref(),
12838 Some(row.state),
12839 None,
12840 )
12841 .await?;
12842 }
12843
12844 for row in &deleted_dlq {
12845 self.sync_unique_claim(
12846 &mut tx,
12847 row.job_id,
12848 &row.unique_key,
12849 row.unique_states.as_deref(),
12850 Some(row.state),
12851 None,
12852 )
12853 .await?;
12854 }
12855
12856 tx.commit().await.map_err(map_sqlx_error)?;
12857 Ok((deleted_done.len() + deleted_dlq.len()) as u64)
12858 }
12859
12860 pub async fn fail_retryable(
12861 &self,
12862 pool: &PgPool,
12863 job_id: i64,
12864 run_lease: i64,
12865 error: &str,
12866 progress: Option<serde_json::Value>,
12867 ) -> Result<Option<JobRow>, AwaError> {
12868 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12869 let result = self
12870 .fail_retryable_in_tx(&mut tx, job_id, run_lease, error, progress)
12871 .await?;
12872 tx.commit().await.map_err(map_sqlx_error)?;
12873 Ok(result)
12874 }
12875
12876 pub async fn fail_retryable_in_tx(
12878 &self,
12879 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12880 job_id: i64,
12881 run_lease: i64,
12882 error: &str,
12883 progress: Option<serde_json::Value>,
12884 ) -> Result<Option<JobRow>, AwaError> {
12885 let Some(moved) = self
12886 .take_running_attempt_tx(tx, job_id, run_lease, "retryable")
12887 .await?
12888 else {
12889 return Ok(None);
12890 };
12891
12892 let mut payload = RuntimePayload::from_json(Self::with_progress(
12893 moved.payload.clone(),
12894 progress.or(moved.progress.clone()),
12895 )?)?;
12896 let exhausted = moved.attempt >= moved.max_attempts;
12897 payload.push_error(lifecycle_error(error, moved.attempt, exhausted));
12898
12899 if exhausted {
12900 let done =
12901 moved
12902 .clone()
12903 .into_done_row(JobState::Failed, Utc::now(), payload.into_json());
12904 self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(moved.state))
12905 .await?;
12906 return Ok(Some(done.into_job_row()?));
12907 }
12908
12909 let deferred = moved.clone().into_deferred_row(
12910 JobState::Retryable,
12911 self.backoff_at_tx(tx, moved.attempt, moved.max_attempts)
12912 .await?,
12913 Some(Utc::now()),
12914 payload.into_json(),
12915 );
12916 self.insert_deferred_rows_tx(tx, vec![deferred.clone()], Some(moved.state))
12917 .await?;
12918 Ok(Some(deferred.into_job_row()?))
12919 }
12920
12921 #[tracing::instrument(skip(self, pool), name = "queue_storage.rescue_stale_heartbeats")]
12922 pub async fn rescue_stale_heartbeats(
12923 &self,
12924 pool: &PgPool,
12925 staleness: Duration,
12926 ) -> Result<Vec<JobRow>, AwaError> {
12927 let schema = self.schema();
12928 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12929 let cutoff = Utc::now()
12930 - TimeDelta::from_std(staleness)
12931 .map_err(|err| AwaError::Validation(format!("invalid staleness: {err}")))?;
12932 let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
12958 r#"
12959 DELETE FROM {schema}.leases AS target
12960 WHERE (target.job_id, target.run_lease) IN (
12961 SELECT lease.job_id, lease.run_lease
12962 FROM {schema}.leases AS lease
12963 LEFT JOIN {schema}.attempt_state AS attempt
12964 ON attempt.job_id = lease.job_id
12965 AND attempt.run_lease = lease.run_lease
12966 WHERE lease.state = 'running'
12967 AND COALESCE(attempt.heartbeat_at, lease.heartbeat_at) < $1
12968 ORDER BY COALESCE(attempt.heartbeat_at, lease.heartbeat_at) ASC
12969 LIMIT 500
12970 FOR UPDATE OF lease SKIP LOCKED
12971 )
12972 RETURNING
12973 ready_slot,
12974 ready_generation,
12975 job_id,
12976 queue,
12977 state,
12978 priority,
12979 attempt,
12980 run_lease,
12981 max_attempts,
12982 lane_seq,
12983 enqueue_shard,
12984 heartbeat_at,
12985 deadline_at,
12986 attempted_at,
12987 callback_id,
12988 callback_timeout_at
12989 "#
12990 ))
12991 .bind(cutoff)
12992 .fetch_all(tx.as_mut())
12993 .await
12994 .map_err(map_sqlx_error)?;
12995
12996 let rescued_receipts = if self.lease_claim_receipts() {
12997 self.rescue_stale_receipt_claims_tx(&mut tx, cutoff).await?
12998 } else {
12999 Vec::new()
13000 };
13001
13002 if deleted.is_empty() && rescued_receipts.is_empty() {
13003 tx.commit().await.map_err(map_sqlx_error)?;
13004 return Ok(Vec::new());
13005 }
13006
13007 let moved_leases = self.hydrate_deleted_leases_tx(&mut tx, deleted).await?;
13008 let moved_receipts = self
13009 .hydrate_deleted_leases_tx(&mut tx, rescued_receipts)
13010 .await?;
13011
13012 let mut rescued = Vec::with_capacity(moved_leases.len() + moved_receipts.len());
13013 for row in moved_leases {
13014 let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state(
13015 row.payload.clone(),
13016 row.progress.clone(),
13017 )?)?;
13018 payload.push_error(lifecycle_error(
13019 "heartbeat stale: worker presumed dead",
13020 row.attempt,
13021 false,
13022 ));
13023 let deferred = row.clone().into_deferred_row(
13024 JobState::Retryable,
13025 self.backoff_at_tx(&mut tx, row.attempt, row.max_attempts)
13026 .await?,
13027 Some(Utc::now()),
13028 payload.into_json(),
13029 );
13030 let rescued_row = self
13031 .insert_rescued_deferred_or_cancel_tx(
13032 &mut tx,
13033 row,
13034 deferred,
13035 "rescued as duplicate: unique claim held by a newer job",
13036 )
13037 .await?;
13038 rescued.push(rescued_row);
13039 }
13040 for row in moved_receipts {
13041 let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state(
13042 row.payload.clone(),
13043 row.progress.clone(),
13044 )?)?;
13045 payload.push_error(lifecycle_error(
13046 "receipt claim stale: worker presumed dead",
13047 row.attempt,
13048 false,
13049 ));
13050 let deferred = row.clone().into_deferred_row(
13051 JobState::Retryable,
13052 self.backoff_at_tx(&mut tx, row.attempt, row.max_attempts)
13053 .await?,
13054 Some(Utc::now()),
13055 payload.into_json(),
13056 );
13057 let rescued_row = self
13058 .insert_rescued_deferred_or_cancel_tx(
13059 &mut tx,
13060 row,
13061 deferred,
13062 "rescued as duplicate: unique claim held by a newer job",
13063 )
13064 .await?;
13065 rescued.push(rescued_row);
13066 }
13067 tx.commit().await.map_err(map_sqlx_error)?;
13068 Ok(rescued)
13069 }
13070
13071 #[tracing::instrument(skip(self, pool), name = "queue_storage.rescue_expired_deadlines")]
13072 pub async fn rescue_expired_deadlines(&self, pool: &PgPool) -> Result<Vec<JobRow>, AwaError> {
13073 let schema = self.schema();
13074 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
13075 let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
13076 r#"
13077 DELETE FROM {schema}.leases
13078 WHERE job_id IN (
13079 SELECT job_id
13080 FROM {schema}.leases
13081 WHERE state = 'running'
13082 AND deadline_at IS NOT NULL
13083 AND deadline_at < clock_timestamp()
13084 ORDER BY deadline_at ASC
13085 LIMIT 500
13086 FOR UPDATE SKIP LOCKED
13087 )
13088 RETURNING
13089 ready_slot,
13090 ready_generation,
13091 job_id,
13092 queue,
13093 state,
13094 priority,
13095 attempt,
13096 run_lease,
13097 max_attempts,
13098 lane_seq,
13099 enqueue_shard,
13100 heartbeat_at,
13101 deadline_at,
13102 attempted_at,
13103 callback_id,
13104 callback_timeout_at
13105 "#
13106 ))
13107 .fetch_all(tx.as_mut())
13108 .await
13109 .map_err(map_sqlx_error)?;
13110
13111 let receipt_deleted = if self.lease_claim_receipts() {
13117 self.rescue_expired_receipt_deadlines_tx(&mut tx).await?
13118 } else {
13119 Vec::new()
13120 };
13121
13122 if deleted.is_empty() && receipt_deleted.is_empty() {
13123 tx.commit().await.map_err(map_sqlx_error)?;
13124 return Ok(Vec::new());
13125 }
13126
13127 let mut moved = self.hydrate_deleted_leases_tx(&mut tx, deleted).await?;
13128 moved.extend(
13129 self.hydrate_deleted_leases_tx(&mut tx, receipt_deleted)
13130 .await?,
13131 );
13132
13133 let mut rescued = Vec::with_capacity(moved.len());
13134 for row in moved {
13135 let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state(
13136 row.payload.clone(),
13137 row.progress.clone(),
13138 )?)?;
13139 payload.push_error(lifecycle_error(
13140 "hard deadline exceeded",
13141 row.attempt,
13142 false,
13143 ));
13144 let deferred = row.clone().into_deferred_row(
13145 JobState::Retryable,
13146 self.backoff_at_tx(&mut tx, row.attempt, row.max_attempts)
13147 .await?,
13148 Some(Utc::now()),
13149 payload.into_json(),
13150 );
13151 let rescued_row = self
13152 .insert_rescued_deferred_or_cancel_tx(
13153 &mut tx,
13154 row,
13155 deferred,
13156 "rescued as duplicate: unique claim held by a newer job",
13157 )
13158 .await?;
13159 rescued.push(rescued_row);
13160 }
13161 tx.commit().await.map_err(map_sqlx_error)?;
13162 Ok(rescued)
13163 }
13164
13165 #[tracing::instrument(skip(self, pool), name = "queue_storage.rescue_expired_callbacks")]
13166 pub async fn rescue_expired_callbacks(&self, pool: &PgPool) -> Result<Vec<JobRow>, AwaError> {
13167 let schema = self.schema();
13168 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
13169 let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
13170 r#"
13171 DELETE FROM {schema}.leases
13172 WHERE job_id IN (
13173 SELECT job_id
13174 FROM {schema}.leases
13175 WHERE state = 'waiting_external'
13176 AND callback_timeout_at IS NOT NULL
13177 AND callback_timeout_at < clock_timestamp()
13178 ORDER BY callback_timeout_at ASC
13179 LIMIT 500
13180 FOR UPDATE SKIP LOCKED
13181 )
13182 RETURNING
13183 ready_slot,
13184 ready_generation,
13185 job_id,
13186 queue,
13187 state,
13188 priority,
13189 attempt,
13190 run_lease,
13191 max_attempts,
13192 lane_seq,
13193 enqueue_shard,
13194 heartbeat_at,
13195 deadline_at,
13196 attempted_at,
13197 callback_id,
13198 callback_timeout_at
13199 "#
13200 ))
13201 .fetch_all(tx.as_mut())
13202 .await
13203 .map_err(map_sqlx_error)?;
13204
13205 if deleted.is_empty() {
13206 tx.commit().await.map_err(map_sqlx_error)?;
13207 return Ok(Vec::new());
13208 }
13209
13210 let moved = self.hydrate_deleted_leases_tx(&mut tx, deleted).await?;
13211
13212 let mut rescued = Vec::with_capacity(moved.len());
13213 for row in moved {
13214 let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state(
13215 row.payload.clone(),
13216 row.progress.clone(),
13217 )?)?;
13218 let exhausted = row.attempt >= row.max_attempts;
13219 payload.push_error(lifecycle_error(
13220 "callback timed out",
13221 row.attempt,
13222 exhausted,
13223 ));
13224 if exhausted {
13225 let done =
13226 row.clone()
13227 .into_done_row(JobState::Failed, Utc::now(), payload.into_json());
13228 let rescued_row = self
13229 .insert_rescued_done_or_cancel_tx(
13230 &mut tx,
13231 row,
13232 done,
13233 "rescued as duplicate: unique claim held by a newer job",
13234 )
13235 .await?;
13236 rescued.push(rescued_row);
13237 } else {
13238 let deferred = row.clone().into_deferred_row(
13239 JobState::Retryable,
13240 self.backoff_at_tx(&mut tx, row.attempt, row.max_attempts)
13241 .await?,
13242 Some(Utc::now()),
13243 payload.into_json(),
13244 );
13245 let rescued_row = self
13246 .insert_rescued_deferred_or_cancel_tx(
13247 &mut tx,
13248 row,
13249 deferred,
13250 "rescued as duplicate: unique claim held by a newer job",
13251 )
13252 .await?;
13253 rescued.push(rescued_row);
13254 }
13255 }
13256 tx.commit().await.map_err(map_sqlx_error)?;
13257 Ok(rescued)
13258 }
13259
13260 pub async fn promote_due(
13261 &self,
13262 pool: &PgPool,
13263 state: JobState,
13264 batch_size: i64,
13265 ) -> Result<usize, AwaError> {
13266 if !matches!(state, JobState::Scheduled | JobState::Retryable) || batch_size <= 0 {
13267 return Ok(0);
13268 }
13269
13270 let schema = self.schema();
13271 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
13272 let moved: Vec<DeferredJobRow> = sqlx::query_as(&format!(
13273 r#"
13274 DELETE FROM {schema}.deferred_jobs
13275 WHERE job_id IN (
13276 SELECT job_id
13277 FROM {schema}.deferred_jobs
13278 WHERE state = $1
13279 AND run_at <= clock_timestamp()
13280 AND NOT EXISTS (
13281 SELECT 1 FROM awa.queue_meta
13282 WHERE queue = {schema}.deferred_jobs.queue AND paused = TRUE
13283 )
13284 ORDER BY run_at ASC, priority ASC, job_id ASC
13285 LIMIT $2
13286 FOR UPDATE SKIP LOCKED
13287 )
13288 RETURNING
13289 job_id,
13290 kind,
13291 queue,
13292 args,
13293 state,
13294 priority,
13295 attempt,
13296 run_lease,
13297 max_attempts,
13298 run_at,
13299 attempted_at,
13300 finalized_at,
13301 created_at,
13302 unique_key,
13303 unique_states,
13304 COALESCE(payload, '{{}}'::jsonb) AS payload
13305 "#
13306 ))
13307 .bind(state)
13308 .bind(batch_size)
13309 .fetch_all(tx.as_mut())
13310 .await
13311 .map_err(map_sqlx_error)?;
13312
13313 if moved.is_empty() {
13314 tx.commit().await.map_err(map_sqlx_error)?;
13315 return Ok(0);
13316 }
13317
13318 let ready_rows: Vec<ExistingReadyRow> = moved
13319 .iter()
13320 .cloned()
13321 .map(|row| ExistingReadyRow {
13322 job_id: row.job_id,
13323 kind: row.kind,
13324 queue: row.queue,
13325 args: row.args,
13326 priority: row.priority,
13327 attempt: row.attempt,
13328 run_lease: row.run_lease,
13329 max_attempts: row.max_attempts,
13330 run_at: Utc::now(),
13331 attempted_at: row.attempted_at,
13332 created_at: row.created_at,
13333 unique_key: row.unique_key,
13334 unique_states: row.unique_states,
13335 payload: row.payload,
13336 })
13337 .collect();
13338 let queues = ready_rows
13339 .iter()
13340 .map(|row| row.queue.clone())
13341 .collect::<Vec<_>>();
13342 self.insert_existing_ready_rows_tx(&mut tx, ready_rows, Some(state))
13343 .await?;
13344 self.notify_queues_tx(&mut tx, queues).await?;
13345 tx.commit().await.map_err(map_sqlx_error)?;
13346 Ok(moved.len())
13347 }
13348
13349 async fn relation_has_rows_tx(
13350 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
13351 relation: &str,
13352 ) -> Result<bool, AwaError> {
13353 sqlx::query_scalar(&format!("SELECT EXISTS (SELECT 1 FROM {relation} LIMIT 1)"))
13354 .fetch_one(tx.as_mut())
13355 .await
13356 .map_err(map_sqlx_error)
13357 }
13358
13359 #[tracing::instrument(skip(self, pool), name = "queue_storage.rotate")]
13360 pub async fn rotate(&self, pool: &PgPool) -> Result<RotateOutcome, AwaError> {
13361 let schema = self.schema();
13362 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
13363
13364 let state: (i32, i64, i32) = sqlx::query_as(&format!(
13365 r#"
13366 SELECT current_slot, generation, slot_count
13367 FROM {schema}.queue_ring_state
13368 WHERE singleton = TRUE
13369 FOR UPDATE
13370 "#
13371 ))
13372 .fetch_one(tx.as_mut())
13373 .await
13374 .map_err(map_sqlx_error)?;
13375
13376 let next_slot = (state.0 + 1).rem_euclid(state.2);
13377 let ready_busy =
13378 Self::relation_has_rows_tx(&mut tx, &ready_child_name(schema, next_slot as usize))
13379 .await?;
13380 let claim_attempt_batch_busy = Self::relation_has_rows_tx(
13381 &mut tx,
13382 &ready_claim_attempt_batch_child_name(schema, next_slot as usize),
13383 )
13384 .await?;
13385 let done_busy =
13386 Self::relation_has_rows_tx(&mut tx, &done_child_name(schema, next_slot as usize))
13387 .await?;
13388 let tombstone_busy = Self::relation_has_rows_tx(
13389 &mut tx,
13390 &ready_tombstone_child_name(schema, next_slot as usize),
13391 )
13392 .await?;
13393 let segment_busy = Self::relation_has_rows_tx(
13394 &mut tx,
13395 &ready_segment_child_name(schema, next_slot as usize),
13396 )
13397 .await?;
13398 let receipt_batch_busy = Self::relation_has_rows_tx(
13399 &mut tx,
13400 &receipt_completion_batch_child_name(schema, next_slot as usize),
13401 )
13402 .await?;
13403 let receipt_tombstone_busy = Self::relation_has_rows_tx(
13404 &mut tx,
13405 &receipt_completion_tombstone_child_name(schema, next_slot as usize),
13406 )
13407 .await?;
13408 let terminal_delta_busy = Self::relation_has_rows_tx(
13409 &mut tx,
13410 &terminal_delta_child_name(schema, next_slot as usize),
13411 )
13412 .await?;
13413
13414 if ready_busy
13415 || claim_attempt_batch_busy
13416 || done_busy
13417 || tombstone_busy
13418 || segment_busy
13419 || receipt_batch_busy
13420 || receipt_tombstone_busy
13421 || terminal_delta_busy
13422 {
13423 tx.commit().await.map_err(map_sqlx_error)?;
13424 return Ok(RotateOutcome::SkippedBusy {
13425 slot: next_slot,
13426 busy: BusyCounts {
13427 queue_ready: busy_indicator(ready_busy),
13428 queue_claim_attempt_batches: busy_indicator(claim_attempt_batch_busy),
13429 queue_done: busy_indicator(done_busy),
13430 queue_tombstones: busy_indicator(tombstone_busy),
13431 queue_ready_segments: busy_indicator(segment_busy),
13432 queue_receipt_completion_batches: busy_indicator(receipt_batch_busy),
13433 queue_receipt_completion_tombstones: busy_indicator(receipt_tombstone_busy),
13434 queue_terminal_deltas: busy_indicator(terminal_delta_busy),
13435 ..Default::default()
13436 },
13437 });
13438 }
13439
13440 let next_generation = state.1 + 1;
13441
13442 sqlx::query(&format!(
13443 r#"
13444 UPDATE {schema}.queue_ring_state
13445 SET current_slot = $1,
13446 generation = $2
13447 WHERE singleton = TRUE
13448 "#
13449 ))
13450 .bind(next_slot)
13451 .bind(next_generation)
13452 .execute(tx.as_mut())
13453 .await
13454 .map_err(map_sqlx_error)?;
13455
13456 sqlx::query(&format!(
13457 r#"
13458 UPDATE {schema}.queue_ring_slots
13459 SET generation = $2
13460 WHERE slot = $1
13461 "#
13462 ))
13463 .bind(next_slot)
13464 .bind(next_generation)
13465 .execute(tx.as_mut())
13466 .await
13467 .map_err(map_sqlx_error)?;
13468
13469 tx.commit().await.map_err(map_sqlx_error)?;
13470 Ok(RotateOutcome::Rotated {
13471 slot: next_slot,
13472 generation: next_generation,
13473 })
13474 }
13475
13476 #[tracing::instrument(skip(self, pool), name = "queue_storage.rotate_leases")]
13477 pub async fn rotate_leases(&self, pool: &PgPool) -> Result<RotateOutcome, AwaError> {
13478 let schema = self.schema();
13479 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
13480
13481 let state: (i32, i64, i32) = sqlx::query_as(&format!(
13488 r#"
13489 SELECT current_slot, generation, slot_count
13490 FROM {schema}.lease_ring_state
13491 WHERE singleton = TRUE
13492 FOR UPDATE
13493 "#
13494 ))
13495 .fetch_one(tx.as_mut())
13496 .await
13497 .map_err(map_sqlx_error)?;
13498
13499 let next_slot = (state.0 + 1).rem_euclid(state.2);
13500 let lease_busy =
13501 Self::relation_has_rows_tx(&mut tx, &lease_child_name(schema, next_slot as usize))
13502 .await?;
13503
13504 if lease_busy {
13505 tx.commit().await.map_err(map_sqlx_error)?;
13506 return Ok(RotateOutcome::SkippedBusy {
13507 slot: next_slot,
13508 busy: BusyCounts {
13509 leases: busy_indicator(lease_busy),
13510 ..Default::default()
13511 },
13512 });
13513 }
13514
13515 let next_generation = state.1 + 1;
13516
13517 let rotated = sqlx::query(&format!(
13518 r#"
13519 UPDATE {schema}.lease_ring_state
13520 SET current_slot = $1,
13521 generation = $2
13522 WHERE singleton = TRUE
13523 AND current_slot = $3
13524 AND generation = $4
13525 "#
13526 ))
13527 .bind(next_slot)
13528 .bind(next_generation)
13529 .bind(state.0)
13530 .bind(state.1)
13531 .execute(tx.as_mut())
13532 .await
13533 .map_err(map_sqlx_error)?;
13534
13535 if rotated.rows_affected() == 0 {
13536 tx.commit().await.map_err(map_sqlx_error)?;
13539 return Ok(RotateOutcome::SkippedBusy {
13540 slot: next_slot,
13541 busy: BusyCounts {
13542 leases: busy_indicator(lease_busy),
13543 ..Default::default()
13544 },
13545 });
13546 }
13547
13548 tx.commit().await.map_err(map_sqlx_error)?;
13549 Ok(RotateOutcome::Rotated {
13550 slot: next_slot,
13551 generation: next_generation,
13552 })
13553 }
13554
13555 #[tracing::instrument(skip(self, pool), name = "queue_storage.rebuild_terminal_counters")]
13579 pub async fn rebuild_terminal_counters(&self, pool: &PgPool) -> Result<i64, AwaError> {
13580 let schema = self.schema();
13581 let rebuild_lock_name = format!("awa.queue_storage.rebuild_terminal_counters:{schema}");
13582 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
13583
13584 sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))")
13585 .bind(&rebuild_lock_name)
13586 .execute(tx.as_mut())
13587 .await
13588 .map_err(map_sqlx_error)?;
13589
13590 sqlx::query(&format!(
13591 "LOCK TABLE {schema}.done_entries, \
13592 {schema}.receipt_completion_batches, \
13593 {schema}.receipt_completion_tombstones, \
13594 {schema}.queue_terminal_count_deltas, \
13595 {schema}.queue_terminal_live_counts \
13596 IN ACCESS EXCLUSIVE MODE"
13597 ))
13598 .execute(tx.as_mut())
13599 .await
13600 .map_err(map_sqlx_error)?;
13601
13602 sqlx::query(&format!(
13603 "TRUNCATE TABLE {schema}.queue_terminal_live_counts, {schema}.queue_terminal_count_deltas"
13604 ))
13605 .execute(tx.as_mut())
13606 .await
13607 .map_err(map_sqlx_error)?;
13608
13609 let inserted: i64 = sqlx::query_scalar(&format!(
13610 r#"
13611 WITH inserted AS (
13612 INSERT INTO {schema}.queue_terminal_live_counts AS counts (
13613 ready_slot, queue, priority, enqueue_shard, counter_bucket, live_terminal_count
13614 )
13615 SELECT
13616 ready_slot,
13617 queue,
13618 priority,
13619 enqueue_shard,
13620 mod(
13621 mod(job_id, {TERMINAL_COUNTER_BUCKETS}::bigint)
13622 + {TERMINAL_COUNTER_BUCKETS}::bigint,
13623 {TERMINAL_COUNTER_BUCKETS}::bigint
13624 )::smallint AS counter_bucket,
13625 count(*)::bigint
13626 FROM {schema}.done_entries
13627 GROUP BY ready_slot, queue, priority, enqueue_shard, counter_bucket
13628 RETURNING 1
13629 )
13630 SELECT COALESCE(count(*), 0)::bigint FROM inserted
13631 "#
13632 ))
13633 .fetch_one(tx.as_mut())
13634 .await
13635 .map_err(map_sqlx_error)?;
13636
13637 sqlx::query(&format!(
13641 r#"
13642 UPDATE {schema}.queue_ring_state
13643 SET terminal_counter_trusted_at = now()
13644 WHERE singleton = TRUE
13645 "#
13646 ))
13647 .execute(tx.as_mut())
13648 .await
13649 .map_err(map_sqlx_error)?;
13650
13651 tx.commit().await.map_err(map_sqlx_error)?;
13652 Ok(inserted)
13653 }
13654
13655 pub async fn terminal_counter_trusted(&self, pool: &PgPool) -> Result<bool, AwaError> {
13667 let schema = self.schema();
13668 let trusted: Option<bool> = sqlx::query_scalar(&format!(
13669 "SELECT terminal_counter_trusted_at IS NOT NULL \
13670 FROM {schema}.queue_ring_state WHERE singleton = TRUE"
13671 ))
13672 .fetch_optional(pool)
13673 .await
13674 .map_err(map_sqlx_error)?;
13675 Ok(trusted.unwrap_or(false))
13679 }
13680
13681 #[tracing::instrument(skip(self, pool), name = "queue_storage.rollup_terminal_count_deltas")]
13697 pub async fn rollup_terminal_count_deltas(
13698 &self,
13699 pool: &PgPool,
13700 max_slots: usize,
13701 ) -> Result<TerminalDeltaRollupOutcome, AwaError> {
13702 if max_slots == 0 {
13703 return Ok(TerminalDeltaRollupOutcome::default());
13704 }
13705
13706 let schema = self.schema();
13707 let current_slot: i32 = sqlx::query_scalar(&format!(
13708 r#"
13709 SELECT current_slot
13710 FROM {schema}.queue_ring_state
13711 WHERE singleton = TRUE
13712 "#
13713 ))
13714 .fetch_one(pool)
13715 .await
13716 .map_err(map_sqlx_error)?;
13717
13718 let slots = self
13719 .terminal_delta_rollup_candidates(pool, current_slot)
13720 .await?;
13721
13722 let mut outcome = TerminalDeltaRollupOutcome::default();
13723 for (slot, generation) in slots {
13724 if outcome.rolled_slots >= max_slots {
13725 break;
13726 }
13727 match self
13728 .rollup_terminal_count_delta_slot(pool, slot, generation)
13729 .await?
13730 {
13731 TerminalDeltaSlotRollup::Empty => {}
13732 TerminalDeltaSlotRollup::Rolled {
13733 delta_rows,
13734 grouped_keys,
13735 } => {
13736 outcome.rolled_slots += 1;
13737 outcome.delta_rows += delta_rows;
13738 outcome.grouped_keys += grouped_keys;
13739 }
13740 TerminalDeltaSlotRollup::SkippedActive => {
13741 outcome.skipped_active_slots += 1;
13742 }
13743 TerminalDeltaSlotRollup::SkippedMvccPinned => {
13744 outcome.skipped_mvcc_pinned = true;
13745 break;
13746 }
13747 TerminalDeltaSlotRollup::Blocked => {
13748 outcome.blocked_slots += 1;
13749 }
13750 }
13751 }
13752
13753 Ok(outcome)
13754 }
13755
13756 async fn terminal_delta_rollup_mvcc_horizon_pinned_tx(
13757 tx: &mut sqlx::Transaction<'_, Postgres>,
13758 ) -> Result<bool, AwaError> {
13759 let pinned: bool = sqlx::query_scalar(
13760 r#"
13761 SELECT EXISTS (
13762 SELECT 1
13763 FROM pg_stat_activity
13764 WHERE datname = current_database()
13765 AND pid <> pg_backend_pid()
13766 AND backend_type = 'client backend'
13767 AND (
13768 backend_xmin IS NOT NULL
13769 OR (
13770 backend_xid IS NOT NULL
13771 AND state LIKE 'idle in transaction%'
13772 )
13773 )
13774 )
13775 "#,
13776 )
13777 .fetch_one(tx.as_mut())
13778 .await
13779 .map_err(map_sqlx_error)?;
13780 Ok(pinned)
13781 }
13782
13783 async fn terminal_delta_rollup_candidates(
13784 &self,
13785 pool: &PgPool,
13786 current_slot: i32,
13787 ) -> Result<Vec<(i32, i64)>, AwaError> {
13788 let schema = self.schema();
13789 let sealed_slots: Vec<(i32, i64)> = sqlx::query_as(&format!(
13790 r#"
13791 SELECT slot, generation
13792 FROM {schema}.queue_ring_slots
13793 WHERE generation >= 0
13794 AND slot <> $1
13795 ORDER BY generation ASC, slot ASC
13796 "#
13797 ))
13798 .bind(current_slot)
13799 .fetch_all(pool)
13800 .await
13801 .map_err(map_sqlx_error)?;
13802
13803 let mut pending_slots = Vec::new();
13804 for (slot, generation) in sealed_slots {
13805 let Ok(slot_index) = usize::try_from(slot) else {
13806 continue;
13807 };
13808 let delta_child = terminal_delta_child_name(schema, slot_index);
13809 let has_pending: bool = sqlx::query_scalar(&format!(
13810 r#"
13811 SELECT EXISTS (
13812 SELECT 1
13813 FROM {delta_child}
13814 WHERE ready_generation = $1
13815 LIMIT 1
13816 )
13817 "#
13818 ))
13819 .bind(generation)
13820 .fetch_one(pool)
13821 .await
13822 .map_err(map_sqlx_error)?;
13823 if has_pending {
13824 pending_slots.push((slot, generation));
13825 }
13826 }
13827
13828 Ok(pending_slots)
13829 }
13830
13831 async fn rollup_terminal_count_delta_slot(
13832 &self,
13833 pool: &PgPool,
13834 slot: i32,
13835 generation: i64,
13836 ) -> Result<TerminalDeltaSlotRollup, AwaError> {
13837 let schema = self.schema();
13838 let delta_child = terminal_delta_child_name(schema, slot as usize);
13839 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
13840
13841 let current_slot: i32 = sqlx::query_scalar(&format!(
13842 r#"
13843 SELECT current_slot
13844 FROM {schema}.queue_ring_state
13845 WHERE singleton = TRUE
13846 FOR UPDATE
13847 "#
13848 ))
13849 .fetch_one(tx.as_mut())
13850 .await
13851 .map_err(map_sqlx_error)?;
13852
13853 let slot_generation: Option<i64> = sqlx::query_scalar(&format!(
13854 r#"
13855 SELECT generation
13856 FROM {schema}.queue_ring_slots
13857 WHERE slot = $1
13858 FOR UPDATE
13859 "#
13860 ))
13861 .bind(slot)
13862 .fetch_optional(tx.as_mut())
13863 .await
13864 .map_err(map_sqlx_error)?;
13865
13866 let Some(slot_generation) = slot_generation else {
13867 tx.commit().await.map_err(map_sqlx_error)?;
13868 return Ok(TerminalDeltaSlotRollup::Empty);
13869 };
13870
13871 if current_slot == slot {
13872 tx.commit().await.map_err(map_sqlx_error)?;
13873 return Ok(TerminalDeltaSlotRollup::SkippedActive);
13874 }
13875
13876 if slot_generation != generation {
13877 tx.commit().await.map_err(map_sqlx_error)?;
13878 return Ok(TerminalDeltaSlotRollup::Empty);
13879 }
13880
13881 if Self::terminal_delta_rollup_mvcc_horizon_pinned_tx(&mut tx).await? {
13882 tx.commit().await.map_err(map_sqlx_error)?;
13883 return Ok(TerminalDeltaSlotRollup::SkippedMvccPinned);
13884 }
13885
13886 let ready_child = ready_child_name(schema, slot as usize);
13887 let pending_ready: bool = sqlx::query_scalar(&format!(
13888 r#"
13889 WITH claim_cursors AS MATERIALIZED (
13890 SELECT
13891 queue,
13892 priority,
13893 enqueue_shard,
13894 {schema}.sequence_next_value(seq_name) AS claim_seq
13895 FROM {schema}.queue_claim_heads
13896 )
13897 SELECT EXISTS (
13898 SELECT 1
13899 FROM claim_cursors AS claims
13900 CROSS JOIN LATERAL (
13901 SELECT 1
13902 FROM {ready_child} AS ready
13903 WHERE ready.ready_generation = $1
13904 AND ready.queue = claims.queue
13905 AND ready.priority = claims.priority
13906 AND ready.enqueue_shard = claims.enqueue_shard
13907 AND ready.lane_seq >= claims.claim_seq
13908 LIMIT 1
13909 ) AS pending_ready
13910 LIMIT 1
13911 )
13912 "#
13913 ))
13914 .bind(generation)
13915 .fetch_one(tx.as_mut())
13916 .await
13917 .map_err(map_sqlx_error)?;
13918
13919 if pending_ready {
13920 tx.commit().await.map_err(map_sqlx_error)?;
13921 return Ok(TerminalDeltaSlotRollup::SkippedActive);
13922 }
13923
13924 set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?;
13925
13926 let lock_delta = sqlx::query(&format!(
13927 "LOCK TABLE {delta_child} IN ACCESS EXCLUSIVE MODE"
13928 ))
13929 .execute(tx.as_mut())
13930 .await;
13931
13932 match lock_delta {
13933 Ok(_) => {}
13934 Err(err) if is_lock_contention_error(&err) => {
13935 let _ = tx.rollback().await;
13936 return Ok(TerminalDeltaSlotRollup::Blocked);
13937 }
13938 Err(err) => {
13939 let _ = tx.rollback().await;
13940 return Err(map_sqlx_error(err));
13941 }
13942 }
13943
13944 let active_leases =
13945 queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?;
13946
13947 if active_leases {
13948 tx.commit().await.map_err(map_sqlx_error)?;
13949 return Ok(TerminalDeltaSlotRollup::SkippedActive);
13950 }
13951
13952 let unclosed_claim_refs =
13953 queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?;
13954
13955 if unclosed_claim_refs {
13956 tx.commit().await.map_err(map_sqlx_error)?;
13957 return Ok(TerminalDeltaSlotRollup::SkippedActive);
13958 }
13959
13960 let delta_rows: i64 = sqlx::query_scalar(&format!(
13961 r#"
13962 SELECT count(*)::bigint
13963 FROM {delta_child}
13964 WHERE ready_generation = $1
13965 "#
13966 ))
13967 .bind(generation)
13968 .fetch_one(tx.as_mut())
13969 .await
13970 .map_err(map_sqlx_error)?;
13971
13972 if delta_rows == 0 {
13973 tx.commit().await.map_err(map_sqlx_error)?;
13974 return Ok(TerminalDeltaSlotRollup::Empty);
13975 }
13976
13977 let grouped_keys: i64 = sqlx::query_scalar(&format!(
13978 r#"
13979 WITH grouped AS MATERIALIZED (
13980 SELECT
13981 ready_slot,
13982 queue,
13983 priority,
13984 enqueue_shard,
13985 counter_bucket,
13986 SUM(terminal_delta)::bigint AS delta
13987 FROM {delta_child}
13988 WHERE ready_generation = $1
13989 GROUP BY ready_slot, queue, priority, enqueue_shard, counter_bucket
13990 HAVING SUM(terminal_delta) <> 0
13991 ),
13992 updated AS (
13993 UPDATE {schema}.queue_terminal_live_counts AS counts
13994 SET live_terminal_count = GREATEST(0, counts.live_terminal_count + grouped.delta)
13995 FROM grouped
13996 WHERE counts.ready_slot = grouped.ready_slot
13997 AND counts.queue = grouped.queue
13998 AND counts.priority = grouped.priority
13999 AND counts.enqueue_shard = grouped.enqueue_shard
14000 AND counts.counter_bucket = grouped.counter_bucket
14001 RETURNING
14002 counts.ready_slot,
14003 counts.queue,
14004 counts.priority,
14005 counts.enqueue_shard,
14006 counts.counter_bucket
14007 ),
14008 inserted AS (
14009 INSERT INTO {schema}.queue_terminal_live_counts AS counts (
14010 ready_slot,
14011 queue,
14012 priority,
14013 enqueue_shard,
14014 counter_bucket,
14015 live_terminal_count
14016 )
14017 SELECT
14018 grouped.ready_slot,
14019 grouped.queue,
14020 grouped.priority,
14021 grouped.enqueue_shard,
14022 grouped.counter_bucket,
14023 grouped.delta
14024 FROM grouped
14025 WHERE grouped.delta > 0
14026 AND NOT EXISTS (
14027 SELECT 1
14028 FROM updated
14029 WHERE updated.ready_slot = grouped.ready_slot
14030 AND updated.queue = grouped.queue
14031 AND updated.priority = grouped.priority
14032 AND updated.enqueue_shard = grouped.enqueue_shard
14033 AND updated.counter_bucket = grouped.counter_bucket
14034 )
14035 ORDER BY ready_slot, queue, priority, enqueue_shard, counter_bucket
14036 ON CONFLICT (ready_slot, queue, priority, enqueue_shard, counter_bucket) DO UPDATE
14037 SET live_terminal_count =
14038 GREATEST(0, counts.live_terminal_count + EXCLUDED.live_terminal_count)
14039 RETURNING 1
14040 )
14041 SELECT count(*)::bigint FROM grouped
14042 "#
14043 ))
14044 .bind(generation)
14045 .fetch_one(tx.as_mut())
14046 .await
14047 .map_err(map_sqlx_error)?;
14048
14049 let truncate_delta = sqlx::query(&format!("TRUNCATE TABLE {delta_child}"))
14050 .execute(tx.as_mut())
14051 .await;
14052
14053 match truncate_delta {
14054 Ok(_) => {
14055 tx.commit().await.map_err(map_sqlx_error)?;
14056 Ok(TerminalDeltaSlotRollup::Rolled {
14057 delta_rows,
14058 grouped_keys,
14059 })
14060 }
14061 Err(err) if is_lock_contention_error(&err) => {
14062 let _ = tx.rollback().await;
14063 Ok(TerminalDeltaSlotRollup::Blocked)
14064 }
14065 Err(err) => {
14066 let _ = tx.rollback().await;
14067 Err(map_sqlx_error(err))
14068 }
14069 }
14070 }
14071
14072 #[tracing::instrument(skip(self, pool), name = "queue_storage.prune_oldest")]
14082 pub async fn prune_oldest(
14083 &self,
14084 pool: &PgPool,
14085 failed_retention: Duration,
14086 ) -> Result<PruneOutcome, AwaError> {
14087 let schema = self.schema();
14088 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
14089
14090 let state: (i32,) = sqlx::query_as(&format!(
14091 r#"
14092 SELECT current_slot
14093 FROM {schema}.queue_ring_state
14094 WHERE singleton = TRUE
14095 FOR UPDATE
14096 "#
14097 ))
14098 .fetch_one(tx.as_mut())
14099 .await
14100 .map_err(map_sqlx_error)?;
14101
14102 let target: Option<(i32, i64)> = sqlx::query_as(&format!(
14103 r#"
14104 SELECT slot, generation
14105 FROM {schema}.queue_ring_slots
14106 WHERE generation >= 0
14107 AND slot <> $1
14108 ORDER BY generation ASC, slot ASC
14109 LIMIT 1
14110 FOR UPDATE
14111 "#
14112 ))
14113 .bind(state.0)
14114 .fetch_optional(tx.as_mut())
14115 .await
14116 .map_err(map_sqlx_error)?;
14117
14118 let Some((slot, generation)) = target else {
14119 tx.commit().await.map_err(map_sqlx_error)?;
14120 return Ok(PruneOutcome::Noop);
14121 };
14122
14123 let ready_child = ready_child_name(schema, slot as usize);
14124 let claim_attempt_child = ready_claim_attempt_batch_child_name(schema, slot as usize);
14125 let done_child = done_child_name(schema, slot as usize);
14126 let tomb_child = ready_tombstone_child_name(schema, slot as usize);
14127 let segment_child = ready_segment_child_name(schema, slot as usize);
14128 let receipt_batch_child = receipt_completion_batch_child_name(schema, slot as usize);
14129 let receipt_tomb_child = receipt_completion_tombstone_child_name(schema, slot as usize);
14130 let delta_child = terminal_delta_child_name(schema, slot as usize);
14131
14132 let active_leases =
14138 queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?;
14139
14140 if active_leases {
14141 tx.commit().await.map_err(map_sqlx_error)?;
14142 return Ok(PruneOutcome::SkippedActive {
14143 slot,
14144 reason: SkipReason::QueueActiveLeases,
14145 count: busy_indicator(active_leases),
14146 });
14147 }
14148
14149 let pending =
14150 queue_prune_has_pending_ready_tx(&mut tx, schema, &ready_child, generation).await?;
14151
14152 if pending {
14153 tx.commit().await.map_err(map_sqlx_error)?;
14154 return Ok(PruneOutcome::SkippedActive {
14155 slot,
14156 reason: SkipReason::QueuePendingReady,
14157 count: busy_indicator(pending),
14158 });
14159 }
14160
14161 let unclosed_claim_refs =
14162 queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?;
14163
14164 if unclosed_claim_refs {
14165 tx.commit().await.map_err(map_sqlx_error)?;
14166 return Ok(PruneOutcome::SkippedActive {
14167 slot,
14168 reason: SkipReason::QueueUnclosedClaimRefs,
14169 count: busy_indicator(unclosed_claim_refs),
14170 });
14171 }
14172
14173 set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?;
14174
14175 let lock_tables = sqlx::query(&format!(
14176 "LOCK TABLE {ready_child}, {claim_attempt_child}, {done_child}, {tomb_child}, {segment_child}, {receipt_batch_child}, {receipt_tomb_child}, {delta_child} IN ACCESS EXCLUSIVE MODE"
14177 ))
14178 .execute(tx.as_mut())
14179 .await;
14180
14181 if let Err(err) = lock_tables {
14182 let _ = tx.rollback().await;
14183 if is_lock_contention_error(&err) {
14184 return Ok(PruneOutcome::Blocked { slot });
14185 }
14186 return Err(map_sqlx_error(err));
14187 }
14188
14189 let active_leases =
14193 queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?;
14194 if active_leases {
14195 tx.commit().await.map_err(map_sqlx_error)?;
14196 return Ok(PruneOutcome::SkippedActive {
14197 slot,
14198 reason: SkipReason::QueueActiveLeases,
14199 count: busy_indicator(active_leases),
14200 });
14201 }
14202
14203 let pending =
14204 queue_prune_has_pending_ready_tx(&mut tx, schema, &ready_child, generation).await?;
14205 if pending {
14206 tx.commit().await.map_err(map_sqlx_error)?;
14207 return Ok(PruneOutcome::SkippedActive {
14208 slot,
14209 reason: SkipReason::QueuePendingReady,
14210 count: busy_indicator(pending),
14211 });
14212 }
14213
14214 let unclosed_claim_refs =
14215 queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?;
14216 if unclosed_claim_refs {
14217 tx.commit().await.map_err(map_sqlx_error)?;
14218 return Ok(PruneOutcome::SkippedActive {
14219 slot,
14220 reason: SkipReason::QueueUnclosedClaimRefs,
14221 count: busy_indicator(unclosed_claim_refs),
14222 });
14223 }
14224
14225 let failed_retention_secs = i64::try_from(failed_retention.as_secs()).unwrap_or(i64::MAX);
14226 let retention_floor = failed_retention_secs > 0;
14227
14228 let carried_failed_rows = if retention_floor {
14234 let (current_generation,): (i64,) = sqlx::query_as(&format!(
14235 "SELECT generation FROM {schema}.queue_ring_slots WHERE slot = $1"
14236 ))
14237 .bind(state.0)
14238 .fetch_one(tx.as_mut())
14239 .await
14240 .map_err(map_sqlx_error)?;
14241
14242 let carried = sqlx::query(&format!(
14251 r#"
14252 INSERT INTO {schema}.done_entries (
14253 ready_slot, ready_generation, job_id, kind, queue,
14254 args, state, priority, attempt, run_lease,
14255 max_attempts, lane_seq, enqueue_shard, run_at,
14256 attempted_at, finalized_at, created_at, unique_key,
14257 unique_states, payload
14258 )
14259 SELECT
14260 $2,
14261 $3,
14262 done.job_id,
14263 done.kind,
14264 done.queue,
14265 COALESCE(done.args, ready.args, '{{}}'::jsonb),
14266 done.state,
14267 done.priority,
14268 done.attempt,
14269 done.run_lease,
14270 COALESCE(done.max_attempts, ready.max_attempts, 25::smallint),
14271 done.lane_seq,
14272 done.enqueue_shard,
14273 COALESCE(done.run_at, ready.run_at, done.finalized_at),
14274 COALESCE(done.attempted_at, ready.attempted_at),
14275 done.finalized_at,
14276 COALESCE(done.created_at, ready.created_at, done.finalized_at),
14277 COALESCE(done.unique_key, ready.unique_key),
14278 COALESCE(done.unique_states, ready.unique_states),
14279 COALESCE(done.payload, ready.payload, '{{}}'::jsonb)
14280 FROM {done_child} AS done
14281 LEFT JOIN {ready_child} AS ready
14282 ON ready.ready_slot = done.ready_slot
14283 AND ready.ready_generation = done.ready_generation
14284 AND ready.queue = done.queue
14285 AND ready.priority = done.priority
14286 AND ready.enqueue_shard = done.enqueue_shard
14287 AND ready.lane_seq = done.lane_seq
14288 WHERE done.state = 'failed'
14289 AND done.finalized_at >= now() - make_interval(secs => $1::bigint)
14290 "#
14291 ))
14292 .bind(failed_retention_secs)
14293 .bind(state.0)
14294 .bind(current_generation)
14295 .execute(tx.as_mut())
14296 .await
14297 .map_err(map_sqlx_error)?
14298 .rows_affected();
14299
14300 if carried > 0 {
14301 sqlx::query(&format!(
14307 r#"
14308 INSERT INTO {schema}.queue_terminal_count_deltas (
14309 ready_slot,
14310 ready_generation,
14311 queue,
14312 priority,
14313 enqueue_shard,
14314 counter_bucket,
14315 terminal_delta
14316 )
14317 SELECT
14318 $2,
14319 $3,
14320 queue,
14321 priority,
14322 enqueue_shard,
14323 mod(
14324 mod(job_id, {TERMINAL_COUNTER_BUCKETS}::bigint)
14325 + {TERMINAL_COUNTER_BUCKETS}::bigint,
14326 {TERMINAL_COUNTER_BUCKETS}::bigint
14327 )::smallint AS counter_bucket,
14328 count(*)::bigint
14329 FROM {done_child}
14330 WHERE state = 'failed'
14331 AND finalized_at >= now() - make_interval(secs => $1::bigint)
14332 GROUP BY queue, priority, enqueue_shard, counter_bucket
14333 ORDER BY queue, priority, enqueue_shard, counter_bucket
14334 "#
14335 ))
14336 .bind(failed_retention_secs)
14337 .bind(state.0)
14338 .bind(current_generation)
14339 .execute(tx.as_mut())
14340 .await
14341 .map_err(map_sqlx_error)?;
14342 }
14343 carried
14344 } else {
14345 0
14346 };
14347
14348 let pruned_terminal_counts: Vec<(String, i16, i64, i64)> = sqlx::query_as(&format!(
14366 r#"
14367 WITH done_counts AS (
14368 SELECT
14369 queue,
14370 priority,
14371 (count(*) FILTER (WHERE state <> 'failed'))::bigint AS completed,
14372 (count(*) FILTER (
14373 WHERE state = 'failed'
14374 AND (
14375 $2::boolean IS FALSE
14376 OR finalized_at < now() - make_interval(secs => $1::bigint)
14377 )
14378 ))::bigint AS failed
14379 FROM {done_child}
14380 GROUP BY queue, priority
14381 ),
14382 -- Scan the entire partition being truncated (all generations),
14383 -- mirroring done_counts' whole-child scan above. The rotate guard
14384 -- keeps a partition to one generation at prune time, so this is the
14385 -- same set today; counting the whole child keeps the rollup fold
14386 -- conservative even if that invariant is ever weakened (a stray
14387 -- generation's batches would otherwise be truncated without being
14388 -- folded into the rollup -> terminal-count undercount).
14389 batch_rows AS (
14390 SELECT
14391 batch.ready_generation,
14392 batch.queue,
14393 batch.priority,
14394 item.job_id,
14395 item.run_lease
14396 FROM {receipt_batch_child} AS batch
14397 CROSS JOIN LATERAL unnest(
14398 batch.job_ids,
14399 batch.run_leases
14400 ) AS item(job_id, run_lease)
14401 ),
14402 batch_counts AS (
14403 SELECT
14404 batch_rows.queue,
14405 batch_rows.priority,
14406 count(*)::bigint AS completed
14407 FROM batch_rows
14408 WHERE NOT EXISTS (
14409 SELECT 1
14410 FROM {receipt_tomb_child} AS tomb
14411 WHERE tomb.ready_generation = batch_rows.ready_generation
14412 AND tomb.job_id = batch_rows.job_id
14413 AND tomb.run_lease = batch_rows.run_lease
14414 )
14415 GROUP BY batch_rows.queue, batch_rows.priority
14416 ),
14417 keys AS (
14418 SELECT queue, priority FROM done_counts
14419 UNION
14420 SELECT queue, priority FROM batch_counts
14421 )
14422 SELECT
14423 keys.queue,
14424 keys.priority,
14425 (
14426 COALESCE(done_counts.completed, 0)
14427 + COALESCE(batch_counts.completed, 0)
14428 )::bigint AS pruned_completed_count,
14429 COALESCE(done_counts.failed, 0)::bigint AS pruned_failed_count
14430 FROM keys
14431 LEFT JOIN done_counts USING (queue, priority)
14432 LEFT JOIN batch_counts USING (queue, priority)
14433 "#
14434 ))
14435 .bind(failed_retention_secs)
14436 .bind(retention_floor)
14437 .fetch_all(tx.as_mut())
14438 .await
14439 .map_err(map_sqlx_error)?;
14440
14441 let truncate = sqlx::query(&format!(
14442 "TRUNCATE TABLE {ready_child}, {claim_attempt_child}, {done_child}, {tomb_child}, {segment_child}, {receipt_batch_child}, {receipt_tomb_child}, {delta_child}"
14443 ))
14444 .execute(tx.as_mut())
14445 .await;
14446
14447 match truncate {
14448 Ok(_) => {
14449 if !pruned_terminal_counts.is_empty() {
14450 self.adjust_terminal_rollups_batch(&mut tx, pruned_terminal_counts.into_iter())
14451 .await?;
14452 }
14453
14454 sqlx::query(&format!(
14462 "DELETE FROM {schema}.queue_terminal_live_counts WHERE ready_slot = $1"
14463 ))
14464 .bind(slot)
14465 .execute(tx.as_mut())
14466 .await
14467 .map_err(map_sqlx_error)?;
14468 tx.commit().await.map_err(map_sqlx_error)?;
14469 if carried_failed_rows > 0 {
14470 tracing::info!(
14471 slot,
14472 carried_failed_rows,
14473 "Carried failed terminal rows inside the retention floor to the live done segment"
14474 );
14475 }
14476 Ok(PruneOutcome::Pruned {
14477 slot,
14478 carried_failed_rows,
14479 })
14480 }
14481 Err(err) if is_lock_contention_error(&err) => {
14482 let _ = tx.rollback().await;
14483 Ok(PruneOutcome::Blocked { slot })
14484 }
14485 Err(err) => {
14486 let _ = tx.rollback().await;
14487 Err(map_sqlx_error(err))
14488 }
14489 }
14490 }
14491
14492 #[tracing::instrument(skip(self, pool), name = "queue_storage.prune_oldest_leases")]
14493 pub async fn prune_oldest_leases(&self, pool: &PgPool) -> Result<PruneOutcome, AwaError> {
14494 let schema = self.schema();
14495 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
14496
14497 let state: (i32, i64, i32) = sqlx::query_as(&format!(
14510 r#"
14511 SELECT current_slot, generation, slot_count
14512 FROM {schema}.lease_ring_state
14513 WHERE singleton = TRUE
14514 FOR UPDATE
14515 "#
14516 ))
14517 .fetch_one(tx.as_mut())
14518 .await
14519 .map_err(map_sqlx_error)?;
14520
14521 let Some((slot, _generation)) = oldest_initialized_ring_slot(state.0, state.1, state.2)
14522 else {
14523 tx.commit().await.map_err(map_sqlx_error)?;
14524 return Ok(PruneOutcome::Noop);
14525 };
14526
14527 let slot_locked: Option<i32> = sqlx::query_scalar(&format!(
14528 r#"
14529 SELECT slot FROM {schema}.lease_ring_slots
14530 WHERE slot = $1
14531 FOR UPDATE
14532 "#
14533 ))
14534 .bind(slot)
14535 .fetch_optional(tx.as_mut())
14536 .await
14537 .map_err(map_sqlx_error)?;
14538
14539 if slot_locked.is_none() {
14540 tx.commit().await.map_err(map_sqlx_error)?;
14541 return Ok(PruneOutcome::Noop);
14542 }
14543
14544 let lease_child = lease_child_name(schema, slot as usize);
14545
14546 set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?;
14547
14548 let lock_table = sqlx::query(&format!(
14549 "LOCK TABLE {lease_child} IN ACCESS EXCLUSIVE MODE"
14550 ))
14551 .execute(tx.as_mut())
14552 .await;
14553
14554 if let Err(err) = lock_table {
14555 let _ = tx.rollback().await;
14556 if is_lock_contention_error(&err) {
14557 return Ok(PruneOutcome::Blocked { slot });
14558 }
14559 return Err(map_sqlx_error(err));
14560 }
14561
14562 let current_slot: i32 = sqlx::query_scalar(&format!(
14563 r#"
14564 SELECT current_slot
14565 FROM {schema}.lease_ring_state
14566 WHERE singleton = TRUE
14567 "#
14568 ))
14569 .fetch_one(tx.as_mut())
14570 .await
14571 .map_err(map_sqlx_error)?;
14572
14573 if current_slot == slot {
14574 tx.commit().await.map_err(map_sqlx_error)?;
14575 return Ok(PruneOutcome::SkippedActive {
14576 slot,
14577 reason: SkipReason::LeaseCurrent,
14578 count: 0,
14579 });
14580 }
14581
14582 let active_leases = Self::relation_has_rows_tx(&mut tx, &lease_child).await?;
14583
14584 if active_leases {
14585 tx.commit().await.map_err(map_sqlx_error)?;
14586 return Ok(PruneOutcome::SkippedActive {
14587 slot,
14588 reason: SkipReason::LeaseActive,
14589 count: busy_indicator(active_leases),
14590 });
14591 }
14592
14593 let truncate = sqlx::query(&format!("TRUNCATE TABLE {lease_child}"))
14594 .execute(tx.as_mut())
14595 .await;
14596
14597 match truncate {
14598 Ok(_) => {
14599 tx.commit().await.map_err(map_sqlx_error)?;
14600 Ok(PruneOutcome::Pruned {
14601 slot,
14602 carried_failed_rows: 0,
14603 })
14604 }
14605 Err(err) if is_lock_contention_error(&err) => {
14606 let _ = tx.rollback().await;
14607 Ok(PruneOutcome::Blocked { slot })
14608 }
14609 Err(err) => {
14610 let _ = tx.rollback().await;
14611 Err(map_sqlx_error(err))
14612 }
14613 }
14614 }
14615
14616 pub async fn vacuum_leases(&self, pool: &PgPool) -> Result<(), AwaError> {
14617 sqlx::query(&format!("VACUUM {}", self.leases_table()))
14618 .execute(pool)
14619 .await
14620 .map_err(map_sqlx_error)?;
14621 Ok(())
14622 }
14623
14624 #[tracing::instrument(skip(self, pool), name = "queue_storage.rotate_claims")]
14635 pub async fn rotate_claims(&self, pool: &PgPool) -> Result<RotateOutcome, AwaError> {
14636 let schema = self.schema();
14637 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
14638
14639 let state: (i32, i64, i32) = sqlx::query_as(&format!(
14640 r#"
14641 SELECT current_slot, generation, slot_count
14642 FROM {schema}.claim_ring_state
14643 WHERE singleton = TRUE
14644 FOR UPDATE
14645 "#
14646 ))
14647 .fetch_one(tx.as_mut())
14648 .await
14649 .map_err(map_sqlx_error)?;
14650
14651 let next_slot = (state.0 + 1).rem_euclid(state.2);
14652
14653 let claim_busy =
14660 Self::relation_has_rows_tx(&mut tx, &claim_child_name(schema, next_slot as usize))
14661 .await?;
14662 let claim_batch_busy = Self::relation_has_rows_tx(
14663 &mut tx,
14664 &claim_batch_child_name(schema, next_slot as usize),
14665 )
14666 .await?;
14667 let closure_busy =
14668 Self::relation_has_rows_tx(&mut tx, &closure_child_name(schema, next_slot as usize))
14669 .await?;
14670 let closure_batch_busy = Self::relation_has_rows_tx(
14671 &mut tx,
14672 &claim_closure_batch_child_name(schema, next_slot as usize),
14673 )
14674 .await?;
14675
14676 if claim_busy || claim_batch_busy || closure_busy || closure_batch_busy {
14677 tx.commit().await.map_err(map_sqlx_error)?;
14678 return Ok(RotateOutcome::SkippedBusy {
14679 slot: next_slot,
14680 busy: BusyCounts {
14681 claims: busy_indicator(claim_busy || claim_batch_busy),
14682 closures: busy_indicator(closure_busy),
14683 closure_batches: busy_indicator(closure_batch_busy),
14684 ..Default::default()
14685 },
14686 });
14687 }
14688
14689 let next_generation = state.1 + 1;
14690
14691 let rotated = sqlx::query(&format!(
14692 r#"
14693 UPDATE {schema}.claim_ring_state
14694 SET current_slot = $1,
14695 generation = $2
14696 WHERE singleton = TRUE
14697 AND current_slot = $3
14698 AND generation = $4
14699 "#
14700 ))
14701 .bind(next_slot)
14702 .bind(next_generation)
14703 .bind(state.0)
14704 .bind(state.1)
14705 .execute(tx.as_mut())
14706 .await
14707 .map_err(map_sqlx_error)?;
14708
14709 if rotated.rows_affected() == 0 {
14710 tx.commit().await.map_err(map_sqlx_error)?;
14713 return Ok(RotateOutcome::SkippedBusy {
14714 slot: next_slot,
14715 busy: BusyCounts {
14716 claims: busy_indicator(claim_busy || claim_batch_busy),
14717 closures: busy_indicator(closure_busy),
14718 closure_batches: busy_indicator(closure_batch_busy),
14719 ..Default::default()
14720 },
14721 });
14722 }
14723
14724 tx.commit().await.map_err(map_sqlx_error)?;
14725 Ok(RotateOutcome::Rotated {
14726 slot: next_slot,
14727 generation: next_generation,
14728 })
14729 }
14730
14731 #[tracing::instrument(skip(self, pool), name = "queue_storage.prune_oldest_claims")]
14756 pub async fn prune_oldest_claims(&self, pool: &PgPool) -> Result<PruneOutcome, AwaError> {
14757 let schema = self.schema();
14758 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
14759
14760 let state: (i32, i64, i32) = sqlx::query_as(&format!(
14761 r#"
14762 SELECT current_slot, generation, slot_count
14763 FROM {schema}.claim_ring_state
14764 WHERE singleton = TRUE
14765 FOR UPDATE
14766 "#
14767 ))
14768 .fetch_one(tx.as_mut())
14769 .await
14770 .map_err(map_sqlx_error)?;
14771
14772 let Some((slot, _generation)) = oldest_initialized_ring_slot(state.0, state.1, state.2)
14773 else {
14774 tx.commit().await.map_err(map_sqlx_error)?;
14775 return Ok(PruneOutcome::Noop);
14776 };
14777
14778 let slot_locked: Option<i32> = sqlx::query_scalar(&format!(
14781 r#"
14782 SELECT slot FROM {schema}.claim_ring_slots
14783 WHERE slot = $1
14784 FOR UPDATE
14785 "#
14786 ))
14787 .bind(slot)
14788 .fetch_optional(tx.as_mut())
14789 .await
14790 .map_err(map_sqlx_error)?;
14791
14792 if slot_locked.is_none() {
14793 tx.commit().await.map_err(map_sqlx_error)?;
14794 return Ok(PruneOutcome::Noop);
14795 }
14796
14797 let claim_child = claim_child_name(schema, slot as usize);
14798 let claim_batch_child = claim_batch_child_name(schema, slot as usize);
14799 let closure_child = closure_child_name(schema, slot as usize);
14800 let closure_batch_child = claim_closure_batch_child_name(schema, slot as usize);
14801
14802 let open_claims = claim_prune_has_open_claims_tx(
14810 &mut tx,
14811 schema,
14812 &claim_child,
14813 &claim_batch_child,
14814 &closure_child,
14815 &closure_batch_child,
14816 )
14817 .await?;
14818
14819 if open_claims {
14820 tx.commit().await.map_err(map_sqlx_error)?;
14821 return Ok(PruneOutcome::SkippedActive {
14822 slot,
14823 reason: SkipReason::ClaimOpen,
14824 count: busy_indicator(open_claims),
14825 });
14826 }
14827
14828 set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?;
14829
14830 let lock_tables = sqlx::query(&format!(
14831 "LOCK TABLE {claim_child}, {claim_batch_child}, {closure_child}, {closure_batch_child} IN ACCESS EXCLUSIVE MODE"
14832 ))
14833 .execute(tx.as_mut())
14834 .await;
14835
14836 if let Err(err) = lock_tables {
14837 let _ = tx.rollback().await;
14838 if is_lock_contention_error(&err) {
14839 return Ok(PruneOutcome::Blocked { slot });
14840 }
14841 return Err(map_sqlx_error(err));
14842 }
14843
14844 let current_slot: i32 = sqlx::query_scalar(&format!(
14849 r#"
14850 SELECT current_slot FROM {schema}.claim_ring_state WHERE singleton = TRUE
14851 "#
14852 ))
14853 .fetch_one(tx.as_mut())
14854 .await
14855 .map_err(map_sqlx_error)?;
14856
14857 if current_slot == slot {
14858 tx.commit().await.map_err(map_sqlx_error)?;
14859 return Ok(PruneOutcome::SkippedActive {
14860 slot,
14861 reason: SkipReason::ClaimCurrent,
14862 count: 0,
14863 });
14864 }
14865
14866 let open_claims = claim_prune_has_open_claims_tx(
14867 &mut tx,
14868 schema,
14869 &claim_child,
14870 &claim_batch_child,
14871 &closure_child,
14872 &closure_batch_child,
14873 )
14874 .await?;
14875 if open_claims {
14876 tx.commit().await.map_err(map_sqlx_error)?;
14877 return Ok(PruneOutcome::SkippedActive {
14878 slot,
14879 reason: SkipReason::ClaimOpen,
14880 count: busy_indicator(open_claims),
14881 });
14882 }
14883
14884 let truncate = sqlx::query(&format!(
14885 "TRUNCATE TABLE {claim_child}, {claim_batch_child}, {closure_child}, {closure_batch_child}"
14886 ))
14887 .execute(tx.as_mut())
14888 .await;
14889
14890 match truncate {
14891 Ok(_) => {
14892 sqlx::query(&format!(
14893 r#"
14894 UPDATE {schema}.claim_ring_slots
14895 SET rescue_cursor_claimed_at = '-infinity'::timestamptz,
14896 rescue_cursor_job_id = 0,
14897 rescue_cursor_run_lease = 0,
14898 deadline_cursor_deadline_at = '-infinity'::timestamptz,
14899 deadline_cursor_job_id = 0,
14900 deadline_cursor_run_lease = 0
14901 WHERE slot = $1
14902 "#
14903 ))
14904 .bind(slot)
14905 .execute(tx.as_mut())
14906 .await
14907 .map_err(map_sqlx_error)?;
14908 tx.commit().await.map_err(map_sqlx_error)?;
14909 Ok(PruneOutcome::Pruned {
14910 slot,
14911 carried_failed_rows: 0,
14912 })
14913 }
14914 Err(err) if is_lock_contention_error(&err) => {
14915 let _ = tx.rollback().await;
14916 Ok(PruneOutcome::Blocked { slot })
14917 }
14918 Err(err) => {
14919 let _ = tx.rollback().await;
14920 Err(map_sqlx_error(err))
14921 }
14922 }
14923 }
14924
14925 fn job_id_sequence(&self) -> String {
14926 format!("{}.job_id_seq", self.schema())
14927 }
14928
14929 fn leases_table(&self) -> String {
14930 format!("{}.{}", self.schema(), self.leases_relname())
14931 }
14932
14933 fn attempt_state_table(&self) -> String {
14934 format!("{}.{}", self.schema(), self.attempt_state_relname())
14935 }
14936}