mod migrations;
mod notify;
mod rows;
use crate::error::Error;
use crate::store::{
CleanupCriteria, HistoryFilter, JobRecord, JournalAppend, JournalRecord, MergePayload, NewJob,
Notifier, Settlement, Snapshot, Store,
};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use deadpool_postgres::{Manager, ManagerConfig, Pool, RecyclingMethod};
use notify::{ListenFactory, PgNotifier, make_listen_factory};
use rows::{JOB_COLUMNS, JOURNAL_COLUMNS, job_from_row, journal_from_row};
use std::time::Duration;
use tokio_postgres::tls::{MakeTlsConnect, TlsConnect};
use tokio_postgres::types::ToSql;
use tokio_postgres::{Client, NoTls, Socket};
use ulid::Ulid;
#[derive(Clone)]
pub struct PostgresStore {
pool: Pool,
prefix: String,
listen_factory: ListenFactory,
}
impl PostgresStore {
pub fn from_config<T>(
config: tokio_postgres::Config,
tls: T,
prefix: impl Into<String>,
) -> Result<Self, Error>
where
T: MakeTlsConnect<Socket> + Clone + Send + Sync + 'static,
T::Stream: Send + Sync,
T::TlsConnect: Send + Sync,
<T::TlsConnect as TlsConnect<Socket>>::Future: Send,
{
let prefix = prefix.into();
validate_prefix(&prefix)?;
let manager_config = ManagerConfig {
recycling_method: RecyclingMethod::Fast,
};
let manager = Manager::from_config(config.clone(), tls.clone(), manager_config);
let pool = Pool::builder(manager).build()?;
let listen_factory = make_listen_factory(config, tls);
Ok(PostgresStore {
pool,
prefix,
listen_factory,
})
}
pub fn connect(dsn: &str, prefix: impl Into<String>) -> Result<Self, Error> {
let config: tokio_postgres::Config = dsn.parse()?;
PostgresStore::from_config(config, NoTls, prefix)
}
#[cfg(feature = "rustls")]
pub fn connect_rustls(
dsn: &str,
prefix: impl Into<String>,
client_config: rustls::ClientConfig,
) -> Result<Self, Error> {
let config: tokio_postgres::Config = dsn.parse()?;
let tls = tokio_postgres_rustls::MakeRustlsConnect::new(client_config);
PostgresStore::from_config(config, tls, prefix)
}
pub fn prefix(&self) -> &str {
&self.prefix
}
pub fn pool(&self) -> &Pool {
&self.pool
}
fn channel(&self) -> String {
format!("{}_jobs", self.prefix)
}
async fn notify(&self, tx: &tokio_postgres::Transaction<'_>) -> Result<(), Error> {
tx.execute("SELECT pg_notify($1, '')", &[&self.channel()])
.await?;
Ok(())
}
}
fn notifies_on_repend(settlement: &Settlement) -> bool {
match settlement {
Settlement::Retry { .. } | Settlement::Pause { .. } | Settlement::Release { .. } => true,
Settlement::Complete { .. } | Settlement::Dead { .. } => false,
}
}
#[async_trait]
impl Store for PostgresStore {
async fn migrate(&self) -> Result<(), Error> {
let mut conn = self.pool.get().await?;
let lock_id = advisory_lock_id(&self.prefix);
let client: &mut Client = &mut conn;
client
.execute("SELECT pg_advisory_lock($1)", &[&lock_id])
.await?;
let result = migrations::apply(client, &self.prefix).await;
let unlock = client
.execute("SELECT pg_advisory_unlock($1)", &[&lock_id])
.await;
result?;
if let Err(error) = unlock {
tracing::warn!(
%error,
"advisory unlock after migration failed; the lock releases on connection close",
);
}
Ok(())
}
async fn enqueue(&self, job: &NewJob) -> Result<(), Error> {
let sql = format!(
"INSERT INTO {prefix}_jobs \
(id, kind, payload, priority, status, created_at, visible_at, \
run_count, failure_count, carry, dedup_key) \
VALUES ($1, $2, $3, $4, 'pending', $5, $6, 0, 0, $7, $8)",
prefix = self.prefix,
);
let mut conn = self.pool.get().await?;
let tx = conn.transaction().await?;
tx.execute(
&sql,
&[
&job.id.to_string(),
&job.kind,
&job.payload,
&job.priority,
&job.created_at,
&job.visible_at,
&job.carry,
&job.dedup_key,
],
)
.await?;
self.notify(&tx).await?;
tx.commit().await?;
Ok(())
}
async fn claim_next(
&self,
kinds: &[String],
priority_floor: i16,
lease: Duration,
claimed_by: &str,
) -> Result<Option<JobRecord>, Error> {
let sql = format!(
"UPDATE {prefix}_jobs SET \
status = 'claimed', \
claimed_by = $1, \
claim_expires_at = now() + interval '1 second' * $2, \
run_count = run_count + 1 \
WHERE id = ( \
SELECT id FROM {prefix}_jobs \
WHERE status = 'pending' \
AND visible_at <= now() \
AND kind = ANY($3) \
AND priority >= $4 \
ORDER BY priority, created_at \
LIMIT 1 \
FOR UPDATE SKIP LOCKED \
) \
RETURNING {columns}",
prefix = self.prefix,
columns = JOB_COLUMNS,
);
let conn = self.pool.get().await?;
let lease_secs = lease.as_secs_f64();
let row = conn
.query_opt(&sql, &[&claimed_by, &lease_secs, &kinds, &priority_floor])
.await?;
row.as_ref().map(job_from_row).transpose()
}
async fn settle(
&self,
id: Ulid,
claimed_by: &str,
settlement: Settlement,
journal: JournalAppend,
) -> Result<bool, Error> {
let guard = "WHERE id = $1 AND claimed_by = $2 AND status = 'claimed'";
let id = id.to_string();
let repends = notifies_on_repend(&settlement);
let mut conn = self.pool.get().await?;
let tx = conn.transaction().await?;
let affected = match settlement {
Settlement::Complete { finished_at } => {
let sql = format!(
"UPDATE {prefix}_jobs SET status = 'completed', finished_at = $3, \
claimed_by = NULL, claim_expires_at = NULL {guard}",
prefix = self.prefix,
);
tx.execute(&sql, &[&id, &claimed_by, &finished_at]).await?
}
Settlement::Retry {
visible_at,
failure_count,
carry,
} => {
let sql = format!(
"UPDATE {prefix}_jobs SET status = 'pending', visible_at = $3, \
failure_count = $4, carry = $5, claimed_by = NULL, \
claim_expires_at = NULL {guard}",
prefix = self.prefix,
);
tx.execute(
&sql,
&[&id, &claimed_by, &visible_at, &failure_count, &carry],
)
.await?
}
Settlement::Pause { visible_at, carry } => {
let sql = format!(
"UPDATE {prefix}_jobs SET status = 'pending', visible_at = $3, \
carry = $4, claimed_by = NULL, claim_expires_at = NULL {guard}",
prefix = self.prefix,
);
tx.execute(&sql, &[&id, &claimed_by, &visible_at, &carry])
.await?
}
Settlement::Dead {
finished_at,
failure_count,
} => {
let sql = format!(
"UPDATE {prefix}_jobs SET status = 'dead', finished_at = $3, \
failure_count = $4, claimed_by = NULL, claim_expires_at = NULL {guard}",
prefix = self.prefix,
);
tx.execute(&sql, &[&id, &claimed_by, &finished_at, &failure_count])
.await?
}
Settlement::Release { visible_at } => {
let sql = format!(
"UPDATE {prefix}_jobs SET status = 'pending', visible_at = $3, \
claimed_by = NULL, claim_expires_at = NULL {guard}",
prefix = self.prefix,
);
tx.execute(&sql, &[&id, &claimed_by, &visible_at]).await?
}
};
if affected > 0 {
let sql = format!(
"INSERT INTO {prefix}_journal \
(job_id, kind, run_no, recorded_at, outcome, note, attachment) \
VALUES ($1, $2, $3, $4, $5, $6, $7)",
prefix = self.prefix,
);
tx.execute(
&sql,
&[
&id,
&journal.kind,
&journal.run_no,
&journal.recorded_at,
&journal.outcome.as_str(),
&journal.note,
&journal.attachment,
],
)
.await?;
}
if affected > 0 && repends {
self.notify(&tx).await?;
}
tx.commit().await?;
Ok(affected > 0)
}
async fn journal(&self, id: Ulid) -> Result<Vec<JournalRecord>, Error> {
let sql = format!(
"SELECT {columns} FROM {prefix}_journal WHERE job_id = $1 ORDER BY id",
columns = JOURNAL_COLUMNS,
prefix = self.prefix,
);
let conn = self.pool.get().await?;
let rows = conn.query(&sql, &[&id.to_string()]).await?;
rows.iter().map(journal_from_row).collect()
}
async fn next_visible_at(&self, kinds: &[String]) -> Result<Option<DateTime<Utc>>, Error> {
let sql = format!(
"SELECT min(visible_at) FROM {prefix}_jobs \
WHERE status = 'pending' AND kind = ANY($1) AND visible_at > now()",
prefix = self.prefix,
);
let conn = self.pool.get().await?;
let row = conn.query_one(&sql, &[&kinds]).await?;
Ok(row.get(0))
}
async fn notifier(&self) -> Result<Box<dyn Notifier>, Error> {
let notifier = PgNotifier::connect(self.listen_factory.clone(), self.channel()).await?;
Ok(Box::new(notifier))
}
async fn find_stale(&self) -> Result<Vec<JobRecord>, Error> {
let sql = format!(
"SELECT {columns} FROM {prefix}_jobs \
WHERE status = 'claimed' AND claim_expires_at < now() \
ORDER BY claim_expires_at \
LIMIT 100",
columns = JOB_COLUMNS,
prefix = self.prefix,
);
let conn = self.pool.get().await?;
let rows = conn.query(&sql, &[]).await?;
rows.iter().map(job_from_row).collect()
}
async fn recover(
&self,
id: Ulid,
visible_at: DateTime<Utc>,
failure_count: i32,
journal: JournalAppend,
) -> Result<bool, Error> {
let id = id.to_string();
let mut conn = self.pool.get().await?;
let tx = conn.transaction().await?;
let update = format!(
"UPDATE {prefix}_jobs SET status = 'pending', visible_at = $2, \
failure_count = $3, claimed_by = NULL, claim_expires_at = NULL \
WHERE id = $1 AND status = 'claimed' AND claim_expires_at < now()",
prefix = self.prefix,
);
let affected = tx
.execute(&update, &[&id, &visible_at, &failure_count])
.await?;
if affected > 0 {
let insert = format!(
"INSERT INTO {prefix}_journal \
(job_id, kind, run_no, recorded_at, outcome, note, attachment) \
VALUES ($1, $2, $3, $4, $5, $6, $7)",
prefix = self.prefix,
);
tx.execute(
&insert,
&[
&id,
&journal.kind,
&journal.run_no,
&journal.recorded_at,
&journal.outcome.as_str(),
&journal.note,
&journal.attachment,
],
)
.await?;
}
if affected > 0 {
self.notify(&tx).await?;
}
tx.commit().await?;
Ok(affected > 0)
}
async fn extend_lease(
&self,
id: Ulid,
claimed_by: &str,
lease: Duration,
) -> Result<bool, Error> {
let sql = format!(
"UPDATE {prefix}_jobs SET claim_expires_at = now() + interval '1 second' * $3 \
WHERE id = $1 AND claimed_by = $2 AND status = 'claimed'",
prefix = self.prefix,
);
let conn = self.pool.get().await?;
let affected = conn
.execute(&sql, &[&id.to_string(), &claimed_by, &lease.as_secs_f64()])
.await?;
Ok(affected > 0)
}
async fn query_jobs(&self, filter: &HistoryFilter) -> Result<Vec<JobRecord>, Error> {
let status = filter.status.map(|s| s.as_str());
let mut params: Vec<&(dyn ToSql + Sync)> = Vec::new();
let mut clauses: Vec<String> = Vec::new();
if let Some(kind) = &filter.kind {
params.push(kind);
clauses.push(format!("kind = ${}", params.len()));
}
if let Some(status) = &status {
params.push(status);
clauses.push(format!("status = ${}", params.len()));
}
if let Some(since) = &filter.finished_since {
params.push(since);
clauses.push(format!("finished_at >= ${}", params.len()));
}
if let Some(until) = &filter.finished_until {
params.push(until);
clauses.push(format!("finished_at < ${}", params.len()));
}
let where_clause = if clauses.is_empty() {
String::new()
} else {
format!("WHERE {}", clauses.join(" AND "))
};
let limit_clause = match &filter.limit {
Some(limit) => {
params.push(limit);
format!("LIMIT ${}", params.len())
}
None => String::new(),
};
let sql = format!(
"SELECT {columns} FROM {prefix}_jobs {where_clause} \
ORDER BY created_at DESC {limit_clause}",
columns = JOB_COLUMNS,
prefix = self.prefix,
);
let conn = self.pool.get().await?;
let rows = conn.query(&sql, ¶ms).await?;
rows.iter().map(job_from_row).collect()
}
async fn cleanup(&self, criteria: &CleanupCriteria) -> Result<u64, Error> {
let status = criteria.status.map(|s| s.as_str());
let mut params: Vec<&(dyn ToSql + Sync)> = vec![&criteria.finished_before];
let mut clauses = vec![
"finished_at IS NOT NULL".to_owned(),
"finished_at < $1".to_owned(),
];
if let Some(kind) = &criteria.kind {
params.push(kind);
clauses.push(format!("kind = ${}", params.len()));
}
if let Some(status) = &status {
params.push(status);
clauses.push(format!("status = ${}", params.len()));
}
let sql = format!(
"DELETE FROM {prefix}_jobs WHERE {}",
clauses.join(" AND "),
prefix = self.prefix,
);
let conn = self.pool.get().await?;
Ok(conn.execute(&sql, ¶ms).await?)
}
async fn stats(&self) -> Result<Snapshot, Error> {
let conn = self.pool.get().await?;
let mut snapshot = Snapshot::default();
let pending_sql = format!(
"SELECT kind, count(*), min(created_at) FROM {prefix}_jobs \
WHERE status = 'pending' GROUP BY kind",
prefix = self.prefix,
);
let now = Utc::now();
for row in conn.query(&pending_sql, &[]).await? {
let kind: String = row.get(0);
let count: i64 = row.get(1);
let oldest: DateTime<Utc> = row.get(2);
snapshot
.pending_by_kind
.insert(kind.clone(), count.max(0) as u64);
let age = (now - oldest).to_std().unwrap_or(Duration::ZERO);
snapshot.oldest_pending_age.insert(kind, age);
}
let claimed_sql = format!(
"SELECT count(*) FROM {prefix}_jobs WHERE status = 'claimed'",
prefix = self.prefix,
);
let claimed: i64 = conn.query_one(&claimed_sql, &[]).await?.get(0);
snapshot.claimed = claimed.max(0) as u64;
let dead_sql = format!(
"SELECT kind, count(*) FROM {prefix}_jobs WHERE status = 'dead' GROUP BY kind",
prefix = self.prefix,
);
for row in conn.query(&dead_sql, &[]).await? {
let kind: String = row.get(0);
let count: i64 = row.get(1);
snapshot.dead_by_kind.insert(kind, count.max(0) as u64);
}
Ok(snapshot)
}
async fn dedup_candidate(
&self,
kind: &str,
dedup_key: &str,
) -> Result<Option<JobRecord>, Error> {
let sql = format!(
"SELECT {columns} FROM {prefix}_jobs \
WHERE kind = $1 AND dedup_key = $2 AND status = 'pending' \
ORDER BY created_at \
LIMIT 1",
columns = JOB_COLUMNS,
prefix = self.prefix,
);
let conn = self.pool.get().await?;
let row = conn.query_opt(&sql, &[&kind, &dedup_key]).await?;
row.as_ref().map(job_from_row).transpose()
}
async fn merge_into(
&self,
id: Ulid,
update: Option<MergePayload>,
journal: JournalAppend,
) -> Result<bool, Error> {
let id = id.to_string();
let mut conn = self.pool.get().await?;
let tx = conn.transaction().await?;
let affected = match update {
Some(MergePayload {
payload,
carry,
priority,
}) => {
let sql = format!(
"UPDATE {prefix}_jobs SET payload = $2, carry = $3, priority = $4 \
WHERE id = $1 AND status = 'pending'",
prefix = self.prefix,
);
tx.execute(&sql, &[&id, &payload, &carry, &priority])
.await?
}
None => {
let sql = format!(
"UPDATE {prefix}_jobs SET dedup_key = dedup_key \
WHERE id = $1 AND status = 'pending'",
prefix = self.prefix,
);
tx.execute(&sql, &[&id]).await?
}
};
if affected > 0 {
let sql = format!(
"INSERT INTO {prefix}_journal \
(job_id, kind, run_no, recorded_at, outcome, note, attachment) \
VALUES ($1, $2, $3, $4, $5, $6, $7)",
prefix = self.prefix,
);
tx.execute(
&sql,
&[
&id,
&journal.kind,
&journal.run_no,
&journal.recorded_at,
&journal.outcome.as_str(),
&journal.note,
&journal.attachment,
],
)
.await?;
}
tx.commit().await?;
Ok(affected > 0)
}
}
const LONGEST_SUFFIX_LEN: usize = "_refinery_schema_history".len();
fn validate_prefix(prefix: &str) -> Result<(), Error> {
if prefix.is_empty() {
return Err(Error::Config("table prefix must not be empty".into()));
}
let max_prefix_len = 63 - LONGEST_SUFFIX_LEN;
if prefix.len() > max_prefix_len {
return Err(Error::Config(format!(
"table prefix {prefix:?} is too long ({} chars); at most {max_prefix_len} are allowed \
so every generated identifier stays within PostgreSQL's 63-character limit",
prefix.len(),
)));
}
let mut chars = prefix.chars();
let first = chars.next().expect("prefix is non-empty");
if !first.is_ascii_lowercase() {
return Err(Error::Config(format!(
"table prefix {prefix:?} must start with a lowercase ASCII letter"
)));
}
if let Some(bad) = prefix
.chars()
.find(|c| !(c.is_ascii_lowercase() || c.is_ascii_digit() || *c == '_'))
{
return Err(Error::Config(format!(
"table prefix {prefix:?} contains the disallowed character {bad:?}; \
only [a-z0-9_] are permitted"
)));
}
Ok(())
}
fn advisory_lock_id(prefix: &str) -> i64 {
const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;
let mut hash = FNV_OFFSET;
for byte in prefix.as_bytes() {
hash ^= u64::from(*byte);
hash = hash.wrapping_mul(FNV_PRIME);
}
hash as i64
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn accepts_a_reasonable_prefix() {
assert!(validate_prefix("venturi").is_ok());
assert!(validate_prefix("jobs_v2").is_ok());
assert!(validate_prefix("a").is_ok());
}
#[test]
fn rejects_empty_prefix() {
assert!(validate_prefix("").is_err());
}
#[test]
fn rejects_prefix_with_uppercase_or_symbols() {
assert!(validate_prefix("Venturi").is_err());
assert!(validate_prefix("jobs-v2").is_err());
assert!(validate_prefix("jobs;drop").is_err());
assert!(validate_prefix("1jobs").is_err());
}
#[test]
fn rejects_overlong_prefix() {
let max = 63 - LONGEST_SUFFIX_LEN;
let ok = "a".repeat(max);
let too_long = "a".repeat(max + 1);
assert!(validate_prefix(&ok).is_ok());
assert!(validate_prefix(&too_long).is_err());
}
#[test]
fn advisory_lock_id_is_stable_and_distinct() {
assert_eq!(advisory_lock_id("venturi"), advisory_lock_id("venturi"));
assert_ne!(advisory_lock_id("venturi"), advisory_lock_id("other"));
}
#[test]
fn repending_settlements_notify_and_terminal_ones_do_not() {
let now = Utc::now();
let carry = serde_json::Value::Null;
assert!(notifies_on_repend(&Settlement::Retry {
visible_at: now,
failure_count: 1,
carry: carry.clone(),
}));
assert!(notifies_on_repend(&Settlement::Pause {
visible_at: now,
carry,
}));
assert!(notifies_on_repend(&Settlement::Release { visible_at: now }));
assert!(!notifies_on_repend(&Settlement::Complete {
finished_at: now
}));
assert!(!notifies_on_repend(&Settlement::Dead {
finished_at: now,
failure_count: 3,
}));
}
}