liminal-server 0.21.1

use std::cell::Cell;
use std::collections::BTreeMap;
use std::sync::{Arc, Mutex, MutexGuard, Weak};
use std::thread::{self, JoinHandle};
use std::time::{SystemTime, UNIX_EPOCH};

use liminal::durability::{StoreCounts, StoreCountsHandle};
use liminal::protocol::{CausalContext, MessageEnvelope, SchemaId};

use crate::ServerError;
use crate::config::types::SelfStreamsConfig;
use crate::server::connection::services::{LiminalConnectionServices, PublishOutcome};
use crate::server::shutdown::ShutdownHandle;

use super::cost::{CostWake, PreparedCost, SelfWork, SelfWorkMark, ServerMetadata};
use super::events::{PublishOrigin, ServerEvent};
use super::publisher::CostSink;
use super::queue::{QueuedRecord, RecordPayload, RecordQueue};

thread_local! {
    static FINAL_FLUSH: Cell<bool> = const { Cell::new(false) };
}

#[derive(Debug)]
pub(in crate::server) struct OriginPublication {
    pub(in crate::server) outcome: PublishOutcome,
    pub(in crate::server) self_row_appended: bool,
    pub(in crate::server) self_row_at: Option<u64>,
}

impl OriginPublication {
    pub(in crate::server) const fn connection(outcome: PublishOutcome) -> Self {
        Self {
            outcome,
            self_row_appended: false,
            self_row_at: None,
        }
    }
}

pub(in crate::server) struct SelfStreams {
    metadata: ServerMetadata,
    counts: StoreCountsHandle,
    queue: Arc<RecordQueue>,
    services: Weak<LiminalConnectionServices>,
    drainer: Mutex<DrainerOwner>,
    wake: Arc<CostWake>,
    retention_work: Arc<Mutex<SelfWork>>,
}

struct DrainerOwner {
    worker: Option<RecordDrainer>,
    thread: Option<JoinHandle<()>>,
}

struct RecordDrainer {
    metadata: ServerMetadata,
    services: Weak<LiminalConnectionServices>,
    queue: Arc<RecordQueue>,
    own_work: Arc<Mutex<SelfWork>>,
    bounds: SelfStreamsConfig,
    retained_from: Mutex<BTreeMap<&'static str, u64>>,
}

impl std::fmt::Debug for SelfStreams {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("SelfStreams")
            .field("metadata", &self.metadata)
            .finish_non_exhaustive()
    }
}

impl SelfStreams {
    pub(in crate::server) fn new(
        metadata: ServerMetadata,
        services: &Arc<LiminalConnectionServices>,
        wake: Arc<CostWake>,
        shutdown: ShutdownHandle,
        bounds: SelfStreamsConfig,
    ) -> Self {
        let counts = services.store_counts();
        let queue = Arc::new(RecordQueue::new(shutdown));
        let retention_work = Arc::new(Mutex::new(SelfWork::default()));
        let worker = RecordDrainer {
            metadata: metadata.clone(),
            services: Arc::downgrade(services),
            queue: Arc::clone(&queue),
            own_work: Arc::clone(&retention_work),
            bounds,
            retained_from: Mutex::new(BTreeMap::new()),
        };
        Self {
            metadata,
            counts,
            queue,
            services: Arc::downgrade(services),
            wake,
            retention_work,
            drainer: Mutex::new(DrainerOwner {
                worker: Some(worker),
                thread: None,
            }),
        }
    }

    pub(in crate::server) fn start(&self) -> Result<(), ServerError> {
        let mut owner = self.drainer.lock().map_err(|error| {
            Self::failure(&format!("record drainer start lock poisoned: {error}"))
        })?;
        if owner.thread.is_some() {
            return Ok(());
        }
        let worker = owner
            .worker
            .take()
            .ok_or_else(|| Self::failure("record drainer unavailable after start failure"))?;
        owner.thread = Some(
            thread::Builder::new()
                .name("liminal-self-records".to_owned())
                .spawn(move || worker.run())
                .map_err(|error| Self::failure(&format!("record drainer construction: {error}")))?,
        );
        drop(owner);
        Ok(())
    }

    pub(in crate::server) fn failure(message: &str) -> ServerError {
        ServerError::ListenerAccept {
            message: format!("SELF_STREAM_EVENTS_APPEND_FAILED: {message}"),
        }
    }

