later 0.0.27

//! Worker-side scanning and execution of topic partition heads.
//!
//! Unlike the default delivery path, a topic-partitioned job is never handed
//! to the shared delivery queue. Instead every server with configured topics
//! runs a bounded number of partition poller tasks that repeatedly try to
//! lease the oldest unfinished job in some `(topic, partition)` and run it.
//! The lease and per-partition head sequence (see [`crate::backend`]) make
//! sure only one worker across every process runs a partition's head job at
//! a time, and only after every earlier sequence in that partition finished.
//!
//! A poller does not re-run the full claim (scan every ready head, win a
//! lease, renew it) before every single job. Once it wins a partition's
//! lease, it holds that assignment — renewed in the background — across
//! every job it finds ready there, and only falls back to a fresh claim once
//! the assignment is lost, rendezvous no longer wants this worker to hold
//! it, or the partition runs dry. Each job still goes through the same
//! completion fencing (owner and lease generation must match) as before, so
//! a stale or reclaimed assignment can never run two jobs at once; this only
//! removes the *redundant* claim work between consecutive jobs in a
//! back-to-back-ready partition.

use crate::{
    backend::PartitionHeadClaim,
    commands::enqueue_continuations,
    core::BgJobHandler,
    models::{FailureTransition, Job, RequeuedStage, Stage},
    UtcDateTime,
};
use std::{sync::Arc, time::Duration};

#[cfg(feature = "prometheus")]
use crate::metrics::{HandlerOutcome, HandlerTimer};

/// A partition lease this poller currently holds, kept warm across however
/// many jobs it finds ready in that partition.
///
/// Dropping this aborts its background renewal task; nothing further tries
/// to keep the lease alive once the assignment is given up.
pub(crate) struct PartitionAssignment {
    topic: String,
    partition: crate::topic::PartitionId,
    lease_epoch: i64,
    /// Signaled once by the renewal task if it ever fails to renew — never
    /// otherwise. `has_changed()` lets the poller check this without
    /// consuming the signal, so the same receiver can later be `.changed()`
    /// awaited for real inside [`handle_partitioned_job`].
    lease_lost: tokio::sync::watch::Receiver<()>,
    renewal: tokio::task::JoinHandle<()>,
}

impl Drop for PartitionAssignment {
    fn drop(&mut self) {
        self.renewal.abort();
    }
}

