TelegramTwitterMediaBot/crates/xmedia-bot/src/queue.rs

//! Generic persistent task queue backed by SQLite (table `tasks`).
//!
//! Concepts kept from the Python `utils/task_queue.py` (untrusted, redesigned):
//! the table schema, the lease/lock/recovery model, and the retry→dead-letter
//! flow. The Python dict-mutation hack (attempts inside the payload) is
//! replaced by dedicated columns.

use crate::db::now_f64;
use parking_lot::Mutex;
use rusqlite::{TransactionBehavior, params};
use serde_json::Value;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::time::Duration;
use tokio::sync::Notify;
use tokio::task::JoinHandle;

pub const MAX_RETRIES: u32 = 2;
pub const LOCK_TTL_SECONDS: f64 = 120.0;

/// Number of concurrent worker loops. Tasks are independent (retries and
/// forward resumes); leases serialize row claims via SQLite transactions, so
/// extra workers drain backlogs faster. Each worker can be mid-send to
/// Telegram at the same time as handler tasks, so keep this modest.
const QUEUE_WORKERS: usize = 4;

/// What a handler returns instead of throwing. The payload it carries is the
/// (possibly updated) task state to persist for the next attempt.
pub enum QueueError {
    /// Reschedule with the given delay; after `MAX_RETRIES` attempts the task
    /// is dead-lettered instead.
    Retryable { delay_seconds: f64, payload: Value },
    /// Give up now.
    Permanent { message: String, payload: Value },
}

type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
type Handler = dyn Fn(Value) -> BoxFuture<'static, Result<(), QueueError>> + Send + Sync;
type DeadLetter = dyn Fn(Value, String) -> BoxFuture<'static, ()> + Send + Sync;

pub struct PersistentTaskQueue {
    pool: std::sync::Arc<crate::db::DbPool>,
    /// Wakes the workers when a row becomes leasable. `notify_one` stores a
    /// permit, so nothing else may share it: a waiter that is not a worker
    /// (the sweep) can consume the permit and leave the due row pending until
    /// the next enqueue.
    notify: Arc<Notify>,
    /// Wakes the lease-expiry sweep; `stop` is the only producer.
    sweep_notify: Arc<Notify>,
    stop: Arc<AtomicBool>,
    worker: Mutex<Vec<JoinHandle<()>>>,
    counter: AtomicU64,
}

struct LeasedRow {
    id: String,
    payload: String,
    attempts: i32,
}

/// Owned worker state so the spawned loop does not borrow the queue handle.
#[derive(Clone)]
struct QueueWorker {
    pool: std::sync::Arc<crate::db::DbPool>,
    notify: Arc<Notify>,
    stop: Arc<AtomicBool>,
    handler: Arc<Handler>,
    dead_letter: Arc<DeadLetter>,
}

/// Resets rows left `in_progress` with an expired lock TTL back to `pending`
/// so they can be leased again (crash/panic recovery).
fn recover_update(conn: &rusqlite::Connection) -> rusqlite::Result<()> {
    conn.execute(
        "UPDATE tasks SET status='pending', locked_until=0 WHERE status='in_progress' AND locked_until < ?1",
        params![now_f64()],
    )?;
    Ok(())
}

/// Base delay × 2^attempts (attempts = retries already done), capped at 300s.
/// Applied at the queue layer so the attempt count actually reaches the
/// backoff computation; Telegram `RetryAfter` delays get the same treatment
/// (conservatively larger wait, no API change needed).
fn scaled_retry_delay(base: f64, attempts: i32) -> f64 {
    (base * 2f64.powi(attempts)).min(300.0)
}

impl PersistentTaskQueue {
    /// Wraps the shared DB pool; the schema is initialized once by
    /// [`crate::db::open_store`] (all three stores share the pool).
    pub fn new(pool: std::sync::Arc<crate::db::DbPool>) -> Self {
        Self {
            pool,
            notify: Arc::new(Notify::new()),
            sweep_notify: Arc::new(Notify::new()),
            stop: Arc::new(AtomicBool::new(false)),
            worker: Mutex::new(Vec::new()),
            counter: AtomicU64::new(0),
        }
    }

