haematite-cli 0.21.7

use std::fmt;
use std::ops::{Deref, DerefMut};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Condvar, Mutex, MutexGuard, TryLockError};
use std::time::{Duration, Instant};

use crate::serve::storage::Failure;

#[derive(Default)]
struct State {
    foreground: bool,
    background: bool,
    #[cfg(test)]
    waiting_background: usize,
}

pub(in crate::serve) struct Admission {
    state: Mutex<State>,
    changed: Condvar,
    wait: Duration,
    foreground: AtomicBool,
}

fn unavailable() -> Failure {
    Failure::new(
        "writer_unavailable",
        "writer admission requires process recovery",
    )
}

impl Admission {
    pub fn new(wait: Duration) -> Result<Self, Failure> {
        if wait.is_zero() || Instant::now().checked_add(wait).is_none() {
            return Err(Failure::new(
                "embed_writer_wait_invalid",
                "embedding writer wait must be positive and representable",
            ));
        }
        Ok(Self {
            state: Mutex::new(State::default()),
            changed: Condvar::new(),
            wait,
            foreground: AtomicBool::new(false),
        })
    }

    /// A waiting foreground writer excludes every later embedding commit.
    /// Only one already-admitted batch, sized by `embed_page_records`, may
    /// finish its one durable commit and full flush before this writer.
    pub fn foreground<'a, T>(
        &'a self,
        storage: &'a Mutex<T>,
        releasing: bool,
    ) -> Result<Guard<'a, T>, Failure> {
        self.foreground_with_deadline(storage, releasing, None)
    }

    /// A served request may wait no longer than its absolute deadline.
    pub fn foreground_until<'a, T>(
        &'a self,
        storage: &'a Mutex<T>,
        deadline: Instant,
    ) -> Result<Guard<'a, T>, Failure> {
        self.foreground_with_deadline(storage, false, Some(deadline))
    }

    fn foreground_with_deadline<'a, T>(
        &'a self,
        storage: &'a Mutex<T>,
        releasing: bool,
        deadline: Option<Instant>,
    ) -> Result<Guard<'a, T>, Failure> {
        let expired = || deadline.is_some_and(|end| Instant::now() >= end);
        let busy = || Failure::new("busy", "embedding writer wait reached the request deadline");
        if expired() {
            return Err(busy());
        }
        let mut reservation = Reservation {
            admission: self,
            foreground: true,
            active: false,
        };
        let mut state = self.state.lock().map_err(|_| unavailable())?;
        if state.foreground {
            return Err(Failure::new(
                "busy",
                "another mutation is being admitted; retry this request identity",
            ));
        }
        state.foreground = true;
        self.foreground.store(true, Ordering::Release);
        reservation.active = true;
        self.changed.notify_all();
        let wait = deadline.map_or(self.wait, |end| {
            self.wait.min(end.saturating_duration_since(Instant::now()))
        });
        let (state, timed) = self
            .changed
            .wait_timeout_while(state, wait, |state| state.background)
            .map_err(|_| unavailable())?;
        if expired() {
            drop(state);
            return Err(busy());
        }
        if timed.timed_out() && state.background {
            drop(state);
            return Err(Failure::new(
                "embed_admission_wait_exceeded",
                "the embedding commit did not release writer admission within embed_writer_wait_ms",
            ));
        }
        drop(state);
        let guard = match storage.try_lock() {
            Ok(guard) => guard,
            Err(TryLockError::WouldBlock) if releasing => {
                storage.lock().map_err(|_| unavailable())?
            }
            Err(TryLockError::WouldBlock) => {
                return Err(Failure::new(
                    "busy",
                    "another mutation is being admitted; retry this request identity",
                ));
            }
            Err(TryLockError::Poisoned(_)) => return Err(unavailable()),
        };
        Ok(Guard {
            guard,
            reservation: Some(reservation),
        })
    }

    pub fn background<'a, T>(&'a self, storage: &'a Mutex<T>) -> Result<Guard<'a, T>, Failure> {
        loop {
            let mut reservation = Reservation {
                admission: self,
                foreground: false,
                active: false,
            };
            let state = self.state.lock().map_err(|_| unavailable())?;
            #[cfg(test)]
            let state = {
                let mut state = state;
                state.waiting_background += 1;
                self.changed.notify_all();
                state
            };
            let mut state = self
                .changed
                .wait_while(state, |state| state.foreground || state.background)
                .map_err(|_| unavailable())?;
            #[cfg(test)]
            {
                state.waiting_background -= 1;
            }
            state.background = true;
            reservation.active = true;
            drop(state);
            let guard = storage.lock().map_err(|_| unavailable())?;
            // A release can own storage before this reservation reaches it.
            // Yield to a writer that arrived in that gap, then wait on its signal.
            if self.foreground.load(Ordering::Acquire) {
                drop(guard);
                drop(reservation);
                continue;
            }
            return Ok(Guard {
                guard,
                reservation: Some(reservation),
            });
        }
    }

    #[cfg(test)]
    pub fn wait_foreground(&self) -> Result<(), Failure> {
        let state = self.state.lock().map_err(|_| unavailable())?;
        let state = self
            .changed
            .wait_while(state, |state| !state.foreground)
            .map_err(|_| unavailable())?;
        drop(state);
        Ok(())
    }

    #[cfg(test)]
    pub fn wait_background(&self) -> Result<(), Failure> {
        let state = self.state.lock().map_err(|_| unavailable())?;
        let state = self
            .changed
            .wait_while(state, |state| state.waiting_background == 0)
            .map_err(|_| unavailable())?;
        drop(state);
        Ok(())
    }
}

struct Reservation<'a> {
    admission: &'a Admission,
    foreground: bool,
    active: bool,
}

impl Drop for Reservation<'_> {
    fn drop(&mut self) {
        if !self.active {
            return;
        }
        // Poison remains a named refusal at every admission boundary.
        if let Ok(mut state) = self.admission.state.lock() {
            if self.foreground {
                state.foreground = false;
                self.admission.foreground.store(false, Ordering::Release);
            } else {
                state.background = false;
            }
        }
        self.admission.changed.notify_all();
    }
}

pub(in crate::serve) struct Guard<'a, T> {
    // Field drop order releases storage before admission acquires its mutex.
    guard: MutexGuard<'a, T>,
    reservation: Option<Reservation<'a>>,
}

impl<'a, T> Guard<'a, T> {
    pub const fn plain(guard: MutexGuard<'a, T>) -> Self {
        Self {
            guard,
            reservation: None,
        }
    }
}

impl<T> fmt::Debug for Guard<'_, T> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("Guard")
            .field(
                "foreground",
                &self
                    .reservation
                    .as_ref()
                    .map(|reservation| reservation.foreground),
            )
            .finish_non_exhaustive()
    }
}

impl<T> Deref for Guard<'_, T> {
    type Target = T;
    fn deref(&self) -> &T {
        &self.guard
    }
}

impl<T> DerefMut for Guard<'_, T> {
    fn deref_mut(&mut self) -> &mut T {
        &mut self.guard
    }
}