    pub(in crate::server) fn drain(&self) -> Result<(), ServerError> {
        self.queue.drain()
    }

    /// Lets the drainer take `takes` more records, then holds it. Test-only.
    #[cfg(test)]
    pub(in crate::server) fn hold_drainer(&self, takes: usize) -> Result<(), ServerError> {
        self.queue.hold(takes)
    }

    /// Lets the drainer take every record again. Test-only.
    #[cfg(test)]
    pub(in crate::server) fn release_drainer(&self) -> Result<(), ServerError> {
        self.queue.release()
    }

    /// Waits until every queued record is written, and every record queued
    /// while those were written. Test-only.
    #[cfg(test)]
    pub(in crate::server) fn drain_all(&self) -> Result<(), ServerError> {
        self.queue.drain_all()
    }

    #[cfg(test)]
    pub(super) fn queue_for_test(&self) -> Arc<RecordQueue> {
        Arc::clone(&self.queue)
    }

    #[cfg(test)]
    pub(super) fn enqueue_event_for_test(&self, event: ServerEvent) -> Result<u64, ServerError> {
        let payload = event
            .encode(&self.metadata, server_time_ms()?)
            .map_err(|error| Self::failure(&error.to_string()))?;
        self.enqueue("liminal.events", payload, false)
    }

    #[cfg(test)]
    pub(super) fn join_drainer_for_test(&self) -> Result<(), ServerError> {
        self.queue.stop()?;
        let thread = self
            .drainer
            .lock()
            .map_err(|error| Self::failure(&format!("drainer join lock: {error}")))?
            .thread
            .take();
        if let Some(thread) = thread {
            thread
                .join()
                .map_err(|_| Self::failure("record drainer join panicked"))?;
        }
        Ok(())
    }

    pub(in crate::server) fn with_flush_barrier<T>(
        &self,
        flush: impl FnOnce() -> Result<T, ServerError>,
    ) -> Result<T, ServerError> {
        let before = self.queue.drain_all();
        if self.queue.was_abandoned() {
            before?;
            return Err(Self::failure("owned durable drain was abandoned"));
        }
        if FINAL_FLUSH.with(|active| active.replace(true)) {
            return Err(Self::failure("nested final flush barrier"));
        }
        let reset = FlushBarrier;
        let result = flush();
        drop(reset);
        let after = self.queue.drain_all();
        if let Err(error) = &before {
            tracing::error!(%error, "self-stream pre-flush drain failed");
        }
        if let Err(error) = &result {
            tracing::error!(%error, "self-stream backing-store flush failed");
        }
        if let Err(error) = &after {
            tracing::error!(%error, "self-stream post-flush drain failed");
        }
        let value = result?;
        before?;
        after?;
        Ok(value)
    }

    fn enqueue(
        &self,
        channel: &'static str,
        payload: Vec<u8>,
        #[cfg(test)] connection_opened: bool,
    ) -> Result<u64, ServerError> {
        let result = (|| {
            self.start()?;
            self.queue.enqueue(QueuedRecord {
                sequence: 0,
                channel,
                payload: RecordPayload::Single(payload),
                #[cfg(test)]
                connection_opened_count: usize::from(connection_opened),
            })
        })();
        if let Err(error) = &result {
            self.queue.fail(error);
        }
        result
    }

    pub(in crate::server) fn lifecycle_batch<'a>(
        &self,
        conversation: u64,
        receipts: impl IntoIterator<Item = (u64, &'a str)>,
    ) -> Result<(), ServerError> {
        let result = (|| {
            let time = server_time_ms()?;
            let payloads = receipts
                .into_iter()
                .map(|(generation, transition)| {
                    ServerEvent::Lifecycle {
                        conversation,
                        generation,
                        transition,
                    }
                    .encode(&self.metadata, time)
                    .map_err(|error| Self::failure(&error.to_string()))
                })
                .collect::<Result<Vec<_>, _>>()?;
            if payloads.is_empty() {
                return Ok(());
            }
            self.start()?;
            let sequence = self.queue.enqueue(QueuedRecord {
                sequence: 0,
                channel: "liminal.events",
                payload: RecordPayload::Batch(payloads),
                #[cfg(test)]
                connection_opened_count: 0,
            })?;
            self.wake.arm_for_work();
            self.queue.await_append(sequence)
        })();
        if let Err(error) = &result {
            self.queue.fail(error);
        }
        result
    }

