dd327d76 Andy McCall 14 days ago ↶
1
2
3
4
5
6
use super :: { now_ms, Pool} ;
7
use sqlx:: FromRow;
8
9
10
pub const MAX_ATTEMPTS : i64 = 8 ;
11
12
# [ derive ( Clone, Debug, FromRow ) ]
13
pub struct Item {
14
pub id : i64 ,
15
pub kind : String,
16
pub payload : String,
17
pub dedupe : String,
18
pub run_after : i64 ,
19
pub attempts : i64 ,
20
pub last_error : String,
21
pub claimed_by : String,
22
pub claimed_until : i64 ,
23
pub dead : bool ,
24
pub created_ms : i64 ,
25
}
26
27
28
29
pub async fn enqueue ( pool : & Pool, kind : & str , payload : & serde_json:: Value, dedupe : & str , delay_ms : i64 ) -> sqlx:: Result< bool > {
30
let now = now_ms ( ) ;
31
let mut tx = super :: begin_write( pool) .await? ;
32
if ! dedupe.is_empty ( ) {
33
let live: i64 = sqlx:: query_scalar( " SELECT COUNT(*) FROM work_queue WHERE kind = $1 AND dedupe = $2 AND dead = $3" )
34
.bind ( kind)
35
.bind ( dedupe)
36
.bind ( false )
37
.fetch_one ( & mut * tx)
38
.await? ;
39
if live > 0 {
40
return Ok ( false ) ;
41
}
42
}
43
sqlx:: query( " INSERT INTO work_queue (kind, payload, dedupe, run_after, created_ms) VALUES ($1, $2, $3, $4, $5)" )
44
.bind ( kind)
45
.bind ( payload.to_string ( ) )
46
.bind ( dedupe)
47
.bind ( now + delay_ms.max ( 0 ) )
48
.bind ( now)
49
.execute ( & mut * tx)
50
.await? ;
51
tx.commit ( ) .await? ;
52
Ok ( true )
53
}
54
55
56
57
pub async fn claim ( pool : & Pool, worker : & str , limit : i64 , lease_ms : i64 ) -> sqlx:: Result< Vec< Item> > {
58
let now = now_ms ( ) ;
59
# [ cfg ( feature = " postgres" ) ]
60
let pick = " SELECT id FROM work_queue WHERE dead = $3 AND run_after <= $4 AND claimed_until <= $4 ORDER BY run_after, id LIMIT $5 FOR UPDATE SKIP LOCKED" ;
61
# [ cfg ( not ( feature = " postgres" ) ) ]
62
let pick = " SELECT id FROM work_queue WHERE dead = $3 AND run_after <= $4 AND claimed_until <= $4 ORDER BY run_after, id LIMIT $5" ;
63
let sql =
64
format! ( " UPDATE work_queue SET claimed_by = $1, claimed_until = $2, attempts = attempts + 1 WHERE id IN ({pick} ) RETURNING *" ) ;
65
let mut tx = pool.begin ( ) .await? ;
66
let items: Vec< Item> =
67
sqlx:: query_as( & sql) .bind ( worker) .bind ( now + lease_ms) .bind ( false ) .bind ( now) .bind ( limit) .fetch_all ( & mut * tx) .await? ;
68
tx.commit ( ) .await? ;
69
Ok ( items)
70
}
71
72
73
pub async fn done ( pool : & Pool, id : i64 ) -> sqlx:: Result< ( ) > {
74
sqlx:: query( " DELETE FROM work_queue WHERE id = $1" ) .bind ( id) .execute ( pool) .await? ;
75
Ok ( ( ) )
76
}
77
78
79
80
pub async fn failed ( pool : & Pool, item : & Item, error : & str ) -> sqlx:: Result< ( ) > {
81
let dead = item.attempts > = MAX_ATTEMPTS ;
82
let backoff = ( 30_000 i64 < < ( item.attempts.clamp ( 1 , 8 ) - 1 ) ) .min ( 3_600_000 ) ;
83
sqlx:: query( " UPDATE work_queue SET last_error = $1, dead = $2, run_after = $3, claimed_by = '', claimed_until = 0 WHERE id = $4" )
84
.bind ( error)
85
.bind ( dead)
86
.bind ( now_ms ( ) + backoff)
87
.bind ( item.id)
88
.execute ( pool)
89
.await? ;
90
Ok ( ( ) )
91
}
92
93
94
pub async fn counts ( pool : & Pool) -> sqlx:: Result< ( i64 , i64 ) > {
95
let waiting: i64 = sqlx:: query_scalar( " SELECT COUNT(*) FROM work_queue WHERE dead = $1" ) .bind ( false ) .fetch_one ( pool) .await? ;
96
let dead: i64 = sqlx:: query_scalar( " SELECT COUNT(*) FROM work_queue WHERE dead = $1" ) .bind ( true ) .fetch_one ( pool) .await? ;
97
Ok ( ( waiting, dead) )
98
}