grommet 0.1.0

//! Building and owning a set of pinned shards.

use crate::clock::{Clock, SystemClock};
use crate::metrics::ShardStats;
use crate::processor::Processor;
use crate::router::Router;
use crate::shard::{self, ShardConfig};
use crate::topology::{Bound, PinPolicy, Plan, ShardPlacement, TopologyReport, Workload};
use crate::work::Envelope;
use std::fmt;
use std::sync::Arc;
use std::thread::JoinHandle;
use tokio::sync::mpsc;

/// What a shard thread knows about itself when it builds its processor.
///
/// This is what makes core-local resources possible: the factory runs on the
/// shard's own thread, after it has been placed, so it can size a connection
/// pool per core and — given [`node`] — pick the offload pool and allocations
/// that are local to the memory it will be touching.
///
/// [`node`]: ShardContext::node
#[derive(Clone, Copy, Debug)]
pub struct ShardContext {
    pub index: usize,
    pub shards: usize,
    /// Where the plan put this shard, if there was one to place it.
    pub placement: Option<ShardPlacement>,
    /// What binding achieved, which is not always what was asked for.
    pub bound: Bound,
}

impl ShardContext {
    /// The memory node this shard should keep its state and its offload work on.
    pub fn node(&self) -> Option<usize> {
        self.placement.map(|placement| placement.node)
    }

    /// The CPU this shard was placed on.
    pub fn cpu(&self) -> Option<usize> {
        self.placement.map(|placement| placement.cpu)
    }
}

#[derive(Debug)]
pub enum BuildError {
    /// [`PinPolicy::Require`] was set and these shard indices could not be
    /// pinned.
    NotPinned(Vec<usize>),
    /// A shard thread died before it reported its placement.
    ShardFailed,
}

impl fmt::Display for BuildError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::NotPinned(shards) => {
                write!(f, "shards {shards:?} could not be pinned under PinPolicy::Require")
            }
            Self::ShardFailed => f.write_str("a shard thread failed during startup"),
        }
    }
}

impl std::error::Error for BuildError {}

pub struct Builder<P: Processor, C: Clock, const CLASSES: usize = 2> {
    shards: usize,
    mailbox: usize,
    shard_config: ShardConfig<CLASSES>,
    pin: PinPolicy,
    plan: Option<Arc<Plan>>,
    clock: C,
    stamp_arrival: bool,
    _processor: std::marker::PhantomData<fn() -> P>,
}

impl<P: Processor, const CLASSES: usize> Builder<P, SystemClock, CLASSES> {
    /// Configure `shards` reactors with the given per-class in-flight budgets.
    ///
    /// Placement is planned from this machine unless [`plan`] supplies one or
    /// [`PinPolicy::Disabled`] turns it off.
    ///
    /// [`plan`]: Builder::plan
    pub fn new(shards: usize, max_inflight: [usize; CLASSES]) -> Self {
        Self::with_clock(shards, max_inflight, SystemClock::new())
    }

    /// One reactor per shard placement in `plan`.
    ///
    /// This is the usual entry point once the layout matters: the plan already
    /// decided how many reactors the machine can carry, after reserving cores
    /// for the offload pool and for the OS, and after honouring any cgroup
    /// bandwidth limit. Choosing a shard count separately is choosing to
    /// disagree with it.
    pub fn for_plan(plan: Arc<Plan>, max_inflight: [usize; CLASSES]) -> Self {
        Self::new(plan.shards.len().max(1), max_inflight).plan(plan)
    }
}

impl<P: Processor, C: Clock, const CLASSES: usize> Builder<P, C, CLASSES> {
    pub fn with_clock(shards: usize, max_inflight: [usize; CLASSES], clock: C) -> Self {
        assert!(shards > 0, "a runtime needs at least one shard");
        Self {
            shards,
            mailbox: 1024,
            shard_config: ShardConfig::new(max_inflight),
            pin: PinPolicy::default(),
            plan: None,
            clock,
            stamp_arrival: true,
            _processor: std::marker::PhantomData,
        }
    }

