//! 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;

/// A lock across every process sharing the hub's database, held while the

/// guard lives (#275): one evaluation of a proposal's auto-land, one bump

/// or landing of a consumer's pin, at a time anywhere. On PostgreSQL a

/// transaction-scoped advisory lock (`pg_advisory_xact_lock`), freed when

/// the guard drops and its transaction ends, or with its connection; on

/// SQLite (one process by definition) a lock in this process.

pub struct Exclusive {

    #[cfg(feature = "postgres")]

    _tx: sqlx::Transaction<'static, sqlx::Postgres>,

    #[cfg(not(feature = "postgres"))]

    _guard: tokio::sync::OwnedMutexGuard<()>,

}

/// Wait for the lock named `name`, then hold it until the guard drops.

#[cfg(feature = "postgres")]

pub async fn exclusive(pool: &crate::db::Pool, name: &str) -> sqlx::Result<Exclusive> {

    let mut tx = pool.begin().await?;

    let schema: String = sqlx::query_scalar("SELECT CAST(current_schema() AS TEXT)").fetch_one(&mut *tx).await?;

    sqlx::query("SELECT pg_advisory_xact_lock($1)").bind(pg::key(&schema, name)).execute(&mut *tx).await?;

    Ok(Exclusive { _tx: tx })

}

/// Wait for the lock named `name`, then hold it until the guard drops.

#[cfg(not(feature = "postgres"))]

pub async fn exclusive(_pool: &crate::db::Pool, name: &str) -> sqlx::Result<Exclusive> {

    static LOCAL: Mutex<Option<std::collections::HashMap<String, std::sync::Arc<tokio::sync::Mutex<()>>>>> = Mutex::new(None);

    let lock = {

        let mut map = LOCAL.lock().unwrap_or_else(|p| p.into_inner());

        let map = map.get_or_insert_with(Default::default);

        // Locks nobody holds or waits for are forgotten.

        map.retain(|_, l| std::sync::Arc::strong_count(l) > 1);

        map.entry(name.to_string()).or_default().clone()

    };

    Ok(Exclusive { _guard: lock.lock_owned().await })

}

/// 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"));

        }

    }

}