    /// Starts the worker loops. Also recovers rows left `in_progress` by a
    /// previous process (lease expired).
    pub async fn start<H, F, D, G>(&self, handler: H, dead_letter: D)
    where
        H: Fn(Value) -> F + Send + Sync + 'static,
        F: Future<Output = Result<(), QueueError>> + Send + 'static,
        D: Fn(Value, String) -> G + Send + Sync + 'static,
        G: Future<Output = ()> + Send + 'static,
    {
        let handler: Arc<Handler> = Arc::new(move |payload| Box::pin(handler(payload)));
        let dead_letter: Arc<DeadLetter> =
            Arc::new(move |payload, message| Box::pin(dead_letter(payload, message)));
        self.recover_stale().await;
        let mut handles = Vec::with_capacity(QUEUE_WORKERS + 1);
        for _ in 0..QUEUE_WORKERS {
            let worker = QueueWorker {
                pool: std::sync::Arc::clone(&self.pool),
                notify: Arc::clone(&self.notify),
                stop: Arc::clone(&self.stop),
                handler: Arc::clone(&handler),
                dead_letter: Arc::clone(&dead_letter),
            };
            handles.push(tokio::spawn(worker.run_loop_supervised()));
        }
        // Periodic lease-expiry sweep: recovers rows a crashed/panicked
        // worker left `in_progress` (the lock TTL bounds the wait). Its own
        // notify (not the workers'): sharing that one let this task consume a
        // `notify_one` permit meant for a worker, which then slept through a
        // due row until some later event. Only `stop` wakes it.
        let sweep_pool = std::sync::Arc::clone(&self.pool);
        let sweep_notify = Arc::clone(&self.sweep_notify);
        let sweep_stop = Arc::clone(&self.stop);
        handles.push(tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(30));
            loop {
                let notified = sweep_notify.notified();
                tokio::pin!(notified);
                tokio::select! {
                    _ = &mut notified => {}
                    _ = interval.tick() => {}
                }
                if sweep_stop.load(Ordering::Relaxed) {
                    break;
                }
                let result = sweep_pool.with_conn(move |conn| recover_update(conn)).await;
                if let Err(e) = result {
                    log::error!("queue sweep failed: {e}");
                }
            }
        }));
        *self.worker.lock() = handles;
    }

    pub async fn stop(&self) {
        self.stop.store(true, Ordering::Relaxed);
        self.notify.notify_waiters();
        self.sweep_notify.notify_waiters();
        let handles = std::mem::take(&mut *self.worker.lock());
        for handle in handles {
            let _ = handle.await;
        }
    }

    /// Persists a task. `run_after` is an absolute unix timestamp (seconds).
    /// Notifies the worker only after the insert has committed, so the worker
    /// never wakes to an invisible row.
    pub async fn enqueue(&self, payload: Value, run_after: f64) -> rusqlite::Result<()> {
        let id = format!(
            "task_{}_{}",
            (now_f64() * 1000.0) as u64,
            self.counter.fetch_add(1, Ordering::Relaxed)
        );
        let payload = payload.to_string();
        log::debug!("enqueued {id} (run_after {run_after:.1})");
        self.pool.with_conn(move |conn| {
            conn.execute(
                "INSERT OR REPLACE INTO tasks (id, payload, run_after, attempts, status, locked_until, created_at) \
                 VALUES (?1, ?2, ?3, 0, 'pending', 0, ?4)",
                params![id, payload, run_after, now_f64()],
            )?;
            Ok(())
        })
        .await?;
        // `notify_one` stores a permit when no worker is registered, so a
        // notification fired between a worker's DB reads and its `notified()`
        // registration is not lost (notify_waiters would drop it). The
        // awakened worker re-leases and finds the new row.
        self.notify.notify_one();
        Ok(())
    }

    async fn recover_stale(&self) {
        self.recover_sweep().await;
    }

    async fn recover_sweep(&self) {
        let result = self.pool.with_conn(move |conn| recover_update(conn)).await;
        if let Err(e) = result {
            log::error!("queue recovery failed: {e}");
        }
    }
}

impl QueueWorker {
    /// Supervised worker: the inner loop runs in its own task so a panic
    /// (e.g. inside a handler or a DB closure) kills only that task; the
    /// supervisor respawns it until stop is set. The row a dead worker had
    /// leased is recovered by the periodic sweep once its lock TTL expires.
    async fn run_loop_supervised(self) {
        while !self.stop.load(Ordering::Relaxed) {
            let worker = self.clone();
            if let Err(e) = tokio::spawn(async move { worker.run_loop().await }).await {
                log::error!("queue worker panicked, restarting: {e}");
            }
        }
    }

