//! Leader locks (docs/roles.md, #272): a background task that must run
//! once across every process sharing the hub's database runs only in the
//! process holding a lock named for it.
//!
//! On PostgreSQL the lock is a session advisory lock
//! (`pg_try_advisory_lock`) on one connection of this process's own, kept
//! apart from the pool for as long as the process runs. A process that does
//! not hold a task's lock asks again every [`RETRY`]. One that holds it
//! checks its connection every [`CHECK`]; if the connection has gone (the
//! database restarted, the network dropped, the session was ended), the
//! server has freed the lock, so the task is stopped (its future dropped)
//! and the process contends again on a new connection. A process that dies
//! closes its connection, freeing its locks, and another takes over within
//! seconds. Each task reads its work from the database, so a successor
//! picking up where it stopped is safe.
//!
//! On SQLite (one machine, one process by definition) every lock is held.

use crate::Shared;
use std::collections::BTreeSet;
use std::future::Future;
use std::sync::Mutex;
use std::time::Duration;

/// How often a process that does not hold a lock asks for it again.
pub const RETRY: Duration = Duration::from_secs(2);

/// How often a holder makes sure its connection, and so its lock, is still
/// there; and how long it waits for the database's answer.
pub const CHECK: Duration = Duration::from_secs(2);

/// This process's background tasks and the locks it holds.
#[derive(Default)]
pub struct Leader {
    /// Every background task this process started, led or not.
    tasks: Mutex<BTreeSet<&'static str>>,
    /// The singleton tasks running here now: the locks this process holds.
    leading: Mutex<BTreeSet<&'static str>>,
    /// The loops contending for locks, stopped by [`Leader::resign`].
    contenders: Mutex<Vec<tokio::task::AbortHandle>>,
    #[cfg(feature = "postgres")]
    session: pg::Session,
}

impl Leader {
    /// The background tasks this process started, by name.
    pub fn tasks(&self) -> Vec<&'static str> {
        self.tasks.lock().unwrap_or_else(|p| p.into_inner()).iter().copied().collect()
    }

    /// The leader locks this process holds: the singleton tasks running here.
    pub fn leading(&self) -> Vec<&'static str> {
        self.leading.lock().unwrap_or_else(|p| p.into_inner()).iter().copied().collect()
    }

    /// Record a background task that needs no lock (it is safe anywhere).
    pub fn note(&self, name: &'static str) {
        self.tasks.lock().unwrap_or_else(|p| p.into_inner()).insert(name);
    }

    /// Stop every singleton task here and give up their locks, as a process
    /// that stops does: another process takes them over.
    pub async fn resign(&self) {
        for c in self.contenders.lock().unwrap_or_else(|p| p.into_inner()).drain(..) {
            c.abort();
        }
        #[cfg(feature = "postgres")]
        self.session.close().await;
    }

    /// The process id of the database session holding this process's
    /// locks, if it has one (for tests that end it).
    #[cfg(feature = "postgres")]
    pub async fn session_pid(&self) -> Option<i32> {
        self.session.pid().await
    }
}

// Taking, checking and giving back a lock: on PostgreSQL through the
// session; on SQLite every lock is this process's.
#[cfg(feature = "postgres")]
impl Leader {
    /// Take `name`'s lock now if it is free: the session's generation it
    /// was taken in, which [`Leader::held`] checks against.
    async fn acquire(&self, app: &Shared, name: &'static str) -> Option<u64> {
        self.session.acquire(&app.db, name).await
    }

    /// Whether a lock taken in generation `term` is still held.
    async fn held(&self, term: u64) -> bool {
        self.session.held(term).await
    }

    /// Give `name`'s lock back, its task done.
    async fn release(&self, name: &'static str, term: u64) {
        self.session.release(name, term).await
    }
}

#[cfg(not(feature = "postgres"))]
impl Leader {
    async fn acquire(&self, _app: &Shared, _name: &'static str) -> Option<u64> {
        Some(0)
    }

    async fn held(&self, _term: u64) -> bool {
        true
    }

