liminal_server/tap/

subscription.rs

//! The tap and its subscriptions.

use liminal::channel::{PublishRecord, Recipient};
use std::collections::HashMap;
use std::num::NonZeroU32;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::Instant;

use super::error::{TapError, lock};
use super::message::{
    Hop, PayloadGrant, PublishFact, Reached, TapBatch, TapFilter, TapIdent, TapMessage, TapSource,
};
use super::ring::BoundedRing;
mod call;
mod publish;
mod reader;

/// Default per-subscriber ring capacity, in messages.
pub const DEFAULT_TAP_CAPACITY: usize = 4096;

/// The largest sampling divisor the server will choose.
pub const MAX_SAMPLE_ONE_IN: NonZeroU32 = match NonZeroU32::new(1 << 16) {
    Some(most) => most,
    None => NonZeroU32::MIN,
};

const TWO: NonZeroU32 = NonZeroU32::MIN.saturating_add(1);

/// Wakes whoever drains the tap; called at most once per publish, after every
/// tap lock is released.
pub type TapWaker = Arc<dyn Fn() + Send + Sync>;

/// One kept message and the divisor it was admitted under, so a drain can
/// close a batch wherever the divisor changed without a second ring.
#[derive(Debug)]
struct Admitted {
    one_in: NonZeroU32,
    message: TapMessage,
}

#[derive(Debug)]
struct SubscriberState {
    ring: BoundedRing<Admitted>,
    one_in: NonZeroU32,
    seen: u64,
    next_seq: u64,
}

impl SubscriberState {
    /// The sampling decision for one matching publish, taken before anything
    /// is allocated for it.
    fn admits(&mut self) -> bool {
        let take = self.seen % u64::from(self.one_in.get()) == 0;
        self.seen = self.seen.wrapping_add(1);
        take
    }

    /// Doubles the divisor once the ring is more than half full.
    fn after_push(&mut self) {
        if self.ring.len().saturating_mul(2) > self.ring.capacity() {
            self.one_in = self.one_in.saturating_mul(TWO).min(MAX_SAMPLE_ONE_IN);
        }
    }
}

struct Subscriber {
    id: u64,
    filter: TapFilter,
    grant: Option<PayloadGrant>,
    state: Mutex<SubscriberState>,
}

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

#[derive(Default)]
struct TapShared {
    active: AtomicUsize,
    observed: AtomicU64,
    /// Of `observed`, the server's own publishes (`liminal.<device>`).
    /// Test-only.
    #[cfg(test)]
    observed_server: AtomicU64,
    dropped_total: AtomicU64,
    unsubscribe_failures: AtomicU64,
    next_id: AtomicU64,
    subscribers: Mutex<Vec<Arc<Subscriber>>>,
    /// Each registered worker's connection and its station id, so a kept
    /// message names its publishing station without the topology lock.
    stations: Mutex<HashMap<u64, Arc<str>>>,
    /// Each live subscription's subscriber pid and its connection, so a
    /// recipient the fan-out recorded names its connection's station.
    receivers: Mutex<HashMap<u64, u64>>,
    /// Recipient buffers between publishes, so a recording publish reuses
    /// an allocation rather than making one.
    buffers: Mutex<Vec<PublishRecord>>,
    /// Recipient buffers lent: publishes that recorded their recipients.
    lent: AtomicU64,
    /// The one empty station list, shared by every kept message whose
    /// recorded fan-out reached no registered worker, so naming none
    /// allocates nothing.
    no_stations: Arc<[Reached]>,
    waker: OnceLock<TapWaker>,
}

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

/// The server's tap. Cheap to clone; every clone is the same tap.
#[derive(Clone, Debug)]
pub struct Tap {
    shared: Arc<TapShared>,
    node: Arc<str>,
    /// This server's depot's station id, when it keeps a persistent store:
    /// a durable message's hop between its publisher and the exchange.
    depot: Option<Arc<str>>,
    capacity: usize,
}

impl Tap {
    /// A tap whose hops name `node`, with per-subscriber rings of `capacity`,
    /// that wakes nobody: for a reader that drains when it chooses, as a test
    /// does.
    #[must_use]
    pub fn new(node: &str, capacity: usize) -> Self {
        Self::build(node, capacity, OnceLock::new())
    }