/// Tries to run one topic partition head, reusing a held assignment when
/// possible.
///
/// Returns `true` when a job was claimed or found ready under an existing
/// assignment (whether or not it could run immediately), so the caller can
/// poll again without sleeping.
pub(crate) async fn poll_partition_head<C, H>(
    handler: &Arc<H>,
    owner: &str,
    live_workers: &[String],
    assignment: &mut Option<PartitionAssignment>,
    #[cfg(feature = "prometheus")] worker_id: &str,
) -> anyhow::Result<bool>
where
    C: Sync + Send,
    H: BgJobHandler<C> + Sync + Send + 'static,
{
    let publisher = handler.get_publisher();

    // Give up a held assignment the moment it's no longer worth keeping:
    // rendezvous moved on, or the renewal task already found the lease
    // taken over by someone else. A job already dispatched under it still
    // runs to completion either way (see handle_partitioned_job); this only
    // stops handing out *new* heads under a stale assignment.
    if let Some(held) = assignment.as_ref() {
        let still_owned = crate::topic::rendezvous_owner(&held.topic, held.partition, live_workers)
            == Some(owner);
        let lease_lost = held.lease_lost.has_changed().unwrap_or(true);
        if !still_owned && !lease_lost {
            let released = publisher
                .release_partition_assignment(owner, &held.topic, held.partition, held.lease_epoch)
                .await?;
            if !released {
                tracing::debug!(
                    topic = %held.topic,
                    partition = held.partition.0,
                    lease_epoch = held.lease_epoch,
                    "Partition assignment was already released before rebalancing"
                );
            }
            *assignment = None;
        } else if lease_lost {
            *assignment = None;
        }
    }

    let claim = match assignment.as_ref() {
        Some(held) => {
            publisher
                .peek_assigned_partition_head(owner, &held.topic, held.partition, held.lease_epoch)
                .await?
        }
        None => None,
    };

    // A held assignment with no ready head must not monopolize one poller for
    // its lease lifetime. Release it before looking for other work. A newly
    // enqueued head is safe either way: it is found by the next scan or by
    // the worker that wins this release's newer epoch.
    if claim.is_none() {
        if let Some(held) = assignment.as_ref() {
            let released = publisher
                .release_partition_assignment(owner, &held.topic, held.partition, held.lease_epoch)
                .await?;
            if !released {
                tracing::debug!(
                    topic = %held.topic,
                    partition = held.partition.0,
                    lease_epoch = held.lease_epoch,
                    "Partition assignment was already released before it became idle"
                );
            }
            // Keep the assignment until the release succeeds. If the database
            // call fails, returning with it still held lets the next poll
            // renew and retry rather than stranding a live lease for 30s.
            *assignment = None;
        }
    }

    let claim = match claim {
        Some(claim) => claim,
        None => {
            // No assignment held, or its next head wasn't ready/existent:
            // fall back to the full scan-and-claim, same as ever.
            let claimed = publisher.claim_partition_head(owner, live_workers).await?;
            #[cfg(feature = "prometheus")]
            crate::metrics::record_partition_claim(
                publisher.metrics_queue(),
                claimed.as_ref().map(|claim| claim.topic.as_str()),
            );
            let Some(claimed) = claimed else {
                return Ok(false);
            };
            let (lease_lost_tx, lease_lost_rx) = tokio::sync::watch::channel(());
            let renewal = spawn_lease_renewal(handler.clone(), claimed.clone(), lease_lost_tx);
            *assignment = Some(PartitionAssignment {
                topic: claimed.topic.clone(),
                partition: claimed.partition,
                lease_epoch: claimed.lease_epoch,
                lease_lost: lease_lost_rx,
                renewal,
            });
            claimed
        }
    };

    let Some(job) = publisher.storage.get_job(claim.job_id.clone()).await? else {
        tracing::warn!(
            job_id = %claim.job_id,
            topic = %claim.topic,
            partition = claim.partition.0,
            "Partition head job is missing; removing orphaned partition row"
        );
        publisher
            .complete_partition_head(&claim, Vec::new())
            .await?;
        return Ok(true);
    };

    let runnable = match &job.stage {
        Stage::Enqueued(_) | Stage::Running(_) => Some(job),
        Stage::Delayed(delayed) if delayed.is_time() => Some(job.transition()),
        Stage::Requeued(requeued) if requeued.is_ready() => Some(job.transition()),
        Stage::Delayed(delayed) => {
            release_not_ready(handler, &claim, delayed.not_before).await?;
            None
        }
        Stage::Requeued(requeued) => {
            release_not_ready(
                handler,
                &claim,
                requeued.not_before.unwrap_or_else(chrono::Utc::now),
            )
            .await?;
            None
        }
        Stage::Waiting(waiting) => {
            // The child usually advances when its parent's completion path
            // calls `enqueue_continuations`. If a worker crashes between
            // completing the parent and enqueuing this continuation, nothing
            // else would ever wake this job up, permanently blocking its
            // whole partition. So each poll independently re-checks the
            // parent, mirroring `handle_job_enqueue_initial`'s Waiting
            // handling, and self-heals by advancing directly.
            let parent_id = waiting.parent_id.clone();
            let parent_settled = match publisher.storage.get_job(parent_id).await? {
                Some(parent) => parent.stage.is_success(),
                None => true,
            };
            if parent_settled {
                Some(job.transition())
            } else {
                let retry_shortly = chrono::Utc::now()
                    .checked_add_signed(chrono::Duration::seconds(1))
                    .unwrap_or_else(chrono::Utc::now);
                release_not_ready(handler, &claim, retry_shortly).await?;
                None
            }
        }
        Stage::Success(_) | Stage::Failed(_) => {
            tracing::warn!(
                job_id = %claim.job_id,
                "Partition head job is already terminal; removing stale partition row"
            );
            publisher
                .complete_partition_head(&claim, Vec::new())
                .await?;
            None
        }
    };

    let Some(job) = runnable else {
        return Ok(true);
    };

    // Safe to unwrap: every path above that reaches here either peeked under
    // `assignment` or just installed a fresh one.
    let lease_lost = &mut assignment
        .as_mut()
        .expect("assignment set above")
        .lease_lost;
    handle_partitioned_job(
        handler,
        &claim,
        job,
        lease_lost,
        #[cfg(feature = "prometheus")]
        worker_id,
    )
    .await?;
    Ok(true)
}

