//! The lock around the replica (`AppState::doc`), counting how long its takers wait
//! (/admin/stats `document.lockWait`): whether the one document is what the node's
//! work queues behind.

use std::sync::Mutex;
use std::time::{Duration, Instant};

use serde::Serialize;
use tokio::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};

use super::replica::Replica;
use crate::locks::MutexExt;

/// Upper bounds (inclusive, microseconds) of the wait buckets; a last bucket takes the
/// rest.
const WAIT_BUCKETS_US: [u64; 7] = [10, 100, 1_000, 10_000, 100_000, 1_000_000, 10_000_000];

/// `tokio::sync::RwLock<Replica>`, with its waits counted.
pub(crate) struct DocumentLock {
    replica: RwLock<Replica>,
    waits: Mutex<LockWaits>,
}

/// How long takers of the lock waited for it, by mode, since start.
#[derive(Debug, Clone, Copy, Default, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct LockWaits {
    /// Readers: wait for a writer only.
    pub(crate) read: Waits,
    /// Writers (applying sync messages, generating them, writes, saves): wait for
    /// everyone.
    pub(crate) write: Waits,
}

#[derive(Debug, Clone, Copy, Default)]
pub(crate) struct Waits {
    /// Per bucket of `WAIT_BUCKETS_US`, then the rest; not cumulative.
    buckets: [u64; WAIT_BUCKETS_US.len() + 1],
    pub(crate) count: u64,
    pub(crate) sum_us: u64,
    pub(crate) max_us: u64,
}

impl Waits {
    fn record(&mut self, waited: Duration) {
        let us = u64::try_from(waited.as_micros()).unwrap_or(u64::MAX);
        let bucket = WAIT_BUCKETS_US
            .iter()
            .position(|&bound| us <= bound)
            .unwrap_or(WAIT_BUCKETS_US.len());
        self.buckets[bucket] += 1;
        self.count += 1;
        self.sum_us = self.sum_us.saturating_add(us);
        self.max_us = self.max_us.max(us);
    }

    /// Each bucket's upper bound in microseconds (None for the last) and count.
    pub(crate) fn buckets(&self) -> impl Iterator<Item = (Option<u64>, u64)> + '_ {
        WAIT_BUCKETS_US
            .iter()
            .copied()
            .map(Some)
            .chain([None])
            .zip(self.buckets.iter().copied())
    }
}

impl Serialize for Waits {
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        #[derive(Serialize)]
        #[serde(rename_all = "camelCase")]
        struct Bucket {
            le_us: Option<u64>,
            count: u64,
        }
        #[derive(Serialize)]
        #[serde(rename_all = "camelCase")]
        struct Out {
            count: u64,
            sum_us: u64,
            max_us: u64,
            /// Not cumulative; the last bucket has no bound.
            buckets: Vec<Bucket>,
        }
        Out {
            count: self.count,
            sum_us: self.sum_us,
            max_us: self.max_us,
            buckets: self
                .buckets()
                .map(|(le_us, count)| Bucket { le_us, count })
                .collect(),
        }
        .serialize(serializer)
    }
}

#[derive(Clone, Copy)]
enum Mode {
    Read,
    Write,
}

impl DocumentLock {
    pub(crate) fn new(replica: Replica) -> Self {
        Self {
            replica: RwLock::new(replica),
            waits: Mutex::default(),
        }
    }

    fn waited(&self, mode: Mode, since: Instant) {
        let waited = since.elapsed();
        let mut waits = self.waits.lock_or_recover();
        match mode {
            Mode::Read => waits.read.record(waited),
            Mode::Write => waits.write.record(waited),
        }
    }

    pub(crate) async fn read(&self) -> RwLockReadGuard<'_, Replica> {
        let since = Instant::now();
        let guard = self.replica.read().await;
        self.waited(Mode::Read, since);
        guard
    }

    pub(crate) async fn write(&self) -> RwLockWriteGuard<'_, Replica> {
        let since = Instant::now();
        let guard = self.replica.write().await;
        self.waited(Mode::Write, since);
        guard
    }

    /// Outside the runtime only (`spawn_blocking`), like `RwLock::blocking_read`.
    pub(crate) fn blocking_read(&self) -> RwLockReadGuard<'_, Replica> {
        let since = Instant::now();
        let guard = self.replica.blocking_read();
        self.waited(Mode::Read, since);
        guard
    }

    /// Outside the runtime only (`spawn_blocking`), like `RwLock::blocking_write`.
    pub(crate) fn blocking_write(&self) -> RwLockWriteGuard<'_, Replica> {
        let since = Instant::now();
        let guard = self.replica.blocking_write();
        self.waited(Mode::Write, since);
        guard
    }

    pub(crate) fn waits(&self) -> LockWaits {
        *self.waits.lock_or_recover()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn waits_fall_in_their_buckets() {
        let mut waits = Waits::default();
        waits.record(Duration::from_micros(5));
        waits.record(Duration::from_micros(10));
        waits.record(Duration::from_millis(2));
        waits.record(Duration::from_secs(60));
        let buckets: Vec<_> = waits.buckets().collect();
        assert_eq!(buckets[0], (Some(10), 2));
        assert_eq!(buckets[3], (Some(10_000), 1));
        assert_eq!(buckets[7], (None, 1));
        assert_eq!(waits.count, 4);
        assert_eq!(waits.sum_us, 5 + 10 + 2_000 + 60_000_000);
        assert_eq!(waits.max_us, 60_000_000);
    }
}