    async fn run_loop(self) {
        while !self.stop.load(Ordering::Relaxed) {
            match self.lease_next().await {
                Ok(Some(row)) => self.process(row).await,
                Ok(None) => {
                    let wait_until = self.earliest_run_after().await;
                    let notified = self.notify.notified();
                    tokio::pin!(notified);
                    match wait_until {
                        Some(until) => {
                            let delay = (until - now_f64()).max(0.0);
                            tokio::select! {
                                _ = &mut notified => {}
                                _ = tokio::time::sleep(Duration::from_secs_f64(delay)) => {}
                            }
                        }
                        None => {
                            notified.await;
                        }
                    }
                }
                // A lease failure while rows are due would otherwise loop
                // with sleep(0) and hammer SQLite; back off briefly.
                Err(e) => {
                    log::error!("queue lease failed: {e}");
                    tokio::time::sleep(Duration::from_secs(1)).await;
                }
            }
        }
    }

    /// Leases the oldest due row (sets it `in_progress` with a lock TTL).
    /// Errors are surfaced so the caller can back off instead of spinning.
    async fn lease_next(&self) -> Result<Option<LeasedRow>, rusqlite::Error> {
        self.pool.with_conn(|conn| {
            // BEGIN IMMEDIATE: with several workers, a deferred transaction
            // that read before another worker's lease commit would fail with
            // SQLITE_BUSY_SNAPSHOT. Taking the write lock up front serializes
            // leases and re-reads the freshest committed state.
            let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
            let now = now_f64();
            let row = tx.query_row(
                "SELECT id, payload, attempts FROM tasks WHERE status='pending' AND run_after <= ?1 AND locked_until <= ?1 \
                 ORDER BY run_after LIMIT 1",
                params![now],
                |r| {
                    Ok((
                        r.get::<_, String>(0)?,
                        r.get::<_, String>(1)?,
                        r.get::<_, i32>(2)?,
                    ))
                },
            );
            let (id, payload, attempts) = match row {
                Ok(row) => row,
                Err(rusqlite::Error::QueryReturnedNoRows) => {
                    tx.commit()?;
                    return Ok(None);
                }
                Err(e) => return Err(e),
            };
            tx.execute(
                "UPDATE tasks SET status='in_progress', locked_until=?1 WHERE id=?2",
                params![now + LOCK_TTL_SECONDS, id],
            )?;
            tx.commit()?;
            Ok(Some(LeasedRow {
                id,
                payload,
                attempts,
            }))
        })
        .await
    }

    async fn earliest_run_after(&self) -> Option<f64> {
        let result = self
            .pool
            .with_conn(|conn| {
                let mut stmt =
                    conn.prepare("SELECT MIN(run_after) FROM tasks WHERE status='pending'")?;
                let mut rows = stmt.query([])?;
                match rows.next()? {
                    Some(row) => Ok(row.get::<_, Option<f64>>(0)?),
                    None => Ok(None),
                }
            })
            .await;
        match result {
            Ok(v) => v,
            Err(e) => {
                log::error!("queue timing query failed: {e}");
                None
            }
        }
    }

    /// Processes one leased row, keeping the lease alive while the handler
    /// runs. Without the heartbeat a task longer than [`LOCK_TTL_SECONDS`]
    /// (slow download, ugoira encode, rate-limited batch forward) would have
    /// its lease expire mid-run; the expiry sweep would flip the row back to
    /// `pending` and another worker would process it again — duplicate sends.
    async fn process(&self, row: LeasedRow) {
        let payload: Value = match serde_json::from_str(&row.payload) {
            Ok(value) => value,
            Err(e) => {
                log::error!("queue: unparseable payload for {}: {e}", row.id);
                self.delete_row(&row.id).await;
                (self.dead_letter)(Value::Null, format!("invalid stored payload: {e}")).await;
                return;
            }
        };
        log::debug!("processing {} (attempt {})", row.id, row.attempts + 1);
        let outcome = self.run_with_lease(&row.id, payload).await;
        match outcome {
            Ok(()) => {
                log::debug!("task {} completed", row.id);
                self.delete_row(&row.id).await;
            }
            Err(QueueError::Retryable {
                delay_seconds,
                payload,
            }) => {
                if row.attempts as u32 >= MAX_RETRIES {
                    let message = format!("task failed after {MAX_RETRIES} retries");
                    log::error!("dead-lettering {}: {message}", row.id);
                    self.delete_row(&row.id).await;
                    (self.dead_letter)(payload, message).await;
                } else {
                    let delay = scaled_retry_delay(delay_seconds, row.attempts);
                    log::debug!(
                        "task {} rescheduled in {delay:.1}s (attempt {})",
                        row.id,
                        row.attempts + 1
                    );
                    self.reschedule(&row.id, payload, delay, row.attempts + 1)
                        .await;
                }
            }
            Err(QueueError::Permanent { message, payload }) => {
                log::error!("dead-lettering {}: {message}", row.id);
                self.delete_row(&row.id).await;
                (self.dead_letter)(payload, message).await;
            }
        }
    }

