add document lock stats

1 files changed, 175 insertions, 0 deletions

diff --git a/src/document/lock.rs b/src/document/lock.rs
new file mode 100644
index 0000000..a394b52
--- /dev/null
+++ b/src/document/lock.rs
@@ -0,0 +1,175 @@
+//! 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);
+    }
+}