    /// A tap that calls `waker` once per publish any subscription kept.
    /// Fixed at construction, so the hot path reads it without a lock and
    /// nothing can replace it.
    #[must_use]
    pub fn with_waker(node: &str, capacity: usize, waker: TapWaker) -> Self {
        Self::build(node, capacity, OnceLock::from(waker))
    }

    fn build(node: &str, capacity: usize, waker: OnceLock<TapWaker>) -> Self {
        Self {
            shared: Arc::new(TapShared {
                waker,
                ..TapShared::default()
            }),
            node: Arc::from(node),
            depot: None,
            capacity: capacity.max(1),
        }
    }

    /// This tap, on a server that keeps a persistent store when `depot`:
    /// its kept durable messages then pass `<node>.depot` (lane brief
    /// ruling 8).
    #[must_use]
    pub fn with_depot(mut self, depot: bool) -> Self {
        self.depot = depot.then(|| Arc::from(format!("{}.depot", self.node)));
        self
    }

    /// The station `id` would collide with: this tap's exchange (its node)
    /// or its depot. A worker with such an identity names no station here.
    #[must_use]
    pub fn reserved(&self, id: &str) -> Option<&'static str> {
        if id == &*self.node {
            Some("exchange")
        } else if self.depot.as_deref() == Some(id) {
            Some("depot")
        } else {
            None
        }
    }

    /// Subscribes without payloads.
    ///
    /// # Errors
    /// Returns [`TapError::Poisoned`] when the subscriber list is poisoned.
    pub fn subscribe(&self, filter: TapFilter) -> Result<TapSubscription, TapError> {
        self.subscribe_with(filter, None)
    }

    /// Subscribes, with payloads when `grant` is present.
    ///
    /// # Errors
    /// Returns [`TapError::Poisoned`] when the subscriber list is poisoned.
    pub fn subscribe_with(
        &self,
        filter: TapFilter,
        grant: Option<PayloadGrant>,
    ) -> Result<TapSubscription, TapError> {
        let id = self.shared.next_id.fetch_add(1, Ordering::Relaxed);
        let subscriber = Arc::new(Subscriber {
            id,
            filter,
            grant,
            state: Mutex::new(SubscriberState {
                ring: BoundedRing::new(self.capacity),
                one_in: NonZeroU32::MIN,
                seen: 0,
                next_seq: 1,
            }),
        });
        let mut subscribers = lock(&self.shared.subscribers, "tap subscribers")?;
        subscribers.push(Arc::clone(&subscriber));
        self.shared
            .active
            .store(subscribers.len(), Ordering::Release);
        drop(subscribers);
        Ok(TapSubscription {
            shared: Arc::clone(&self.shared),
            subscriber,
        })
    }

    /// Publishes the hook has shown the tap while a subscription was live (a
    /// count pin: N ordinary publishes give N, the tap's own records give
    /// zero).
    #[must_use]
    pub fn observations(&self) -> u64 {
        self.shared.observed.load(Ordering::Relaxed)
    }

    /// Of [`Self::observations`], those of the server's own publishes (its
    /// `liminal.events` and `liminal.cost` rows), which its publisher owner
    /// writes on its own schedule. Test-only.
    #[cfg(test)]
    #[must_use]
    pub(crate) fn server_observations(&self) -> u64 {
        self.shared.observed_server.load(Ordering::Relaxed)
    }

    /// Live subscriptions.
    #[must_use]
    pub fn subscriber_count(&self) -> usize {
        self.shared.active.load(Ordering::Acquire)
    }

    /// Messages evicted from every subscription's ring since the tap was made.
    #[must_use]
    pub fn dropped_total(&self) -> u64 {
        self.shared.dropped_total.load(Ordering::Relaxed)
    }

    /// Poisons the subscriber list, as a panic on a thread holding it would,
    /// and answers whether it is now poisoned. Test-only.
    #[cfg(test)]
    pub(crate) fn poison_subscribers(&self) -> bool {
        std::thread::scope(|scope| {
            scope
                .spawn(|| {
                    let _held = self.shared.subscribers.lock();
                    std::panic::resume_unwind(Box::new("poisoned for a test"));
                })
                .join()
                .is_err()
        }) && self.shared.subscribers.is_poisoned()
    }

    /// Names `station` as the station `connection` publishes as, from its
    /// worker registration until [`Self::forget_station`].
    ///
    /// # Errors
    /// Returns [`TapError::Poisoned`] when the station table is poisoned.
    pub fn name_station(&self, connection: u64, station: &str) -> Result<(), TapError> {
        lock(&self.shared.stations, "tap stations")?.insert(connection, Arc::from(station));
        Ok(())
    }

    /// Forgets `connection`'s station, as its connection closes.
    ///
    /// # Errors
    /// Returns [`TapError::Poisoned`] when the station table is poisoned.
    pub fn forget_station(&self, connection: u64) -> Result<(), TapError> {
        lock(&self.shared.stations, "tap stations")?.remove(&connection);
        Ok(())
    }

    /// Names `connection` as the owner of the subscription whose subscriber
    /// pid is `subscriber`, until [`Self::forget_receiver`].
    ///
    /// # Errors
    /// Returns [`TapError::Poisoned`] when the receiver table is poisoned.
    pub fn name_receiver(&self, subscriber: u64, connection: u64) -> Result<(), TapError> {
        lock(&self.shared.receivers, "tap receivers")?.insert(subscriber, connection);
        Ok(())
    }

    /// Forgets the subscription whose subscriber pid is `subscriber`, as it
    /// leaves its channel.
    ///
    /// # Errors
    /// Returns [`TapError::Poisoned`] when the receiver table is poisoned.
    pub fn forget_receiver(&self, subscriber: u64) -> Result<(), TapError> {
        lock(&self.shared.receivers, "tap receivers")?.remove(&subscriber);
        Ok(())
    }

    /// Forgets every subscription of `connection`, as it closes.
    ///
    /// # Errors
    /// Returns [`TapError::Poisoned`] when the receiver table is poisoned.
    pub fn forget_receivers_of(&self, connection: u64) -> Result<(), TapError> {
        lock(&self.shared.receivers, "tap receivers")?.retain(|_, owner| *owner != connection);
        Ok(())
    }

    /// A buffer for a publish to record its recipients in, only while a
    /// subscription is live: with none it is one atomic load, no lock and no
    /// allocation, and the publish records nothing. A buffer lent comes from
    /// those given back, so it allocates only while the pool and each buffer
    /// grow to the widest concurrent fan-out.
    ///
    /// # Errors
    /// Returns [`TapError::Poisoned`] when the buffer pool is poisoned.
    pub fn lend_recipients(&self) -> Result<Option<PublishRecord>, TapError> {
        if self.shared.active.load(Ordering::Acquire) == 0 {
            return Ok(None);
        }
        self.shared.lent.fetch_add(1, Ordering::Relaxed);
        Ok(Some(
            lock(&self.shared.buffers, "tap buffers")?
                .pop()
                .unwrap_or_default(),
        ))
    }

    /// Takes back a buffer [`Self::lend_recipients`] lent, emptied, for the
    /// next recording publish.
    ///
    /// # Errors
    /// Returns [`TapError::Poisoned`] when the buffer pool is poisoned.
    pub fn give_back(&self, mut buffer: PublishRecord) -> Result<(), TapError> {
        buffer.clear();
        lock(&self.shared.buffers, "tap buffers")?.push(buffer);
        Ok(())
    }

    /// Recipient buffers lent since the tap was made: publishes that asked
    /// the fan-out for their recipients.
    #[must_use]
    pub fn recipients_lent(&self) -> u64 {
        self.shared.lent.load(Ordering::Relaxed)
    }

    /// Subscriptions whose drop found the subscriber list poisoned and so
    /// could not leave it.
    #[must_use]
    pub fn unsubscribe_failures(&self) -> u64 {
        self.shared.unsubscribe_failures.load(Ordering::Relaxed)
    }
}

