use async_trait::async_trait;
use awa::model::{insert_with, migrations, InsertOpts};
use awa::{Client, JobArgs, JobContext, JobError, JobResult, QueueConfig, Worker};
use chrono::{Duration as ChronoDuration, Utc};
use serde::{Deserialize, Serialize};
use sqlx::postgres::PgPoolOptions;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::time::{Duration, Instant};
use uuid::Uuid;
fn repo_root() -> PathBuf {
Path::new(env!("CARGO_MANIFEST_DIR"))
.parent()
.expect("workspace root should exist")
.to_path_buf()
}
fn compose_file() -> PathBuf {
repo_root().join("docker/failover-smoke/compose.yml")
}
fn project_name() -> String {
format!("awa-failover-{}", Uuid::new_v4().simple())
}
fn database_url(port: u16) -> String {
format!("postgres://postgres:test@127.0.0.1:{port}/awa_failover_test")
}
fn run_command(mut command: Command, context: &str) -> String {
let output = command
.output()
.unwrap_or_else(|err| panic!("{context} failed to start: {err}"));
assert!(
output.status.success(),
"{context} failed with status {}.\nstdout:\n{}\nstderr:\n{}",
output.status,
String::from_utf8_lossy(&output.stdout),
String::from_utf8_lossy(&output.stderr),
);
String::from_utf8(output.stdout).expect("command output should be utf-8")
}
fn try_command(mut command: Command) -> Result<String, String> {
let output = command
.output()
.map_err(|err| format!("failed to start command: {err}"))?;
if !output.status.success() {
return Err(format!(
"status {}.\nstdout:\n{}\nstderr:\n{}",
output.status,
String::from_utf8_lossy(&output.stdout),
String::from_utf8_lossy(&output.stderr),
));
}
String::from_utf8(output.stdout).map_err(|err| format!("utf8 decode failed: {err}"))
}
fn docker_compose(project: &str, args: &[&str]) -> Command {
let mut command = Command::new("docker");
command.arg("compose");
command.arg("-f").arg(compose_file());
command.arg("-p").arg(project);
command.args(args);
command.current_dir(repo_root());
command
}
struct ComposeStack {
project: String,
}
impl ComposeStack {
fn up() -> Self {
let project = project_name();
run_command(
docker_compose(&project, &["up", "-d", "--build"]),
"docker compose up",
);
Self { project }
}
fn proxy_port(&self) -> u16 {
let output = run_command(
docker_compose(&self.project, &["port", "haproxy", "5432"]),
"docker compose port haproxy",
);
let binding = output.trim();
let port = binding
.rsplit(':')
.next()
.expect("compose port output should contain a port");
port.parse::<u16>()
.expect("compose port output should end in a valid port")
}
fn stop_primary(&self) {
run_command(
docker_compose(&self.project, &["stop", "primary"]),
"docker compose stop primary",
);
}
fn primary_wal_lsn(&self) -> String {
run_command(
docker_compose(
&self.project,
&[
"exec",
"-T",
"primary",
"sh",
"-lc",
"PGPASSWORD=test psql -U postgres -d awa_failover_test -At -c \"SELECT pg_current_wal_lsn()\"",
],
),
"docker compose exec primary pg_current_wal_lsn",
)
.trim()
.to_string()
}
fn replica_has_replayed(&self, lsn: &str) -> bool {
let output = try_command(
docker_compose(
&self.project,
&[
"exec",
"-T",
"replica",
"sh",
"-lc",
&format!(
"PGPASSWORD=test psql -U postgres -d awa_failover_test -At -c \"SELECT pg_last_wal_replay_lsn() >= '{lsn}'::pg_lsn\""
),
],
),
);
matches!(output.as_deref(), Ok("t\n") | Ok("t"))
}
fn promote_replica(&self) {
run_command(
docker_compose(
&self.project,
&[
"exec",
"-T",
"replica",
"sh",
"-lc",
"PGPASSWORD=test psql -U postgres -d awa_failover_test -c \"SELECT pg_promote(wait_seconds => 30);\"",
],
),
"docker compose exec replica pg_promote",
);
}
}
impl Drop for ComposeStack {
fn drop(&mut self) {
let _ = Command::new("docker")
.arg("compose")
.arg("-f")
.arg(compose_file())
.arg("-p")
.arg(&self.project)
.args(["down", "-v", "--remove-orphans"])
.current_dir(repo_root())
.status();
}
}
async fn connect_pool_with(database_url: &str, max_connections: u32) -> sqlx::PgPool {
PgPoolOptions::new()
.max_connections(max_connections)
.acquire_timeout(Duration::from_secs(5))
.connect(database_url)
.await
.expect("failed to connect to failover test database")
}
async fn connect_pool(database_url: &str) -> sqlx::PgPool {
connect_pool_with(database_url, 20).await
}
async fn wait_for_pool(database_url: &str, timeout: Duration) -> sqlx::PgPool {
let start = Instant::now();
loop {
match PgPoolOptions::new()
.max_connections(5)
.acquire_timeout(Duration::from_secs(2))
.connect(database_url)
.await
{
Ok(pool) => return pool,
Err(err) => {
assert!(
start.elapsed() < timeout,
"timed out connecting to failover test database: {err}"
);
tokio::time::sleep(Duration::from_millis(250)).await;
}
}
}
}
async fn wait_for_replica_replay(stack: &ComposeStack, lsn: &str, timeout: Duration) {
let start = Instant::now();
loop {
if stack.replica_has_replayed(lsn) {
return;
}
assert!(
start.elapsed() < timeout,
"timed out waiting for replica to replay primary LSN {lsn}"
);
tokio::time::sleep(Duration::from_millis(250)).await;
}
}
async fn wait_for_writable(database_url: &str, timeout: Duration) {
let start = Instant::now();
loop {
if let Ok(pool) = PgPoolOptions::new()
.max_connections(2)
.acquire_timeout(Duration::from_secs(2))
.connect(database_url)
.await
{
let writable = sqlx::query_scalar::<_, bool>("SELECT NOT pg_is_in_recovery()")
.fetch_one(&pool)
.await
.unwrap_or(false);
if writable {
pool.close().await;
return;
}
pool.close().await;
}
assert!(
start.elapsed() < timeout,
"timed out waiting for promoted writable database"
);
tokio::time::sleep(Duration::from_millis(250)).await;
}
}
async fn queue_state_counts(pool: &sqlx::PgPool, queue: &str) -> HashMap<String, i64> {
if let Some(schema) = queue_storage_schema_for_counts(pool).await {
let sql = format!(
"WITH live_counts AS ( \
SELECT state::text, count(*)::bigint AS count \
FROM ( \
SELECT 'available'::awa.job_state AS state \
FROM {schema}.ready_entries AS ready \
JOIN {schema}.queue_claim_heads AS claims \
ON claims.queue = ready.queue \
AND claims.priority = ready.priority \
AND claims.enqueue_shard = ready.enqueue_shard \
WHERE ready.queue = $1 \
AND ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name) \
UNION ALL \
SELECT state FROM {schema}.deferred_jobs WHERE queue = $1 \
UNION ALL \
SELECT state FROM {schema}.leases WHERE queue = $1 \
UNION ALL \
SELECT 'running'::awa.job_state AS state \
FROM {schema}.lease_claims AS lc \
WHERE lc.queue = $1 \
AND lc.closed_at IS NULL \
AND NOT EXISTS ( \
SELECT 1 FROM {schema}.lease_claim_closures AS cx \
WHERE cx.claim_slot = lc.claim_slot \
AND cx.job_id = lc.job_id \
AND cx.run_lease = lc.run_lease \
) \
AND NOT EXISTS ( \
SELECT 1 FROM {schema}.lease_claim_closure_batches AS cb \
WHERE cb.receipt_ranges @> lc.receipt_id \
) \
AND NOT EXISTS ( \
SELECT 1 FROM {schema}.leases AS lease \
WHERE lease.job_id = lc.job_id \
AND lease.run_lease = lc.run_lease \
) \
AND NOT EXISTS ( \
SELECT 1 FROM {schema}.deferred_jobs AS deferred \
WHERE deferred.job_id = lc.job_id \
AND deferred.run_lease = lc.run_lease \
) \
AND NOT EXISTS ( \
SELECT 1 FROM {schema}.terminal_jobs AS terminal \
WHERE terminal.job_id = lc.job_id \
AND terminal.run_lease = lc.run_lease \
) \
AND NOT EXISTS ( \
SELECT 1 FROM {schema}.dlq_entries AS dlq \
WHERE dlq.job_id = lc.job_id \
AND dlq.run_lease = lc.run_lease \
) \
UNION ALL \
SELECT state FROM {schema}.terminal_jobs WHERE queue = $1 \
UNION ALL \
SELECT state FROM {schema}.dlq_entries WHERE queue = $1 \
) AS jobs \
GROUP BY state \
), \
pruned_terminal AS ( \
SELECT 'completed'::text AS state, \
COALESCE( \
sum( \
GREATEST( \
COALESCE(lanes.pruned_completed_count, 0), \
COALESCE(rollups.pruned_completed_count, 0) \
) \
), \
0 \
)::bigint AS count \
FROM ( \
SELECT queue, priority, pruned_completed_count \
FROM {schema}.queue_lanes \
WHERE queue = $1 \
) AS lanes \
FULL OUTER JOIN ( \
SELECT queue, priority, pruned_completed_count \
FROM {schema}.queue_terminal_rollups \
WHERE queue = $1 \
) AS rollups \
USING (queue, priority) \
) \
SELECT state, sum(count)::bigint AS count \
FROM ( \
SELECT state, count FROM live_counts \
UNION ALL \
SELECT state, count FROM pruned_terminal WHERE count > 0 \
) AS counts \
GROUP BY state"
);
let rows: Vec<(String, i64)> = sqlx::query_as(&sql)
.bind(queue)
.fetch_all(pool)
.await
.expect("failed to query queue-storage state counts");
return rows.into_iter().collect();
}
let rows: Vec<(String, i64)> = sqlx::query_as(
r#"
SELECT state::text, count(*)::bigint
FROM awa.jobs
WHERE queue = $1
GROUP BY state
"#,
)
.bind(queue)
.fetch_all(pool)
.await
.expect("failed to query queue state counts");
rows.into_iter().collect()
}
async fn active_queue_storage_schema(pool: &sqlx::PgPool) -> Option<String> {
sqlx::query_scalar("SELECT awa.active_queue_storage_schema()")
.fetch_one(pool)
.await
.expect("failed to resolve active queue storage schema")
}
async fn queue_storage_schema_for_counts(pool: &sqlx::PgPool) -> Option<String> {
active_queue_storage_schema(pool).await
}
fn state_count(counts: &HashMap<String, i64>, state: &str) -> i64 {
counts.get(state).copied().unwrap_or(0)
}
async fn wait_for_counts(
pool: &sqlx::PgPool,
queue: &str,
predicate: impl Fn(&HashMap<String, i64>) -> bool,
timeout: Duration,
) -> HashMap<String, i64> {
let start = Instant::now();
loop {
let counts = queue_state_counts(pool, queue).await;
if predicate(&counts) {
return counts;
}
assert!(
start.elapsed() < timeout,
"timed out waiting for queue {queue}; last counts: {counts:?}; storage: {}",
storage_debug(pool, queue).await
);
tokio::time::sleep(Duration::from_millis(100)).await;
}
}
async fn storage_debug(pool: &sqlx::PgPool, queue: &str) -> String {
let storage: Option<(String, String, Option<String>)> = sqlx::query_as(
"SELECT current_engine, state, awa.active_queue_storage_schema() \
FROM awa.storage_transition_state WHERE singleton",
)
.fetch_optional(pool)
.await
.ok()
.flatten();
let canonical_count: i64 =
sqlx::query_scalar("SELECT count(*)::bigint FROM awa.jobs WHERE queue = $1")
.bind(queue)
.fetch_one(pool)
.await
.unwrap_or(-1);
let queue_storage_count = if let Some((_, _, Some(schema))) = &storage {
let sql = format!(
"SELECT \
(SELECT count(*)::bigint FROM {schema}.ready_entries WHERE queue = $1) + \
(SELECT count(*)::bigint FROM {schema}.deferred_jobs WHERE queue = $1) + \
(SELECT count(*)::bigint FROM {schema}.leases WHERE queue = $1) + \
(SELECT count(*)::bigint FROM {schema}.lease_claims WHERE queue = $1) + \
(SELECT count(*)::bigint FROM {schema}.lease_claim_closures AS cx \
JOIN {schema}.lease_claims AS lc \
ON lc.claim_slot = cx.claim_slot \
AND lc.job_id = cx.job_id \
AND lc.run_lease = cx.run_lease \
WHERE lc.queue = $1) + \
(SELECT count(*)::bigint FROM {schema}.lease_claim_closure_batches AS cb \
JOIN {schema}.lease_claims AS lc \
ON cb.receipt_ranges @> lc.receipt_id \
WHERE lc.queue = $1) + \
(SELECT count(*)::bigint FROM {schema}.terminal_jobs WHERE queue = $1) + \
(SELECT COALESCE( \
sum( \
GREATEST( \
COALESCE(lanes.pruned_completed_count, 0), \
COALESCE(rollups.pruned_completed_count, 0) \
) \
), \
0 \
)::bigint \
FROM ( \
SELECT queue, priority, pruned_completed_count \
FROM {schema}.queue_lanes \
WHERE queue = $1 \
) AS lanes \
FULL OUTER JOIN ( \
SELECT queue, priority, pruned_completed_count \
FROM {schema}.queue_terminal_rollups \
WHERE queue = $1 \
) AS rollups \
USING (queue, priority))"
);
sqlx::query_scalar::<_, i64>(&sql)
.bind(queue)
.fetch_one(pool)
.await
.unwrap_or(-1)
} else {
-1
};
format!(
"transition={storage:?}, canonical_rows={canonical_count}, queue_storage_rows={queue_storage_count}"
)
}
async fn wait_for_client_postgres_recovery(client: &Client, timeout: Duration) {
let start = Instant::now();
loop {
let health = client.health_check().await;
if health.postgres_connected
&& health.poll_loop_alive
&& health.heartbeat_alive
&& health.maintenance_alive
{
return;
}
assert!(
start.elapsed() < timeout,
"timed out waiting for client recovery; last health: {:?}",
health
);
tokio::time::sleep(Duration::from_millis(100)).await;
}
}
#[derive(Debug, Serialize, Deserialize, JobArgs)]
struct FailoverJob {
seq: i64,
}
struct CompleteWorker;
#[async_trait]
impl Worker for CompleteWorker {
fn kind(&self) -> &'static str {
"failover_job"
}
async fn perform(&self, _ctx: &JobContext) -> Result<JobResult, JobError> {
Ok(JobResult::Completed)
}
}
fn failover_client(pool: sqlx::PgPool, queue: &str) -> Client {
Client::builder(pool)
.queue(
queue,
QueueConfig {
max_workers: 4,
poll_interval: Duration::from_millis(25),
..QueueConfig::default()
},
)
.heartbeat_interval(Duration::from_millis(50))
.promote_interval(Duration::from_millis(50))
.heartbeat_rescue_interval(Duration::from_millis(100))
.leader_election_interval(Duration::from_millis(100))
.leader_check_interval(Duration::from_millis(100))
.register_worker(CompleteWorker)
.build()
.expect("failed to build failover smoke client")
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
#[ignore = "requires docker compose"]
async fn test_postgres_hot_standby_promotion_keeps_awa_working() {
let stack = ComposeStack::up();
let proxy_port = stack.proxy_port();
let url = database_url(proxy_port);
let mut app_pool = wait_for_pool(&url, Duration::from_secs(30)).await;
migrations::run(&app_pool)
.await
.expect("migrations should succeed through proxy");
let queue = format!("failover_smoke_{}", Uuid::new_v4().simple());
let client_pool = connect_pool(&url).await;
let client = failover_client(client_pool, &queue);
client.start().await.expect("client should start");
app_pool.close().await;
app_pool = connect_pool_with(&url, 1).await;
for seq in 0..8_i64 {
insert_with(
&app_pool,
&FailoverJob { seq },
InsertOpts {
queue: queue.clone(),
..Default::default()
},
)
.await
.expect("initial insert should succeed");
}
wait_for_counts(
&app_pool,
&queue,
|counts| state_count(counts, "completed") == 8,
Duration::from_secs(10),
)
.await;
let synced_lsn = stack.primary_wal_lsn();
wait_for_replica_replay(&stack, &synced_lsn, Duration::from_secs(30)).await;
stack.stop_primary();
stack.promote_replica();
wait_for_writable(&url, Duration::from_secs(30)).await;
wait_for_client_postgres_recovery(&client, Duration::from_secs(30)).await;
app_pool.close().await;
app_pool = connect_pool_with(&url, 1).await;
for seq in 8..16_i64 {
insert_with(
&app_pool,
&FailoverJob { seq },
InsertOpts {
queue: queue.clone(),
run_at: Some(Utc::now() + ChronoDuration::milliseconds(200)),
..Default::default()
},
)
.await
.expect("post-promotion insert should succeed");
}
let counts = wait_for_counts(
&app_pool,
&queue,
|counts| {
state_count(counts, "completed") == 16
&& state_count(counts, "scheduled") == 0
&& state_count(counts, "running") == 0
&& state_count(counts, "available") == 0
},
Duration::from_secs(20),
)
.await;
assert_eq!(state_count(&counts, "completed"), 16);
client.shutdown(Duration::from_secs(5)).await;
}