improvise/crates/improvise-protocol/src/delivery/client.rs

//! The client end of an edge: watermarks, the held window, and when to ask (ADR 009 clause 9).
//! Why each choice: `context/code/delivery.md` §4.

use std::collections::{BTreeMap, BTreeSet};

use zset::ZSet;

use super::exchange::{Ask, Correction, CorrectionBody, Edge, EdgeError, Reply, Request};
use super::exchange::{SessionToken, Window};
use super::notice::{AxisChange, Flags, Moved, Notice, Relabel, ShapeKey, Stamp};
use super::wire::{Cell, WireKey, decode};

/// What the shape's consumer must do with the structure around the cells.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Event {
    Axis(AxisChange),
    Relabel(Relabel),
    Flags(Flags),
    /// Positions and labels are unknown: re-derive them from what is held.
    Resync,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Want {
    Correction {
        watermark: Option<Stamp>,
    },
    Fetch {
        watermark: Option<Stamp>,
        keys: BTreeSet<WireKey>,
        through: Stamp,
        window: bool,
    },
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Outcome {
    /// Applied; `delta` is what changed in the held cells.
    Applied { delta: ZSet<Cell> },
    /// Already covered by the watermark: nothing was done.
    Duplicate,
    /// Starts after the watermark, or in another epoch: ask again from the watermark.
    Gap,
    /// Names a shape this client does not read.
    Untracked,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Synced {
    Settled,
    Failed,
    /// The rounds ran out with work left: a defect in the edge, not a state to wait out.
    Exhausted,
}

/// How many abandoned epochs a shape remembers.
const LEFT_EPOCHS: usize = 64;

#[derive(Debug, Clone)]
struct Tracked {
    window: Window,
    held: ZSet<Cell>,
    watermark: Option<Stamp>,
    seen: Option<Stamp>,
    /// Notices received and not yet fetched: `(publication, what moved)`.
    unapplied: Vec<(Stamp, Moved)>,
    /// The last notice that moved an axis: until a window is read past it, positions are unknown.
    restructured: Option<Stamp>,
    /// Epochs the watermark has moved out of: nothing from them applies again.
    left: Vec<(u64, u64)>,
    stale: bool,
}

#[derive(Debug, Clone, Default)]
pub struct EdgeClient {
    session: Option<SessionToken>,
    shapes: BTreeMap<ShapeKey, Tracked>,
    events_through: Option<Stamp>,
}

impl EdgeClient {
    pub fn new(session: Option<SessionToken>) -> Self {
        Self {
            session,
            ..Self::default()
        }
    }

    /// Start reading `shape`; the first thing it wants is a full refetch.
    pub fn track(&mut self, shape: ShapeKey, window: Window) {
        self.shapes.insert(
            shape,
            Tracked {
                window,
                held: ZSet::new(),
                watermark: None,
                seen: None,
                unapplied: Vec::new(),
                restructured: None,
                left: Vec::new(),
                stale: true,
            },
        );
    }

    pub fn held(&self, shape: &ShapeKey) -> Option<&ZSet<Cell>> {
        self.shapes.get(shape).map(|t| &t.held)
    }

    pub fn watermark(&self, shape: &ShapeKey) -> Option<Stamp> {
        self.shapes.get(shape).and_then(|t| t.watermark)
    }

    pub fn wants(&self, shape: &ShapeKey) -> Option<Want> {
        let t = self.shapes.get(shape)?;
        let needs_all = t.stale || t.unapplied.iter().any(|(_, m)| *m == Moved::Correct);
        let keys: BTreeSet<WireKey> = t
            .unapplied
            .iter()
            .flat_map(|(_, m)| match m {
                Moved::Keys(keys) => keys.iter().cloned().collect::<Vec<_>>(),
                Moved::Correct => Vec::new(),
            })
            .collect();
        let owed = !t.unapplied.is_empty() || t.restructured.is_some();
        match (needs_all, t.seen) {
            (true, _) => Some(Want::Correction {
                watermark: t.watermark,
            }),
            (false, Some(through)) if owed => Some(Want::Fetch {
                watermark: t.watermark,
                keys,
                through,
                window: t.restructured.is_some(),
            }),
            _ => None,
        }
    }

    pub fn request(&self, shape: &ShapeKey, want: Want) -> Option<Request> {
        let t = self.shapes.get(shape)?;
        Some(Request {
            session: self.session.clone(),
            shape: shape.clone(),
            window: t.window,
            ask: match want {
                Want::Correction { watermark } => Ask::Correction { watermark },
                Want::Fetch {
                    watermark,
                    keys,
                    through,
                    window,
                } => Ask::Fetch {
                    watermark,
                    keys,
                    through,
                    window,
                },
            },
        })
    }

    /// A pushed notice. Duplicates are ignored, a gap or a new epoch makes every shape ask.
    pub fn on_notice(&mut self, notice: &Notice) -> Vec<Event> {
        self.accept(notice, None, false)
    }

    /// `only` limits the notice to one shape: a pulled notice carries no other shape's keys.
    fn accept(&mut self, notice: &Notice, only: Option<&ShapeKey>, answer: bool) -> Vec<Event> {
        let end = notice.stamp();
        let (mut read, mut ignored) = (0, 0);
        for (shape, t) in &mut self.shapes {
            if only.is_some_and(|o| o != shape) {
                continue;
            }
            read += 1;
            let older = |s: Option<Stamp>| {
                s.is_some_and(|s| s.host == end.host && end.generation < s.generation)
            };
            let left = !answer && t.left.contains(&(end.host, end.generation));
            if older(t.watermark) || older(t.seen) || left {
                ignored += 1;
                continue;
            }
            let contiguous = match t.seen {
                Some(seen) if seen.same_epoch(&end) => match notice.range.to <= seen.publication {
                    true => continue,
                    false => notice.range.from <= seen.publication,
                },
                _ => false,
            };
            t.seen = Some(end);
            if !contiguous || notice.evicted.contains(shape) {
                t.unapplied.push((end, Moved::Correct));
                continue;
            }
            if !notice.axis.is_empty() || !notice.relabels.is_empty() || notice.flags.views {
                t.restructured = Some(end);
            }
            match notice.moved.get(shape) {
                Some(Moved::Keys(keys)) if keys.is_empty() => {}
                Some(moved) => t.unapplied.push((end, moved.clone())),
                None => {}
            }
            let current = t.watermark.is_some_and(|w| w.same_epoch(&end));
            if t.unapplied.is_empty() && t.restructured.is_none() && !t.stale && current {
                t.watermark = Some(end);
            }
        }
        match read > 0 && ignored == read {
            true => Vec::new(),
            false => self.events_of(notice),
        }
    }

    /// The connection dropped: every shape asks for a correction from its watermark.
    pub fn connection_lost(&mut self) {
        self.shapes.values_mut().for_each(|t| t.stale = true);
    }

    pub fn request_failed(&mut self, shape: &ShapeKey) {
        if let Some(t) = self.shapes.get_mut(shape) {
            t.stale = true;
        }
    }

    /// A notice that could not be read: what it said is lost, so everything asks.
    pub fn unreadable(&mut self) {
        self.connection_lost();
        self.events_through = None;
    }

    fn events_of(&mut self, n: &Notice) -> Vec<Event> {
        let end = n.stamp();
        let resync = vec![Event::Resync];
        let usable = match self.events_through {
            Some(t) if t.same_epoch(&end) => {
                if n.range.to <= t.publication {
                    return Vec::new();
                }
                n.range.from == t.publication
            }
            _ => false,
        };
        self.events_through = Some(end);
        match usable {
            false => resync,
            true => n
                .axis
                .iter()
                .cloned()
                .map(Event::Axis)
                .chain(n.relabels.iter().cloned().map(Event::Relabel))
                .chain(n.flags.any().then_some(Event::Flags(n.flags)))
                .collect(),
        }
    }

    /// A correction that arrives on its own: it applies only at the watermark it starts from,
    /// as `held + (target - held)`, and never from an epoch the client has left.
    pub fn apply(&mut self, c: &Correction) -> (Outcome, Vec<Event>) {
        self.apply_with(c, false)
    }

    /// The reply to a request this client just made. The host that answers is the live one, so
    /// an epoch the client had written off is live again (`context/code/delivery.md` section 4).
    pub fn apply_answer(&mut self, c: &Correction) -> (Outcome, Vec<Event>) {
        self.apply_with(c, true)
    }

    fn apply_with(&mut self, c: &Correction, answer: bool) -> (Outcome, Vec<Event>) {
        let Some(t) = self.shapes.get_mut(&c.shape) else {
            return (Outcome::Untracked, Vec::new());
        };
        if answer {
            t.left
                .retain(|e| *e != (c.through.host, c.through.generation));
        }
        let target = decode(&c.target);
        let older = |s: Option<Stamp>| {
            s.is_some_and(|s| s.host == c.through.host && c.through.generation < s.generation)
        };
        if older(t.watermark)
            || older(t.seen)
            || t.left.contains(&(c.through.host, c.through.generation))
        {
            return (Outcome::Duplicate, Vec::new());
        }
        let covered = t
            .watermark
            .is_some_and(|w| w.same_epoch(&c.through) && w.publication >= c.through.publication);
        let starts_here = match (&c.from, t.watermark) {
            (Some(from), Some(w)) => w.same_epoch(from) && w.publication >= from.publication,
            _ => false,
        };
        let delta = match (&c.body, covered) {
            (_, true) => {
                t.stale &= t.watermark != Some(c.through);
                let live = c.through;
                t.unapplied
                    .retain(|(at, _)| at.same_epoch(&live) && at.publication > live.publication);
                if t.restructured.is_some_and(|r| !r.same_epoch(&live)) {
                    t.restructured = None;
                }
                if t.seen.is_some_and(|s| !s.same_epoch(&live)) {
                    t.seen = Some(live);
                }
                return (Outcome::Duplicate, Vec::new());
            }
            (CorrectionBody::Window, false) => &target - &t.held,
            (CorrectionBody::Moved { .. }, false) if !starts_here => {
                return (Outcome::Gap, Vec::new());
            }
            (CorrectionBody::Moved { keys }, false) => {
                let held_there: ZSet<Cell> = t
                    .held
                    .iter()
                    .filter(|((k, _), _)| keys.contains(k))
                    .map(|(cell, w)| (cell.clone(), w))
                    .collect();
                &target - &held_there
            }
        };
        t.held += &delta;
        if let Some(w) = t.watermark.filter(|w| !w.same_epoch(&c.through)) {
            t.left.push((w.host, w.generation));
            let excess = t.left.len().saturating_sub(LEFT_EPOCHS);
            t.left.drain(..excess);
        }
        t.watermark = Some(c.through);
        let window = matches!(c.body, CorrectionBody::Window);
        t.stale &= !window;
        if window
            && t.restructured.is_some_and(|r| {
                !r.same_epoch(&c.through) || r.publication <= c.through.publication
            })
        {
            t.restructured = None;
        }
        let seen_epoch = t.seen.map(|s| (s.host, s.generation));
        t.unapplied.retain(|(at, _)| {
            Some((at.host, at.generation)) == seen_epoch
                && !(at.same_epoch(&c.through) && at.publication <= c.through.publication)
        });
        let seen_ahead = t
            .seen
            .is_some_and(|s| s.same_epoch(&c.through) && s.publication >= c.through.publication);
        if !seen_ahead {
            t.seen = Some(c.through);
        }
        let events = match (&c.carried, window) {
            (Some(carried), _) => self.events_of(carried),
            (None, true) => {
                self.events_through = Some(c.through);
                vec![Event::Resync]
            }
            (None, false) => Vec::new(),
        };
        (Outcome::Applied { delta }, events)
    }

    /// Ask until `shape` wants nothing, at most `rounds` times.
    pub fn sync(
        &mut self,
        edge: &mut impl Edge,
        shape: &ShapeKey,
        rounds: usize,
    ) -> (Synced, Vec<Event>) {
        let mut events = Vec::new();
        for _ in 0..rounds {
            let Some(want) = self.wants(shape) else {
                return (Synced::Settled, events);
            };
            let Some(request) = self.request(shape, want) else {
                return (Synced::Settled, events);
            };
            match edge.call(request) {
                Ok(Reply::Correction(c)) => events.extend(self.apply_answer(&c).1),
                Ok(Reply::Notices(list)) => {
                    list.iter()
                        .for_each(|n| events.extend(self.accept(n, Some(shape), true)));
                }
                Err(EdgeError::Failed) => {
                    self.request_failed(shape);
                    return (Synced::Failed, events);
                }
                Err(EdgeError::Unreadable) => {
                    self.unreadable();
                    return (Synced::Failed, events);
                }
            }
        }
        match self.wants(shape) {
            None => (Synced::Settled, events),
            Some(_) => (Synced::Exhausted, events),
        }
    }

    /// Pull instead of listening: ask for the notices after the horizon seen.
    pub fn poll(&mut self, edge: &mut impl Edge, shape: &ShapeKey) -> (Synced, Vec<Event>) {
        let (Some(t), Some(after)) = (
            self.shapes.get(shape),
            self.shapes.get(shape).and_then(|t| t.seen),
        ) else {
            return self.sync(edge, shape, 4);
        };
        let request = Request {
            session: self.session.clone(),
            shape: shape.clone(),
            window: t.window,
            ask: Ask::Notices { after },
        };
        let mut events = Vec::new();
        match edge.call(request) {
            Ok(Reply::Notices(list)) => list
                .iter()
                .for_each(|n| events.extend(self.accept(n, Some(shape), true))),
            Ok(Reply::Correction(c)) => events.extend(self.apply_answer(&c).1),
            Err(EdgeError::Failed) => {
                self.request_failed(shape);
                return (Synced::Failed, events);
            }
            Err(EdgeError::Unreadable) => {
                self.unreadable();
                return (Synced::Failed, events);
            }
        }
        let (synced, more) = self.sync(edge, shape, 4);
        events.extend(more);
        (synced, events)
    }
}