async fn release_not_ready<C, H>(
    handler: &Arc<H>,
    claim: &PartitionHeadClaim,
    earliest_run_at: UtcDateTime,
) -> anyhow::Result<()>
where
    C: Sync + Send,
    H: BgJobHandler<C> + Sync + Send + 'static,
{
    let released = handler
        .get_publisher()
        .release_partition_head(claim, earliest_run_at, Vec::new())
        .await?;
    if !released {
        tracing::debug!(
            job_id = %claim.job_id,
            topic = %claim.topic,
            partition = claim.partition.0,
            lease_epoch = claim.lease_epoch,
            "Partition lease was lost before an unready head could be released"
        );
    }
    Ok(())
}

/// Spawns a task that renews `claim`'s lease every 10 seconds until aborted.
///
/// Runs for as long as this poller keeps its assignment to `claim`'s
/// partition — across every job it finds ready there, not just one —
/// regardless of whether this worker is still the rendezvous-assigned
/// owner: preempting a job that is already executing would risk two workers
/// running it at once, so a reassignment made while a job is in flight only
/// takes effect once that job finishes and this worker gives up the
/// assignment (see [`poll_partition_head`]).
fn spawn_lease_renewal<C, H>(
    handler: Arc<H>,
    claim: PartitionHeadClaim,
    lease_lost: tokio::sync::watch::Sender<()>,
) -> tokio::task::JoinHandle<()>
where
    C: Sync + Send,
    H: BgJobHandler<C> + Sync + Send + 'static,
{
    tokio::spawn(async move {
        loop {
            tokio::time::sleep(Duration::from_secs(10)).await;
            let renewed = handler.get_publisher().renew_partition_lease(&claim).await;
            #[cfg(feature = "prometheus")]
            if let Ok(renewed) = &renewed {
                crate::metrics::record_partition_lease_renewal(
                    handler.get_publisher().metrics_queue(),
                    &claim.topic,
                    *renewed,
                );
            }
            match renewed {
                Ok(true) => {}
                Ok(false) => {
                    tracing::warn!(
                        job_id = %claim.job_id,
                        topic = %claim.topic,
                        partition = claim.partition.0,
                        "Partition lease was lost"
                    );
                    let _ = lease_lost.send(());
                    return;
                }
                Err(error) => {
                    tracing::warn!(%error, job_id = %claim.job_id, "Failed to renew partition lease");
                }
            }
        }
    })
}