    async fn release(&self, _name: &'static str, _term: u64) {}
}

/// Marks a task as led here while it runs; dropping it (the task ended,
/// was stopped, or the loop was aborted) unmarks it.
struct Leading<'a>(&'a Leader, &'static str);

impl<'a> Leading<'a> {
    fn new(leader: &'a Leader, name: &'static str) -> Self {
        leader.leading.lock().unwrap_or_else(|p| p.into_inner()).insert(name);
        Leading(leader, name)
    }
}

impl Drop for Leading<'_> {
    fn drop(&mut self) {
        self.0.leading.lock().unwrap_or_else(|p| p.into_inner()).remove(self.1);
    }
}

/// Run the singleton task `name` in the background, only while this
/// process holds its lock. `task` makes the task's future; it is made
/// again each time the lock is taken after being lost. A task that ends by
/// itself gives its lock back and is not run again.
pub fn run<F, Fut>(app: &Shared, name: &'static str, task: F)
where
    F: Fn(Shared) -> Fut + Send + 'static,
    Fut: Future<Output = ()> + Send + 'static,
{
    app.leader.note(name);
    let handle = tokio::spawn(contend(app.clone(), name, task));
    app.leader.contenders.lock().unwrap_or_else(|p| p.into_inner()).push(handle.abort_handle());
}

async fn contend<F, Fut>(app: Shared, name: &'static str, task: F)
where
    F: Fn(Shared) -> Fut,
    Fut: Future<Output = ()>,
{
    let leader = &app.leader;
    let mut waiting = false;
    loop {
        let Some(term) = leader.acquire(&app, name).await else {
            if !waiting {
                tracing::info!("leader: {name} runs elsewhere; standing by");
                waiting = true;
            }
            tokio::time::sleep(RETRY).await;
            continue;
        };
        if waiting {
            tracing::info!("leader: took over {name}");
            waiting = false;
        }
        let _leading = Leading::new(leader, name);
        let work = task(app.clone());
        tokio::pin!(work);
        let finished = loop {
            tokio::select! {
                _ = &mut work => break true,
                _ = tokio::time::sleep(CHECK) => {
                    if !leader.held(term).await {
                        break false;
                    }
                }
            }
        };
        if finished {
            leader.release(name, term).await;
            return;
        }
        // The future is dropped here: the task stops where it was.
        tracing::warn!("leader: lost the lock for {name}; stopped it, contending again");
    }
}

#[cfg(feature = "postgres")]
mod pg {
    use super::CHECK;
    use crate::db::Pool;
    use sqlx::postgres::PgConnection;
    use sqlx::Connection;
    use std::time::Instant;

    /// The connection holding this process's locks.
    struct Conn {
        conn: PgConnection,
        /// The schema the hub's tables are in: two hubs in two schemas of
        /// one database are two hubs, and do not share locks.
        schema: String,
        pid: i32,
        /// When the database last answered on it.
        checked: Instant,
    }

    #[derive(Default)]
    pub(super) struct Session {
        conn: tokio::sync::Mutex<Option<Conn>>,
        /// Counts connections lost: a lock taken in an earlier generation
        /// went with its connection.
        term: std::sync::atomic::AtomicU64,
        closed: std::sync::atomic::AtomicBool,
    }

    /// The lock's key: a stable 64-bit hash of the schema and the name.
    pub(super) fn key(schema: &str, name: &str) -> i64 {
        let h = blake3::hash(format!("oxygenhub leader\0{schema}\0{name}").as_bytes());
        i64::from_le_bytes(h.as_bytes()[..8].try_into().expect("eight bytes"))
    }

    impl Session {
        fn term(&self) -> u64 {
            self.term.load(std::sync::atomic::Ordering::SeqCst)
        }

        /// The connection is gone, and every lock with it.
        fn lose(&self, slot: &mut Option<Conn>) {
            if slot.take().is_some() {
                self.term.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
            }
        }

        async fn connect(db: &Pool) -> anyhow::Result<Conn> {
            let options = db.connect_options();
            let mut conn = PgConnection::connect_with(&options).await?;
            let (schema, pid): (String, i32) =
                sqlx::query_as("SELECT CAST(current_schema() AS TEXT), pg_backend_pid()").fetch_one(&mut conn).await?;
            Ok(Conn { conn, schema, pid, checked: Instant::now() })
        }

        pub(super) async fn acquire(&self, db: &Pool, name: &'static str) -> Option<u64> {
            let mut slot = self.conn.lock().await;
            if self.closed.load(std::sync::atomic::Ordering::SeqCst) {
                return None;
            }
            if slot.is_none() {
                match tokio::time::timeout(CHECK * 5, Self::connect(db)).await {
                    Ok(Ok(c)) => *slot = Some(c),
                    Ok(Err(e)) => {
                        tracing::warn!("leader: connecting for {name}'s lock: {e}");
                        return None;
                    }
                    Err(_) => {
                        tracing::warn!("leader: connecting for {name}'s lock: timed out");
                        return None;
                    }
                }
            }
            let term = self.term();
            let c = slot.as_mut()?;
            let ask = sqlx::query_scalar::<_, bool>("SELECT pg_try_advisory_lock($1)").bind(key(&c.schema, name));
            match tokio::time::timeout(CHECK, ask.fetch_one(&mut c.conn)).await {
                Ok(Ok(got)) => {
                    c.checked = Instant::now();
                    got.then_some(term)
                }
                Ok(Err(e)) => {
                    tracing::warn!("leader: asking for {name}'s lock: {e}");
                    self.lose(&mut slot);
                    None
                }
                Err(_) => {
                    tracing::warn!("leader: asking for {name}'s lock: timed out");
                    self.lose(&mut slot);
                    None
                }
            }
        }

        pub(super) async fn held(&self, term: u64) -> bool {
            let mut slot = self.conn.lock().await;
            if self.term() != term {
                return false;
            }
            let Some(c) = slot.as_mut() else { return false };
            // One answer serves every task's check made soon after it.
            if c.checked.elapsed() < CHECK / 2 {
                return true;
            }
            let ping = sqlx::query_scalar::<_, i32>("SELECT 1").fetch_one(&mut c.conn);
            match tokio::time::timeout(CHECK, ping).await {
                Ok(Ok(_)) => {
                    c.checked = Instant::now();
                    true
                }
                _ => {
                    tracing::warn!("leader: the connection holding the locks is gone");
                    self.lose(&mut slot);
                    false
                }
            }
        }

        pub(super) async fn release(&self, name: &'static str, term: u64) {
            let mut slot = self.conn.lock().await;
            if self.term() != term {
                return;
            }
            let Some(c) = slot.as_mut() else { return };
            let unlock = sqlx::query_scalar::<_, bool>("SELECT pg_advisory_unlock($1)").bind(key(&c.schema, name));
            if !matches!(tokio::time::timeout(CHECK, unlock.fetch_one(&mut c.conn)).await, Ok(Ok(_))) {
                self.lose(&mut slot);
            }
        }

        pub(super) async fn close(&self) {
            self.closed.store(true, std::sync::atomic::Ordering::SeqCst);
            let mut slot = self.conn.lock().await;
            if let Some(c) = slot.take() {
                self.term.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
                let _ = c.conn.close().await;
            }
        }

        pub(super) async fn pid(&self) -> Option<i32> {
            self.conn.lock().await.as_ref().map(|c| c.pid)
        }
    }

    #[cfg(test)]
    mod tests {
        #[test]
        fn keys_are_stable_and_per_schema() {
            assert_eq!(super::key("hub_a", "flow.scheduler"), super::key("hub_a", "flow.scheduler"));
            assert_ne!(super::key("hub_a", "flow.scheduler"), super::key("hub_b", "flow.scheduler"));
            assert_ne!(super::key("hub_a", "flow.scheduler"), super::key("hub_a", "flow.reaper"));
        }
    }
}