    /// Mailbox depth per shard. This is the queue that absorbs bursts before
    /// submitters feel backpressure.
    pub fn mailbox(mut self, capacity: usize) -> Self {
        assert!(capacity > 0, "a mailbox needs capacity");
        self.mailbox = capacity;
        self
    }

    pub fn shard_config(mut self, config: ShardConfig<CLASSES>) -> Self {
        self.shard_config = config;
        self
    }

    pub fn pin(mut self, policy: PinPolicy) -> Self {
        self.pin = policy;
        self
    }

    /// Place shards according to `plan`, round-robin if there are more shards
    /// than the plan has placements for.
    ///
    /// The same plan should be given to the offload pools, so that a shard and
    /// the workers it submits to agree about which memory node they are on.
    pub fn plan(mut self, plan: Arc<Plan>) -> Self {
        self.plan = Some(plan);
        self
    }

    /// Suppress a retry whose request id is already queued or in flight for
    /// the same key. See [`ShardConfig::coalesce_duplicates`].
    pub fn coalesce_duplicates(mut self, coalesce: bool) -> Self {
        self.shard_config.coalesce_duplicates = coalesce;
        self
    }

    /// See [`Router::with_options`].
    pub fn stamp_arrival(mut self, stamp: bool) -> Self {
        self.stamp_arrival = stamp;
        self
    }

    /// Start every shard, building one processor per shard on its own thread.
    ///
    /// The factory runs inside the shard's runtime, which is what lets each
    /// shard own core-local resources — connection pools, caches, buffers —
    /// rather than sharing one set across cores.
    pub fn spawn<F>(self, factory: F) -> Result<Runtime<P, C, CLASSES>, BuildError>
    where
        F: Fn(&ShardContext) -> P + Send + Sync + 'static,
    {
        // Reading the machine is deferred to here rather than done in `new`, so
        // that a runtime which never starts never pays for it, and so a caller
        // who supplies a plan never reads the machine twice.
        let plan = match (self.plan, self.pin) {
            (plan @ Some(_), _) => plan,
            (None, PinPolicy::Disabled) => None,
            (None, _) => crate::topology::detect(&Workload::default()).ok().map(Arc::new),
        };
        let placements: &[ShardPlacement] =
            plan.as_ref().map(|plan| plan.shards.as_slice()).unwrap_or_default();
        let placement_for =
            |index: usize| (!placements.is_empty()).then(|| placements[index % placements.len()]);

        let mut cpus: Vec<usize> =
            (0..self.shards).filter_map(|index| placement_for(index).map(|at| at.cpu)).collect();
        cpus.sort_unstable();
        cpus.dedup();
        let distinct_cores = cpus.len();

        let factory = Arc::new(factory);
        let (report, reports) = std::sync::mpsc::channel();

        let mut senders = Vec::with_capacity(self.shards);
        let mut workers = Vec::with_capacity(self.shards);
        let mut stats = Vec::with_capacity(self.shards);

        for index in 0..self.shards {
            let (tx, rx) = mpsc::channel::<Envelope<P::Work>>(self.mailbox);
            senders.push(tx);
            let shard_stats = Arc::new(ShardStats::<CLASSES>::default());
            stats.push(shard_stats.clone());

            let placement = placement_for(index);
            let context_shards = self.shards;
            let clock = self.clock.clone();
            let config = self.shard_config;
            let policy = self.pin;
            let factory = factory.clone();
            let report = report.clone();
            let plan = plan.clone();

            workers.push(
                std::thread::Builder::new()
                    .name(format!("shard-{index}"))
                    .spawn(move || {
                        // Bind first. Memory binding only governs pages touched
                        // afterwards, and everything this thread allocates from
                        // here on — the runtime, the processor, the key states —
                        // should come from its own node.
                        let bound = match (policy, placement, &plan) {
                            (PinPolicy::Disabled, _, _) | (_, None, _) | (_, _, None) => {
                                Bound::default()
                            }
                            (_, Some(placement), Some(plan)) => plan.bind_shard(&placement),
                        };
                        // Report before blocking forever, so the builder can
                        // fail fast rather than wait on a shard that started.
                        let _ = report.send((index, bound));

                        let runtime = tokio::runtime::Builder::new_current_thread()
                            .enable_all()
                            .build()
                            .expect("shard runtime");
                        runtime.block_on(async move {
                            let context =
                                ShardContext { index, shards: context_shards, placement, bound };
                            let processor = factory(&context);
                            shard::run(rx, processor, clock, shard_stats, config).await;
                        });
                    })
                    .expect("spawn shard thread"),
            );
        }
        drop(report);

        let mut pinned = 0;
        let mut memory_bound = 0;
        let mut unpinned = Vec::new();
        for _ in 0..self.shards {
            let (index, bound) = reports.recv().map_err(|_| BuildError::ShardFailed)?;
            if bound.cpu {
                pinned += 1;
            } else {
                unpinned.push(index);
            }
            if bound.memory {
                memory_bound += 1;
            }
        }

        if self.pin == PinPolicy::Require && !unpinned.is_empty() {
            unpinned.sort_unstable();
            // Closing every mailbox tells the shards to drain and exit.
            drop(senders);
            for worker in workers {
                let _ = worker.join();
            }
            return Err(BuildError::NotPinned(unpinned));
        }

        Ok(Runtime {
            router: Some(Arc::new(Router::with_options(senders, self.clock, self.stamp_arrival))),
            workers,
            stats,
            report: TopologyReport {
                shards: self.shards,
                distinct_cores,
                pinned,
                memory_bound,
                policy: self.pin,
            },
        })
    }
}