    pub(in crate::server) fn event(&self, event: ServerEvent<'_>) {
        if !matches!(event, ServerEvent::Flush { .. }) {
            self.wake.arm_for_work();
        }
        let wait = matches!(event, ServerEvent::Flush { .. }) && FINAL_FLUSH.with(Cell::get);
        let result = (|| {
            let payload = event
                .encode(&self.metadata, server_time_ms()?)
                .map_err(|error| Self::failure(&error.to_string()))?;
            let sequence = self.enqueue(
                "liminal.events",
                payload,
                #[cfg(test)]
                matches!(event, ServerEvent::ConnectionOpened { .. }),
            )?;
            if wait {
                self.queue.await_append(sequence)?;
                let services = self
                    .services
                    .upgrade()
                    .ok_or_else(|| Self::failure("head checkpoint services owner is gone"))?;
                services.checkpoint_self_stream_heads(self.queue.heads()?)?;
            }
            Ok(())
        })();
        if let Err(error) = result {
            self.queue.fail(&error);
        }
    }

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

    #[cfg(test)]
    pub(super) fn own_work_for_test(&self) -> Result<SelfWork, ServerError> {
        self.queue.own_work()
    }

    #[cfg(test)]
    pub(super) fn queue_measurements(&self) -> Result<(usize, usize, Option<u64>), ServerError> {
        self.queue.measurements()
    }
}

impl RecordDrainer {
    fn run(self) {
        // Every store call this thread makes is the server's own, for its
        // whole life, so none of it wakes the cost tick.
        let _self_work = SelfWorkMark::enter();
        let panic_guard = DrainerPanicGuard(Arc::clone(&self.queue));
        let result = (|| {
            while let Some(record) = self.queue.next()? {
                let completion = self.queue.append_guard();
                let head = match &record.payload {
                    RecordPayload::Single(payload) => self.publish(record.channel, payload)?,
                    RecordPayload::Batch(payloads) => self.publish_batch(payloads)?,
                };
                let sequence = head.checked_sub(1).ok_or_else(|| {
                    SelfStreams::failure("self-stream append returned an empty head")
                })?;
                self.retain(record.channel, sequence);
                self.queue.appended(&record, head)?;
                drop(completion);
            }
            Ok::<(), ServerError>(())
        })();
        if let Err(error) = result {
            self.queue.fail(&error);
        }
        drop(panic_guard);
    }

    fn publish_batch(&self, payloads: &[Vec<u8>]) -> Result<u64, ServerError> {
        let services = self
            .services
            .upgrade()
            .ok_or_else(|| SelfStreams::failure("receipt batch services owner is gone"))?;
        services.publish_self_stream_batch(payloads, self.metadata.publisher())
    }

    fn publish(&self, channel: &str, payload: &[u8]) -> Result<u64, ServerError> {
        let services = self
            .services
            .upgrade()
            .ok_or_else(|| SelfStreams::failure("connection services owner is gone"))?;
        let envelope = MessageEnvelope::new(
            SchemaId::from_schema_bytes(b"{}"),
            CausalContext::default(),
            payload.to_vec(),
        );
        let mut published = services.publish_with_origin(
            channel,
            &envelope,
            None,
            PublishOrigin::Server {
                device: self.metadata.device(),
            },
        );
        let mut confirmed_head = None;
        let appended = match &mut published {
            Ok(publication) => {
                confirmed_head = publication
                    .self_row_at
                    .map(|sequence| {
                        sequence
                            .checked_add(1)
                            .ok_or_else(|| SelfStreams::failure("self-stream sequence overflow"))
                    })
                    .transpose()?;
                publication.self_row_appended
            }
            Err(ServerError::PublishUnresolved {
                channel: addressed,
                sequence: Some(sequence),
                message,
                ..
            }) if addressed.as_str() == channel => {
                match self.confirm_unresolved_row(&services, channel, *sequence, &envelope) {
                    Ok(stored) => {
                        if stored {
                            confirmed_head = Some(sequence.checked_add(1).ok_or_else(|| {
                                SelfStreams::failure("confirmed self-stream sequence overflow")
                            })?);
                        }
                        stored
                    }
                    Err(error) => {
                        *message = format!("{message}; self-row confirmation failed: {error}");
                        false
                    }
                }
            }
            Err(_) => false,
        };
        if appended {
            return confirmed_head.map_or_else(|| services.self_stream_head(channel), Ok);
        }
        published?;
        Err(SelfStreams::failure(
            "self-stream publication was refused before append",
        ))
    }

