liminal-server 0.19.0

use std::collections::BTreeMap;
use std::future::Future;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, OnceLock, Weak};
use std::time::Duration;

use liminal::durability::{StoreCountObserver, StoreCounts};
use serde::Serialize;
use tokio::sync::Notify;

use super::events::{PublishOrigin, ServerEvent};
use super::sink::SelfStreams;
use crate::ServerError;
use crate::server::connection::services::LiminalConnectionServices;

pub(in crate::server) const fn build_commit() -> &'static str {
    env!("LIMINAL_BUILD_COMMIT")
}

/// The metadata shared by both records from one server incarnation.
#[derive(Clone, Debug)]
pub(in crate::server) struct ServerMetadata {
    device: String,
    publisher: String,
    server_incarnation: u64,
}

#[derive(Debug, thiserror::Error)]
pub(in crate::server) enum RecordError {
    #[error("self-stream device is empty")]
    EmptyDevice,
    #[error("self-stream lifecycle transition is empty")]
    EmptyLifecycleTransition,
    #[error("self-stream pressure response is empty")]
    EmptyPressureResponse,
    #[error("self-stream flush store is empty")]
    EmptyFlushStore,
    #[error("self-stream JSON encoding failed: {0}")]
    Encoding(#[from] serde_json::Error),
}

impl ServerMetadata {
    pub(in crate::server) fn from_config(
        config: &crate::config::types::ServerConfig,
        server_incarnation: u64,
    ) -> Result<Self, RecordError> {
        let device = config
            .cluster
            .as_ref()
            .map_or_else(|| "local".to_owned(), |cluster| cluster.node_name.clone());
        Self::new(device, server_incarnation)
    }

    pub(in crate::server) fn new(
        device: String,
        server_incarnation: u64,
    ) -> Result<Self, RecordError> {
        if device.trim().is_empty() {
            return Err(RecordError::EmptyDevice);
        }
        Ok(Self {
            publisher: format!("liminal.{device}"),
            device,
            server_incarnation,
        })
    }

    pub(in crate::server) fn device(&self) -> &str {
        &self.device
    }

    pub(in crate::server) fn publisher(&self) -> &str {
        &self.publisher
    }

    pub(in crate::server) const fn server_incarnation(&self) -> u64 {
        self.server_incarnation
    }
}

/// Cumulative work attributed to the server's own publications.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub(in crate::server) struct SelfWork {
    pub(in crate::server) rows: u64,
    pub(in crate::server) flushes: u64,
    /// Truncate calls the sink made to keep a self stream within its bound.
    pub(in crate::server) prunes: u64,
    /// Rows those truncate calls deleted.
    pub(in crate::server) rows_pruned: u64,
    /// Truncate calls that failed (each logged as `SELF_STREAM_PRUNE_FAILED`).
    pub(in crate::server) prune_failures: u64,
}

impl SelfWork {
    const fn since(self, earlier: Self) -> Self {
        Self {
            rows: self.rows.saturating_sub(earlier.rows),
            flushes: self.flushes.saturating_sub(earlier.flushes),
            prunes: self.prunes.saturating_sub(earlier.prunes),
            rows_pruned: self.rows_pruned.saturating_sub(earlier.rows_pruned),
            prune_failures: self.prune_failures.saturating_sub(earlier.prune_failures),
        }
    }
}

/// A prepared record retains the exact relaxed snapshot it describes.
#[derive(Debug)]
pub(in crate::server) struct PreparedCost {
    previous_snapshot: StoreCounts,
    snapshot: StoreCounts,
    self_work: SelfWork,
    delta: StoreCounts,
    self_delta: SelfWork,
}

impl PreparedCost {
    #[cfg(test)]
    pub(in crate::server) const fn delta(&self) -> StoreCounts {
        self.delta
    }