#[tracing::instrument(skip_all, fields(job_id = %job.id, job_type = %job.payload_type, topic = %claim.topic, partition = claim.partition.0))]
async fn handle_partitioned_job<C, H>(
    handler: &Arc<H>,
    claim: &PartitionHeadClaim,
    job: Job,
    lease_lost: &mut tokio::sync::watch::Receiver<()>,
    #[cfg(feature = "prometheus")] worker_id: &str,
) -> anyhow::Result<()>
where
    C: Sync + Send,
    H: BgJobHandler<C> + Sync + Send + 'static,
{
    let ptype = job.payload_type.clone();
    let payload = job.payload.clone();
    let job_id = job.id.clone();
    let publisher = handler.get_publisher();

    let running_job = job.transition();
    publisher.save(&running_job).await?;

    #[cfg(feature = "prometheus")]
    let _partition_active_guard =
        crate::metrics::PartitionActiveGuard::start(publisher.metrics_queue(), &claim.topic);
    #[cfg(feature = "prometheus")]
    let handler_timer = HandlerTimer::start(publisher.metrics_queue(), worker_id, &ptype);
    let handler_result = tokio::select! {
        result = handler.dispatch(ptype.clone(), &payload, job_id) => result,
        changed = lease_lost.changed() => match changed {
            Ok(()) => Err(anyhow::anyhow!(
                "partition lease lost while handler was running"
            )),
            Err(_) => Err(anyhow::anyhow!(
                "partition lease renewal task stopped while handler was running"
            )),
        },
    };
    #[cfg(feature = "prometheus")]
    handler_timer.finish(match &handler_result {
        Ok(()) => HandlerOutcome::Success,
        Err(_) => HandlerOutcome::Error,
    });

    match handler_result {
        Ok(_) => {
            let success_job = running_job.transition_success()?;
            let success_job_id = success_job.id.clone();
            let operations = publisher
                .job_save_and_expire_operations(&success_job, Duration::from_secs(3600))?;
            if publisher.complete_partition_head(claim, operations).await? {
                enqueue_continuations(handler.clone(), &success_job, &success_job_id).await?;
            } else {
                tracing::warn!(
                    job_id = %success_job_id,
                    topic = %claim.topic,
                    partition = claim.partition.0,
                    lease_epoch = claim.lease_epoch,
                    "Partition lease was lost before successful completion"
                );
            }
        }
        Err(error) => {
            tracing::warn!("Failed partitioned job {}: {}", running_job.id, error);

            let mut running_job = running_job;
            if running_job.config.needs_retry_policy() {
                match handler.retry_policy(&ptype, &payload).await {
                    Ok(retry_policy) => running_job
                        .config
                        .resolve_retry_policy(retry_policy.as_ref()),
                    Err(policy_error) => tracing::warn!(
                        "Could not resolve retry policy for job {}: {}",
                        running_job.id,
                        policy_error
                    ),
                }
            }
            let failed_job_id = running_job.id.clone();
            match running_job.transition_failure(error.to_string())? {
                FailureTransition::Retry(job) => {
                    let not_before = requeued_not_before(&job).unwrap_or_else(chrono::Utc::now);
                    let operations = publisher.job_save_operations(&job)?;
                    let released = publisher
                        .release_partition_head(claim, not_before, operations)
                        .await?;
                    if !released {
                        tracing::warn!(
                            job_id = %failed_job_id,
                            topic = %claim.topic,
                            partition = claim.partition.0,
                            lease_epoch = claim.lease_epoch,
                            "Partition lease was lost before retry release"
                        );
                    }
                }
                FailureTransition::Failed(job) => {
                    let operations = publisher
                        .job_save_and_expire_operations(&job, Duration::from_secs(3600))?;
                    if !publisher.complete_partition_head(claim, operations).await? {
                        tracing::warn!(
                            job_id = %failed_job_id,
                            topic = %claim.topic,
                            partition = claim.partition.0,
                            lease_epoch = claim.lease_epoch,
                            "Partition lease was lost before final failure completion"
                        );
                    }
                }
            }
        }
    }

    Ok(())
}

fn requeued_not_before(job: &Job) -> Option<UtcDateTime> {
    match &job.stage {
        Stage::Requeued(RequeuedStage { not_before, .. }) => *not_before,
        _ => None,
    }
}