    /// SRV-012 R2: keeps `channel` within its configured row bound B. Once
    /// the rows retained since the last prune this process knows of exceed
    /// B, one truncate keeps the newest B - B / 8, so a prune runs at most
    /// once per B / 8 + 1 appends. Runs after the append is durable, outside
    /// the own-work lock; a failure is logged by name, counted, and retried
    /// on the next append, and never fails the append that triggered it.
    fn retain(&self, channel: &'static str, sequence: u64) {
        let bound = match channel {
            "liminal.cost" => self.bounds.cost_rows,
            _ => self.bounds.events_rows,
        };
        let keep = bound - bound / 8;
        let cut = match self.retained_from.lock() {
            Ok(retained_from) => {
                let first = retained_from.get(channel).copied().unwrap_or(0);
                let retained = sequence.saturating_add(1).saturating_sub(first);
                if retained <= bound {
                    return;
                }
                sequence.saturating_add(1).saturating_sub(keep)
            }
            Err(error) => {
                self.prune_failed(channel, &format!("retention lock poisoned: {error}"));
                return;
            }
        };
        let pruned = self
            .services
            .upgrade()
            .ok_or_else(|| "connection services owner is gone".to_owned())
            .and_then(|services| {
                let key = format!("{channel}:0");
                liminal::durability::bridge::block_on(
                    services.durable_store().truncate_before(&key, cut),
                )
                .map_err(|error| error.to_string())?
                .map_err(|error| error.to_string())
            });
        match pruned {
            Ok(rows) => {
                let recorded = self
                    .retained_from
                    .lock()
                    .map(|mut retained_from| {
                        retained_from.insert(channel, cut);
                    })
                    .map_err(|error| format!("retention lock poisoned: {error}"))
                    .and_then(|()| {
                        self.own_work()
                            .map(|mut own_work| {
                                own_work.prunes = own_work.prunes.saturating_add(1);
                                own_work.rows_pruned = own_work.rows_pruned.saturating_add(rows);
                            })
                            .map_err(|error| error.to_string())
                    });
                if let Err(error) = recorded {
                    self.prune_failed(channel, &error);
                }
            }
            Err(error) => self.prune_failed(channel, &error),
        }
    }

    fn prune_failed(&self, channel: &str, error: &str) {
        tracing::warn!(stream = channel, %error, "SELF_STREAM_PRUNE_FAILED");
        match self.own_work() {
            Ok(mut own_work) => {
                own_work.prune_failures = own_work.prune_failures.saturating_add(1);
            }
            Err(error) => {
                tracing::error!(stream = channel, %error, "SELF_STREAM_PRUNE_FAILED not counted");
            }
        }
    }

    fn own_work(&self) -> Result<MutexGuard<'_, SelfWork>, ServerError> {
        self.own_work
            .lock()
            .map_err(|error| SelfStreams::failure(&format!("own-work lock poisoned: {error}")))
    }

    fn confirm_unresolved_row(
        &self,
        services: &LiminalConnectionServices,
        channel: &str,
        sequence: u64,
        expected: &MessageEnvelope,
    ) -> Result<bool, ServerError> {
        let store = services.durable_store();
        let key = format!("{channel}:0");
        let row = liminal::durability::bridge::block_on(store.read_at(&key, sequence))
            .map_err(|error| {
                SelfStreams::failure(&format!("self-row confirmation bridge: {error}"))
            })?
            .map_err(|error| {
                SelfStreams::failure(&format!("self-row confirmation read: {error}"))
            })?;
        let Some(row) = row else {
            return Ok(false);
        };
        if row.sequence != sequence {
            return Err(SelfStreams::failure(
                "self-row confirmation returned a different sequence",
            ));
        }
        let stored =
            liminal::durability::MessageEnvelope::deserialize(&row.payload).map_err(|error| {
                SelfStreams::failure(&format!("self-row confirmation envelope: {error}"))
            })?;
        Ok(stored.payload == expected.payload
            && stored.publisher_id == self.metadata.publisher()
            && stored.idempotency_key.is_none())
    }
}