    pub(in crate::server) fn encode(
        &self,
        metadata: &ServerMetadata,
        server_time_ms: u64,
    ) -> Result<Vec<u8>, RecordError> {
        let record = CostRecord {
            v: 1,
            build_commit: build_commit(),
            counts: self.delta.named().into_iter().collect(),
            crate_version: env!("CARGO_PKG_VERSION"),
            device: metadata.device(),
            previous_snapshot: self.previous_snapshot.named().into_iter().collect(),
            publisher: metadata.publisher(),
            self_flushes: self.self_delta.flushes,
            self_prune_failures: self.self_delta.prune_failures,
            self_prunes: self.self_delta.prunes,
            self_rows: self.self_delta.rows,
            self_rows_pruned: self.self_delta.rows_pruned,
            server_incarnation: metadata.server_incarnation(),
            server_time_ms,
            snapshot: self.snapshot.named().into_iter().collect(),
        };
        Ok(serde_json::to_vec(&record)?)
    }
}

#[derive(Serialize)]
struct CostRecord<'a> {
    v: u8,
    build_commit: &'static str,
    counts: BTreeMap<&'static str, u64>,
    crate_version: &'static str,
    device: &'a str,
    previous_snapshot: BTreeMap<&'static str, u64>,
    publisher: &'a str,
    self_flushes: u64,
    self_prune_failures: u64,
    self_prunes: u64,
    self_rows: u64,
    self_rows_pruned: u64,
    server_incarnation: u64,
    server_time_ms: u64,
    snapshot: BTreeMap<&'static str, u64>,
}

/// Publication success, rather than encoding, advances the reporting boundary.
#[derive(Debug)]
pub(in crate::server) struct CostWindow {
    previous: StoreCounts,
    previous_self: SelfWork,
}

impl CostWindow {
    pub(in crate::server) const fn new(snapshot: StoreCounts, self_work: SelfWork) -> Self {
        Self {
            previous: snapshot,
            previous_self: self_work,
        }
    }

    pub(in crate::server) fn prepare(
        &self,
        snapshot: StoreCounts,
        self_work: SelfWork,
    ) -> Option<PreparedCost> {
        let delta = snapshot.since(&self.previous);
        if delta.named().into_iter().all(|(_, value)| value == 0) {
            return None;
        }
        Some(PreparedCost {
            previous_snapshot: self.previous,
            snapshot,
            self_work,
            delta,
            self_delta: self_work.since(self.previous_self),
        })
    }

    pub(in crate::server) const fn published(&mut self, record: &PreparedCost) {
        self.previous = record.snapshot;
        self.previous_self = record.self_work;
    }
}

/// Coalesces external work into one timer, with no timer while idle.
#[derive(Debug, Default)]
pub(in crate::server) struct CostWake {
    wake: Notify,
    armed: AtomicBool,
    /// Every arm for connection-caused work, counted. Test-only.
    #[cfg(test)]
    arms: std::sync::atomic::AtomicU64,
    external_publication: AtomicBool,
    source: OnceLock<WakeSource>,
    #[cfg(test)]
    notifications: std::sync::atomic::AtomicU64,
    #[cfg(test)]
    timer_pending: AtomicBool,
    #[cfg(test)]
    waiting: Notify,
    #[cfg(test)]
    timer_started: Notify,
    /// While set, a due tick is never served (the owner's stop still ends
    /// it): a test holds the server's cost rows out of its windows. Test-only.
    #[cfg(test)]
    held: AtomicBool,
}

#[derive(Debug)]
struct WakeSource {
    services: Weak<LiminalConnectionServices>,
    streams: Weak<SelfStreams>,
    baseline: Mutex<StoreCounts>,
}

impl CostWake {
    pub(in crate::server) fn install_source(
        &self,
        services: &Arc<LiminalConnectionServices>,
        streams: &Arc<SelfStreams>,
    ) -> Result<(), ServerError> {
        self.source
            .set(WakeSource {
                services: Arc::downgrade(services),
                streams: Arc::downgrade(streams),
                baseline: Mutex::new(services.cost_snapshot()),
            })
            .map_err(|_| SelfStreams::failure("cost wake source already installed"))
    }

    pub(in crate::server) fn self_append_total(
        &self,
        snapshot: StoreCounts,
    ) -> Result<(), ServerError> {
        if let Some(source) = self.source.get() {
            *source.baseline.lock().map_err(|error| {
                SelfStreams::failure(&format!("cost wake baseline lock poisoned: {error}"))
            })? = snapshot;
        }
        Ok(())
    }