/// A running set of shards. Dropping it closes every mailbox and waits for the
/// shards to drain.
pub struct Runtime<P: Processor, C: Clock, const CLASSES: usize = 2> {
    router: Option<Arc<Router<P::Work, C, CLASSES>>>,
    workers: Vec<JoinHandle<()>>,
    stats: Vec<Arc<ShardStats<CLASSES>>>,
    report: TopologyReport,
}

impl<P: Processor, const CLASSES: usize> Runtime<P, SystemClock, CLASSES> {
    /// Start configuring a runtime on the system clock. Use
    /// [`Builder::with_clock`] directly for a different one.
    pub fn builder(
        shards: usize,
        max_inflight: [usize; CLASSES],
    ) -> Builder<P, SystemClock, CLASSES> {
        Builder::new(shards, max_inflight)
    }

    /// Start configuring a runtime laid out by `plan`, one shard per placement.
    pub fn for_plan(
        plan: Arc<Plan>,
        max_inflight: [usize; CLASSES],
    ) -> Builder<P, SystemClock, CLASSES> {
        Builder::for_plan(plan, max_inflight)
    }
}

impl<P: Processor, C: Clock, const CLASSES: usize> Runtime<P, C, CLASSES> {
    pub fn router(&self) -> &Arc<Router<P::Work, C, CLASSES>> {
        self.router.as_ref().expect("router is present until shutdown")
    }

    pub fn stats(&self) -> &[Arc<ShardStats<CLASSES>>] {
        &self.stats
    }

    pub fn topology(&self) -> &TopologyReport {
        &self.report
    }

    /// Close the mailboxes and wait for every shard to finish draining.
    ///
    /// Shutdown is driven by dropping the router, so any clone of it that you
    /// are still holding will keep the shards alive. Drop those first.
    pub fn shutdown(mut self) {
        self.close();
    }

    fn close(&mut self) {
        drop(self.router.take());
        for worker in self.workers.drain(..) {
            let _ = worker.join();
        }
    }
}

impl<P: Processor, C: Clock, const CLASSES: usize> Drop for Runtime<P, C, CLASSES> {
    fn drop(&mut self) {
        self.close();
    }
}