use chrono::{DateTime, Utc};
use es_entity::AtomicOperation;
use es_entity::clock::ClockHandle;
use serde_json::Value as JsonValue;
use sqlx::postgres::{PgPool, types::PgInterval};
use tracing::{Instrument, Span, instrument};
use std::{
sync::{
Arc,
atomic::{AtomicBool, Ordering},
},
time::Duration,
};
use super::{
JobId,
config::JobPollerConfig,
dispatcher::*,
entity::{Job, JobType},
error::JobError,
handle::OwnedTaskHandle,
registry::JobRegistry,
repo::JobRepo,
tracker::JobTracker,
};
#[cfg(all(feature = "tokio-task-names", tokio_unstable))]
macro_rules! spawn_named_task {
($name:expr, $future:expr) => {
tokio::task::Builder::new()
.name($name)
.spawn($future)
.expect("failed to spawn task")
};
}
#[cfg(not(all(feature = "tokio-task-names", tokio_unstable)))]
macro_rules! spawn_named_task {
($name:expr, $future:expr) => {
tokio::spawn($future)
};
}
pub(crate) struct JobPoller {
config: JobPollerConfig,
repo: Arc<JobRepo>,
registry: JobRegistry,
tracker: Arc<JobTracker>,
instance_id: uuid::Uuid,
shutdown_tx: tokio::sync::broadcast::Sender<
tokio::sync::mpsc::Sender<tokio::sync::oneshot::Receiver<()>>,
>,
clock: ClockHandle,
}
pub(crate) struct JobPollerHandle {
#[allow(dead_code)]
poller: Arc<JobPoller>,
#[allow(dead_code)]
handle: OwnedTaskHandle,
#[allow(dead_code)]
router_listener_handle: OwnedTaskHandle,
#[allow(dead_code)]
router_waiter_handle: OwnedTaskHandle,
shutdown_tx: tokio::sync::broadcast::Sender<
tokio::sync::mpsc::Sender<tokio::sync::oneshot::Receiver<()>>,
>,
shutdown_called: Arc<AtomicBool>,
shutdown_timeout: Duration,
max_jobs_per_process: usize,
repo: Arc<JobRepo>,
instance_id: uuid::Uuid,
clock: ClockHandle,
}
const MAX_WAIT: Duration = Duration::from_secs(60);
impl JobPoller {
pub fn new(
config: JobPollerConfig,
repo: Arc<JobRepo>,
registry: JobRegistry,
tracker: Arc<JobTracker>,
clock: ClockHandle,
) -> Self {
let (shutdown_tx, _) = tokio::sync::broadcast::channel::<
tokio::sync::mpsc::Sender<tokio::sync::oneshot::Receiver<()>>,
>(1);
Self {
tracker,
repo,
config,
registry,
instance_id: uuid::Uuid::now_v7(),
shutdown_tx,
clock,
}
}
pub fn registered_job_types(&self) -> Vec<JobType> {
self.registry.registered_job_types()
}
pub fn start(
self,
router_listener_handle: OwnedTaskHandle,
router_waiter_handle: OwnedTaskHandle,
) -> JobPollerHandle {
let lost_handle = self.start_lost_handler();
let keep_alive_handle = self.start_keep_alive_handler();
let stale_jobs_handle = self.start_stale_jobs_handler();
let shutdown_tx = self.shutdown_tx.clone();
let repo = Arc::clone(&self.repo);
let instance_id = self.instance_id;
let shutdown_timeout = self.config.shutdown_timeout;
let max_jobs_per_process = self.config.max_jobs_per_process;
let clock = self.clock.clone();
let executor = Arc::new(self);
let handle = OwnedTaskHandle::new(spawn_named_task!(
"job-poller-main-loop",
Self::main_loop(
Arc::clone(&executor),
lost_handle,
keep_alive_handle,
stale_jobs_handle,
)
));
JobPollerHandle {
poller: executor,
handle,
router_listener_handle,
router_waiter_handle,
shutdown_tx,
shutdown_called: Arc::new(AtomicBool::new(false)),
repo,
instance_id,
shutdown_timeout,
max_jobs_per_process,
clock,
}
}
async fn main_loop(
self: Arc<Self>,
_lost_task: OwnedTaskHandle,
_keep_alive_task: OwnedTaskHandle,
_stale_jobs_task: OwnedTaskHandle,
) {
let mut failures = 0;
let mut woken_up = false;
let mut shutdown_rx = self.shutdown_tx.subscribe();
loop {
let timeout = match self.poll_and_dispatch(woken_up).await {
Ok(duration) => {
failures = 0;
duration
}
Err(e) => {
failures += 1;
tracing::error!(
exception.message = %e,
exception.type = std::any::type_name_of_val(&e),
failures,
"main loop error"
);
Duration::from_millis(50 << failures)
}
};
tokio::select! {
biased;
_ = shutdown_rx.recv() => {
break;
}
result = self.clock.timeout(timeout, self.tracker.notified()) => {
woken_up = result.is_ok();
}
}
}
}
#[instrument(
name = "job.poll_and_dispatch",
level = "debug",
skip(self),
fields(poller_id, n_jobs_running, n_jobs_to_start, now, next_poll_in)
)]
async fn poll_and_dispatch(self: &Arc<Self>, woken_up: bool) -> Result<Duration, JobError> {
let span = Span::current();
span.record("poller_id", tracing::field::display(self.instance_id));
let Some(n_jobs_to_poll) = self.tracker.next_batch_size() else {
span.record("next_poll_in", tracing::field::debug(MAX_WAIT));
span.record("n_jobs_to_start", 0);
return Ok(MAX_WAIT);
};
let supported_job_types = self.registry.registered_job_types();
let rows = match poll_jobs(
self.repo.pool(),
n_jobs_to_poll,
self.instance_id,
&supported_job_types,
&self.clock,
)
.await?
{
JobPollResult::WaitTillNextJob(duration) => {
span.record("next_poll_in", tracing::field::debug(duration));
span.record("n_jobs_to_start", 0);
return Ok(duration);
}
JobPollResult::Jobs(jobs) => jobs,
};
span.record("n_jobs_to_start", rows.len());
if !rows.is_empty() {
let ids: Vec<JobId> = rows.iter().map(|row| row.id).collect();
let mut entities = self.repo.find_all::<Job>(&ids).await?;
for row in rows {
let Some(job) = entities.remove(&row.id) else {
tracing::error!(
job_id = %row.id,
"claimed job row has no entity; skipping dispatch"
);
continue;
};
self.dispatch_job(job, row).await?;
}
}
span.record("next_poll_in", tracing::field::debug(Duration::ZERO));
Ok(Duration::ZERO)
}
fn start_lost_handler(&self) -> OwnedTaskHandle {
let job_lost_interval = self.config.job_lost_interval;
let pool = self.repo.pool().clone();
let clock = self.clock.clone();
let supported_job_types = self.registry.registered_job_types();
let instance_id = self.instance_id;
let tracker = Arc::clone(&self.tracker);
OwnedTaskHandle::new(spawn_named_task!("job-poller-lost-handler", async move {
loop {
tokio::time::sleep(job_lost_interval / 2).await;
let alive_threshold = chrono::Utc::now() - job_lost_interval;
let reschedule_at = clock.now();
let self_live_ids = tracker.live_job_ids();
let span = tracing::debug_span!(
parent: None,
"job.detect_lost_jobs",
alive_threshold = %alive_threshold,
reschedule_at = %reschedule_at,
instance_id = %instance_id,
n_live_jobs = self_live_ids.len(),
n_lost_jobs = tracing::field::Empty,
);
async {
match reclaim_lost_jobs(
&pool,
instance_id,
&supported_job_types,
alive_threshold,
reschedule_at,
&self_live_ids,
)
.await
{
Ok(ids) => {
Span::current().record("n_lost_jobs", ids.len());
for id in ids {
tracing::error!(job_id = %id, "lost job");
}
}
Err(e) => {
tracing::error!(
exception.message = %e,
exception.type = std::any::type_name_of_val(&e),
"lost-handler failed to reclaim lost jobs"
);
Span::current().record("n_lost_jobs", 0);
}
}
}
.instrument(span)
.await;
}
}))
}
fn start_keep_alive_handler(&self) -> OwnedTaskHandle {
let job_lost_interval = self.config.job_lost_interval;
let pool = self.repo.pool().clone();
let instance_id = self.instance_id;
let tracker = Arc::clone(&self.tracker);
OwnedTaskHandle::new(spawn_named_task!(
"job-poller-keep-alive-handler",
async move {
let mut failures = 0;
loop {
let live_ids = tracker.live_job_ids();
let now = chrono::Utc::now();
let span = tracing::debug_span!(
parent: None,
"job.keep_alive",
instance_id = %instance_id,
now = %now,
n_live_jobs = live_ids.len(),
failures
);
let timeout = async {
if live_ids.is_empty() {
failures = 0;
return job_lost_interval / 4;
}
match sqlx::query!(
r#"
UPDATE job_executions
SET alive_at = $1
WHERE poller_instance_id = $2
AND state = 'running'
AND id = ANY($3)
"#,
now,
instance_id,
&live_ids,
)
.execute(&pool)
.await
{
Ok(_) => {
failures = 0;
job_lost_interval / 4
}
Err(e) => {
failures += 1;
tracing::error!(
instance_id = %instance_id,
exception.message = %e,
exception.type = std::any::type_name_of_val(&e),
"keep alive error"
);
Duration::from_millis(50 << failures)
}
}
}
.instrument(span)
.await;
tokio::time::sleep(timeout).await;
}
}
))
}
fn start_stale_jobs_handler(&self) -> OwnedTaskHandle {
let pending_jobs_check_interval = self.config.pending_jobs_check_interval;
let pool = self.repo.pool().clone();
let clock = self.clock.clone();
let supported_job_types = self.registry.registered_job_types();
OwnedTaskHandle::new(spawn_named_task!(
"job-poller-stale-jobs-handler",
async move {
loop {
tokio::time::sleep(pending_jobs_check_interval).await;
let now = clock.now();
let span = tracing::info_span!(
parent: None,
"job.check_stale_pending_jobs",
n_stale_pending = tracing::field::Empty,
max_pending_duration_secs = tracing::field::Empty,
);
async {
match sqlx::query!(
r#"
SELECT
job_type,
COUNT(*)::INT4 AS "count!: i32",
EXTRACT(EPOCH FROM ($1::timestamptz - MIN(execute_at)))::FLOAT8
AS "max_pending_duration_secs!: f64"
FROM job_executions
WHERE state = 'pending'
AND execute_at <= $1::timestamptz
AND job_type = ANY($2)
GROUP BY job_type
"#,
now,
&supported_job_types as _,
)
.fetch_all(&pool)
.await
{
Ok(rows) => {
let mut total_stale: i64 = 0;
let mut max_pending_secs: f64 = 0.0;
for row in &rows {
total_stale += row.count as i64;
if row.max_pending_duration_secs > max_pending_secs {
max_pending_secs = row.max_pending_duration_secs;
}
tracing::warn!(
job_type = %row.job_type,
count = row.count,
max_pending_duration_secs = row.max_pending_duration_secs,
"stale pending jobs detected"
);
}
Span::current().record("n_stale_pending", total_stale);
Span::current()
.record("max_pending_duration_secs", max_pending_secs);
}
Err(e) => {
tracing::error!(
exception.message = %e,
exception.type = std::any::type_name_of_val(&e),
"failed to check stale pending jobs"
);
}
}
}
.instrument(span)
.await;
}
}
))
}
#[instrument(
name = "job.dispatch_job",
skip(self, job, polled_job),
fields(job_id, job_type, poller_id, attempt, now)
)]
async fn dispatch_job(&self, job: Job, polled_job: PolledJob) -> Result<(), JobError> {
let span = Span::current();
span.record("attempt", polled_job.attempt);
span.record("job_id", tracing::field::display(job.id));
span.record("job_type", tracing::field::display(&job.job_type));
let runner = self
.registry
.init_job(&job, Arc::clone(&self.repo), self.clock.clone())?;
let retry_settings = self.registry.retry_settings(&job.job_type).clone();
let repo = Arc::clone(&self.repo);
let tracker = self.tracker.clone();
let instance_id = self.instance_id;
let clock = self.clock.clone();
span.record("now", tracing::field::display(clock.now()));
span.record("poller_id", tracing::field::display(instance_id));
let shutdown_rx_job = self.shutdown_tx.subscribe();
let mut shutdown_rx_monitor = self.shutdown_tx.subscribe();
let shutdown_timeout = self.config.shutdown_timeout;
let job_id = job.id;
let job_type = job.job_type.clone();
#[cfg_attr(
not(all(feature = "tokio-task-names", tokio_unstable)),
allow(unused_variables)
)]
let task_name = format!("job-{}-{}", job_type, job_id);
spawn_named_task!(&task_name, async move {
use tracing::Instrument;
let attempt = polled_job.attempt;
let job_fut = JobDispatcher::new(
repo,
tracker,
retry_settings,
job_id,
runner,
instance_id,
clock,
)
.execute_job(job, polled_job, shutdown_rx_job);
tokio::pin!(job_fut);
tokio::select! {
res = &mut job_fut => {
if let Err(e) = res {
tracing::error!(
job_id = %job_id,
attempt,
exception.message = %e,
exception.type = std::any::type_name_of_val(&e),
"job dispatcher error"
);
}
}
Ok(shutdown_notifier) = shutdown_rx_monitor.recv() => {
let (send, recv) = tokio::sync::oneshot::channel();
async {
match shutdown_notifier.send(recv).await {
Ok(()) => {
tracing::Span::current().record("ack_sent", true);
tracing::info!("Acknowledgement sent, waiting for job completion");
drop(shutdown_notifier);
match tokio::time::timeout(shutdown_timeout, &mut job_fut).await {
Ok(res) => {
tracing::Span::current().record("job_completed", true);
tracing::info!("Job completed gracefully");
if let Err(e) = res {
tracing::error!(
job_id = %job_id,
attempt,
exception.message = %e,
exception.type = std::any::type_name_of_val(&e),
"job dispatcher error"
);
}
}
Err(_) => {
tracing::Span::current().record("job_completed", false);
tracing::warn!("Job exceeded timeout, aborting");
}
}
let _ = send.send(());
tracing::info!("Final completion signal sent");
}
Err(_) => {
tracing::Span::current().record("ack_sent", false);
tracing::error!("Failed to send acknowledgement - stopped listening");
}
}
}.instrument(tracing::info_span!(
parent: None,
"job.shutdown_coordination",
job_id = %job_id,
job_type = %job_type,
coordination_path = "shutdown_first",
ack_sent = tracing::field::Empty,
job_completed = tracing::field::Empty,
)
).await;
}
}
});
Ok(())
}
}
async fn reclaim_lost_jobs(
pool: &PgPool,
instance_id: uuid::Uuid,
supported_job_types: &[JobType],
alive_threshold: DateTime<Utc>,
reschedule_at: DateTime<Utc>,
self_live_ids: &[uuid::Uuid],
) -> Result<Vec<JobId>, sqlx::Error> {
let rows = sqlx::query!(
r#"
UPDATE job_executions
SET state = 'pending', execute_at = $3, attempt_index = attempt_index + 1, poller_instance_id = NULL
WHERE state = 'running'
AND alive_at < $1::timestamptz
AND job_type = ANY($2)
AND (poller_instance_id IS DISTINCT FROM $4 OR id <> ALL($5))
RETURNING id AS "id!: JobId"
"#,
alive_threshold,
supported_job_types as _,
reschedule_at,
instance_id,
self_live_ids,
)
.fetch_all(pool)
.await?;
Ok(rows.into_iter().map(|r| r.id).collect())
}
#[instrument(name = "job.poll_jobs", level = "debug", skip(pool, supported_job_types, clock), fields(n_jobs_to_poll, instance_id = %instance_id, n_jobs_found = tracing::field::Empty))]
async fn poll_jobs(
pool: &PgPool,
n_jobs_to_poll: usize,
instance_id: uuid::Uuid,
supported_job_types: &[super::entity::JobType],
clock: &ClockHandle,
) -> Result<JobPollResult, sqlx::Error> {
let sim_now = clock.now();
let wall_now = chrono::Utc::now();
Span::current().record("now", tracing::field::display(sim_now));
let rows = sqlx::query_as!(
JobPollRow,
r#"
-- Narrowest set first: only due jobs, in execute_at order,
-- bounded by a small overscan of the poll limit.
-- idx_job_executions_pending_job_type_execute_at serves this as
-- an ordered index scan; future-scheduled jobs never reach the
-- anti-join / window-function work below.
WITH due AS (
SELECT id, queue_id, execute_at
FROM job_executions
WHERE state = 'pending'
AND job_type = ANY($4)
AND execute_at <= $2::timestamptz
ORDER BY execute_at
LIMIT $1 * 4
),
candidates AS (
-- The 1-at-a-time-per-queue anti-join and the queue-dedup
-- window run only over the bounded due set instead of every
-- pending row.
SELECT id, execute_at,
ROW_NUMBER() OVER (
PARTITION BY COALESCE(queue_id, id::text)
ORDER BY execute_at
) AS rn
FROM due
WHERE NOT EXISTS (
SELECT 1 FROM job_executions AS running
WHERE running.state = 'running'
AND running.queue_id IS NOT NULL
AND running.queue_id = due.queue_id
)
),
selected_jobs AS (
-- The wide execution_state_json is fetched only for the
-- ~$1 winners, not carried through the CTEs and the window
-- sort for every pending job.
SELECT je.id, je.execution_state_json AS data_json, je.attempt_index
FROM candidates c
JOIN job_executions je ON je.id = c.id
WHERE c.rn = 1
ORDER BY je.execute_at ASC
LIMIT $1
FOR UPDATE OF je
),
updated AS (
UPDATE job_executions AS je
SET state = 'running', alive_at = $5, execute_at = NULL, poller_instance_id = $3
FROM selected_jobs
WHERE je.id = selected_jobs.id
AND je.state = 'pending'
RETURNING je.id, selected_jobs.data_json, je.attempt_index
),
min_wait AS (
-- Index-only scan over the pending partial index, no
-- anti-join. May wake slightly early when the nearest
-- future job is queue-blocked — one wasted wake-up at
-- worst, only while its queue stays busy.
SELECT MIN(execute_at) - $2::timestamptz AS wait_time
FROM job_executions
WHERE state = 'pending'
AND job_type = ANY($4)
AND execute_at > $2::timestamptz
)
SELECT * FROM (
SELECT
u.id AS "id?: JobId",
u.data_json AS "data_json?: JsonValue",
u.attempt_index AS "attempt_index?",
NULL::INTERVAL AS "max_wait?: PgInterval"
FROM updated u
UNION ALL
SELECT
NULL::UUID AS "id?: JobId",
NULL::JSONB AS "data_json?: JsonValue",
NULL::INT AS "attempt_index?",
mw.wait_time AS "max_wait?: PgInterval"
FROM min_wait mw
WHERE NOT EXISTS (SELECT 1 FROM updated)
) AS result
"#,
n_jobs_to_poll as i32,
sim_now,
instance_id,
supported_job_types as _,
wall_now,
)
.fetch_all(pool)
.await?;
Span::current().record("n_jobs_found", rows.len());
Ok(JobPollResult::from_rows(rows))
}
#[derive(Debug)]
enum JobPollResult {
Jobs(Vec<PolledJob>),
WaitTillNextJob(Duration),
}
#[derive(Debug)]
struct JobPollRow {
id: Option<JobId>,
data_json: Option<JsonValue>,
attempt_index: Option<i32>,
max_wait: Option<PgInterval>,
}
impl JobPollResult {
pub fn from_rows(rows: Vec<JobPollRow>) -> Self {
if rows.is_empty() {
JobPollResult::WaitTillNextJob(MAX_WAIT)
} else if rows.len() == 1 && rows[0].id.is_none() {
if let Some(interval) = &rows[0].max_wait {
JobPollResult::WaitTillNextJob(pg_interval_to_duration(interval))
} else {
JobPollResult::WaitTillNextJob(MAX_WAIT)
}
} else {
let jobs = rows
.into_iter()
.filter_map(|row| {
if let (Some(id), Some(attempt_index)) = (row.id, row.attempt_index) {
Some(PolledJob {
id,
data_json: row.data_json,
attempt: attempt_index as u32,
})
} else {
None
}
})
.collect();
JobPollResult::Jobs(jobs)
}
}
}
fn pg_interval_to_duration(interval: &PgInterval) -> Duration {
const SECONDS_PER_DAY: u64 = 24 * 60 * 60;
if interval.microseconds < 0 || interval.days < 0 || interval.months < 0 {
Duration::default()
} else {
let days = (interval.days as u64) + (interval.months as u64) * 30;
Duration::from_micros(interval.microseconds as u64)
+ Duration::from_secs(days * SECONDS_PER_DAY)
}
}
impl JobPollerHandle {
pub async fn shutdown(&self) -> Result<(), JobError> {
perform_shutdown(
self.shutdown_tx.clone(),
Arc::clone(&self.repo),
self.instance_id,
self.shutdown_called.clone(),
self.shutdown_timeout,
self.max_jobs_per_process,
self.clock.clone(),
)
.await
}
}
impl Drop for JobPollerHandle {
fn drop(&mut self) {
let shutdown_tx = self.shutdown_tx.clone();
let repo = Arc::clone(&self.repo);
let instance_id = self.instance_id;
let shutdown_called = self.shutdown_called.clone();
let shutdown_timeout = self.shutdown_timeout;
let max_jobs_per_process = self.max_jobs_per_process;
let clock = self.clock.clone();
spawn_named_task!("job-poller-shutdown-on-drop", async move {
let _ = perform_shutdown(
shutdown_tx,
repo,
instance_id,
shutdown_called,
shutdown_timeout,
max_jobs_per_process,
clock,
)
.await;
});
}
}
#[instrument(
name = "jobs.perform_shutdown",
skip(shutdown_tx, repo, clock),
fields(n_jobs, instance_id = %instance_id, broadcast_ok, n_responses)
)]
async fn perform_shutdown(
shutdown_tx: tokio::sync::broadcast::Sender<
tokio::sync::mpsc::Sender<tokio::sync::oneshot::Receiver<()>>,
>,
repo: Arc<JobRepo>,
instance_id: uuid::Uuid,
shutdown_called: Arc<AtomicBool>,
shutdown_timeout: Duration,
max_jobs_per_process: usize,
clock: ClockHandle,
) -> Result<(), JobError> {
if shutdown_called
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
.is_err()
{
return Ok(());
}
let (send, mut recv) =
tokio::sync::mpsc::channel::<tokio::sync::oneshot::Receiver<()>>(max_jobs_per_process);
let broadcast_ok = shutdown_tx.send(send).is_ok();
tracing::Span::current().record("broadcast_ok", broadcast_ok);
if broadcast_ok {
let mut receivers = Vec::with_capacity(max_jobs_per_process);
let receive_timeout = Duration::from_millis(100);
tracing::info!("Starting to collect shutdown acknowledgements from job monitors");
loop {
match tokio::time::timeout(receive_timeout, recv.recv()).await {
Ok(Some(oneshot_rx)) => {
receivers.push(oneshot_rx);
tracing::info!(
n_collected = receivers.len(),
"Received acknowledgement from monitor task"
);
}
Ok(None) => {
tracing::info!(
n_collected = receivers.len(),
"Channel closed, all monitors responded"
);
break;
}
Err(_) => {
tracing::warn!(
n_collected = receivers.len(),
"Receive timeout expired, moving on with collected responses"
);
break;
}
}
}
tracing::Span::current().record("n_responses", receivers.len());
tracing::info!(
n_responses = receivers.len(),
"Waiting for all acknowledged jobs to complete"
);
if tokio::time::timeout(shutdown_timeout, futures::future::join_all(receivers))
.await
.is_err()
{
tracing::warn!("Some jobs did not signal completion within shutdown timeout");
} else {
tracing::info!("All acknowledged jobs completed");
}
} else {
tracing::warn!("No active shutdown subscribers, waiting for shutdown timeout");
tokio::time::sleep(shutdown_timeout).await;
}
kill_remaining_jobs(repo, instance_id, clock).await
}
#[instrument(name = "jobs.kill_remaining_jobs", skip(repo, clock), fields(instance_id = %instance_id, n_killed = tracing::field::Empty))]
async fn kill_remaining_jobs(
repo: Arc<JobRepo>,
instance_id: uuid::Uuid,
clock: ClockHandle,
) -> Result<(), JobError> {
let mut op = repo.begin_op_with_clock(&clock).await?;
let now = clock.now();
let rows = sqlx::query!(
r#"
UPDATE job_executions
SET state = 'pending',
execute_at = $1,
poller_instance_id = NULL
WHERE poller_instance_id = $2 AND state = 'running'
RETURNING id as "id!: JobId", attempt_index
"#,
now,
instance_id
)
.fetch_all(op.as_executor())
.await?;
let n_killed = rows.len();
tracing::Span::current().record("n_killed", n_killed);
if n_killed == 0 {
return Ok(());
}
let attempt_map: std::collections::HashMap<JobId, u32> = rows
.into_iter()
.map(|r| (r.id, r.attempt_index as u32))
.collect();
let ids: Vec<JobId> = attempt_map.keys().copied().collect();
let entities = repo.find_all::<crate::Job>(&ids).await?;
for (job_id, mut job) in entities {
let attempt_index = attempt_map[&job_id];
tracing::warn!(
job_id = %job_id,
job_type = %job.job_type,
attempt = attempt_index,
"Job still running after shutdown timeout, forcing reschedule"
);
job.abort_execution("killed job".to_string(), now, attempt_index);
repo.update_in_op(&mut op, &mut job).await?;
}
op.commit().await?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
async fn init_pool() -> anyhow::Result<PgPool> {
let pg_con = std::env::var("PG_CON").unwrap();
Ok(sqlx::PgPool::connect(&pg_con).await?)
}
async fn seed_running_job(
pool: &PgPool,
job_type: &str,
instance_id: uuid::Uuid,
alive_at: DateTime<Utc>,
) -> anyhow::Result<JobId> {
let id = JobId::new();
let uuid = uuid::Uuid::from(id);
let now = chrono::Utc::now();
sqlx::query(
"INSERT INTO jobs (id, unique_per_type, job_type, created_at) VALUES ($1, false, $2, $3)",
)
.bind(uuid)
.bind(job_type)
.bind(now)
.execute(pool)
.await?;
sqlx::query(
"INSERT INTO job_executions \
(id, job_type, state, poller_instance_id, attempt_index, alive_at, created_at) \
VALUES ($1, $2, 'running', $3, 1, $4, $5)",
)
.bind(uuid)
.bind(job_type)
.bind(instance_id)
.bind(alive_at)
.bind(now)
.execute(pool)
.await?;
Ok(id)
}
#[tokio::test]
async fn self_reclaim_skips_live_jobs() -> anyhow::Result<()> {
let pool = init_pool().await?;
let self_id = uuid::Uuid::now_v7();
let other_id = uuid::Uuid::now_v7();
let job_type = format!("reclaim-gate-{}", uuid::Uuid::now_v7());
let stale = chrono::Utc::now() - chrono::Duration::seconds(600);
let live_self = seed_running_job(&pool, &job_type, self_id, stale).await?;
let orphan_self = seed_running_job(&pool, &job_type, self_id, stale).await?;
let other_instance = seed_running_job(&pool, &job_type, other_id, stale).await?;
let threshold = chrono::Utc::now() - chrono::Duration::seconds(60);
let reschedule_at = chrono::Utc::now();
let self_live_ids = vec![uuid::Uuid::from(live_self)];
let types = vec![JobType::from_owned(job_type.clone())];
let reclaimed: std::collections::HashSet<JobId> = reclaim_lost_jobs(
&pool,
self_id,
&types,
threshold,
reschedule_at,
&self_live_ids,
)
.await?
.into_iter()
.collect();
assert!(
reclaimed.contains(&orphan_self),
"self-owned orphan (no live future) must be reclaimed"
);
assert!(
reclaimed.contains(&other_instance),
"another instance's stale row must be reclaimed"
);
assert!(
!reclaimed.contains(&live_self),
"self-owned row with a live runner must NOT be reclaimed"
);
let row: (String, Option<uuid::Uuid>, i32) = sqlx::query_as(
"SELECT state::text, poller_instance_id, attempt_index \
FROM job_executions WHERE id = $1",
)
.bind(uuid::Uuid::from(live_self))
.fetch_one(&pool)
.await?;
assert_eq!(row.0, "running");
assert_eq!(row.1, Some(self_id));
assert_eq!(row.2, 1);
Ok(())
}
}