/// One subscriber's handle. Dropping it unsubscribes.
#[derive(Debug)]
pub struct TapSubscription {
    shared: Arc<TapShared>,
    subscriber: Arc<Subscriber>,
}

/// `elapsed` in whole milliseconds; `None` beyond what a u64 holds, which
/// the caller names rather than writing the most a u64 holds.
fn elapsed_ms(elapsed: std::time::Duration) -> Option<u64> {
    u64::try_from(elapsed.as_millis()).ok()
}

/// The `named` first-named stations of `recipients` recipients, in order, in
/// one allocation: `first_at(index)` is the station the recipient at `index`
/// names first, if it names one, with the instant that recipient was queued. `None` when none is named.
///
/// # Errors
/// Returns [`TapError::StationsDisagree`] when the fill finds other than
/// `named` stations.
fn gather<'station>(
    recipients: usize,
    named: usize,
    first_at: impl Fn(usize) -> Option<(&'station Arc<str>, Instant)>,
) -> Result<Option<Arc<[Reached]>>, TapError> {
    let Some(first) = (0..recipients).find_map(&first_at) else {
        return if named == 0 {
            Ok(None)
        } else {
            Err(TapError::StationsDisagree { named, found: 0 })
        };
    };
    // A range mapped is of known length, so the collect allocates once.
    // Each step finds the next first-named station. A step that finds none
    // fills its slot with `first` only so the slice is whole; the count of
    // those found then refuses the list, so that slot never leaves here.
    let mut next = 0;
    let mut found = 0_usize;
    let reached = |(station, queued_at): (&Arc<str>, Instant)| Reached {
        station: Arc::clone(station),
        queued_at,
    };
    let list: Arc<[Reached]> = (0..named)
        .map(|_| {
            match (next..recipients)
                .find_map(|index| first_at(index).map(|station| (index, station)))
            {
                Some((index, station)) => {
                    next = index + 1;
                    found += 1;
                    reached(station)
                }
                None => reached(first),
            }
        })
        .collect();
    let beyond = (next..recipients)
        .filter(|index| first_at(*index).is_some())
        .count();
    if found != named || beyond != 0 {
        return Err(TapError::StationsDisagree {
            named,
            found: found + beyond,
        });
    }
    Ok(Some(list))
}