    /// Drives the handler to completion, refreshing the row's `locked_until`
    /// every 30 s so the expiry sweep never re-leases a still-running task.
    /// The heartbeat is part of this future, not a separate spawned task: if
    /// the worker task dies (panic) the heartbeat dies with it and the sweep
    /// recovers the row exactly as before.
    async fn run_with_lease(&self, id: &str, payload: Value) -> Result<(), QueueError> {
        let fut = (self.handler)(payload);
        tokio::pin!(fut);
        let mut interval = tokio::time::interval(Duration::from_secs(30));
        // The first interval tick fires immediately; skip it (the lease was
        // just set by lease_next).
        interval.tick().await;
        let id_owned = id.to_string();
        loop {
            tokio::select! {
                result = &mut fut => return result,
                _ = interval.tick() => {
                    let now = now_f64();
                    let id = id_owned.clone();
                    let result = self
                        .pool
                        .with_conn(move |conn| {
                            conn.execute(
                                "UPDATE tasks SET locked_until=?1 WHERE id=?2 AND status='in_progress'",
                                params![now + LOCK_TTL_SECONDS, id],
                            )
                        })
                        .await;
                    if let Err(e) = result {
                        log::error!("queue lease heartbeat failed: {e}");
                    }
                }
            }
        }
    }

    async fn delete_row(&self, id: &str) {
        let id = id.to_string();
        let result = self
            .pool
            .with_conn(move |conn| {
                conn.execute("DELETE FROM tasks WHERE id = ?1", params![id])?;
                Ok(())
            })
            .await;
        if let Err(e) = result {
            log::error!("queue delete failed: {e}");
        }
    }

    async fn reschedule(&self, id: &str, payload: Value, delay_seconds: f64, attempts: i32) {
        let id = id.to_string();
        let payload = payload.to_string();
        let result = self.pool.with_conn(move |conn| {
            conn.execute(
                "UPDATE tasks SET payload=?1, run_after=?2, attempts=?3, status='pending', locked_until=0 WHERE id=?4",
                params![payload, now_f64() + delay_seconds, attempts, id],
            )?;
            Ok(())
        })
        .await;
        if let Err(e) = result {
            log::error!("queue reschedule failed: {e}");
        }
        // Same permit semantics as enqueue: never lose the wakeup.
        self.notify.notify_one();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};

    #[test]
    fn scaled_retry_delay_scales_and_caps() {
        assert_eq!(scaled_retry_delay(1.0, 0), 1.0);
        assert_eq!(scaled_retry_delay(1.0, 1), 2.0);
        assert_eq!(scaled_retry_delay(1.0, 2), 4.0);
        assert_eq!(scaled_retry_delay(1.5, 1), 3.0);
        assert_eq!(scaled_retry_delay(1.0, 10), 300.0, "capped at 300s");
        assert_eq!(scaled_retry_delay(300.0, 0), 300.0);
    }

    async fn new_queue() -> (PersistentTaskQueue, tempfile::TempDir) {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("queue.db");
        let pool = crate::db::open_store(path.to_str().unwrap()).unwrap();
        let queue = PersistentTaskQueue::new(pool);
        (queue, dir)
    }

    #[tokio::test]
    async fn enqueue_runs_handler_once() {
        let (queue, _dir) = new_queue().await;
        let calls = Arc::new(AtomicUsize::new(0));
        let calls_worker = calls.clone();
        queue
            .start(
                move |payload| {
                    assert_eq!(payload["n"], 42);
                    calls_worker.fetch_add(1, AtomicOrdering::SeqCst);
                    async { Ok(()) }
                },
                |_payload, _message| async {},
            )
            .await;
        queue
            .enqueue(serde_json::json!({"n": 42}), now_f64())
            .await
            .unwrap();
        tokio::time::sleep(Duration::from_millis(300)).await;
        assert_eq!(calls.load(AtomicOrdering::SeqCst), 1);
        queue.stop().await;
    }

