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