#[cfg(test)]
mod gather_tests {
    use std::sync::Arc;
    use std::time::{Duration, Instant};

    use super::{Reached, TapError, elapsed_ms, gather};

    /// Buttercup's V2, Hermes 04:27: a hop further out than u64
    /// milliseconds is named, never written as the most a u64 holds.
    #[test]
    fn an_elapsed_time_beyond_u64_milliseconds_has_no_reading() {
        assert_eq!(elapsed_ms(Duration::from_millis(7)), Some(7));
        assert_eq!(elapsed_ms(Duration::MAX), None);
    }

    /// Buttercup's finding 2 (Hermes 03:05): when the filling pass finds
    /// fewer stations than the count named, the list is refused by name,
    /// never filled in with a station that stands in for the missing one.
    /// The count named two; the fill finds only the ledger, as if the audit
    /// worker had left between the passes.
    #[test]
    fn a_fill_that_disagrees_with_the_count_is_refused() {
        let ledger: Arc<str> = Arc::from("orders.ledger");
        let queued = Instant::now();
        let gathered = gather(2, 2, |index| (index == 0).then_some((&ledger, queued)));
        assert!(
            matches!(
                gathered,
                Err(TapError::StationsDisagree { named: 2, found: 1 })
            ),
            "{gathered:?}"
        );
    }

    /// Positive control: passes that agree build the list, and none named
    /// builds none.
    #[test]
    fn a_fill_that_agrees_builds_the_list() -> Result<(), TapError> {
        let ledger: Arc<str> = Arc::from("orders.ledger");
        let audit: Arc<str> = Arc::from("orders.audit");
        let queued = Instant::now();
        let first_at = |index: usize| match index {
            0 => Some((&ledger, queued)),
            2 => Some((&audit, queued)),
            _ => None,
        };
        let reached = |station: &Arc<str>| Reached {
            station: Arc::clone(station),
            queued_at: queued,
        };
        assert_eq!(
            gather(3, 2, first_at)?,
            Some(Arc::from(vec![reached(&ledger), reached(&audit)]))
        );
        assert_eq!(gather(3, 0, |_| None)?, None);
        Ok(())
    }
}