liminal-server 0.25.4

use std::cell::Cell;
use std::collections::BTreeMap;
use std::future::Future;
use std::marker::PhantomData;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, 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;

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,
    /// Work the server did not do for itself, since the tick last began:
    /// a connection's publication, or store work made off the drainer's
    /// thread.
    external_work: 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 {
    streams: Weak<SelfStreams>,
}

thread_local! {
    /// Held on the self-record drainer's thread for its whole life: every
    /// store call made there is the server's own
    /// (docs/design/SELF-STREAMS-COST-WAKE-2026-10-07.md).
    static SELF_WORK: Cell<bool> = const { Cell::new(false) };
    static READER_WORK: Cell<bool> = const { Cell::new(false) };
}

/// Marks the store work of the thread that holds it as the server's own,
/// until it is dropped. It is `!Send`, so it never leaves that thread.
pub(in crate::server) struct SelfWorkMark {
    _thread: PhantomData<*const ()>,
}

impl SelfWorkMark {
    pub(in crate::server) fn enter() -> Self {
        SELF_WORK.with(|mark| mark.set(true));
        Self {
            _thread: PhantomData,
        }
    }
}

impl Drop for SelfWorkMark {
    fn drop(&mut self) {
        SELF_WORK.with(|mark| mark.set(false));
    }
}

/// Keeps a self-stream observer's store work from producing more self-stream work.
pub struct ReaderWork {
    previous: bool,
    _thread: PhantomData<*const ()>,
}

impl ReaderWork {
    pub(crate) fn enter() -> Self {
        Self {
            previous: READER_WORK.with(|mark| mark.replace(true)),
            _thread: PhantomData,
        }
    }
}

impl Drop for ReaderWork {
    fn drop(&mut self) {
        READER_WORK.with(|mark| mark.set(self.previous));
    }
}

pub fn install_bridge_reader(
    services: &crate::server::connection::services::LiminalConnectionServices,
    feeds: &Arc<crate::link_bridge::subscriptions::ExportSubscriptions>,
) -> Result<(), ServerError> {
    let wake = services.signals().wake();
    if let Some(streams) = wake
        .source
        .get()
        .and_then(|source| source.streams.upgrade())
    {
        streams.observe_bridge(services, feeds)?;
    }
    Ok(())
}

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

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

    /// Connection-caused work: external by its caller's word, wherever
    /// it runs.
    pub(in crate::server) fn arm_for_work(&self) {
        #[cfg(test)]
        self.arms.fetch_add(1, Ordering::Relaxed);
        self.external();
    }

    /// Records external work and wakes the tick, unless a tick is already
    /// armed and will count it.
    fn external(&self) {
        self.external_work.store(true, Ordering::Relaxed);
        if !self.armed.swap(true, Ordering::Relaxed) {
            #[cfg(test)]
            self.notifications.fetch_add(1, Ordering::Relaxed);
            self.wake.notify_one();
        }
    }

    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.external_work.swap(false, Ordering::Relaxed) {
                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_work(&self) -> bool {
        self.external_work.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 {
    /// Store work wakes the tick unless the drainer made it, on its own
    /// thread, for the server's own streams.
    fn changed(&self) {
        if !SELF_WORK.with(Cell::get) && !READER_WORK.with(Cell::get) {
            self.external();
        }
    }

    fn flushed(&self, store: &str, batch_rows: Option<u64>) {
        if READER_WORK.with(Cell::get) {
            return;
        }
        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);
    }
}