    #[tokio::test]
    async fn retryable_reschedules_then_dead_letters() {
        let (queue, _dir) = new_queue().await;
        let calls = Arc::new(AtomicUsize::new(0));
        let dead_calls = Arc::new(AtomicUsize::new(0));
        let c = calls.clone();
        let d = dead_calls.clone();
        queue
            .start(
                move |payload| {
                    c.fetch_add(1, AtomicOrdering::SeqCst);
                    let payload = payload.clone();
                    async move {
                        Err(QueueError::Retryable {
                            delay_seconds: 0.001,
                            payload,
                        })
                    }
                },
                move |_payload, _message| {
                    d.fetch_add(1, AtomicOrdering::SeqCst);
                    async {}
                },
            )
            .await;
        queue
            .enqueue(serde_json::json!({"a": 1}), now_f64())
            .await
            .unwrap();
        tokio::time::sleep(Duration::from_millis(600)).await;
        assert_eq!(
            calls.load(AtomicOrdering::SeqCst),
            MAX_RETRIES as usize + 1,
            "handler should run once per attempt"
        );
        assert_eq!(dead_calls.load(AtomicOrdering::SeqCst), 1);
        queue.stop().await;
    }

    #[tokio::test]
    async fn permanent_error_dead_letters_immediately() {
        let (queue, _dir) = new_queue().await;
        let calls = Arc::new(AtomicUsize::new(0));
        let dead_calls = Arc::new(AtomicUsize::new(0));
        let c = calls.clone();
        let d = dead_calls.clone();
        queue
            .start(
                move |payload| {
                    c.fetch_add(1, AtomicOrdering::SeqCst);
                    let payload = payload.clone();
                    async move {
                        Err(QueueError::Permanent {
                            message: "nope".into(),
                            payload,
                        })
                    }
                },
                move |_payload, message| {
                    assert_eq!(message, "nope");
                    d.fetch_add(1, AtomicOrdering::SeqCst);
                    async {}
                },
            )
            .await;
        queue
            .enqueue(serde_json::json!({"a": 1}), now_f64())
            .await
            .unwrap();
        tokio::time::sleep(Duration::from_millis(300)).await;
        assert_eq!(calls.load(AtomicOrdering::SeqCst), 1);
        assert_eq!(dead_calls.load(AtomicOrdering::SeqCst), 1);
        queue.stop().await;
    }

    #[tokio::test]
    async fn stale_in_progress_row_is_recovered_on_start() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("queue.db");
        // Insert a stale leased row directly (lease expired). open_store runs
        // the schema; the queue below shares the same pool.
        let pool = crate::db::open_store(path.to_str().unwrap()).unwrap();
        {
            let conn = rusqlite::Connection::open(&path).unwrap();
            conn.execute(
                "INSERT INTO tasks (id, payload, run_after, attempts, status, locked_until, created_at) \
                 VALUES ('task_stale', '{\"s\":1}', 0, 0, 'in_progress', ?1, 0)",
                params![now_f64() - 10.0],
            )
            .unwrap();
        }
        let queue = PersistentTaskQueue::new(pool);
        let calls = Arc::new(AtomicUsize::new(0));
        let c = calls.clone();
        queue
            .start(
                move |payload| {
                    assert_eq!(payload["s"], 1);
                    c.fetch_add(1, AtomicOrdering::SeqCst);
                    async { Ok(()) }
                },
                |_payload, _message| async {},
            )
            .await;
        tokio::time::sleep(Duration::from_millis(300)).await;
        assert_eq!(calls.load(AtomicOrdering::SeqCst), 1);
        queue.stop().await;
    }

    #[tokio::test]
    async fn runtime_sweep_recovers_expired_lease() {
        let (queue, _dir) = new_queue().await;
        let calls = Arc::new(AtomicUsize::new(0));
        let c = calls.clone();
        queue
            .start(
                move |payload| {
                    assert_eq!(payload["s"], 1);
                    c.fetch_add(1, AtomicOrdering::SeqCst);
                    async { Ok(()) }
                },
                |_payload, _message| async {},
            )
            .await;
        // Insert a stale leased row AFTER startup: without a runtime sweep it
        // would stay `in_progress` forever (only start() used to recover).
        {
            let conn = rusqlite::Connection::open(queue.pool.path()).unwrap();
            conn.execute(
                "INSERT INTO tasks (id, payload, run_after, attempts, status, locked_until, created_at) \
                 VALUES ('task_stale_runtime', '{\"s\":1}', 0, 0, 'in_progress', ?1, 0)",
                params![now_f64() - 1000.0],
            )
            .unwrap();
        }
        queue.recover_sweep().await;
        tokio::time::sleep(Duration::from_millis(300)).await;
        assert_eq!(
            calls.load(AtomicOrdering::SeqCst),
            1,
            "expired lease must be recovered and processed exactly once"
        );
        queue.stop().await;
    }
}