    pub(in crate::server) fn arm(&self, origin: &PublishOrigin<'_>) {
        if !origin.is_server() {
            self.arm_for_work();
        }
    }

    pub(in crate::server) fn arm_for_work(&self) {
        #[cfg(test)]
        self.arms.fetch_add(1, Ordering::Relaxed);
        self.external_publication.store(true, Ordering::Relaxed);
        self.changed();
    }

    fn changed_since_self_append(&self) -> Result<bool, ServerError> {
        let external_publication = self.external_publication.swap(false, Ordering::Relaxed);
        let Some(source) = self.source.get() else {
            return Ok(true);
        };
        let services = source
            .services
            .upgrade()
            .ok_or_else(|| SelfStreams::failure("cost wake services owner is gone"))?;
        let snapshot = services.cost_snapshot();
        let mut baseline = source.baseline.lock().map_err(|error| {
            SelfStreams::failure(&format!("cost wake baseline lock poisoned: {error}"))
        })?;
        let changed = *baseline != snapshot;
        *baseline = snapshot;
        drop(baseline);
        Ok(external_publication || changed)
    }

    pub(in crate::server) async fn next_tick<Sleep, SleepFuture>(
        &self,
        sleep: Sleep,
    ) -> Result<(), ServerError>
    where
        Sleep: FnOnce(Duration) -> SleepFuture + Send,
        SleepFuture: Future<Output = ()> + Send,
    {
        loop {
            self.armed.store(false, Ordering::Relaxed);
            #[cfg(test)]
            self.waiting.notify_one();
            self.wake.notified().await;
            if self.changed_since_self_append()? {
                break;
            }
        }
        self.armed.store(true, Ordering::Relaxed);
        let timer = ArmedTimer { wake: self };
        #[cfg(test)]
        self.timer_pending.store(true, Ordering::Release);
        #[cfg(test)]
        self.timer_started.notify_one();
        sleep(Duration::from_secs(10)).await;
        #[cfg(test)]
        if self.held.load(Ordering::Acquire) {
            std::future::pending::<()>().await;
        }
        #[cfg(test)]
        self.timer_pending.store(false, Ordering::Release);
        drop(timer);
        Ok(())
    }

    #[cfg(test)]
    pub(super) fn external_publication(&self) -> bool {
        self.external_publication.load(Ordering::Relaxed)
    }

    #[cfg(test)]
    pub(super) fn notifications(&self) -> u64 {
        self.notifications.load(Ordering::Relaxed)
    }

    /// Holds every due tick for the rest of this wake's life. Test-only.
    #[cfg(test)]
    pub(in crate::server) fn hold(&self) {
        self.held.store(true, Ordering::Release);
    }

    #[cfg(test)]
    pub(in crate::server) fn armed(&self) -> bool {
        self.armed.load(Ordering::Acquire)
    }

    /// How many times connection-caused work armed the tick.
    #[cfg(test)]
    pub(in crate::server) fn arms(&self) -> u64 {
        self.arms.load(Ordering::Relaxed)
    }

    #[cfg(test)]
    pub(in crate::server) fn timer_pending(&self) -> bool {
        self.timer_pending.load(Ordering::Acquire)
    }

    #[cfg(test)]
    pub(super) async fn waiting(&self) {
        self.waiting.notified().await;
    }

    #[cfg(test)]
    pub(super) async fn timer_started(&self) {
        self.timer_started.notified().await;
    }
}

impl StoreCountObserver for CostWake {
    fn changed(&self) {
        if !self.armed.swap(true, Ordering::Relaxed) {
            #[cfg(test)]
            self.notifications.fetch_add(1, Ordering::Relaxed);
            self.wake.notify_one();
        }
    }

    fn flushed(&self, store: &str, batch_rows: Option<u64>) {
        if let Some(streams) = self
            .source
            .get()
            .and_then(|source| source.streams.upgrade())
        {
            streams.event(ServerEvent::Flush { store, batch_rows });
        }
    }
}

struct ArmedTimer<'a> {
    wake: &'a CostWake,
}

impl Drop for ArmedTimer<'_> {
    fn drop(&mut self) {
        self.wake.armed.store(false, Ordering::Release);
        #[cfg(test)]
        self.wake.timer_pending.store(false, Ordering::Release);
    }
}