use std::{
collections::{HashMap, HashSet},
fmt,
future::Future,
pin::Pin,
sync::Arc,
time::Duration,
};
use sqlx::{PgConnection, PgExecutor, PgPool, Postgres, Row, Transaction, postgres::PgPoolOptions};
use tokio::sync::{OwnedSemaphorePermit, Semaphore, watch};
use uuid::Uuid;
pub use sqlx::types::chrono::{DateTime, Utc};
pub type Result<T, E = Error> = std::result::Result<T, E>;
#[derive(Debug)]
#[non_exhaustive]
pub enum Error {
Database(sqlx::Error),
Encode(serde_json::Error),
Job(String),
NotACode {
what: &'static str,
},
}
impl fmt::Display for Error {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Database(error) => error.fmt(f),
Self::Encode(error) => error.fmt(f),
Self::Job(text) => write!(f, "worker error: {text}"),
Self::NotACode { what } => write!(
f,
"`{what}` must be 1 to {MAX_REASON_CHARS} characters with no line breaks"
),
}
}
}
impl std::error::Error for Error {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Database(error) => Some(error),
Self::Encode(error) => Some(error),
Self::Job(_) | Self::NotACode { .. } => None,
}
}
}
impl From<sqlx::Error> for Error {
fn from(error: sqlx::Error) -> Self {
Self::Database(error)
}
}
impl From<serde_json::Error> for Error {
fn from(error: serde_json::Error) -> Self {
Self::Encode(error)
}
}
type JobFuture = Pin<Box<dyn Future<Output = Result<(), String>> + Send>>;
type Handler = Arc<dyn Fn(serde_json::Value) -> JobFuture + Send + Sync>;
const SETUP: &str = "CREATE TABLE IF NOT EXISTS axa_queue (
id UUID PRIMARY KEY,
name TEXT NOT NULL,
args JSONB NOT NULL,
status TEXT NOT NULL DEFAULT 'queued',
error TEXT,
run_at TIMESTAMPTZ NOT NULL DEFAULT now(),
started_at TIMESTAMPTZ,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
)";
const SETUP_INDEX: &str =
"CREATE INDEX IF NOT EXISTS axa_queue_ready ON axa_queue (run_at) WHERE status = 'queued'";
const SETUP_ATTEMPTS: &str =
"ALTER TABLE axa_queue ADD COLUMN IF NOT EXISTS attempts INTEGER NOT NULL DEFAULT 0";
const SETUP_MAX_ATTEMPTS: &str =
"ALTER TABLE axa_queue ADD COLUMN IF NOT EXISTS max_attempts INTEGER";
const SETUP_PRIORITY: &str =
"ALTER TABLE axa_queue ADD COLUMN IF NOT EXISTS priority SMALLINT NOT NULL DEFAULT 0";
const SETUP_TIMEOUT: &str = "ALTER TABLE axa_queue ADD COLUMN IF NOT EXISTS timeout_ms BIGINT";
const SETUP_CLAIM_INDEX: &str = "CREATE INDEX IF NOT EXISTS axa_queue_claim
ON axa_queue (priority DESC, run_at) WHERE status = 'queued'";
const SETUP_SCHEDULE: &str = "CREATE TABLE IF NOT EXISTS axa_schedule (
name TEXT PRIMARY KEY,
due_at TIMESTAMPTZ NOT NULL,
last_job UUID,
interval_ms BIGINT,
seen_at TIMESTAMPTZ NOT NULL DEFAULT now(),
archived_at TIMESTAMPTZ,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
)";
const SETUP_ACTS: &str = "CREATE TABLE IF NOT EXISTS axa_job_acts (
id BIGINT GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
job_id UUID NOT NULL,
job_name TEXT NOT NULL,
act TEXT NOT NULL CHECK (act IN ('cancelled', 'retried', 'released')),
attempts INTEGER NOT NULL,
by TEXT NOT NULL CHECK (length(by) BETWEEN 1 AND 200 AND by !~ '[[:cntrl:]]'),
reason TEXT NOT NULL CHECK (length(reason) BETWEEN 1 AND 200 AND reason !~ '[[:cntrl:]]'),
at TIMESTAMPTZ NOT NULL DEFAULT now()
)";
const SETUP_ACTS_INDEX: &str =
"CREATE INDEX IF NOT EXISTS axa_job_acts_job ON axa_job_acts (job_id, id)";
const SETUP_ACTS_REFUSE: &str = "CREATE OR REPLACE FUNCTION axa_job_acts_refuse() RETURNS trigger
LANGUAGE plpgsql AS $$
BEGIN
RAISE EXCEPTION '% on % is refused: acts on jobs are only ever added to',
TG_OP, TG_TABLE_NAME;
END
$$";
const SETUP_ACTS_ROWS: &str = "CREATE OR REPLACE TRIGGER axa_job_acts_refuse
BEFORE UPDATE OR DELETE ON axa_job_acts
FOR EACH ROW EXECUTE FUNCTION axa_job_acts_refuse()";
const SETUP_ACTS_TRUNCATE: &str = "CREATE OR REPLACE TRIGGER axa_job_acts_refuse_truncate
BEFORE TRUNCATE ON axa_job_acts
FOR EACH STATEMENT EXECUTE FUNCTION axa_job_acts_refuse()";
const SETUP_ACTS_ALWAYS: &str = "ALTER TABLE axa_job_acts
ENABLE ALWAYS TRIGGER axa_job_acts_refuse,
ENABLE ALWAYS TRIGGER axa_job_acts_refuse_truncate";
const ACTS_TRIGGERS: [&str; 2] = ["axa_job_acts_refuse", "axa_job_acts_refuse_truncate"];
const MISSING_TRIGGERS: &str = "SELECT want FROM unnest($1::text[]) AS want
WHERE NOT EXISTS (
SELECT 1 FROM pg_trigger t
WHERE t.tgrelid = to_regclass('axa_job_acts') AND t.tgname = want
AND t.tgenabled::text = 'A'
)";
const SETUP_LOCK: i64 = 0x0061_7861_5F71;
const ENQUEUE: &str = "INSERT INTO axa_queue
(id, name, args, max_attempts, priority, run_at, timeout_ms)
VALUES ($1, $2, $3, $4, $5, COALESCE($6, now() + make_interval(secs => $7)), $8)";
const CLAIM: &str = "UPDATE axa_queue
SET status = 'running', started_at = now(), updated_at = now(), attempts = attempts + 1
WHERE id = (
SELECT id FROM axa_queue
WHERE status = 'queued' AND run_at <= now() AND name = ANY($1)
ORDER BY priority DESC, run_at
FOR UPDATE SKIP LOCKED
LIMIT 1
)
RETURNING id, name, args, attempts, max_attempts, timeout_ms";
const HEARTBEAT: &str = "UPDATE axa_queue SET updated_at = now()
WHERE id = $1 AND attempts = $2 AND status = 'running'";
const FINISH: &str = "UPDATE axa_queue SET status = $3, error = $4, updated_at = now()
WHERE id = $1 AND attempts = $2 AND status = 'running'";
const RETRY: &str = "UPDATE axa_queue
SET status = 'queued', error = $3, started_at = NULL, updated_at = now(),
run_at = now() + make_interval(secs => $4)
WHERE id = $1 AND attempts = $2 AND status = 'running'";
const REQUEUE_STALE: &str = "UPDATE axa_queue
SET status = 'queued', started_at = NULL, updated_at = now()
WHERE status = 'running' AND updated_at < now() - make_interval(mins => $1)
AND attempts < COALESCE(max_attempts, $2)";
const FAIL_STALE: &str = "UPDATE axa_queue
SET status = 'failed', error = 'abandoned by its worker', updated_at = now()
WHERE status = 'running' AND updated_at < now() - make_interval(mins => $1)
AND attempts >= COALESCE(max_attempts, $2)";
const PURGE: &str = "DELETE FROM axa_queue
WHERE status IN ('completed', 'failed', 'cancelled')
AND updated_at < now() - make_interval(days => $1)";
const STATUS: &str = "SELECT id, name, status, attempts, max_attempts, priority, error,
timeout_ms, run_at, started_at, created_at, updated_at FROM axa_queue WHERE id = $1";
const MARK_LOST: &str =
"UPDATE axa_queue SET status = 'needs_review', error = $2, updated_at = now()
WHERE id = $1 AND status IN ('queued', 'running', 'failed')";
const CANCEL: &str = "WITH changed AS (
UPDATE axa_queue SET status = 'cancelled', updated_at = now()
WHERE id = $1 AND status = 'queued'
RETURNING id, name, attempts
)
INSERT INTO axa_job_acts (job_id, job_name, act, attempts, by, reason)
SELECT id, name, 'cancelled', attempts, $2, $3 FROM changed";
const RETRY_FAILED: &str = "WITH target AS (
SELECT id, status FROM axa_queue
WHERE id = $1 AND status IN ('failed', 'needs_review')
FOR UPDATE
), changed AS (
UPDATE axa_queue q
SET status = 'queued', run_at = now(), started_at = NULL, updated_at = now(),
max_attempts = q.attempts + 1
FROM target WHERE q.id = target.id
RETURNING q.id, q.name, q.attempts, target.status AS was
)
INSERT INTO axa_job_acts (job_id, job_name, act, attempts, by, reason)
SELECT id, name, CASE was WHEN 'needs_review' THEN 'released' ELSE 'retried' END,
attempts, $2, $3
FROM changed";
const ACTS: &str = "SELECT id, job_id, job_name, act, attempts, by, reason, at
FROM axa_job_acts WHERE job_id = $1 ORDER BY id";
const COUNTS: &str = "SELECT
count(*) FILTER (WHERE status = 'queued' AND run_at <= now()) AS ready,
count(*) FILTER (WHERE status = 'queued' AND run_at > now()) AS waiting,
count(*) FILTER (WHERE status = 'running') AS running,
count(*) FILTER (WHERE status = 'completed') AS completed,
count(*) FILTER (WHERE status = 'failed') AS failed,
count(*) FILTER (WHERE status = 'cancelled') AS cancelled,
count(*) FILTER (WHERE status = 'needs_review') AS needs_review,
EXTRACT(EPOCH FROM now() - min(run_at) FILTER (WHERE status = 'queued' AND run_at <= now()))::float8 AS oldest
FROM axa_queue";
const SCHEDULE_CREATE: &str = "INSERT INTO axa_schedule (name, due_at, interval_ms)
VALUES ($1, COALESCE($2, now()), $3)
ON CONFLICT (name) DO UPDATE
SET seen_at = now(), archived_at = NULL, interval_ms = EXCLUDED.interval_ms
WHERE axa_schedule.archived_at IS NOT NULL
OR axa_schedule.interval_ms IS DISTINCT FROM EXCLUDED.interval_ms
OR axa_schedule.seen_at < now() - interval '1 minute'";
const SCHEDULE_LOCK: &str = "SELECT last_job, due_at <= now() AS due FROM axa_schedule
WHERE name = $1 AND archived_at IS NULL FOR UPDATE SKIP LOCKED";
const SCHEDULE_ARCHIVE: &str = "UPDATE axa_schedule SET archived_at = now(), updated_at = now()
WHERE archived_at IS NULL AND name <> ALL($1)
AND seen_at < now() - make_interval(mins => $2)";
const SCHEDULES: &str = "SELECT s.name, s.interval_ms, s.due_at, s.last_job, s.seen_at,
s.archived_at, s.created_at, q.status AS last_status, q.error AS last_error
FROM axa_schedule s LEFT JOIN axa_queue q ON q.id = s.last_job ORDER BY s.name";
const SCHEDULE_BUSY: &str = "SELECT EXISTS (SELECT 1 FROM axa_queue
WHERE id = $1 AND status IN ('queued', 'running'))";
const SCHEDULE_ADVANCE: &str = "UPDATE axa_schedule
SET due_at = due_at + make_interval(secs => $2)
* (floor(EXTRACT(EPOCH FROM now() - due_at)::float8 / $2) + 1),
last_job = COALESCE($3, last_job), updated_at = now()
WHERE name = $1";
pub const SETUP_STATEMENTS: [&str; 14] = [
SETUP,
SETUP_ATTEMPTS,
SETUP_MAX_ATTEMPTS,
SETUP_INDEX,
SETUP_PRIORITY,
SETUP_TIMEOUT,
SETUP_CLAIM_INDEX,
SETUP_SCHEDULE,
SETUP_ACTS,
SETUP_ACTS_INDEX,
SETUP_ACTS_REFUSE,
SETUP_ACTS_ROWS,
SETUP_ACTS_TRUNCATE,
SETUP_ACTS_ALWAYS,
];
pub const MAX_REASON_CHARS: usize = 200;
const MIN_HEARTBEAT: Duration = Duration::from_millis(20);
const MAX_BACKOFF: Duration = Duration::from_hours(1);
const MAX_ERROR: usize = 1024;
#[must_use]
pub fn backoff(base: Duration, attempts: i32) -> Duration {
let doublings = u32::try_from(attempts.saturating_sub(1))
.unwrap_or(0)
.min(16);
base.saturating_mul(1 << doublings).min(MAX_BACKOFF)
}
fn cut(mut text: String) -> String {
if text.len() > MAX_ERROR {
let mut end = MAX_ERROR;
while !text.is_char_boundary(end) {
end -= 1;
}
text.truncate(end);
}
text
}
fn millis(duration: Duration) -> i64 {
i64::try_from(duration.as_millis()).unwrap_or(i64::MAX)
}
fn from_millis(millis: i64) -> Duration {
Duration::from_millis(u64::try_from(millis).unwrap_or(0))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum When {
Now,
At(DateTime<Utc>),
In(Duration),
}
#[derive(Clone)]
pub struct NewJob {
name: String,
args: serde_json::Value,
max_attempts: Option<i32>,
priority: i16,
when: When,
timeout: Option<Duration>,
}
impl fmt::Debug for NewJob {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("NewJob")
.field("name", &self.name)
.field("max_attempts", &self.max_attempts)
.field("priority", &self.priority)
.field("when", &self.when)
.field("timeout", &self.timeout)
.finish_non_exhaustive()
}
}
impl NewJob {
pub fn new(name: impl Into<String>, args: serde_json::Value) -> Self {
Self {
name: name.into(),
args,
max_attempts: None,
priority: 0,
when: When::Now,
timeout: None,
}
}
#[must_use]
pub const fn timeout(mut self, limit: Duration) -> Self {
self.timeout = Some(limit);
self
}
#[must_use]
pub const fn max_attempts(mut self, max_attempts: Option<i32>) -> Self {
self.max_attempts = max_attempts;
self
}
#[must_use]
pub const fn priority(mut self, priority: i16) -> Self {
self.priority = priority;
self
}
#[must_use]
pub const fn run_in(mut self, delay: Duration) -> Self {
self.when = When::In(delay);
self
}
#[must_use]
pub const fn run_at(mut self, at: DateTime<Utc>) -> Self {
self.when = When::At(at);
self
}
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
}
pub async fn enqueue_on<'e, E: PgExecutor<'e>>(executor: E, job: &NewJob) -> Result<Uuid> {
let id = Uuid::new_v4();
let (at, delay) = match job.when {
When::Now => (None, 0.0),
When::At(at) => (Some(at), 0.0),
When::In(delay) => (None, delay.as_secs_f64()),
};
sqlx::query(ENQUEUE)
.bind(id)
.bind(&job.name)
.bind(&job.args)
.bind(job.max_attempts)
.bind(job.priority)
.bind(at)
.bind(delay)
.bind(job.timeout.map(millis))
.execute(executor)
.await?;
Ok(id)
}
pub async fn enqueue_all_on(connection: &mut PgConnection, jobs: &[NewJob]) -> Result<Vec<Uuid>> {
let mut ids = Vec::with_capacity(jobs.len());
for job in jobs {
ids.push(enqueue_on(&mut *connection, job).await?);
}
Ok(ids)
}
fn check_code(text: &str, what: &'static str) -> Result<()> {
if text.is_empty()
|| text.chars().count() > MAX_REASON_CHARS
|| text.chars().any(char::is_control)
{
Err(Error::NotACode { what })
} else {
Ok(())
}
}
pub async fn cancel(
transaction: &mut Transaction<'_, Postgres>,
id: Uuid,
by: &str,
reason: &str,
) -> Result<bool> {
act(transaction, CANCEL, id, by, reason).await
}
pub async fn retry(
transaction: &mut Transaction<'_, Postgres>,
id: Uuid,
by: &str,
reason: &str,
) -> Result<bool> {
act(transaction, RETRY_FAILED, id, by, reason).await
}
async fn act(
transaction: &mut Transaction<'_, Postgres>,
statement: &'static str,
id: Uuid,
by: &str,
reason: &str,
) -> Result<bool> {
check_code(by, "by")?;
check_code(reason, "reason")?;
let kept = sqlx::query(statement)
.bind(id)
.bind(by)
.bind(reason)
.execute(&mut **transaction)
.await?
.rows_affected();
Ok(kept == 1)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum ActKind {
Cancelled,
Retried,
Released,
}
impl ActKind {
fn parse(text: &str) -> Result<Self> {
match text {
"cancelled" => Ok(Self::Cancelled),
"retried" => Ok(Self::Retried),
"released" => Ok(Self::Released),
_ => Err(Error::Job("an act with an unknown kind".to_string())),
}
}
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct Act {
pub id: i64,
pub job_id: Uuid,
pub job_name: String,
pub kind: ActKind,
pub attempts: i32,
pub by: String,
pub reason: String,
pub at: DateTime<Utc>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum JobState {
Queued,
Running,
Completed,
Failed,
Cancelled,
NeedsReview,
}
impl JobState {
fn parse(text: &str) -> Result<Self> {
match text {
"queued" => Ok(Self::Queued),
"running" => Ok(Self::Running),
"completed" => Ok(Self::Completed),
"failed" => Ok(Self::Failed),
"cancelled" => Ok(Self::Cancelled),
"needs_review" => Ok(Self::NeedsReview),
_ => Err(Error::Job("a job row with an unknown status".to_string())),
}
}
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct JobInfo {
pub id: Uuid,
pub name: String,
pub state: JobState,
pub attempts: i32,
pub max_attempts: Option<i32>,
pub priority: i16,
pub timeout: Option<Duration>,
pub error: Option<String>,
pub run_at: DateTime<Utc>,
pub started_at: Option<DateTime<Utc>>,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
}
#[derive(Debug, Clone, Copy, Default, PartialEq)]
#[non_exhaustive]
pub struct Counts {
pub ready: i64,
pub waiting: i64,
pub running: i64,
pub completed: i64,
pub failed: i64,
pub cancelled: i64,
pub needs_review: i64,
pub oldest_ready_seconds: f64,
}
#[derive(Clone)]
pub struct Schedule {
name: String,
interval: Duration,
first: Option<DateTime<Utc>>,
job: NewJob,
}
impl fmt::Debug for Schedule {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Schedule")
.field("name", &self.name)
.field("interval", &self.interval)
.field("first", &self.first)
.field("job", &self.job)
.finish()
}
}
impl Schedule {
pub fn every(name: impl Into<String>, interval: Duration, job: NewJob) -> Self {
let name = name.into();
assert!(
!name.is_empty() && name.len() <= 200,
"a schedule name must be 1 to 200 characters"
);
assert!(
interval >= Duration::from_secs(1),
"the schedule {name} needs an interval of at least a second"
);
Self {
name,
interval,
first: None,
job: NewJob {
when: When::Now,
..job
},
}
}
#[must_use]
pub const fn first_at(mut self, first: DateTime<Utc>) -> Self {
self.first = Some(first);
self
}
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct ScheduleInfo {
pub name: String,
pub interval: Option<Duration>,
pub due_at: DateTime<Utc>,
pub last_job: Option<Uuid>,
pub last_state: Option<JobState>,
pub last_error: Option<String>,
pub seen_at: DateTime<Utc>,
pub archived_at: Option<DateTime<Utc>>,
pub created_at: DateTime<Utc>,
}
#[derive(Debug, Clone)]
pub struct QueueOptions {
pub workers: usize,
pub poll_interval: Duration,
pub stale_minutes: i32,
pub heartbeat_interval: Duration,
pub max_attempts: i32,
pub retry_backoff: Duration,
pub retention_days: i32,
pub maintenance_interval: Duration,
}
impl Default for QueueOptions {
fn default() -> Self {
Self {
workers: 2,
poll_interval: Duration::from_secs(1),
stale_minutes: 60,
heartbeat_interval: Duration::from_mins(1),
max_attempts: 1,
retry_backoff: Duration::from_secs(30),
retention_days: 7,
maintenance_interval: Duration::from_mins(5),
}
}
}
pub struct Queue {
pool: PgPool,
handlers: HashMap<String, Handler>,
limits: HashMap<String, Arc<Semaphore>>,
schedules: Vec<Schedule>,
options: QueueOptions,
}
impl fmt::Debug for Queue {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Queue")
.field("handlers", &self.handlers.len())
.field("limits", &self.limits.len())
.field("schedules", &self.schedules.len())
.field("options", &self.options)
.finish_non_exhaustive()
}
}
struct Claimed {
id: Uuid,
name: String,
args: serde_json::Value,
attempts: i32,
limit: i32,
timeout: Option<Duration>,
}
enum Outcome {
Done(Result<(), String>),
Lost(String),
}
impl Queue {
#[must_use]
pub fn from_pool(pool: PgPool, options: QueueOptions) -> Self {
Self {
pool,
handlers: HashMap::new(),
limits: HashMap::new(),
schedules: Vec::new(),
options,
}
}
async fn open(uri: &str, options: QueueOptions) -> Result<Self> {
let loops = u32::try_from(options.workers).unwrap_or(u32::MAX);
let pool = PgPoolOptions::new()
.max_connections(loops.saturating_add(2))
.connect(uri)
.await?;
Ok(Self::from_pool(pool, options))
}
pub async fn connect_existing(uri: &str, options: QueueOptions) -> Result<Self> {
let queue = Self::open(uri, options).await?;
queue.check().await?;
Ok(queue)
}
pub async fn connect(uri: &str, options: QueueOptions) -> Result<Self> {
let queue = Self::open(uri, options).await?;
let mut setup = queue.pool.begin().await?;
sqlx::query("SELECT pg_advisory_xact_lock($1)")
.bind(SETUP_LOCK)
.execute(&mut *setup)
.await?;
for statement in SETUP_STATEMENTS {
sqlx::query(statement).execute(&mut *setup).await?;
}
setup.commit().await?;
Ok(queue)
}
pub async fn check(&self) -> Result<()> {
sqlx::query("SELECT attempts, max_attempts, priority, timeout_ms FROM axa_queue LIMIT 0")
.execute(&self.pool)
.await?;
sqlx::query(
"SELECT name, due_at, last_job, interval_ms, seen_at, archived_at FROM axa_schedule LIMIT 0",
)
.execute(&self.pool)
.await?;
sqlx::query(
"SELECT id, job_id, job_name, act, attempts, by, reason, at FROM axa_job_acts LIMIT 0",
)
.execute(&self.pool)
.await?;
let missing: Vec<String> = sqlx::query_scalar(MISSING_TRIGGERS)
.bind(ACTS_TRIGGERS.to_vec())
.fetch_all(&self.pool)
.await?;
if missing.is_empty() {
Ok(())
} else {
Err(Error::Database(sqlx::Error::Configuration(
format!(
"the job queue is missing, or has turned off, {}: run SETUP_STATEMENTS",
missing.join(", ")
)
.into(),
)))
}
}
#[must_use]
pub const fn pool(&self) -> &PgPool {
&self.pool
}
#[must_use]
pub const fn options(&self) -> &QueueOptions {
&self.options
}
pub fn handle<F, Fut, E>(&mut self, name: impl Into<String>, handler: F)
where
F: Fn(serde_json::Value) -> Fut + Send + Sync + 'static,
Fut: Future<Output = std::result::Result<(), E>> + Send + 'static,
E: fmt::Display + 'static,
{
let handler = Arc::new(handler);
let run: Handler = Arc::new(move |args| {
let job = handler(args);
Box::pin(async move { job.await.map_err(|error| error.to_string()) })
});
self.handlers.insert(name.into(), run);
}
pub fn limit(&mut self, name: impl Into<String>, jobs: usize) {
self.limits
.insert(name.into(), Arc::new(Semaphore::new(jobs.max(1))));
}
pub fn schedule(&mut self, schedule: Schedule) {
assert!(
self.schedules.iter().all(|s| s.name != schedule.name),
"the schedule {} was added twice",
schedule.name
);
self.schedules.push(schedule);
}
pub async fn enqueue(&self, name: &str, args: serde_json::Value) -> Result<Uuid> {
enqueue_on(&self.pool, &NewJob::new(name, args)).await
}
pub async fn enqueue_limited(
&self,
name: &str,
args: serde_json::Value,
max_attempts: Option<i32>,
) -> Result<Uuid> {
enqueue_on(
&self.pool,
&NewJob::new(name, args).max_attempts(max_attempts),
)
.await
}
pub async fn enqueue_job(&self, job: &NewJob) -> Result<Uuid> {
enqueue_on(&self.pool, job).await
}
pub async fn enqueue_all(&self, jobs: &[NewJob]) -> Result<Vec<Uuid>> {
let mut transaction = self.pool.begin().await?;
let ids = enqueue_all_on(&mut transaction, jobs).await?;
transaction.commit().await?;
Ok(ids)
}
pub async fn job(&self, id: Uuid) -> Result<Option<JobInfo>> {
let Some(row) = sqlx::query(STATUS)
.bind(id)
.fetch_optional(&self.pool)
.await?
else {
return Ok(None);
};
Ok(Some(JobInfo {
id: row.try_get("id")?,
name: row.try_get("name")?,
state: JobState::parse(row.try_get::<&str, _>("status")?)?,
attempts: row.try_get("attempts")?,
max_attempts: row.try_get("max_attempts")?,
priority: row.try_get("priority")?,
timeout: row
.try_get::<Option<i64>, _>("timeout_ms")?
.map(from_millis),
error: row.try_get("error")?,
run_at: row.try_get("run_at")?,
started_at: row.try_get("started_at")?,
created_at: row.try_get("created_at")?,
updated_at: row.try_get("updated_at")?,
}))
}
pub async fn acts(&self, id: Uuid) -> Result<Vec<Act>> {
let rows = sqlx::query(ACTS).bind(id).fetch_all(&self.pool).await?;
rows.iter()
.map(|row| {
Ok(Act {
id: row.try_get("id")?,
job_id: row.try_get("job_id")?,
job_name: row.try_get("job_name")?,
kind: ActKind::parse(row.try_get::<&str, _>("act")?)?,
attempts: row.try_get("attempts")?,
by: row.try_get("by")?,
reason: row.try_get("reason")?,
at: row.try_get("at")?,
})
})
.collect()
}
pub async fn counts(&self) -> Result<Counts> {
let row = sqlx::query(COUNTS).fetch_one(&self.pool).await?;
Ok(Counts {
ready: row.try_get("ready")?,
waiting: row.try_get("waiting")?,
running: row.try_get("running")?,
completed: row.try_get("completed")?,
failed: row.try_get("failed")?,
cancelled: row.try_get("cancelled")?,
needs_review: row.try_get("needs_review")?,
oldest_ready_seconds: row.try_get::<Option<f64>, _>("oldest")?.unwrap_or(0.0),
})
}
pub async fn ping(&self) -> Result<()> {
sqlx::query("SELECT 1").execute(&self.pool).await?;
Ok(())
}
pub async fn requeue_stale(&self) -> Result<u64> {
sqlx::query(FAIL_STALE)
.bind(self.options.stale_minutes)
.bind(self.options.max_attempts)
.execute(&self.pool)
.await?;
let result = sqlx::query(REQUEUE_STALE)
.bind(self.options.stale_minutes)
.bind(self.options.max_attempts)
.execute(&self.pool)
.await?;
Ok(result.rows_affected())
}
pub async fn purge_finished(&self) -> Result<u64> {
let result = sqlx::query(PURGE)
.bind(self.options.retention_days)
.execute(&self.pool)
.await?;
Ok(result.rows_affected())
}
pub async fn tick_schedules(&self) -> Result<usize> {
let mut queued = 0;
for schedule in &self.schedules {
if self.tick(schedule).await? {
queued += 1;
}
}
Ok(queued)
}
async fn tick(&self, schedule: &Schedule) -> Result<bool> {
sqlx::query(SCHEDULE_CREATE)
.bind(&schedule.name)
.bind(schedule.first)
.bind(millis(schedule.interval))
.execute(&self.pool)
.await?;
let mut transaction = self.pool.begin().await?;
let Some(row) = sqlx::query(SCHEDULE_LOCK)
.bind(&schedule.name)
.fetch_optional(&mut *transaction)
.await?
else {
return Ok(false);
};
if !row.try_get::<bool, _>("due")? {
transaction.commit().await?;
return Ok(false);
}
let busy = match row.try_get::<Option<Uuid>, _>("last_job")? {
Some(last) => {
sqlx::query_scalar::<_, bool>(SCHEDULE_BUSY)
.bind(last)
.fetch_one(&mut *transaction)
.await?
}
None => false,
};
let job = if busy {
tracing::info!(schedule = %schedule.name, "schedule skipped, its last job has not finished");
None
} else {
Some(enqueue_on(&mut *transaction, &schedule.job).await?)
};
sqlx::query(SCHEDULE_ADVANCE)
.bind(&schedule.name)
.bind(schedule.interval.as_secs_f64())
.bind(job)
.execute(&mut *transaction)
.await?;
transaction.commit().await?;
if job.is_some() {
tracing::debug!(schedule = %schedule.name, "schedule queued its job");
}
Ok(job.is_some())
}
pub async fn archive_schedules(&self) -> Result<u64> {
let names: Vec<&str> = self.schedules.iter().map(|s| s.name.as_str()).collect();
let result = sqlx::query(SCHEDULE_ARCHIVE)
.bind(&names)
.bind(self.options.stale_minutes)
.execute(&self.pool)
.await?;
Ok(result.rows_affected())
}
pub async fn schedules(&self) -> Result<Vec<ScheduleInfo>> {
let rows = sqlx::query(SCHEDULES).fetch_all(&self.pool).await?;
rows.iter()
.map(|row| {
Ok(ScheduleInfo {
name: row.try_get("name")?,
interval: row
.try_get::<Option<i64>, _>("interval_ms")?
.map(from_millis),
due_at: row.try_get("due_at")?,
last_job: row.try_get("last_job")?,
last_state: row
.try_get::<Option<&str>, _>("last_status")?
.map(JobState::parse)
.transpose()?,
last_error: row.try_get("last_error")?,
seen_at: row.try_get("seen_at")?,
archived_at: row.try_get("archived_at")?,
created_at: row.try_get("created_at")?,
})
})
.collect()
}
async fn maintain(&self) {
match self.archive_schedules().await {
Ok(0) => {}
Ok(archived) => tracing::info!(archived, "archived schedules no longer in code"),
Err(error) => tracing::error!(%error, "could not archive schedules"),
}
if let Err(error) = self.requeue_stale().await {
tracing::error!(%error, "could not recover stale jobs");
}
if let Err(error) = self.purge_finished().await {
tracing::error!(%error, "could not delete old jobs");
}
}
async fn supervise(&self, job: JobFuture, claim: &Claimed) -> Outcome {
let mut job = AbortOnDrop(tokio::spawn(job));
let every = self.options.heartbeat_interval.max(MIN_HEARTBEAT);
let stale =
Duration::from_secs(u64::try_from(self.options.stale_minutes).unwrap_or(0) * 60);
let mut last_heard = tokio::time::Instant::now();
let limit = claim
.timeout
.map(|timeout| tokio::time::Instant::now() + timeout);
let deadline = async {
match limit {
Some(at) => tokio::time::sleep_until(at).await,
None => std::future::pending().await,
}
};
tokio::pin!(deadline);
loop {
tokio::select! {
finished = &mut job.0 => {
return Outcome::Done(finished.unwrap_or_else(|_| {
Err(Error::Job("the job panicked".to_string()).to_string())
}));
}
() = &mut deadline => {
let limit = claim.timeout.unwrap_or_default();
return Outcome::Done(Err(
Error::Job(format!("timed out after {}s", limit.as_secs_f64())).to_string(),
));
}
() = tokio::time::sleep(every) => {
let beat = sqlx::query(HEARTBEAT)
.bind(claim.id)
.bind(claim.attempts)
.execute(&self.pool)
.await;
match beat {
Ok(done) if done.rows_affected() == 1 => last_heard = tokio::time::Instant::now(),
Ok(_) => return Outcome::Lost("another worker or recovery took its claim".to_string()),
Err(error) => {
tracing::warn!(id = %claim.id, %error, "job heartbeat failed");
if last_heard.elapsed() >= stale {
return Outcome::Lost(format!(
"it could not report in for {} minutes",
self.options.stale_minutes
));
}
}
}
}
}
}
}
async fn mark_lost(&self, claim: &Claimed, reason: &str) {
let text = Error::Job(format!("stopped: {reason}")).to_string();
let every = self.options.heartbeat_interval.max(MIN_HEARTBEAT);
loop {
match sqlx::query(MARK_LOST)
.bind(claim.id)
.bind(&text)
.execute(&self.pool)
.await
{
Ok(_) => return,
Err(error) => {
tracing::error!(id = %claim.id, %error, "could not mark a lost job for review");
tokio::time::sleep(every).await;
}
}
}
}
fn claimable(&self) -> (Vec<String>, HashMap<String, OwnedSemaphorePermit>) {
let mut places = HashMap::new();
let mut names = Vec::new();
for name in self.handlers.keys() {
match self.limits.get(name) {
None => names.push(name.clone()),
Some(limit) => {
if let Ok(place) = limit.clone().try_acquire_owned() {
places.insert(name.clone(), place);
names.push(name.clone());
}
}
}
}
(names, places)
}
async fn claim(&self, names: &[String]) -> Result<Option<Claimed>> {
let Some(row) = sqlx::query(CLAIM)
.bind(names)
.fetch_optional(&self.pool)
.await?
else {
return Ok(None);
};
Ok(Some(Claimed {
id: row.try_get("id")?,
name: row.try_get("name")?,
args: row.try_get("args")?,
attempts: row.try_get("attempts")?,
limit: row
.try_get::<Option<i32>, _>("max_attempts")?
.unwrap_or(self.options.max_attempts),
timeout: row
.try_get::<Option<i64>, _>("timeout_ms")?
.map(from_millis),
}))
}
pub async fn run_one(&self) -> Result<bool> {
let (names, mut places) = self.claimable();
if names.is_empty() {
return Ok(false);
}
let Some(job) = self.claim(&names).await? else {
return Ok(false);
};
let _place = places.remove(&job.name);
drop(places);
let outcome = match self.handlers.get(&job.name) {
Some(handler) => self.supervise(handler(job.args.clone()), &job).await,
None => Outcome::Done(Err(
Error::Job(format!("no handler for job {}", job.name)).to_string()
)),
};
let (id, name, attempts) = (job.id, job.name.as_str(), job.attempts);
let finished = match outcome {
Outcome::Lost(reason) => {
tracing::error!(job = %name, %id, attempts, %reason, "job stopped, it lost its claim");
self.mark_lost(&job, &reason).await;
return Ok(true);
}
Outcome::Done(Ok(())) => {
sqlx::query(FINISH)
.bind(id)
.bind(attempts)
.bind("completed")
.bind(None::<String>)
.execute(&self.pool)
.await?
}
Outcome::Done(Err(error)) if attempts < job.limit => {
tracing::warn!(job = %name, %id, attempts, %error, "job failed, will retry");
let wait = backoff(self.options.retry_backoff, attempts);
sqlx::query(RETRY)
.bind(id)
.bind(attempts)
.bind(cut(error))
.bind(wait.as_secs_f64())
.execute(&self.pool)
.await?
}
Outcome::Done(Err(error)) => {
tracing::error!(job = %name, %id, attempts, %error, "job failed");
sqlx::query(FINISH)
.bind(id)
.bind(attempts)
.bind("failed")
.bind(cut(error))
.execute(&self.pool)
.await?
}
};
if finished.rows_affected() == 0 {
tracing::error!(job = %name, %id, "job finished after losing its claim");
self.mark_lost(
&job,
"it finished after another worker or recovery took its claim",
)
.await;
}
Ok(true)
}
pub async fn run(self: Arc<Self>, shutdown: watch::Receiver<bool>) {
let mut loops = Vec::new();
loops.push(AbortOnDrop(tokio::spawn(every_until(
shutdown.clone(),
self.options.maintenance_interval,
{
let queue = self.clone();
move || {
let queue = queue.clone();
async move { queue.maintain().await }
}
},
))));
if !self.schedules.is_empty() {
let queue = self.clone();
loops.push(AbortOnDrop(tokio::spawn(every_until(
shutdown.clone(),
self.options.poll_interval,
move || {
let queue = queue.clone();
async move {
if let Err(error) = queue.tick_schedules().await {
tracing::error!(%error, "could not queue scheduled jobs");
}
}
},
))));
}
for _ in 0..self.options.workers.max(1) {
let queue = self.clone();
let mut shutdown = shutdown.clone();
loops.push(AbortOnDrop(tokio::spawn(async move {
while !*shutdown.borrow() {
let worked = queue.run_one().await.unwrap_or_else(|error| {
tracing::error!(%error, "job queue error");
false
});
if !worked {
tokio::select! {
() = tokio::time::sleep(queue.options.poll_interval) => {},
_ = shutdown.changed() => {},
}
}
}
})));
}
for job_loop in &mut loops {
let _ = (&mut job_loop.0).await;
}
}
#[must_use]
pub fn handled(&self) -> HashSet<&str> {
self.handlers.keys().map(String::as_str).collect()
}
}
async fn every_until<F, Fut>(mut shutdown: watch::Receiver<bool>, gap: Duration, round: F)
where
F: Fn() -> Fut,
Fut: Future<Output = ()>,
{
while !*shutdown.borrow() {
round().await;
tokio::select! {
() = tokio::time::sleep(gap) => {},
_ = shutdown.changed() => {},
}
}
}
struct AbortOnDrop<T>(tokio::task::JoinHandle<T>);
impl<T> Drop for AbortOnDrop<T> {
fn drop(&mut self) {
self.0.abort();
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn each_retry_waits_twice_as_long_up_to_a_limit() {
let base = Duration::from_secs(30);
assert_eq!(backoff(base, 1), base);
assert_eq!(backoff(base, 2), base * 2);
assert_eq!(backoff(base, 3), base * 4);
assert_eq!(backoff(base, 1_000), MAX_BACKOFF);
}
#[test]
fn the_debug_text_leaves_out_arguments() {
let job = NewJob::new("Report", serde_json::json!({ "nhs": "943 476 5919" }));
let schedule = Schedule::every("nightly", Duration::from_hours(1), job.clone());
for text in [format!("{job:?}"), format!("{schedule:?}")] {
assert!(text.contains("Report"), "{text}");
assert!(!text.contains("943"), "{text}");
}
}
#[test]
fn a_by_or_reason_that_is_not_a_code_is_refused() {
assert!(check_code("ops:alice", "by").is_ok());
assert!(check_code(&"é".repeat(MAX_REASON_CHARS), "reason").is_ok());
for bad in [
String::new(),
"é".repeat(MAX_REASON_CHARS + 1),
"two\nlines".to_string(),
"a\ttab".to_string(),
] {
let error = check_code(&bad, "reason").unwrap_err();
assert!(matches!(error, Error::NotACode { what: "reason" }));
assert!(!error.to_string().contains("lines"), "the text is left out");
}
}
#[test]
fn a_long_error_is_cut_on_a_character_boundary() {
let text = "é".repeat(MAX_ERROR);
let kept = cut(text);
assert!(kept.len() <= MAX_ERROR);
assert!(kept.chars().all(|c| c == 'é'));
assert_eq!(cut("short".to_string()), "short");
}
#[test]
#[should_panic(expected = "at least a second")]
fn a_schedule_under_a_second_is_refused_at_boot() {
let _ = Schedule::every(
"fast",
Duration::from_millis(10),
NewJob::new("J", serde_json::Value::Null),
);
}
#[test]
#[should_panic(expected = "1 to 200 characters")]
fn a_schedule_with_no_name_is_refused_at_boot() {
let _ = Schedule::every(
"",
Duration::from_secs(1),
NewJob::new("J", serde_json::Value::Null),
);
}
#[tokio::test]
#[should_panic(expected = "added twice")]
async fn a_schedule_added_twice_is_refused_at_boot() {
let pool = PgPoolOptions::new()
.connect_lazy("postgres://u@localhost/none")
.unwrap();
let mut queue = Queue::from_pool(pool, QueueOptions::default());
let job = NewJob::new("J", serde_json::Value::Null);
queue.schedule(Schedule::every("a", Duration::from_secs(5), job.clone()));
queue.schedule(Schedule::every("a", Duration::from_secs(9), job));
}
#[test]
fn a_schedule_ignores_a_delay_on_its_job() {
let job = NewJob::new("J", serde_json::Value::Null).run_in(Duration::from_hours(1));
let schedule = Schedule::every("s", Duration::from_secs(5), job);
assert_eq!(schedule.job.when, When::Now);
}
}