impl CostSink for Arc<SelfStreams> {
    fn completion_failure_signal(&self) -> Option<super::publisher::CompletionFailureSignal> {
        let queue = Arc::clone(&self.queue);
        Some(Arc::new(move |error| queue.abandon(error)))
    }

    fn before_stop(&mut self) -> Result<(), ServerError> {
        let before = self.counts.snapshot();
        self.queue.drain_all()?;
        tracing::debug!(counts = ?self.counts.snapshot().since(&before), "self-stream final queued drain store calls");
        Ok(())
    }
    fn snapshot(&self) -> Result<(StoreCounts, SelfWork), ServerError> {
        self.drain()?;
        let mut own_work = self.queue.own_work()?;
        let retention = self.retention_work.lock().map_err(|error| {
            SelfStreams::failure(&format!("retention work lock poisoned: {error}"))
        })?;
        own_work.prunes = retention.prunes;
        own_work.rows_pruned = retention.rows_pruned;
        own_work.prune_failures = retention.prune_failures;
        drop(retention);
        Ok((self.counts.snapshot(), own_work))
    }

    fn publish(&mut self, record: &PreparedCost) -> Result<(), ServerError> {
        let before = self.counts.snapshot();
        let payload = record
            .encode(&self.metadata, server_time_ms()?)
            .map_err(|error| SelfStreams::failure(&error.to_string()))?;
        let sequence = self.enqueue(
            "liminal.cost",
            payload,
            #[cfg(test)]
            false,
        )?;
        tracing::debug!(sequence, "self-stream cost record awaits append");
        self.queue.await_append(sequence)?;
        tracing::debug!(sequence, "self-stream cost record append completed");
        tracing::debug!(counts = ?self.counts.snapshot().since(&before), "self-stream final cost store calls");
        Ok(())
    }
    fn console_wake(&self) -> Option<Arc<tokio::sync::Notify>> {
        self.services.upgrade()?.signals().console_wake()
    }

    /// SRV-012 R4: the console pump on this loop's thread. A pump failure is
    /// a lock or encoding failure the hub has already counted; it is logged
    /// by name and the loop goes on, since the console is an observer and
    /// never a reason to stop the server's own streams.
    fn pump_console(&mut self) {
        let Some(services) = self.services.upgrade() else {
            return;
        };
        if let Err(error) = services.console().pump(&*services) {
            tracing::warn!(%error, "CONSOLE_PUMP_FAILED");
        }
    }
}

struct FlushBarrier;

impl Drop for FlushBarrier {
    fn drop(&mut self) {
        FINAL_FLUSH.with(|active| active.set(false));
    }
}

struct DrainerPanicGuard(Arc<RecordQueue>);

impl Drop for DrainerPanicGuard {
    fn drop(&mut self) {
        if thread::panicking() {
            self.0
                .fail(&SelfStreams::failure("record drainer panicked"));
        }
    }
}

impl Drop for SelfStreams {
    fn drop(&mut self) {
        if let Err(error) = self.queue.stop() {
            self.queue.fail(&error);
        }
        let drainer = match self.drainer.get_mut() {
            Ok(drainer) => drainer.thread.take(),
            Err(error) => {
                self.queue.fail(&Self::failure(&format!(
                    "record drainer join lock poisoned: {error}"
                )));
                None
            }
        };
        if let Some(drainer) = drainer {
            if drainer.thread().id() == thread::current().id() {
                tracing::debug!(
                    "self-stream drainer drops its final services owner and is stopping"
                );
            } else if self.queue.was_abandoned() && !drainer.is_finished() {
                tracing::error!(
                    "self-stream drainer completion remains unconfirmed after owned drain abandonment"
                );
            } else if drainer.join().is_err() {
                self.queue
                    .fail(&Self::failure("record drainer join panicked"));
            }
        }
    }
}

fn server_time_ms() -> Result<u64, ServerError> {
    let elapsed = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map_err(|error| ServerError::ListenerAccept {
            message: format!("self-stream server clock is before the epoch: {error}"),
        })?;
    u64::try_from(elapsed.as_millis()).map_err(|error| ServerError::ListenerAccept {
        message: format!("self-stream server time cannot fit in milliseconds: {error}"),
    })
}