liminal-server 0.25.4

use std::future::{Future, poll_fn};
use std::pin::pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::task::Poll;
use std::time::Duration;

use liminal::durability::StoreCounts;
use tokio::sync::Notify;

use crate::ServerError;

use super::cost::{CostWake, CostWindow, PreparedCost, SelfWork};

/// A blocked publisher cannot report its own drain abandonment to the sink.
pub(in crate::server) type CompletionFailureSignal = Arc<dyn Fn(&ServerError) + Send + Sync>;

/// The acknowledged boundary keeps failed publications in the next delta.
pub(in crate::server) trait CostSink {
    fn completion_failure_signal(&self) -> Option<CompletionFailureSignal> {
        None
    }
    fn before_stop(&mut self) -> Result<(), ServerError> {
        Ok(())
    }
    fn snapshot(&self) -> Result<(StoreCounts, SelfWork), ServerError>;
    fn publish(&mut self, record: &PreparedCost) -> Result<(), ServerError>;

    /// SRV-012 R4: the console hub's wake, when this sink also publishes the
    /// console streams. The loop below then waits on it beside the cost
    /// tick, so the console needs no thread, timer or poll of its own.
    fn console_wake(&self) -> Option<Arc<Notify>> {
        None
    }

    /// Publishes what the console hub has due. Called on this loop's thread
    /// when [`Self::console_wake`] fires.
    fn pump_console(&mut self) {}
}

#[derive(Debug, Default)]
pub(in crate::server) struct PublisherStop {
    requested: AtomicBool,
    notify: Notify,
}

impl PublisherStop {
    pub(in crate::server) fn request(&self) {
        self.requested.store(true, Ordering::Release);
        self.notify.notify_one();
    }

    async fn wait(&self) {
        let notified = self.notify.notified();
        if self.requested.load(Ordering::Acquire) {
            return;
        }
        notified.await;
    }
}

/// Both production and the paused-clock proof drive this single loop.
pub(in crate::server) async fn run_cost_publisher<Sleep, SleepFuture, Sink>(
    wake: &CostWake,
    stop: &PublisherStop,
    baseline: (StoreCounts, SelfWork),
    sleep: Sleep,
    sink: &mut Sink,
) -> Result<(), ServerError>
where
    Sleep: Fn(Duration) -> SleepFuture + Send + Sync,
    SleepFuture: Future<Output = ()> + Send,
    Sink: CostSink + Send,
{
    let (snapshot, self_work) = baseline;
    let mut window = CostWindow::new(snapshot, self_work);
    let console = sink.console_wake();
    loop {
        let mut tick = pin!(wake.next_tick(&sleep));
        let mut cancelled = pin!(stop.wait());
        // A console wake is served and the same tick kept pinned, so a busy
        // console never restarts (and so never starves) the cost timer.
        loop {
            let mut console_due = pin!(async {
                match &console {
                    Some(console) => console.notified().await,
                    None => std::future::pending().await,
                }
            });
            let woken = poll_fn(|context| {
                if cancelled.as_mut().poll(context).is_ready() {
                    return Poll::Ready(Woken::Stop);
                }
                if console_due.as_mut().poll(context).is_ready() {
                    return Poll::Ready(Woken::Console);
                }
                tick.as_mut().poll(context).map(Woken::Tick)
            })
            .await;
            match woken {
                Woken::Stop => {
                    tracing::debug!("self-stream publisher stop drains queued records");
                    sink.before_stop()?;
                    tracing::debug!("self-stream publisher stop publishes final cost window");
                    publish_window(sink, &mut window)?;
                    tracing::debug!("self-stream publisher stop finished final cost window");
                    return Ok(());
                }
                Woken::Console => sink.pump_console(),
                Woken::Tick(ticked) => {
                    ticked?;
                    break;
                }
            }
        }
        publish_window(sink, &mut window)?;
    }
}

fn publish_window<Sink: CostSink>(
    sink: &mut Sink,
    window: &mut CostWindow,
) -> Result<(), ServerError> {
    let (snapshot, self_work) = sink.snapshot()?;
    if let Some(record) = window.prepare(snapshot, self_work) {
        sink.publish(&record)?;
        window.published(&record);
    }
    Ok(())
}

/// What woke the publisher loop.
enum Woken {
    Stop,
    Console,
    Tick(Result<(), ServerError>),
}