From 56eaec41f881b7637d3e0608c5a676f93c64f89b Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Mon, 14 Sep 2026 15:34:23 +0200 Subject: [PATCH 1/4] feat: add the atlas delta pipeline (feed, placement, projector, task) --- .../graph/atlas/src/serve/delta/feed/mod.rs | 608 ++++++++++++ .../atlas/src/serve/delta/feed/pending/mod.rs | 403 ++++++++ .../src/serve/delta/feed/pending/tests.rs | 868 ++++++++++++++++++ .../graph/atlas/src/serve/delta/feed/tests.rs | 524 +++++++++++ .../@local/graph/atlas/src/serve/delta/mod.rs | 9 + .../atlas/src/serve/delta/placement/mod.rs | 506 ++++++++++ .../src/serve/delta/placement/pending/mod.rs | 180 ++++ .../serve/delta/placement/pending/tests.rs | 460 ++++++++++ .../atlas/src/serve/delta/placement/tests.rs | 466 ++++++++++ .../atlas/src/serve/delta/projector/mod.rs | 448 +++++++++ .../atlas/src/serve/delta/projector/tests.rs | 329 +++++++ .../graph/atlas/src/serve/delta/task/mod.rs | 311 +++++++ .../graph/atlas/src/serve/delta/task/tests.rs | 272 ++++++ libs/@local/graph/atlas/src/serve/mod.rs | 1 + libs/@local/graph/atlas/src/serve/walk/mod.rs | 114 +++ .../graph/atlas/src/serve/walk/tests.rs | 129 +++ 16 files changed, 5628 insertions(+) create mode 100644 libs/@local/graph/atlas/src/serve/delta/feed/mod.rs create mode 100644 libs/@local/graph/atlas/src/serve/delta/feed/pending/mod.rs create mode 100644 libs/@local/graph/atlas/src/serve/delta/feed/pending/tests.rs create mode 100644 libs/@local/graph/atlas/src/serve/delta/feed/tests.rs create mode 100644 libs/@local/graph/atlas/src/serve/delta/placement/mod.rs create mode 100644 libs/@local/graph/atlas/src/serve/delta/placement/pending/mod.rs create mode 100644 libs/@local/graph/atlas/src/serve/delta/placement/pending/tests.rs create mode 100644 libs/@local/graph/atlas/src/serve/delta/placement/tests.rs create mode 100644 libs/@local/graph/atlas/src/serve/delta/projector/mod.rs create mode 100644 libs/@local/graph/atlas/src/serve/delta/projector/tests.rs create mode 100644 libs/@local/graph/atlas/src/serve/delta/task/mod.rs create mode 100644 libs/@local/graph/atlas/src/serve/delta/task/tests.rs create mode 100644 libs/@local/graph/atlas/src/serve/walk/mod.rs create mode 100644 libs/@local/graph/atlas/src/serve/walk/tests.rs diff --git a/libs/@local/graph/atlas/src/serve/delta/feed/mod.rs b/libs/@local/graph/atlas/src/serve/delta/feed/mod.rs new file mode 100644 index 00000000000..68b8dedaf38 --- /dev/null +++ b/libs/@local/graph/atlas/src/serve/delta/feed/mod.rs @@ -0,0 +1,608 @@ +//! Event coalescing and placement delivery for one mutable delta. +//! +//! Each entity retains its newest event version across overlapping feed reads. Placements match +//! completions to pending versions under [`EventId`]'s range conditions. A withdrawal discards its +//! pending update, preventing a later result from reviving it. Store poll failures retry the poll. +//! Classification and display failures retry without advancing the failed stage. Missing +//! classification or display rows also retry. Within one feed task, a delivered placement failure +//! is terminal for that event version, whose overlap record remains until an event with a later +//! transaction time replaces it. + +mod pending; + +#[cfg(test)] +mod tests; + +use alloc::sync::Arc; +use core::{fmt, future::Future, ops::ControlFlow, pin::pin, time::Duration}; + +use arc_swap::ArcSwap; +use error_stack::{Report, ReportSink, ResultExt as _}; +use futures::StreamExt as _; +use hash_graph_postgres_store::store::{EntityEvent, PostgresStorePool}; +use hash_graph_store::{error::QueryError, pool::StorePool as _}; +use hash_graph_temporal_versioning::{Timestamp, TransactionTime}; +use hashql_core::{collections::FastHashMap, id::Id as _}; +use tokio::{ + sync::{ + Notify, + mpsc::{self, Permit, error::TrySendError}, + oneshot, + }, + time::Interval, +}; +use type_system::knowledge::entity::id::EntityEditionId; + +use self::pending::{DeltaAction, Pending}; +use super::{ + Delta, + placement::{Completed, Initial, PendingEntry}, +}; +use crate::{ + math::Vec2, + postgres::{self, edition_display::DisplayParts, id::ArchivedEntityId}, +}; + +/// A feed configuration or execution failure. +#[derive(Debug)] +pub(super) enum DeltaFeedError { + /// The polling interval is zero or outside the monotonic clock's range. + InvalidInterval, + /// The replay safety lag is unrepresentable or precedes the transaction-time range. + InvalidSafetyLag, + /// The publication exchange closed before answering. + Closed, + /// Acquiring a store connection failed. + Connect, + /// Reading the entity-event stream failed. + Event, + /// Classifying pending entities failed. + Classification, + /// Reading pending editions' display metadata failed. + Display, + /// The placement request or completion channel closed. + PlacementClosed, +} + +impl fmt::Display for DeltaFeedError { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::InvalidInterval => fmt.write_str( + "the feed polling interval must be non-zero and fit the monotonic clock", + ), + Self::InvalidSafetyLag => { + fmt.write_str("the safety lag precedes the representable transaction-time range") + } + Self::Closed => fmt.write_str("the feed task stopped before publication"), + Self::Connect => fmt.write_str("failed to connect to the database"), + Self::Event => fmt.write_str("failed to read entity events"), + Self::Classification => fmt.write_str("failed to classify entities"), + Self::Display => fmt.write_str("failed to read edition displays"), + Self::PlacementClosed => fmt.write_str("the placement task stopped"), + } + } +} + +impl core::error::Error for DeltaFeedError {} + +hashql_core::id::newtype! { + /// A per-delta sequence number for a coalesced entity update. + /// + /// New unarchived updates advance the counter through `usize` arithmetic. + /// + /// # Panics + /// + /// At `u32::MAX`, the next increment panics on a 64-bit target or with overflow checking on a 32-bit target. + /// + /// # Warning + /// + /// On a 32-bit target without overflow checking, the counter wraps after 2³² increments. A delayed placement result can then match a newer update of the same entity if its event ID repeats. + pub(crate) struct EventId(u32) +} + +/// Submits placement requests for entities awaiting a wire position, in encounter order. +/// +/// A supplied `permit` covers exactly one request. Once spent, further requests use `tx` directly +/// and stop at its first full or closed channel. +fn pump( + pending: &mut Pending, + tx: &mpsc::Sender>, + mut permit: Option>>, +) -> ControlFlow<()> { + let mut placements = pending.placements(); + while let Some(placement) = placements.next_placement() { + let request = placement.request(); + if let Some(permit) = permit.take() { + permit.send(request); + } else { + match tx.try_send(request) { + Ok(()) => {} + Err(TrySendError::Full(_)) => break, + Err(TrySendError::Closed(_)) => return ControlFlow::Break(()), + } + } + placement.submitted(); + } + ControlFlow::Continue(()) +} + +/// The feed's replay lag and polling cadence. +#[derive(Copy, Clone)] +pub(crate) struct DeltaFeedTaskOptions { + /// How far behind the transaction-time watermark the feed replays events from. + pub safety_lag: Duration, + /// How often the feed polls for new events and re-evaluates pending work. + pub tick_rate: Duration, +} + +/// The feed's ends of the placement request and completion channels. +pub(super) struct Placement { + /// Placement requests sent after classification identifies an unplaced node. + pub requests: mpsc::Sender>, + /// Placement outcomes returned for version-checked application. + pub completed: mpsc::Receiver>, +} + +/// A single publisher's change notification and delta exchange. +/// +/// The notification stores at most one wake-up. Changes coalesce while the publisher is busy. The +/// next exchange transfers the feed's complete accumulated working delta. +pub(super) struct Publication { + notify: Arc, + update: mpsc::Sender<(Delta, oneshot::Sender)>, +} + +impl Publication { + /// Waits for the next change notification, then exchanges the caller's delta for the task's. + /// + /// # Errors + /// + /// Returns [`DeltaFeedError::Closed`] if the feed task stops before completing the exchange. + #[expect( + clippy::integer_division_remainder_used, + reason = "Tokio select uses a remainder to traverse its branch set" + )] + pub(super) async fn next(&self, previous: Arc) -> Result> { + tokio::select! { + biased; + () = self.update.closed() => return Err(Report::new(DeltaFeedError::Closed)), + () = self.notify.notified() => {} + } + + let (tx, rx) = oneshot::channel(); + self.update + .send((Arc::unwrap_or_clone(previous), tx)) + .await + .change_context(DeltaFeedError::Closed)?; + + rx.await.change_context(DeltaFeedError::Closed) + } + + /// Exchanges deltas after change notifications until `shutdown` resolves. + /// + /// Multiple notifications may coalesce before one exchange. Shutdown wins over a pending + /// notification. This leaves the last already-exchanged delta published rather than forcing + /// the feed's working copy into the reader. + /// + /// # Errors + /// + /// Returns [`DeltaFeedError::Closed`] if the feed task stops before shutdown. + #[expect( + clippy::integer_division_remainder_used, + reason = "Tokio select uses a remainder to traverse its branch set" + )] + pub(super) async fn run( + self, + current: Arc>, + mut previous: Arc, + shutdown: impl Future, + ) -> Result<(), Report> { + let mut shutdown = pin!(shutdown); + + // `next` waits for notification before `Arc::unwrap_or_clone(previous)`. The wait allows + // holders of `previous` to drop their strong references. At one strong reference, the + // operation moves the `Delta` value without cloning it. + loop { + let next = tokio::select! { + biased; + () = &mut shutdown => return Ok(()), + next = self.next(previous) => next?, + }; + + previous = current.swap(Arc::new(next)); + } + } +} + +/// Reused buffers for one tick's database, classification and display work. +#[derive(Default)] +struct DeltaFeedTaskScratch { + events: Vec, + positions: Vec, + entities: Vec, + editions: Vec, + displays: FastHashMap>, +} + +/// One generation's entity-feed task for polling, event coalescing and delta application. +pub(super) struct DeltaFeedTask { + delta: Delta, + pool: Arc, + options: DeltaFeedTaskOptions, + + watermark: Timestamp, + safety_lag: ::time::Duration, + replayed: bool, + + update: mpsc::Receiver<(Delta, oneshot::Sender)>, + notify: Arc, + + placement: Option, + + pending: Pending, + scratch: DeltaFeedTaskScratch, +} + +impl DeltaFeedTask { + /// Prepares replay from `delta` and the base snapshot's transaction-time `watermark`. + /// + /// The working revision increments once under [`DeltaRevision`](super::DeltaRevision)'s + /// arithmetic conditions. Without placement channels, new node placements remain pending while + /// withdrawals and metadata updates continue. Construction starts no task and opens no database + /// connection. + /// + /// # Errors + /// + /// Returns [`DeltaFeedError::InvalidInterval`] for a zero or unrepresentable polling interval. + /// Returns [`DeltaFeedError::InvalidSafetyLag`] when the timestamp types cannot represent the + /// safety lag or the resulting replay window precedes their range. + /// + /// # Panics + /// + /// With overflow checking, panics if option validation succeeds and the revision's `usize` + /// conversion equals `usize::MAX`. + pub(super) fn new( + mut delta: Delta, + pool: Arc, + options: DeltaFeedTaskOptions, + watermark: Timestamp, + placement: Option, + ) -> Result<(Self, Publication), Report> { + if options.tick_rate.is_zero() + || tokio::time::Instant::now() + .checked_add(options.tick_rate) + .is_none() + { + return Err(Report::new(DeltaFeedError::InvalidInterval)); + } + + let safety_lag = ::time::Duration::try_from(options.safety_lag) + .change_context(DeltaFeedError::InvalidSafetyLag)?; + + let earliest = Timestamp::from_unix_timestamp( + ::time::Date::MIN.midnight().assume_utc().unix_timestamp(), + ); + + if safety_lag > watermark - earliest { + return Err(Report::new(DeltaFeedError::InvalidSafetyLag)); + } + + let (tx, rx) = mpsc::channel(1); + + let notify = Arc::new(Notify::new()); + let publication = Publication { + notify: Arc::clone(¬ify), + update: tx, + }; + + delta.revision.increment_by(1); + + let this = Self { + delta, + pool, + options, + watermark, + safety_lag, + replayed: false, + update: rx, + notify, + placement, + pending: Pending::default(), + scratch: DeltaFeedTaskScratch::default(), + }; + + Ok((this, publication)) + } + + /// Runs the feed's tick loop until its update channel closes, notifying on every change. + /// + /// The task logs store polling, classification and display failures and retries them on later + /// ticks. They do not end this loop. A closed update channel is normal publication shutdown. + /// + /// # Errors + /// + /// Returns [`DeltaFeedError::PlacementClosed`] if the placement channels close. + /// + /// # Panics + /// + /// Panics unless polled within a Tokio runtime with time enabled. A sufficiently late tick can + /// also panic when adding [`DeltaFeedTaskOptions::tick_rate`] to the current instant would + /// exceed Tokio's representable deadline. The panic conditions on [`Self::step`] also apply. + pub(super) async fn run(mut self) -> Result<(), Report> { + let mut interval = tokio::time::interval(self.options.tick_rate); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + + while let ControlFlow::Continue(changed) = self.step(&mut interval).await? { + if changed { + self.notify.notify_one(); + } + } + + Ok(()) + } + + /// Advances one feed event. + /// + /// An event is a poll, a placement submission, a placement result or an exchange. + /// + /// Returns whether the delta changed, or [`ControlFlow::Break`] once the update channel closes. + /// + /// # Errors + /// + /// Returns [`DeltaFeedError::PlacementClosed`] if the placement channels close. + /// + /// # Panics + /// + /// Panics under the conditions on [`Self::tick`] or [`Self::exchange`]. A late tick can also + /// panic if `interval` computes an unrepresentable next deadline. + #[expect( + clippy::integer_division_remainder_used, + reason = "Tokio select uses a remainder to randomize its first branch" + )] + async fn step( + &mut self, + interval: &mut Interval, + ) -> Result, Report> { + let tx = self + .placement + .as_ref() + .map(|placement| placement.requests.clone()); + + let changed = tokio::select! { + _ = interval.tick() => match self.tick().await { + Ok(changed) => changed, + Err(error) => { + tracing::warn!(?error, "Retry the entity feed on the next tick"); + false + } + }, + permit = async { + match &tx { + Some(tx) => tx.reserve().await, + None => core::future::pending().await, + } + }, if self.pending.has_placements() => { + let permit = permit.change_context(DeltaFeedError::PlacementClosed)?; + let tx = tx.as_ref().expect("should have a placement sender after reserving capacity"); + if pump(&mut self.pending, tx, Some(permit)).is_break() { + return Err(Report::new(DeltaFeedError::PlacementClosed)); + } + + false + }, + result = async { + match &mut self.placement { + Some(placement) => placement.completed.recv().await, + None => core::future::pending().await, + } + } => { + let result = result.ok_or_else(|| Report::new(DeltaFeedError::PlacementClosed))?; + let placement = self.placement.as_mut().expect("should have placement channels after receiving a result"); + self.pending.receive(result); + while let Ok(result) = placement.completed.try_recv() { + self.pending.receive(result); + } + + self.pending.normalize(self.delta.world.bounds(), &mut self.scratch.positions); + self.pending.apply(&mut self.delta) + }, + request = self.update.recv() => { + let Some((previous, reply)) = request else { + return Ok(ControlFlow::Break(())); + }; + + self.exchange(previous, reply); + false + }, + }; + + if let Some(placement) = &self.placement + && pump(&mut self.pending, &placement.requests, None).is_break() + { + return Err(Report::new(DeltaFeedError::PlacementClosed)); + } + + Ok(ControlFlow::Continue(changed)) + } + + /// Exchanges the working delta and reuses `previous` for further changes. + /// + /// Revision arithmetic follows [`DeltaRevision`](super::DeltaRevision). If the current revision + /// fits `usize`, a cancelled reply restores that revision. Narrowing on a 32-bit target can + /// instead discard its upper bits before the increment and rollback. + /// + /// # Panics + /// + /// With overflow checking, panics if the current revision's `usize` conversion equals + /// `usize::MAX`. + fn exchange(&mut self, mut previous: Delta, reply: oneshot::Sender) { + previous.clone_from(&self.delta); + previous.revision.increment_by(1); + + let requested = core::mem::replace(&mut self.delta, previous); + if reply.send(requested).is_err() { + self.delta.revision.decrement_by(1); + } + } + + /// Reads entity events since the safety-lagged watermark and applies their withdrawals. + /// + /// Returns whether the delta changed. + /// + /// # Errors + /// + /// Returns [`DeltaFeedError::Connect`] or [`DeltaFeedError::Event`] for a store failure. + /// + /// # Panics + /// + /// Panics under [`Self::apply_events`]' counter and history-order conditions after a successful + /// stream read. + async fn poll_database(&mut self) -> Result> { + self.scratch.events.clear(); + let pool = Arc::clone(&self.pool); + let store = pool + .acquire(None) + .await + .change_context(DeltaFeedError::Connect)?; + + let mut sink = ReportSink::::new_armed(); + + let stream = store.entity_events_since(self.watermark - self.safety_lag); + let mut stream = pin!(stream); + + while let Some(result) = stream.next().await { + if let Some(event) = sink.attempt(result) { + self.scratch.events.push(event); + } + } + + sink.finish().change_context(DeltaFeedError::Event)?; + Ok(self.apply_events()) + } + + /// Coalesces drained events and advances the transaction-time high-water mark. + /// + /// Each event raises the watermark to the greater of its current value and the event's + /// transaction time. Input order cannot move it backward. + /// + /// Returns whether the delta changed, including the first call after construction. + /// + /// # Panics + /// + /// Panics if a newer unarchived update exhausts the event counter under [`Pending::observe`]'s + /// target-width and overflow-checking conditions. Withdrawals also have [`Delta::withdraw`]'s + /// history-order condition. + fn apply_events(&mut self) -> bool { + let mut changed = false; + + for event in self.scratch.events.drain(..) { + self.watermark = self.watermark.max(Pending::identity_and_time(&event).1); + match self.pending.observe(event) { + DeltaAction::None => {} + DeltaAction::Withdraw(entity) => changed |= self.delta.withdraw(entity), + } + } + + changed |= !self.replayed; + self.replayed = true; + changed + } + + /// Classifies entities awaiting a node/edge decision and advances their pending stage. + /// + /// # Errors + /// + /// Returns [`DeltaFeedError::Connect`] or [`DeltaFeedError::Classification`] for a store + /// failure. + async fn classify(&mut self) -> Result<(), Report> { + self.scratch.entities.clear(); + self.scratch.entities.extend(self.pending.classifications()); + + if self.scratch.entities.is_empty() { + return Ok(()); + } + + let pool = Arc::clone(&self.pool); + let store = pool + .acquire(None) + .await + .change_context(DeltaFeedError::Connect)?; + let classifications = postgres::classify_entities(&store, self.scratch.entities.drain(..)) + .await + .change_context(DeltaFeedError::Classification)?; + + self.pending.classify(&self.delta, classifications); + Ok(()) + } + + /// Reads edition displays for entities awaiting one and advances their pending stage. + /// + /// Returns whether the delta changed. + /// + /// # Errors + /// + /// Returns [`DeltaFeedError::Connect`] or [`DeltaFeedError::Display`] for a store failure. + /// + /// # Panics + /// + /// Panics under [`Pending::capture`]'s revision and allocation conditions. + async fn capture_displays(&mut self) -> Result> { + self.scratch.editions.clear(); + self.scratch.editions.extend(self.pending.editions()); + + if self.scratch.editions.is_empty() { + return Ok(false); + } + let total_editions = self.scratch.editions.len(); + + let store = self + .pool + .acquire(None) + .await + .change_context(DeltaFeedError::Connect)?; + let displays = postgres::read_edition_displays(&store, self.scratch.editions.drain(..)) + .await + .change_context(DeltaFeedError::Display)?; + + if displays.len() != total_editions { + tracing::warn!( + answers = displays.len(), + requests = total_editions, + "The display read returned a different edition count" + ); + } + + self.scratch.displays.extend(displays); + Ok(self + .pending + .capture(&mut self.delta, &mut self.scratch.displays)) + } + + /// Runs one poll/classify/capture/apply cycle, retrying a failed stage on the next tick. + /// + /// Returns whether the delta changed. + /// + /// # Errors + /// + /// Returns [`DeltaFeedError`] only for a database poll failure. Classification and display + /// failures are logged and retried. + /// + /// # Panics + /// + /// The panic conditions on [`Self::poll_database`], [`Self::capture_displays`] and + /// [`Pending::apply`] also apply. + async fn tick(&mut self) -> Result> { + let mut changed = self.poll_database().await?; + if let Err(error) = self.classify().await { + tracing::warn!(?error, "Retry entity classification on the next tick"); + } + + match self.capture_displays().await { + Ok(captured) => changed |= captured, + Err(error) => tracing::warn!(?error, "Retry edition displays on the next tick"), + } + + changed |= self.pending.apply(&mut self.delta); + Ok(changed) + } +} diff --git a/libs/@local/graph/atlas/src/serve/delta/feed/pending/mod.rs b/libs/@local/graph/atlas/src/serve/delta/feed/pending/mod.rs new file mode 100644 index 00000000000..f6a703d4db9 --- /dev/null +++ b/libs/@local/graph/atlas/src/serve/delta/feed/pending/mod.rs @@ -0,0 +1,403 @@ +//! Versioned per-entity work from feed event through delta application. + +use hash_graph_postgres_store::store::{EntityEvent, EntityUpdate}; +use hash_graph_temporal_versioning::{Timestamp, TransactionTime}; +use hashbrown::hash_map::IterMut; +use hashql_core::{collections::FastHashMap, id::Id as _}; +use type_system::knowledge::entity::id::EntityEditionId; + +use super::EventId; +use crate::{ + dataset::auxiliary::OwnedLegend, + math::{Bounds2, Vec2}, + postgres::{Classification, edition_display::DisplayParts, id::ArchivedEntityId}, + salt::lod::stage::WIRE_FRAME, + serve::delta::{ + Delta, + placement::{Completed, Initial, PendingEntry}, + projector::Position, + }, +}; + +#[cfg(test)] +mod tests; + +/// A classified entity's wire-frame position or unresolved edge endpoints. +#[derive(Debug, Clone, Copy)] +enum Geometry { + Node(Vec2), + Edge(Option<[ArchivedEntityId; 2]>), +} + +/// One entity update's progress from classification to capture. +#[derive(Debug)] +enum Stage { + Classify, + Place, + Placing, + Normalize(Position), + Capture(Geometry), + Ready { + geometry: Geometry, + legend: OwnedLegend, + }, +} + +/// The newest coalesced update for one entity, with its staged progress. +struct Update { + id: EventId, + event: EntityUpdate, + stage: Stage, +} + +/// A queued placement held until submission succeeds. +pub(super) struct Placement<'pending> { + entity: ArchivedEntityId, + update: &'pending mut Update, + placements: &'pending mut FastHashMap, +} + +impl Placement<'_> { + /// Builds the placement request for this entity's newest coalesced event. + pub(super) const fn request(&self) -> PendingEntry { + PendingEntry { + event: self.update.id, + entity: self.update.event.entity, + phase: Initial, + } + } + + /// Marks this entity's placement request as outstanding. + pub(super) fn submitted(self) { + self.placements.insert(self.entity, self.update.id); + self.update.stage = Stage::Placing; + } +} + +/// Entities awaiting a placement request, excluding those already outstanding. +pub(super) struct Placements<'pending> { + updates: IterMut<'pending, ArchivedEntityId, Update>, + outstanding: &'pending mut FastHashMap, +} + +impl Placements<'_> { + /// Returns the next entity awaiting a placement request, marking none as submitted yet. + pub(super) fn next_placement(&mut self) -> Option> { + let (&entity, update) = self.updates.find(|(entity, update)| { + matches!(update.stage, Stage::Place) && !self.outstanding.contains_key(*entity) + })?; + Some(Placement { + entity, + update, + placements: self.outstanding, + }) + } +} + +/// An immediate entity-visibility decision. +#[derive(Debug)] +pub(super) enum DeltaAction { + /// An ignored event or an update awaiting later application. + None, + Withdraw(ArchivedEntityId), +} + +/// The newest event time and at most one placement request per entity. +/// +/// A newer event replaces the entity's retained transaction-time high-water mark and unfinished +/// work. Its placement waits until the preceding request completes, even after a withdrawal. +/// Completion or failure leaves that newest time recorded until this feed task ends. Consequently, +/// an overlapping replay at an equal or older time does not retry a delivered placement failure in +/// the same task. Only a later event version starts new work. +pub(super) struct Pending { + counter: EventId, + versions: FastHashMap>, + placements: FastHashMap, + updates: FastHashMap, +} + +impl Pending { + /// Returns an event's entity identity and the transaction time it takes effect. + pub(super) fn identity_and_time( + event: &EntityEvent, + ) -> (ArchivedEntityId, Timestamp) { + match event { + EntityEvent::Updated(update) => { + (ArchivedEntityId::from(update.entity), update.changed_at) + } + EntityEvent::Ended(end) => (ArchivedEntityId::from(end.entity), end.ended_at), + EntityEvent::Deleted(deletion) => ( + ArchivedEntityId::from(deletion.entity), + deletion.provenance.deleted_at_transaction_time, + ), + } + } + + /// Coalesces a newer event and selects its immediate delta action. + /// + /// Equal or older timestamps leave both pending work and its placement request unchanged. + /// + /// # Panics + /// + /// Panics when a newer, unarchived update must advance the counter past `u32::MAX`, either on a + /// 64-bit target or on a 32-bit target with overflow checking. A 32-bit target without overflow + /// checking wraps instead, with the reuse limitations documented on [`EventId`]. + #[expect( + clippy::needless_pass_by_value, + reason = "The feed transfers ownership of buffered events to the coalescer" + )] + pub(super) fn observe(&mut self, event: EntityEvent) -> DeltaAction { + let (entity, time) = Self::identity_and_time(&event); + if self.versions.get(&entity).is_some_and(|&held| held >= time) { + return DeltaAction::None; + } + + match event { + EntityEvent::Updated(event) if !event.archived => { + let id = self.counter; + self.counter.increment_by(1); + self.versions.insert(entity, time); + self.updates.insert( + entity, + Update { + id, + event, + stage: Stage::Classify, + }, + ); + DeltaAction::None + } + EntityEvent::Updated(_) | EntityEvent::Ended(_) | EntityEvent::Deleted(_) => { + self.versions.insert(entity, time); + self.updates.remove(&entity); + DeltaAction::Withdraw(entity) + } + } + } + + /// Returns entities awaiting classification. + pub(super) fn classifications(&self) -> impl Iterator { + self.updates.iter().filter_map(|(&entity, update)| { + matches!(update.stage, Stage::Classify).then_some(entity) + }) + } + + /// Advances classified entities to placement or capture, by their resolved kind. + /// + /// An entity no longer awaiting classification, or missing from `classifications`, is left + /// unchanged. + pub(super) fn classify( + &mut self, + delta: &Delta, + classifications: impl IntoIterator, + ) { + for (entity, classification) in classifications { + let Some(update) = self.updates.get_mut(&entity) else { + continue; + }; + if !matches!(update.stage, Stage::Classify) { + continue; + } + + update.stage = match classification { + Classification::Node => delta + .node_position(entity) + .map_or(Stage::Place, |position| { + Stage::Capture(Geometry::Node(position)) + }), + Classification::Edge { source, target } => { + let endpoints = source.zip(target).map(<[_; 2]>::from); + if endpoints.is_none() { + tracing::warn!(?entity, "The link has an incomplete endpoint pair"); + } + Stage::Capture(Geometry::Edge(endpoints)) + } + }; + } + } + + /// Returns whether an entity awaits a placement request that is not yet outstanding. + pub(super) fn has_placements(&self) -> bool { + self.updates.iter().any(|(entity, update)| { + matches!(update.stage, Stage::Place) && !self.placements.contains_key(entity) + }) + } + + /// Returns a view over entities awaiting a placement request. + pub(super) fn placements(&mut self) -> Placements<'_> { + Placements { + updates: self.updates.iter_mut(), + outstanding: &mut self.placements, + } + } + + /// Matches a placement result against its outstanding request and current update. + /// + /// The task drops a reply for a since-removed entity without logging an error or applying it. + /// Superseded-event rejection relies on distinct [`EventId`] values: counter wrap can make an + /// old reply match the same entity's newer update. A failed placement discards the pending + /// update but leaves the entity's newest observed transaction time unchanged for the rest of + /// this feed task. An overlapping replay at an equal or older time does not retry the failed + /// placement. + pub(super) fn receive( + &mut self, + PendingEntry { + event, + entity, + phase: Completed(result), + }: PendingEntry, + ) { + let entity = ArchivedEntityId::from(entity); + if self.placements.get(&entity) != Some(&event) { + return; + } + self.placements.remove(&entity); + + let Some(update) = self + .updates + .get_mut(&entity) + .filter(|update| update.id == event) + else { + return; + }; + + match result { + Ok(position) => update.stage = Stage::Normalize(position), + Err(error) => { + self.updates.remove(&entity); + tracing::warn!(?entity, ?error, "Discard the failed placement"); + } + } + } + + /// Normalizes placed node positions into the wire frame and advances them to capture. + /// + /// `scratch` is cleared on entry and on return. + pub(super) fn normalize(&mut self, bounds: Bounds2, scratch: &mut Vec) { + scratch.clear(); + scratch.extend( + self.updates + .values() + .filter_map(|update| match &update.stage { + Stage::Normalize(position) => Some(position.get()), + Stage::Classify + | Stage::Place + | Stage::Placing + | Stage::Capture(_) + | Stage::Ready { .. } => None, + }), + ); + if scratch.is_empty() { + return; + } + + let positions = bounds.normalize_into(WIRE_FRAME, scratch); + // Both traversals use the same table order without inserting or removing entries. + let updates = self + .updates + .values_mut() + .filter(|update| matches!(update.stage, Stage::Normalize(_))); + for (update, position) in updates.zip(positions) { + update.stage = Stage::Capture(Geometry::Node(position)); + } + scratch.clear(); + } + + /// Returns editions awaiting a display read. + pub(super) fn editions(&self) -> impl Iterator { + self.updates.values().filter_map(|update| { + matches!(update.stage, Stage::Capture(_)).then_some(update.event.edition) + }) + } + + /// Advances entities with a captured display to ready, registering their ontology icon. + /// + /// Returns whether `delta` changed. An entity missing from `displays`, or with no display + /// found, is left unchanged and retried on the next tick. + /// + /// # Panics + /// + /// Panics under [`Delta::register_ontology`]'s revision condition or [`OwnedLegend::new`]'s + /// allocation conditions. + pub(super) fn capture( + &mut self, + delta: &mut Delta, + displays: &mut FastHashMap>, + ) -> bool { + let mut changed = false; + for update in self.updates.values_mut() { + let Stage::Capture(geometry) = update.stage else { + continue; + }; + let Some(Some(DisplayParts { + label, + icon, + representative, + })) = displays.remove(&update.event.edition) + else { + continue; + }; + let Some((representative, registered)) = delta.register_ontology(representative, icon) + else { + tracing::warn!(entity = ?update.event.entity, "No ontology row remains for the display"); + continue; + }; + + changed |= registered; + update.stage = Stage::Ready { + geometry, + legend: OwnedLegend::new(representative, &label), + }; + } + displays.clear(); + changed + } + + /// Applies every ready update to `delta`, removing entities whose row remains available. + /// + /// Returns whether `delta` changed. An entity with no row available, or an edge still missing + /// an endpoint, is retried on a later call. + /// + /// # Panics + /// + /// Panics under the revision conditions on [`Delta::update_node`] and [`Delta::update_edge`]. + pub(super) fn apply(&mut self, delta: &mut Delta) -> bool { + let mut changed = false; + self.updates.retain(|&entity, update| { + let Stage::Ready { geometry, legend } = &update.stage else { + return true; + }; + let outcome = match *geometry { + Geometry::Node(position) => delta.update_node(entity, legend.clone(), position), + Geometry::Edge(Some([source, target])) => { + let (Some(source), Some(target)) = + (delta.node_row(source), delta.node_row(target)) + else { + return true; + }; + delta.update_edge(entity, legend.clone(), Some([source, target])) + } + Geometry::Edge(None) => return true, + }; + + let Some(applied) = outcome else { + tracing::warn!(?entity, "No entity row remains for the update"); + return true; + }; + changed |= applied; + false + }); + changed + } +} + +impl Default for Pending { + fn default() -> Self { + Self { + counter: EventId::MIN, + versions: FastHashMap::default(), + placements: FastHashMap::default(), + updates: FastHashMap::default(), + } + } +} diff --git a/libs/@local/graph/atlas/src/serve/delta/feed/pending/tests.rs b/libs/@local/graph/atlas/src/serve/delta/feed/pending/tests.rs new file mode 100644 index 00000000000..fd3082b1159 --- /dev/null +++ b/libs/@local/graph/atlas/src/serve/delta/feed/pending/tests.rs @@ -0,0 +1,868 @@ +use alloc::sync::Arc; +use core::assert_matches; + +use arc_swap::Guard; +use error_stack::Report; +use hash_graph_postgres_store::store::{EntityDeletion, EntityEnd, EntityEvent}; +use hash_graph_temporal_versioning::Timestamp; +use hashql_core::{collections::FastHashMap, id::Id as _}; +use rand::{SeedableRng as _, rngs::StdRng}; +use tokio::sync::mpsc; +use type_system::{ + knowledge::entity::{EntityId, id::EntityEditionId, provenance::EntityDeletionProvenance}, + principal::actor::ActorEntityUuid, +}; +use uuid::Uuid; + +use super::{DeltaAction, Geometry, Pending, Stage}; +use crate::{ + dataset::{ + PROJECTOR_DIMENSIONS, + auxiliary::{Label, OwnedIcon, OwnedLabel, OwnedLegend}, + }, + identity::{NodeRowId, OntologyRowId}, + math::{Bounds2, BoxedVecN, Vec2}, + postgres::{Classification, edition_display::DisplayParts, id::ArchivedEntityId}, + salt::{fit::prepare::IdentityProvider as _, lod::stage::WIRE_FRAME}, + serve::{ + delta::{ + Delta, + epoch::Epoch, + feed::{ + pump, + tests::{entity, update}, + }, + overlay::{DeltaIdentityProvider, NaiveIdentityProvider}, + placement::{Completed, PendingEntry, PlacementError}, + projector::{Position, projector}, + }, + tests::fixture::{TamperFixture, secret}, + world::World, + }, +}; + +/// A synthetic generation and the delta opened over it, retaining the generation's temp files. +struct Fixture { + _generation: TamperFixture, + delta: Delta, +} + +/// Opens a synthetic world under `name` and creates a delta over it, retaining its temp files. +/// +/// # Panics +/// +/// Panics on failure during generation publication or world opening. +fn fixture(name: &str) -> Fixture { + let generation = TamperFixture::publish(name); + let world = World::open(generation.generation().clone(), &secret()) + .expect("should open the synthetic world"); + let delta = + Delta::new(Arc::new(world), StdRng::seed_from_u64(12)).expect("should create a delta"); + Fixture { + _generation: generation, + delta, + } +} + +/// Observes an update for `seed` and classifies it as a node, queuing it for placement. +fn queue_node(pending: &mut Pending, delta: &Delta, seed: u128, seconds: i64) { + pending.observe(update(seed, seconds, false)); + pending.classify( + delta, + [(ArchivedEntityId::from(entity(seed)), Classification::Node)], + ); +} + +/// Builds display parts carrying `label`, a fixed icon, and a fixed representative id. +fn display(label: &str) -> DisplayParts { + DisplayParts { + label: OwnedLabel::from(label), + icon: OwnedIcon::from("icon"), + representative: Uuid::from_u128(1000).into(), + } +} + +/// Projects the zero vector to a single finite position, reused across placement fixtures. +/// +/// # Panics +/// +/// Panics if the projector returns no row or rejects the projected position. +fn projected() -> Position { + projector(None) + .project([&BoxedVecN::zero()]) + .next() + .expect("should return a projection") + .expect("should project a finite position") +} + +/// Snapshots `delta` into an epoch usable for read-side queries. +fn epoch(delta: &Delta) -> Epoch { + Epoch::from(Guard::from_inner(Arc::new(delta.clone()))) +} + +/// Collects `pairs` into a display-answer map keyed by edition id. +fn displays( + pairs: impl IntoIterator)>, +) -> FastHashMap> { + let mut map = FastHashMap::default(); + map.extend(pairs); + map +} + +/// Keeps equal and older overlaps in their pending phase. +/// +/// Only a later transaction time replaces the retained edition and resets its stage to +/// classification. +#[test] +fn observe_overlap() { + let mut pending = Pending::default(); + let entity = ArchivedEntityId::from(entity(1)); + assert_matches!(pending.observe(update(1, 10, false)), DeltaAction::None); + let first = pending.updates[&entity].id; + pending + .updates + .get_mut(&entity) + .expect("should retain the update") + .stage = Stage::Place; + assert_matches!(pending.observe(update(1, 10, false)), DeltaAction::None); + assert_matches!(pending.observe(update(1, 9, true)), DeltaAction::None); + assert_eq!(pending.updates[&entity].id, first); + assert_matches!(pending.updates[&entity].stage, Stage::Place); + assert_eq!(pending.updates.len(), 1); + assert_matches!(pending.observe(update(1, 11, false)), DeltaAction::None); + assert_ne!(pending.updates[&entity].id, first); + assert_matches!(pending.updates[&entity].stage, Stage::Classify); + assert_eq!(pending.updates.len(), 1); +} + +/// Keeps a replacement waiting for the prior request, whose result cannot apply to it. +#[test] +fn receive_superseded() { + let mut pending = Pending::default(); + let raw = entity(1); + let entity = ArchivedEntityId::from(raw); + + pending.observe(update(1, 1, false)); + pending + .updates + .get_mut(&entity) + .expect("should retain the node") + .stage = Stage::Place; + let (tx, mut requests) = mpsc::channel(1); + assert!(pump(&mut pending, &tx, None).is_continue()); + let first = requests + .try_recv() + .expect("should submit the first update") + .event; + + pending.observe(update(1, 2, false)); + let second = pending.updates[&entity].id; + pending + .updates + .get_mut(&entity) + .expect("should retain the replacement") + .stage = Stage::Place; + assert!( + !pending.has_placements(), + "should wait for the outstanding request" + ); + assert!(pump(&mut pending, &tx, None).is_continue()); + assert!(requests.try_recv().is_err()); + + pending.receive(PendingEntry { + event: second, + entity: raw, + phase: Completed(Ok(projected())), + }); + assert!( + pending.placements.contains_key(&entity), + "should ignore an unsubmitted ID" + ); + assert!(!pending.has_placements()); + + pending.receive(PendingEntry { + event: first, + entity: raw, + phase: Completed(Ok(projected())), + }); + assert!( + !pending.placements.contains_key(&entity), + "should discard the superseded result, but still free the outstanding slot" + ); + assert_matches!(pending.updates[&entity].stage, Stage::Place); + assert!( + pending.has_placements(), + "should release the waiting replacement" + ); + + assert!(pump(&mut pending, &tx, None).is_continue()); + assert_eq!( + requests + .try_recv() + .expect("should submit the replacement") + .event, + second + ); + pending.receive(PendingEntry { + event: first, + entity: raw, + phase: Completed(Ok(projected())), + }); + assert!( + pending.placements.contains_key(&entity), + "should ignore a repeated old completion" + ); + + pending.receive(PendingEntry { + event: second, + entity: raw, + phase: Completed(Ok(projected())), + }); + assert!(!pending.placements.contains_key(&entity)); + assert_matches!(pending.updates[&entity].stage, Stage::Normalize(_)); +} + +/// Retains outstanding placement state across every removal. +/// +/// Each removal cancels its update, but the submitted placement remains until completion. +#[test] +fn observe_removals() { + let id = entity(1); + let entity = ArchivedEntityId::from(id); + let removals = [ + update(1, 2, true), + EntityEvent::Ended(EntityEnd { + entity: id, + ended_at: Timestamp::from_unix_timestamp(2), + }), + EntityEvent::Deleted(EntityDeletion { + entity: id, + provenance: EntityDeletionProvenance { + deleted_by_id: ActorEntityUuid::new(Uuid::from_u128(4)), + deleted_at_transaction_time: Timestamp::from_unix_timestamp(2), + deleted_at_decision_time: Timestamp::from_unix_timestamp(2), + }, + }), + ]; + for removal in removals { + let mut pending = Pending::default(); + pending.observe(update(1, 1, false)); + pending + .updates + .get_mut(&entity) + .expect("should retain the node") + .stage = Stage::Place; + let (tx, mut requests) = mpsc::channel(1); + assert!(pump(&mut pending, &tx, None).is_continue()); + let first = requests.try_recv().expect("should submit the node").event; + + assert_matches!(pending.observe(removal), DeltaAction::Withdraw(withdrawn) if withdrawn == entity); + assert!(pending.updates.is_empty()); + assert!( + pending.placements.contains_key(&entity), + "should retain the outstanding request after withdrawal" + ); + + pending.observe(update(1, 1, false)); + assert!(pending.updates.is_empty(), "should reject the stale repeat"); + + pending.receive(PendingEntry { + event: first, + entity: id, + phase: Completed(Ok(projected())), + }); + assert!( + pending.updates.is_empty(), + "should leave removed work absent" + ); + assert!(pending.placements.is_empty()); + } +} + +/// Keeps versions and outstanding requests independent for equal UUIDs in different webs. +#[test] +fn observe_other_web() { + let mut pending = Pending::default(); + let first_id = entity(1); + let mut second_id = first_id; + second_id.web_id = type_system::principal::actor_group::WebId::new(Uuid::from_u128(2)); + + pending.observe(update(1, 10, false)); + let EntityEvent::Updated(mut other) = update(1, 1, false) else { + unreachable!() + }; + other.entity = second_id; + pending.observe(EntityEvent::Updated(other)); + assert_eq!(pending.updates.len(), 2); + + let first = ArchivedEntityId::from(first_id); + let second = ArchivedEntityId::from(second_id); + pending + .updates + .get_mut(&first) + .expect("should retain the node") + .stage = Stage::Place; + let (tx, mut requests) = mpsc::channel(1); + assert!(pump(&mut pending, &tx, None).is_continue()); + let request = requests + .try_recv() + .expect("should submit the first web's update") + .event; + + pending.receive(PendingEntry { + event: request, + entity: second_id, + phase: Completed(Ok(projected())), + }); + assert!( + pending.placements.contains_key(&first), + "should leave the other web's request outstanding" + ); + + pending.receive(PendingEntry { + event: request, + entity: first_id, + phase: Completed(Ok(projected())), + }); + assert!(pending.placements.is_empty()); + assert_eq!(pending.updates[&second].event, other); +} + +/// Reuses fitted-node coordinates for metadata updates and revival. +#[test] +fn classify_fitted_revival() { + let mut fixture = fixture("pending-fitted-revival"); + let mut pending = Pending::default(); + + let row = NodeRowId::MIN; + let entity = fixture + .delta + .world + .layout + .index + .identity + .key_of(row) + .expect("should resolve the fitted identity"); + let position = fixture + .delta + .node_position(entity) + .expect("should retain fitted coordinates"); + assert!(fixture.delta.withdraw(entity)); + + let EntityEvent::Updated(mut event) = update(100, 1, false) else { + unreachable!() + }; + event.entity = EntityId::from(entity); + pending.observe(EntityEvent::Updated(event)); + pending.classify(&fixture.delta, [(entity, Classification::Node)]); + + assert!(!pending.has_placements()); + assert_matches!( + pending.updates[&entity].stage, + Stage::Capture(Geometry::Node(held)) if held == position + ); + + let mut answers = displays([(event.edition, Some(display("revived")))]); + assert!(pending.capture(&mut fixture.delta, &mut answers)); + assert!(pending.apply(&mut fixture.delta)); + assert!(pending.updates.is_empty()); + + let epoch = epoch(&fixture.delta); + assert_eq!( + fixture.delta.world.layout.position(&epoch, row), + Some(position) + ); +} + +/// Preserves queued input under admission backpressure. +/// +/// A full input retains unsent nodes and drains all admitted requests without duplication. +#[test] +fn pump_backpressure() { + let fixture = fixture("pending-pump-backpressure"); + let mut pending = Pending::default(); + for seed in 100..103 { + queue_node(&mut pending, &fixture.delta, seed, 1); + } + let (tx, mut requests) = mpsc::channel(1); + assert!(pump(&mut pending, &tx, None).is_continue()); + assert!(pump(&mut pending, &tx, None).is_continue()); + let mut admitted = Vec::new(); + for _ in 0..3 { + admitted.push( + requests + .try_recv() + .expect("should admit one request at capacity") + .entity, + ); + assert!(pump(&mut pending, &tx, None).is_continue()); + } + admitted.sort_unstable_by_key(|entity| Uuid::from(entity.entity_uuid)); + assert_eq!(admitted, [entity(100), entity(101), entity(102)]); + assert!(requests.try_recv().is_err()); + assert!(!pending.has_placements()); +} + +/// Ends input admission on closure without marking an unsent update as running. +#[test] +fn pump_closed() { + let fixture = fixture("pending-pump-closed"); + let mut pending = Pending::default(); + queue_node(&mut pending, &fixture.delta, 100, 1); + let (tx, mut requests) = mpsc::channel(1); + requests.close(); + assert!(pump(&mut pending, &tx, None).is_break()); + assert!(pending.has_placements()); +} + +/// Cancels queued placement when an update archives its entity. +/// +/// No allocation or revival can follow the archived result. +#[test] +fn receive_archived() { + let fixture = fixture("pending-receive-archived"); + let mut pending = Pending::default(); + queue_node(&mut pending, &fixture.delta, 100, 1); + let (tx, mut requests) = mpsc::channel(1); + assert!(pump(&mut pending, &tx, None).is_continue()); + let request = requests.try_recv().expect("should submit the node"); + + assert_matches!( + pending.observe(update(100, 2, true)), + DeltaAction::Withdraw(withdrawn) if withdrawn == ArchivedEntityId::from(entity(100)) + ); + pending.receive(PendingEntry { + event: request.event, + entity: request.entity, + phase: Completed(Ok(projected())), + }); + + assert!(pending.updates.is_empty()); + assert_eq!( + fixture.delta.node_row(ArchivedEntityId::from(entity(100))), + None + ); +} + +/// Prevents overlap from retrying a failed placement budget. +/// +/// Failure removes pending work but retains the event's transaction-time high-water mark. +#[test] +fn receive_failed() { + let fixture = fixture("pending-receive-failed"); + let mut pending = Pending::default(); + queue_node(&mut pending, &fixture.delta, 100, 1); + let (tx, mut requests) = mpsc::channel(1); + assert!(pump(&mut pending, &tx, None).is_continue()); + let request = requests.try_recv().expect("should submit the node"); + + pending.receive(PendingEntry { + event: request.event, + entity: request.entity, + phase: Completed(Err(Report::new(PlacementError::Exhaustion))), + }); + assert!(pending.updates.is_empty()); + + pending.observe(update(100, 1, false)); + assert!(pending.updates.is_empty(), "should reject the stale repeat"); + + queue_node(&mut pending, &fixture.delta, 100, 2); + assert!(pending.has_placements()); +} + +/// Normalizes successful placements before display capture. +#[test] +fn receive_wire_coordinates() { + let mut fixture = fixture("pending-receive-wire-coordinates"); + let mut pending = Pending::default(); + queue_node(&mut pending, &fixture.delta, 100, 1); + let (tx, mut requests) = mpsc::channel(1); + assert!(pump(&mut pending, &tx, None).is_continue()); + let request = requests.try_recv().expect("should submit the node"); + + let position = projected(); + let expected = fixture + .delta + .world + .bounds() + .normalize_into(WIRE_FRAME, &[position.get()])[0]; + assert_ne!( + expected, + position.get(), + "should use a non-identity normalization case" + ); + + pending.receive(PendingEntry { + event: request.event, + entity: request.entity, + phase: Completed(Ok(position)), + }); + let mut scratch = Vec::new(); + pending.normalize(fixture.delta.world.bounds(), &mut scratch); + assert!(scratch.is_empty()); + + let entity = ArchivedEntityId::from(request.entity); + let edition = pending.updates[&entity].event.edition; + let mut answers = displays([(edition, Some(display("placed")))]); + pending.capture(&mut fixture.delta, &mut answers); + assert!(pending.apply(&mut fixture.delta)); + + let node = fixture + .delta + .node_row(entity) + .expect("should allocate the node"); + let epoch = epoch(&fixture.delta); + assert_eq!( + fixture.delta.world.layout.position(&epoch, node), + Some(expected) + ); +} + +/// Requires endpoint rows for link updates and never requests an embedding. +#[test] +fn edge_missing_endpoint() { + let mut fixture = fixture("pending-edge-missing-endpoint"); + let mut pending = Pending::default(); + let edge = ArchivedEntityId::from(entity(100)); + let node = ArchivedEntityId::from(entity(101)); + + pending.observe(update(100, 1, false)); + pending.classify( + &fixture.delta, + [( + edge, + Classification::Edge { + source: Some(node), + target: Some(node), + }, + )], + ); + assert!(!pending.has_placements()); + + let edition = pending.updates[&edge].event.edition; + let mut answers = displays([(edition, Some(display("link")))]); + pending.capture(&mut fixture.delta, &mut answers); + assert!(!pending.apply(&mut fixture.delta)); + assert_eq!(pending.updates.len(), 1); + + fixture + .delta + .update_node( + node, + OwnedLegend::new(OntologyRowId::MIN, Label::new("node")), + Vec2::ZERO, + ) + .expect("should allocate the endpoint"); + assert!(pending.apply(&mut fixture.delta)); + assert!(pending.updates.is_empty()); +} + +/// Ignores a superseded edition's display response. +/// +/// Only a response for the current edition advances capture. +#[test] +fn display_replaced_edition() { + let mut fixture = fixture("pending-display-replaced-edition"); + let mut pending = Pending::default(); + let entity = ArchivedEntityId::from(entity(100)); + fixture + .delta + .update_node( + entity, + OwnedLegend::new(OntologyRowId::MIN, Label::new("initial")), + Vec2::ZERO, + ) + .expect("should allocate the node"); + + queue_node(&mut pending, &fixture.delta, 100, 1); + let old = pending.updates[&entity].event.edition; + queue_node(&mut pending, &fixture.delta, 100, 2); + let current = pending.updates[&entity].event.edition; + + let row = fixture + .delta + .node_row(entity) + .expect("should retain the existing row"); + let mut answers = displays([(old, Some(display("old")))]); + assert!(!pending.capture(&mut fixture.delta, &mut answers)); + assert!(answers.is_empty()); + assert_matches!( + pending.updates[&entity].stage, + Stage::Capture(_), + "should keep the current request waiting after a stale answer" + ); + + let mut answers = displays([(current, None)]); + assert!(!pending.capture(&mut fixture.delta, &mut answers)); + assert_matches!( + pending.updates[&entity].stage, + Stage::Capture(_), + "should keep the request waiting after a missing display" + ); + + let mut answers = displays([(current, Some(display("new")))]); + assert!(pending.capture(&mut fixture.delta, &mut answers)); + let Stage::Ready { legend, .. } = &pending.updates[&entity].stage else { + panic!("should capture the current display"); + }; + assert_eq!(legend.label(), "new"); + assert!(pending.apply(&mut fixture.delta)); + assert_eq!(fixture.delta.node_row(entity), Some(row)); + let identities = DeltaIdentityProvider::from_parts( + &fixture.delta.node, + NaiveIdentityProvider::from_ref(&fixture.delta.world.layout.index.identity), + ); + assert_eq!( + identities + .payload_of_row(row) + .expect("should replace the row's legend") + .label(), + "new" + ); +} + +/// Releases admission after a failed superseded request without discarding its replacement. +#[test] +fn receive_superseded_failure() { + let fixture = fixture("pending-receive-superseded-failure"); + let mut pending = Pending::default(); + queue_node(&mut pending, &fixture.delta, 100, 1); + let (tx, mut requests) = mpsc::channel(1); + assert!(pump(&mut pending, &tx, None).is_continue()); + let request = requests.try_recv().expect("should submit the first update"); + + queue_node(&mut pending, &fixture.delta, 100, 2); + assert!( + !pending.has_placements(), + "should wait for the outstanding request" + ); + assert!(pump(&mut pending, &tx, None).is_continue()); + assert!(requests.try_recv().is_err()); + + pending.receive(PendingEntry { + event: request.event, + entity: request.entity, + phase: Completed(Err(Report::new(PlacementError::Exhaustion))), + }); + assert!( + pending.has_placements(), + "should release the replacement after the failed completion" + ); + + assert!(pump(&mut pending, &tx, None).is_continue()); + let replacement = requests.try_recv().expect("should submit the replacement"); + assert_ne!(replacement.event, request.event); +} + +/// Reuses an added node's first position after revival. +/// +/// Older epochs remain withdrawn. +#[test] +fn classify_added_revival() { + let mut fixture = fixture("pending-classify-added-revival"); + let mut pending = Pending::default(); + let entity = ArchivedEntityId::from(entity(100)); + let position = Vec2::new(0.5, -0.25); + + fixture + .delta + .update_node( + entity, + OwnedLegend::new(OntologyRowId::MIN, Label::new("original")), + position, + ) + .expect("should allocate a node"); + fixture.delta.revision.increment_by(1); + fixture.delta.withdraw(entity); + let withdrawn = epoch(&fixture.delta); + fixture.delta.revision.increment_by(1); + + queue_node(&mut pending, &fixture.delta, 100, 1); + assert_matches!( + pending.updates[&entity].stage, + Stage::Capture(Geometry::Node(held)) if held == position + ); + + let edition = pending.updates[&entity].event.edition; + let mut answers = displays([(edition, Some(display("revived")))]); + pending.capture(&mut fixture.delta, &mut answers); + assert!(pending.apply(&mut fixture.delta)); + + let row = fixture + .delta + .node_row(entity) + .expect("should resolve the retained row"); + assert_eq!(fixture.delta.world.layout.position(&withdrawn, row), None); + let epoch = epoch(&fixture.delta); + assert_eq!( + fixture.delta.world.layout.position(&epoch, row), + Some(position) + ); +} +/// Restores request order and drains normalization scratch state. +/// +/// The classification and dispatched-edition counts equal the queued counts. Each captured +/// position differs from the expected wire-frame normalization by no more than the tolerance. +#[test] +fn normalize_batch() { + let mut pending = Pending::default(); + let inputs: [BoxedVecN; 5] = core::array::from_fn(|axis| { + let mut input = BoxedVecN::zero(); + input.as_array_mut()[axis] = 1.0; + input + }); + let positions: Vec<_> = projector(None) + .project(&inputs) + .try_collect() + .expect("should project finite inputs"); + assert!(positions.windows(2).all(|pair| { + let [first, second] = pair else { + unreachable!("windows(2) should yield pairs") + }; + first != second + })); + + for seed in 100..105 { + pending.observe(update(seed, 1, false)); + pending + .updates + .get_mut(&ArchivedEntityId::from(entity(seed))) + .expect("should retain the node") + .stage = Stage::Place; + } + let (tx, mut rx) = mpsc::channel(5); + assert!(pump(&mut pending, &tx, None).is_continue()); + let mut entities = Vec::new(); + let mut raw = Vec::new(); + for position in positions.into_iter().rev() { + let request = rx.try_recv().expect("should submit each node"); + entities.push(ArchivedEntityId::from(request.entity)); + raw.push(position.get()); + pending.receive(PendingEntry { + event: request.event, + entity: request.entity, + phase: Completed(Ok(position)), + }); + } + assert!(rx.try_recv().is_err()); + for seed in 200..210 { + pending.observe(update(seed, 1, false)); + } + + let bounds = Bounds2::new(Vec2::splat(-100.0), Vec2::splat(100.0)) + .expect("should have ordered finite bounds"); + let expected = bounds.normalize_into(WIRE_FRAME, &raw); + let mut scratch = vec![Vec2::splat(99.0)]; + pending.normalize(bounds, &mut scratch); + assert!(scratch.is_empty()); + assert_eq!(pending.classifications().count(), 10); + assert_eq!(pending.editions().count(), 5); + for (entity, expected) in entities.into_iter().zip(expected) { + let Stage::Capture(Geometry::Node(actual)) = pending.updates[&entity].stage else { + panic!("should await the node's display"); + }; + assert!(actual.distance_squared_wide(expected) < 1e-10); + } + pending.normalize(bounds, &mut scratch); + assert!(scratch.is_empty()); + assert_eq!(pending.editions().count(), 5); +} + +/// Discards a completed placement superseded before normalization. +/// +/// The replacement returns to classification and dispatches no edition. +#[test] +fn normalize_superseded() { + let mut pending = Pending::default(); + let key = ArchivedEntityId::from(entity(100)); + pending.observe(update(100, 1, false)); + pending + .updates + .get_mut(&key) + .expect("should retain the node") + .stage = Stage::Place; + let (tx, mut rx) = mpsc::channel(1); + assert!(pump(&mut pending, &tx, None).is_continue()); + let request = rx.try_recv().expect("should submit the node"); + pending.receive(PendingEntry { + event: request.event, + entity: request.entity, + phase: Completed(Ok(projected())), + }); + assert_matches!(pending.updates[&key].stage, Stage::Normalize(_)); + pending.observe(update(100, 2, false)); + let bounds = Bounds2::new(Vec2::splat(-100.0), Vec2::splat(100.0)) + .expect("should have ordered finite bounds"); + let mut scratch = Vec::new(); + pending.normalize(bounds, &mut scratch); + assert!(scratch.is_empty()); + assert_ne!(pending.updates[&key].id, request.event); + assert_matches!(pending.updates[&key].stage, Stage::Classify); + assert_eq!(pending.editions().count(), 0); +} + +/// Defers an edge until both endpoint identities resolve. +/// +/// An edge with only one resolved endpoint captures its display but cannot apply. Once both +/// endpoints resolve, the edge applies and its identity resolves through the topology provider. +#[test] +fn edge_incomplete_pair() { + let mut fixture = fixture("pending-edge-incomplete-pair"); + let node = ArchivedEntityId::from(entity(101)); + fixture + .delta + .update_node( + node, + OwnedLegend::new(OntologyRowId::MIN, Label::new("node")), + Vec2::ZERO, + ) + .expect("should allocate the endpoint"); + let edge = ArchivedEntityId::from(entity(100)); + for (source, target) in [(Some(node), None), (None, Some(node)), (None, None)] { + let mut pending = Pending::default(); + pending.observe(update(100, 1, false)); + pending.classify( + &fixture.delta, + [(edge, Classification::Edge { source, target })], + ); + assert!(!pending.has_placements()); + let edition = pending.updates[&edge].event.edition; + pending.capture( + &mut fixture.delta, + &mut displays([(edition, Some(display("incomplete")))]), + ); + assert_matches!( + pending.updates[&edge].stage, + Stage::Ready { + geometry: Geometry::Edge(None), + .. + } + ); + assert!(!pending.apply(&mut fixture.delta)); + assert_eq!(pending.updates.len(), 1); + + pending.observe(update(100, 2, false)); + pending.classify( + &fixture.delta, + [( + edge, + Classification::Edge { + source: Some(node), + target: Some(node), + }, + )], + ); + let edition = pending.updates[&edge].event.edition; + pending.capture( + &mut fixture.delta, + &mut displays([(edition, Some(display("complete")))]), + ); + pending.apply(&mut fixture.delta); + assert!(pending.updates.is_empty()); + let identities = DeltaIdentityProvider::from_parts( + &fixture.delta.edge, + NaiveIdentityProvider::from_ref(&fixture.delta.world.topology.identity), + ); + assert!(identities.row_of(edge).is_some()); + } +} diff --git a/libs/@local/graph/atlas/src/serve/delta/feed/tests.rs b/libs/@local/graph/atlas/src/serve/delta/feed/tests.rs new file mode 100644 index 00000000000..2a33921a4bb --- /dev/null +++ b/libs/@local/graph/atlas/src/serve/delta/feed/tests.rs @@ -0,0 +1,524 @@ +use alloc::sync::Arc; +use core::{assert_matches, future, ops::ControlFlow, pin::pin, time::Duration}; + +use arc_swap::{ArcSwap, Guard}; +use futures::FutureExt as _; +use hash_graph_postgres_store::store::{ + DatabaseConnectionInfo, DatabasePoolConfig, DatabaseType, EntityEvent, EntityUpdate, + PostgresStorePool, PostgresStoreSettings, +}; +use hash_graph_temporal_versioning::Timestamp; +use hashql_core::id::Id as _; +use rand::{SeedableRng as _, rngs::StdRng}; +use tokio::{ + sync::{mpsc, oneshot}, + time::Instant, +}; +use tokio_postgres::NoTls; +use type_system::knowledge::entity::{EntityId, id::EntityEditionId}; +use uuid::Uuid; + +use super::{DeltaFeedError, DeltaFeedTask, DeltaFeedTaskOptions, Placement, Publication, pump}; +use crate::{ + dataset::auxiliary::{Label, OwnedLegend}, + identity::{NodeRowId, OntologyRowId}, + math::{Vec2, nz}, + postgres::{Classification, id::ArchivedEntityId}, + serve::{ + delta::{ + Delta, + epoch::Epoch, + placement::{Completed, Initial, PendingEntry}, + }, + tests::fixture::{TamperFixture, secret}, + world::World, + }, +}; + +/// A feed task opened over a synthetic generation, with its placement channel ends. +struct Fixture { + _generation: TamperFixture, + task: DeltaFeedTask, + requests: mpsc::Receiver>, + completed: mpsc::Sender>, + publication: Publication, +} + +/// Builds a feed task over a synthetic world published under `name`. +/// +/// The task has disconnected placement channels and a fresh publication handle. +/// +/// # Panics +/// +/// Panics on failure during generation publication, world opening, store pool construction, or +/// feed initialization. +async fn fixture(name: &str) -> Fixture { + let generation = TamperFixture::publish(name); + let world = World::open(generation.generation().clone(), &secret()) + .expect("should open the synthetic world"); + let delta = + Delta::new(Arc::new(world), StdRng::seed_from_u64(12)).expect("should create a delta"); + // Pool construction opens no connection. These cases supply read results directly. + let pool = PostgresStorePool::new( + &DatabaseConnectionInfo::new( + DatabaseType::Postgres, + "feed-test".to_owned(), + String::new(), + "/no-feed-test-postgres".to_owned(), + 5432, + "feed-test".to_owned(), + ), + &DatabasePoolConfig { + max_connections: nz!(1), + }, + NoTls, + PostgresStoreSettings::default(), + ) + .await + .expect("should construct an unconnected pool"); + let (tx, requests) = mpsc::channel(1); + let (completed, rx) = mpsc::channel(4); + let (task, publication) = DeltaFeedTask::new( + delta, + Arc::new(pool), + DeltaFeedTaskOptions { + safety_lag: Duration::from_secs(60), + tick_rate: Duration::from_secs(5), + }, + Timestamp::UNIX_EPOCH, + Some(Placement { + requests: tx, + completed: rx, + }), + ) + .expect("should construct the feed"); + Fixture { + _generation: generation, + task, + requests, + completed, + publication, + } +} + +/// Rejects a zero tick rate regardless of the safety lag. +#[tokio::test] +async fn new_interval_zero() { + let fixture = fixture("feed-new-interval-zero").await; + let result = DeltaFeedTask::new( + fixture.task.delta, + fixture.task.pool, + DeltaFeedTaskOptions { + tick_rate: Duration::ZERO, + safety_lag: Duration::ZERO, + }, + Timestamp::UNIX_EPOCH, + None, + ); + let error = result + .map(|_| ()) + .expect_err("should refuse a zero interval"); + assert_matches!(error.current_context(), DeltaFeedError::InvalidInterval); +} + +/// Rejects replay windows outside the representable transaction-time range. +/// +/// An unrepresentable replay window - an infinite safety lag at the epoch, or a lag exceeding +/// the earliest representable timestamp - is refused. A zero lag at the earliest timestamp is +/// admitted and sets the watermark exactly there. +#[tokio::test] +async fn new_replay_window() { + let fixture = fixture("feed-new-replay-window").await; + let earliest = + Timestamp::from_unix_timestamp(time::Date::MIN.midnight().assume_utc().unix_timestamp()); + for (watermark, lag) in [ + (Timestamp::UNIX_EPOCH, Duration::MAX), + (earliest, Duration::from_secs(1)), + ] { + let result = DeltaFeedTask::new( + fixture.task.delta.clone(), + Arc::clone(&fixture.task.pool), + DeltaFeedTaskOptions { + tick_rate: Duration::from_secs(5), + safety_lag: lag, + }, + watermark, + None, + ); + let error = result + .map(|_| ()) + .expect_err("should refuse an unrepresentable replay window"); + assert_matches!(error.current_context(), DeltaFeedError::InvalidSafetyLag); + } + let (task, _publication) = DeltaFeedTask::new( + fixture.task.delta, + fixture.task.pool, + DeltaFeedTaskOptions { + tick_rate: Duration::from_secs(5), + safety_lag: Duration::ZERO, + }, + earliest, + None, + ) + .expect("should admit a zero lag at the earliest timestamp"); + assert_eq!(task.watermark - task.safety_lag, earliest); +} + +/// Keeps a publication captured before startup independent of subsequent writes. +#[tokio::test] +async fn new_captured_revision() { + let fixture = fixture("feed-new-captured-revision").await; + let captured = Epoch::from(Guard::from_inner(Arc::new(fixture.task.delta.clone()))); + let world = Arc::clone(&fixture.task.delta.world); + let base = NodeRowId::MIN; + let identity = world + .layout + .index + .key_of(&captured, base) + .expect("should resolve the base row"); + let (mut task, _publication) = DeltaFeedTask::new( + fixture.task.delta, + fixture.task.pool, + fixture.task.options, + Timestamp::UNIX_EPOCH, + None, + ) + .expect("should construct the feed"); + assert_eq!(task.delta.revision, captured.revision().plus(1)); + assert!(task.delta.withdraw(identity)); + assert!(world.layout.position(&captured, base).is_some()); + let current = Epoch::from(Guard::from_inner(Arc::new(task.delta.clone()))); + assert_eq!(world.layout.position(¤t, base), None); +} + +/// Exchanges a queued publication without waiting for a database tick. +#[tokio::test] +async fn publication_exchange() { + let mut fixture = fixture("feed-publication-exchange").await; + let revision = fixture.task.delta.revision; + let mut publication = pin!( + fixture + .publication + .next(Arc::new(fixture.task.delta.clone())) + ); + assert!(publication.as_mut().now_or_never().is_none()); + fixture.task.notify.notify_one(); + assert!(publication.as_mut().now_or_never().is_none()); + let period = Duration::from_hours(24); + let mut interval = tokio::time::interval_at(Instant::now() + period, period); + assert_eq!( + fixture + .task + .step(&mut interval) + .now_or_never() + .expect("should receive the exchange") + .expect("should publish"), + ControlFlow::Continue(false) + ); + let published = publication.await.expect("should receive the publication"); + assert_eq!(published.revision, revision); + assert_eq!(fixture.task.delta.revision, revision.plus(1)); +} + +/// Swaps complete deltas while preserving an earlier capture. +#[tokio::test] +async fn publication_captured() { + let fixture = fixture("feed-publication-captured").await; + let current = Arc::new(ArcSwap::from_pointee(fixture.task.delta.clone())); + let previous = current.load_full(); + let captured = Epoch::from(current.load()); + let world = Arc::clone(&fixture.task.delta.world); + let row = NodeRowId::MIN; + let identity = world + .layout + .index + .key_of(&captured, row) + .expect("should resolve the base identity"); + let (mut task, publication) = DeltaFeedTask::new( + fixture.task.delta, + fixture.task.pool, + fixture.task.options, + Timestamp::UNIX_EPOCH, + None, + ) + .expect("should construct the feed"); + let mut publishing = pin!(publication.run(Arc::clone(¤t), previous, future::pending())); + assert!(publishing.as_mut().now_or_never().is_none()); + let period = Duration::from_hours(24); + let mut interval = tokio::time::interval_at(Instant::now() + period, period); + for withdraw in [true, false, true] { + if withdraw { + assert!(task.delta.withdraw(identity)); + } else { + assert_eq!( + task.delta.update_node( + identity, + OwnedLegend::new(OntologyRowId::MIN, Label::new("revived")), + Vec2::ZERO + ), + Some(true) + ); + } + let revision = task.delta.revision; + task.notify.notify_one(); + assert!(publishing.as_mut().now_or_never().is_none()); + assert_eq!( + task.step(&mut interval) + .now_or_never() + .expect("should accept the exchange") + .expect("should publish"), + ControlFlow::Continue(false) + ); + assert!(publishing.as_mut().now_or_never().is_none()); + let latest = Epoch::from(current.load()); + assert_eq!(latest.revision(), revision); + assert_eq!(world.layout.position(&latest, row).is_none(), withdraw); + assert!(world.layout.position(&captured, row).is_some()); + } +} + +/// Closes exchange on shutdown and both placement directions on feed exit. +#[tokio::test] +async fn publication_shutdown() { + let mut fixture = fixture("feed-publication-shutdown").await; + let current = Arc::new(ArcSwap::from_pointee(fixture.task.delta.clone())); + let previous = current.load_full(); + let (stop, shutdown) = oneshot::channel(); + let mut publishing = pin!(fixture.publication.run(current, previous, async move { + shutdown.await.expect("should receive the shutdown request"); + })); + assert!(publishing.as_mut().now_or_never().is_none()); + stop.send(()).expect("should request shutdown"); + publishing + .as_mut() + .now_or_never() + .expect("should observe shutdown") + .expect("should stop publication normally"); + assert!(fixture.task.update.is_closed()); + let period = Duration::from_hours(24); + let mut interval = tokio::time::interval_at(Instant::now() + period, period); + assert_eq!( + fixture + .task + .step(&mut interval) + .now_or_never() + .expect("should observe exchange closure") + .expect("should stop the feed normally"), + ControlFlow::Break(()) + ); + drop(fixture.task); + assert!(fixture.completed.is_closed()); + assert!(fixture.requests.recv().await.is_none()); +} + +/// Reports the feed as closed when a requester asks for publication after task drop. +#[tokio::test] +async fn publication_closed() { + let fixture = fixture("feed-publication-closed").await; + let previous = Arc::new(fixture.task.delta.clone()); + drop(fixture.task); + let error = fixture + .publication + .next(previous) + .await + .err() + .expect("should report the stopped task"); + assert_matches!(error.current_context(), DeltaFeedError::Closed); +} + +/// Keeps placement pending when projector channels are absent. +/// +/// Publication shutdown still completes. +#[tokio::test] +async fn placement_disabled() { + let fixture = fixture("feed-placement-disabled").await; + let (mut task, publication) = DeltaFeedTask::new( + fixture.task.delta, + fixture.task.pool, + fixture.task.options, + Timestamp::UNIX_EPOCH, + None, + ) + .expect("should construct a feed without placement"); + queue_node(&mut task, 100, 1); + let period = Duration::from_hours(24); + let mut interval = tokio::time::interval_at(Instant::now() + period, period); + assert!(task.step(&mut interval).now_or_never().is_none()); + assert!(task.pending.has_placements()); + drop(publication); + assert_eq!( + task.step(&mut interval) + .now_or_never() + .expect("should observe publication closure") + .expect("should stop the feed"), + ControlFlow::Break(()) + ); +} + +/// Builds an entity id from a fixed web and `seed`, for compact literal test identities. +pub(super) fn entity(seed: u128) -> EntityId { + EntityId::from(ArchivedEntityId { + web_id: Uuid::from_u128(1).into(), + entity_uuid: Uuid::from_u128(seed).into(), + }) +} + +/// Builds an updated-entity event for `seed` at `seconds`, optionally marked archived. +pub(super) fn update(seed: u128, seconds: i64, archived: bool) -> EntityEvent { + EntityEvent::Updated(EntityUpdate { + entity: entity(seed), + edition: EntityEditionId::new(Uuid::from_u128( + seed * 1_000 + u128::from(seconds.unsigned_abs()), + )), + archived, + changed_at: Timestamp::from_unix_timestamp(seconds), + }) +} + +/// Observes an update for `seed` and classifies it as a node, queuing it for placement. +fn queue_node(task: &mut DeltaFeedTask, seed: u128, seconds: i64) { + task.pending.observe(update(seed, seconds, false)); + task.pending.classify( + &task.delta, + [(ArchivedEntityId::from(entity(seed)), Classification::Node)], + ); +} + +/// Publishes once after quiet replay and never decreases the watermark during overlap. +#[tokio::test] +async fn replay_watermark() { + let mut fixture = fixture("feed-replay-watermark").await; + let task = &mut fixture.task; + assert!(task.apply_events()); + assert!(!task.apply_events()); + task.scratch.events.push(update(100, 10, false)); + assert!(!task.apply_events()); + assert_eq!(task.watermark, Timestamp::from_unix_timestamp(10)); + task.scratch.events.push(update(100, 9, false)); + assert!(!task.apply_events()); + assert_eq!(task.watermark, Timestamp::from_unix_timestamp(10)); + assert_eq!( + task.pending.classifications().collect::>(), + [ArchivedEntityId::from(entity(100))] + ); +} + +/// Preserves allocation and revision when the requester drops an exchange reply. +#[tokio::test] +async fn exchange_cancelled_reply() { + let mut fixture = fixture("feed-exchange-cancelled-reply").await; + let task = &mut fixture.task; + let previous = task.delta.clone(); + let revision = task.delta.revision; + let (reply, receive) = oneshot::channel(); + task.exchange(previous, reply); + let published = receive.await.expect("should receive the publication"); + assert_eq!(published.revision, revision); + assert_eq!(task.delta.revision, revision.plus(1)); + let (reply, receive) = oneshot::channel(); + drop(receive); + task.exchange(published, reply); + assert_eq!(task.delta.revision, revision.plus(1)); +} + +/// Resumes admission when capacity opens without waiting for a database tick. +#[tokio::test] +async fn step_capacity_ready() { + let mut fixture = fixture("feed-step-capacity-ready").await; + queue_node(&mut fixture.task, 100, 1); + queue_node(&mut fixture.task, 101, 1); + let placement = fixture + .task + .placement + .as_ref() + .expect("should have placement channels"); + assert!(pump(&mut fixture.task.pending, &placement.requests, None).is_continue()); + let period = Duration::from_hours(24); + let mut interval = tokio::time::interval_at(Instant::now() + period, period); + assert!(fixture.task.step(&mut interval).now_or_never().is_none()); + let first = fixture + .requests + .try_recv() + .expect("should release input capacity"); + let result = fixture + .task + .step(&mut interval) + .now_or_never() + .expect("should use the available capacity") + .expect("should keep the placement channel open"); + assert_eq!(result, ControlFlow::Continue(false)); + let second = fixture + .requests + .try_recv() + .expect("should submit the next node"); + assert_ne!(first.entity, second.entity); + assert_eq!(fixture.task.watermark, Timestamp::from_unix_timestamp(0)); + assert!(!fixture.task.replayed); +} + +/// Reports task failure on result-channel closure without waiting for a database tick. +#[tokio::test] +async fn step_result_closed() { + let mut fixture = fixture("feed-step-result-closed").await; + fixture + .task + .placement + .as_mut() + .expect("should have placement channels") + .completed + .close(); + let period = Duration::from_hours(24); + let mut interval = tokio::time::interval_at(Instant::now() + period, period); + let error = fixture + .task + .step(&mut interval) + .now_or_never() + .expect("should observe channel closure") + .expect_err("should report the missing placement task"); + assert_matches!(error.current_context(), DeltaFeedError::PlacementClosed); +} + +/// Ends on publication closure without advancing the database watermark. +#[tokio::test] +async fn step_publication_closed() { + let mut fixture = fixture("feed-step-publication-closed").await; + fixture.task.update.close(); + let period = Duration::from_hours(24); + let mut interval = tokio::time::interval_at(Instant::now() + period, period); + assert_eq!( + fixture + .task + .step(&mut interval) + .now_or_never() + .expect("should observe publication closure") + .expect("should stop normally"), + ControlFlow::Break(()) + ); + assert_eq!(fixture.task.watermark, Timestamp::from_unix_timestamp(0)); +} + +/// Preserves an allocated row while replaying its withdrawal. +/// +/// Replay hides the node's position but leaves its row resolvable. +#[tokio::test] +async fn replay_withdrawal() { + let mut fixture = fixture("feed-replay-withdrawal").await; + let task = &mut fixture.task; + let key = ArchivedEntityId::from(entity(100)); + task.delta + .update_node( + key, + OwnedLegend::new(OntologyRowId::MIN, Label::new("visible")), + Vec2::ZERO, + ) + .expect("should allocate the node"); + let row = task + .delta + .node_row(key) + .expect("should retain the node's row"); + task.scratch.events.push(update(100, 1, true)); + assert!(task.apply_events()); + let epoch = Epoch::from(Guard::from_inner(Arc::new(task.delta.clone()))); + assert_eq!(task.delta.world.layout.position(&epoch, row), None); + assert_eq!(task.delta.node_row(key), Some(row)); +} diff --git a/libs/@local/graph/atlas/src/serve/delta/mod.rs b/libs/@local/graph/atlas/src/serve/delta/mod.rs index 341fb6c0b17..fe884bdec6a 100644 --- a/libs/@local/graph/atlas/src/serve/delta/mod.rs +++ b/libs/@local/graph/atlas/src/serve/delta/mod.rs @@ -10,11 +10,15 @@ )] pub(crate) mod epoch; +mod feed; mod history; mod id; mod importance; pub(crate) mod layout; pub(crate) mod overlay; +pub(crate) mod placement; +mod projector; +mod task; pub(crate) mod topology; use alloc::sync::Arc; @@ -23,6 +27,11 @@ use hashql_core::id::Id as _; use rand::TryCryptoRng; use zerocopy::{NativeEndian, U64}; +pub(crate) use self::{ + feed::DeltaFeedTaskOptions, + placement::{DeltaPlacementTaskOptions, EmbeddingWorkflow}, + task::{DeltaReader, DeltaTask, DeltaTaskError, DeltaTaskOptions}, +}; use self::{ layout::{LayoutDelta, provider::NaiveLayoutProvider}, overlay::{ diff --git a/libs/@local/graph/atlas/src/serve/delta/placement/mod.rs b/libs/@local/graph/atlas/src/serve/delta/placement/mod.rs new file mode 100644 index 00000000000..33701d4cccb --- /dev/null +++ b/libs/@local/graph/atlas/src/serve/delta/placement/mod.rs @@ -0,0 +1,506 @@ +//! Embedding retries and owned projection batches for incoming events. +//! +//! [`DeltaPlacementTask`] keeps database and workflow operations on Tokio. A failed store lookup +//! preserves that stage's embedding-miss budget for a later tick, while a successful miss spends +//! it. An unavailable workflow actor, a failed workflow start or a per-row projection error +//! completes the affected event with a terminal placement error. Each projection job owns the +//! model, scratch storage and input batch on Rayon until it returns them. + +mod pending; + +use alloc::sync::Arc; +use core::{ + error::Error, fmt, future::Future, num::NonZero, ops::ControlFlow, panic::AssertUnwindSafe, + time::Duration, +}; +use std::collections::HashMap; + +use error_stack::{Report, ResultExt as _}; +use hash_graph_authorization::policies::store::PrincipalStore; +use hash_graph_postgres_store::store::{AsClient, PostgresStorePool}; +use hash_graph_store::pool::StorePool as _; +use hash_temporal_client::TemporalClient; +use hashql_core::{collections::FastHashMap, id::Id as _}; +use tokio::{sync::mpsc, time::Interval}; +use type_system::{ontology::id::BaseUrl, principal::actor::ActorId}; + +use self::pending::{Pending, PollWorkflow, Project}; +use super::projector::DeltaProjector; +use crate::{ + dataset::PROJECTOR_DIMENSIONS, + math::BoxedVecN, + postgres::{id::ArchivedEntityId, read_projector_embeddings}, +}; + +#[cfg(test)] +mod tests; + +pub(crate) use self::pending::{Completed, Initial, PendingEntry}; + +/// The workflow client and property exclusions used to produce missing embeddings. +#[derive(Debug)] +pub struct EmbeddingWorkflow { + /// The client used to start or find each entity's embedding workflow. + pub temporal: TemporalClient, + /// The per-base-URL exclusions supplied to the embedding workflow. + pub exclusions: HashMap>, +} + +/// A failed event placement. +#[derive(Debug)] +pub(crate) enum PlacementError { + /// The configured reads found no embedding, or workflow submission is unavailable. + Exhaustion, + /// Starting the embedding workflow failed. + Workflow, + /// Projecting the embedding failed. + Projection, +} + +impl fmt::Display for PlacementError { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Exhaustion => fmt.write_str("the embedding lookup budget ran out"), + Self::Workflow => fmt.write_str("could not start the embedding workflow"), + Self::Projection => fmt.write_str("could not project the embedding"), + } + } +} + +impl Error for PlacementError {} + +/// A placement task configuration or execution failure. +#[derive(Debug)] +pub(crate) enum DeltaPlacementError { + /// The polling interval is zero. + InvalidInterval, + /// Acquiring a store connection failed. + Connect, + /// Reading the projector embeddings failed. + Read, + /// The projection worker panicked or disappeared without returning its owned state. + Offload, +} + +impl fmt::Display for DeltaPlacementError { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::InvalidInterval => { + fmt.write_str("the placement polling interval must be non-zero") + } + Self::Connect => fmt.write_str("could not connect to the store"), + Self::Read => fmt.write_str("could not read the projector embeddings"), + Self::Offload => fmt.write_str("the projection worker failed"), + } + } +} + +impl Error for DeltaPlacementError {} + +/// Polling cadence and bounds supplied by the placement task's owner. +#[derive(Copy, Clone)] +pub(crate) struct DeltaPlacementTaskOptions { + /// How often the task polls for embeddings and re-evaluates pending work. + pub tick_rate: Duration, + /// Workflow-embedding misses allowed before an entry exhausts. + /// + /// A lookup still runs at zero: a hit succeeds, while a miss exhausts the entry. + pub tries_workflow: u16, + /// Database-embedding misses allowed before an entry submits a workflow. + /// + /// A lookup still runs at zero: a hit succeeds, while a miss advances to submission. + pub tries_database: u16, + /// Minimum ticks between non-empty projection batches. + /// + /// The previous batch's tick plus this interval must fit `usize` for the spacing comparison to + /// preserve this bound. + pub minimum_projection_interval: usize, + /// Maximum admitted events, including results awaiting delivery. + pub max_pending: NonZero, +} + +/// Reused embedding-lookup results for one poll. +#[derive(Default)] +struct Scratch { + embeddings: FastHashMap>, +} + +hashql_core::id::newtype! { + /// A placement task's own tick counter, used to space projection batches. + /// + /// Unit increments and spacing sums convert the stored value to `usize` before addition. A 32-bit target discards upper bits during that conversion. + /// + /// # Warning + /// + /// Without overflow checking, overflowing additions wrap. Repeated counter values or wrapped spacing thresholds can violate the minimum interval between projection batches. + struct Tick(u64) +} + +/// One generation's embedding polling and projection task. +pub(crate) struct DeltaPlacementTask { + pool: Arc, + tick: Tick, + options: DeltaPlacementTaskOptions, + actor: Option>, + pending: Pending, + scratch: Scratch, + workflow: Option>, + last_projection_at: Option, +} + +impl DeltaPlacementTask { + /// Builds a task with bounded admission and no cached actor. + /// + /// # Errors + /// + /// Returns [`DeltaPlacementError::InvalidInterval`] for a zero polling interval. + pub(crate) fn new( + pool: Arc, + options: DeltaPlacementTaskOptions, + workflow: Option>, + ) -> Result> { + if options.tick_rate.is_zero() { + return Err(Report::new(DeltaPlacementError::InvalidInterval)); + } + + Ok(Self { + pool, + tick: Tick::new(0), + options, + actor: None, + pending: Pending::default(), + scratch: Scratch::default(), + workflow, + last_projection_at: None, + }) + } + + /// Drains accepted requests after input closure, or stops when the result receiver closes. + /// + /// Results retain their event IDs and use fitted-world coordinates. A store failure retries the + /// failed lookup stage without spending that stage's miss budget. Stages already completed in + /// the same poll remain applied. An unavailable actor exhausts workflow submissions. One + /// projection batch runs at a time. + /// + /// Dropping this future cancels its asynchronous polling and delivery. A projection already + /// submitted to Rayon continues after its result receiver drops. Completion releases the owned + /// model, scratch storage and undeliverable results. + /// + /// # Errors + /// + /// Returns [`DeltaPlacementError::Offload`] if the projection worker panics or disappears + /// without returning its owned state. The task logs and retries store failures. It delivers + /// workflow and per-row projection failures as [`PlacementError`] values without ending the + /// task. + /// + /// # Panics + /// + /// Panics unless polled within a Tokio runtime with time enabled. A sufficiently late tick can + /// also panic when adding [`DeltaPlacementTaskOptions::tick_rate`] to the current instant would + /// exceed Tokio's representable deadline. With overflow checking, a unit tick increment at + /// `usize::MAX` or an evaluated spacing sum above that bound also panics, as + /// [`Self::projection_due`] describes. + pub(crate) async fn run( + mut self, + mut projector: DeltaProjector, + mut rx: mpsc::Receiver>, + tx: mpsc::Sender>, + ) -> Result<(), Report> { + let mut interval = tokio::time::interval(self.options.tick_rate); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + let mut input_open = true; + + while self + .wait_for_tick(&mut rx, &tx, &mut input_open, &mut interval) + .await + .is_continue() + { + self.tick.increment_by(1); + if let Err(error) = self.poll().await { + tracing::warn!(?error, "Retry placement polling after a store failure"); + } + + if self.projection_due() { + let batch = core::mem::take(&mut self.pending.project); + let (returned, outcomes) = Self::project(projector, batch).await?; + projector = returned; + + self.pending.completed.extend(outcomes); + self.last_projection_at = Some(self.tick); + } + + if self.pending.flush(&tx).is_break() { + return Ok(()); + } + } + + Ok(()) + } + + /// Processes placement work until a tick, completion-channel closure or complete input drain. + /// + /// Returns [`ControlFlow::Continue`] whenever a tick fires, including while pending work fills + /// the admission limit. Only input admission observes that limit. A full task can keep polling + /// and projecting its admitted work. Returns [`ControlFlow::Break`] when the completion channel + /// closes. After input closure, it also breaks once all pending work drains and the task + /// delivers every result. + /// + /// # Panics + /// + /// A late tick can panic if `interval` computes an unrepresentable next deadline. + #[expect( + clippy::integer_division_remainder_used, + reason = "Tokio select uses a remainder to randomize its first branch" + )] + async fn wait_for_tick( + &mut self, + rx: &mut mpsc::Receiver>, + tx: &mpsc::Sender>, + input_open: &mut bool, + interval: &mut Interval, + ) -> ControlFlow<()> { + while *input_open || self.pending.len() > 0 { + tokio::select! { + () = tx.closed() => return ControlFlow::Break(()), + permit = tx.reserve(), if !self.pending.completed.is_empty() => { + let Ok(permit) = permit else { + return ControlFlow::Break(()); + }; + if let Some(completed) = self.pending.completed.pop_front() { + permit.send(completed); + } + if self.pending.flush(tx).is_break() { + return ControlFlow::Break(()); + } + } + request = rx.recv(), if *input_open && self.pending.len() < self.options.max_pending.get() => { + match request { + Some(request) => self.pending.enqueue(request, self.options.tries_database), + None => *input_open = false, + } + } + _ = interval.tick() => return ControlFlow::Continue(()), + } + } + + ControlFlow::Break(()) + } + + /// Returns whether a non-empty projection batch is due at the current tick. + /// + /// The current tick and the previous projection tick plus + /// [`DeltaPlacementTaskOptions::minimum_projection_interval`] must fit `usize` for this + /// comparison to preserve the requested spacing. A first batch needs no spacing comparison. + /// + /// # Panics + /// + /// With overflow checking, panics when a non-empty pending batch and a previous projection + /// require a spacing sum above `usize::MAX`. + fn projection_due(&self) -> bool { + !self.pending.project.is_empty() + && self + .last_projection_at + .is_none_or(|last| self.tick >= last.plus(self.options.minimum_projection_interval)) + } + + /// Reads `ids`' projector embeddings into `scratch`, clearing it first. + /// + /// # Errors + /// + /// Returns [`DeltaPlacementError::Read`] if the store read fails. + async fn fetch_embeddings( + store: &impl AsClient, + ids: impl Iterator, + scratch: &mut Scratch, + ) -> Result<(), Report> { + scratch.embeddings.clear(); + + let embeddings = read_projector_embeddings(store, ids) + .await + .change_context(DeltaPlacementError::Read)?; + scratch.embeddings.extend(embeddings); + + Ok(()) + } + + /// Polls outstanding workflow and database lookups, then submits exhausted database entries. + /// + /// Returns without connecting to the store when no entry awaits a poll or submission. + /// + /// # Errors + /// + /// Returns [`DeltaPlacementError::Connect`] or [`DeltaPlacementError::Read`] for a store + /// failure. + async fn poll(&mut self) -> Result<(), Report> { + if self.pending.poll_database.is_empty() + && self.pending.poll_workflow.is_empty() + && self.pending.submit_workflow.is_empty() + { + return Ok(()); + } + + let pool = Arc::clone(&self.pool); + let mut store = pool + .acquire(None) + .await + .change_context(DeltaPlacementError::Connect)?; + + if !self.pending.poll_workflow.is_empty() { + Self::fetch_embeddings( + &store, + self.pending + .poll_workflow + .iter() + .map(|entry| ArchivedEntityId::from(entry.entity)), + &mut self.scratch, + ) + .await?; + + self.pending + .transition_workflow(&mut self.scratch.embeddings); + } + + if !self.pending.poll_database.is_empty() { + Self::fetch_embeddings( + &store, + self.pending + .poll_database + .iter() + .map(|entry| ArchivedEntityId::from(entry.entity)), + &mut self.scratch, + ) + .await?; + + self.pending + .transition_database(&mut self.scratch.embeddings); + } + + self.scratch.embeddings.clear(); + self.submit_workflow(&mut store).await; + + Ok(()) + } + + /// Starts an embedding workflow for entries awaiting submission, resolving the actor lazily. + /// + /// An unavailable workflow client or actor exhausts every waiting entry instead of submitting. + async fn submit_workflow(&mut self, store: &mut impl PrincipalStore) { + if self.pending.submit_workflow.is_empty() { + return; + } + + let Some(workflow) = &self.workflow else { + self.pending.exhaust_submit_workflow(); + return; + }; + + let ControlFlow::Continue(actor) = + Self::resolve_actor(&mut self.actor, &mut self.pending, async { + store + .get_or_create_system_machine("h") + .await + .map(ActorId::from) + }) + .await + else { + return; + }; + + for entry in self.pending.submit_workflow.drain(..) { + let workflow_id = format!("atlas-embedding-{}", entry.entity); + let start = workflow + .temporal + .ensure_update_entity_embeddings_workflow( + workflow_id, + actor.into(), + entry.entity, + &workflow.exclusions, + ) + .await; + + match start { + Ok(_) => self + .pending + .poll_workflow + .push(entry.transition(PollWorkflow { + tries: self.options.tries_workflow, + })), + Err(error) => self + .pending + .completed + .push_back(entry.transition(Completed(Err( + error.change_context(PlacementError::Workflow), + )))), + } + } + } + + /// Returns the cached actor lookup outcome, running and caching `lookup` on the first call. + /// + /// A failed lookup is cached as [`ControlFlow::Break`] and exhausts `pending`'s waiting + /// workflow submissions. + async fn resolve_actor( + cached: &mut Option>, + pending: &mut Pending, + lookup: impl Future>, + ) -> ControlFlow<(), ActorId> { + let actor = if let Some(actor) = *cached { + actor + } else { + let actor = match lookup.await { + Ok(actor) => ControlFlow::Continue(actor), + Err(error) => { + tracing::warn!( + ?error, + "Disable embedding workflows after actor lookup failure" + ); + ControlFlow::Break(()) + } + }; + + *cached = Some(actor); + actor + }; + + if actor.is_break() { + pending.exhaust_submit_workflow(); + } + actor + } + + /// Projects `batch` on a blocking offload thread, returning the projector for reuse. + /// + /// # Errors + /// + /// Returns [`DeltaPlacementError::Offload`] if the projection worker panics or disappears + /// without returning, discarding the projector. + #[tracing::instrument(skip_all, fields(rows = batch.len()))] + async fn project( + projector: DeltaProjector, + batch: Vec>, + ) -> Result< + ( + DeltaProjector, + impl ExactSizeIterator>, + ), + Report, + > { + // The job owns the projector and discards it on panic. Partially mutated scratch is never + // reused. + let mut projector = AssertUnwindSafe(projector); + + crate::offload::run(move || { + let results: Vec<_> = projector + .project(batch.iter().map(|entry| &entry.phase.embedding)) + .collect(); + + let completed = batch.into_iter().zip(results).map(|(entry, result)| { + entry.transition(Completed(result.change_context(PlacementError::Projection))) + }); + + (projector.0, completed) + }) + .await + .change_context(DeltaPlacementError::Offload) + } +} diff --git a/libs/@local/graph/atlas/src/serve/delta/placement/pending/mod.rs b/libs/@local/graph/atlas/src/serve/delta/placement/pending/mod.rs new file mode 100644 index 00000000000..e4118b1dad3 --- /dev/null +++ b/libs/@local/graph/atlas/src/serve/delta/placement/pending/mod.rs @@ -0,0 +1,180 @@ +//! Placement requests partitioned by embedding and projection phase. + +use alloc::collections::VecDeque; +use core::ops::ControlFlow; + +use error_stack::Report; +use hashql_core::collections::FastHashMap; +use tokio::sync::mpsc::{self, error::TrySendError}; +use type_system::knowledge::entity::EntityId; + +use super::PlacementError; +use crate::{ + dataset::PROJECTOR_DIMENSIONS, + math::BoxedVecN, + postgres::id::ArchivedEntityId, + serve::delta::{feed::EventId, projector::Position}, +}; + +#[cfg(test)] +mod tests; + +/// An event awaiting its first embedding lookup. +pub(crate) struct Initial; + +/// An event awaiting a database embedding lookup, with its remaining miss budget. +/// +/// The zero value still permits one lookup: a hit succeeds, while a miss leaves this phase. +pub(super) struct PollDatabase { + pub tries: u16, +} + +/// An event awaiting workflow submission after its database tries are exhausted. +pub(super) struct SubmitWorkflow; + +/// An event awaiting a workflow embedding lookup, with its remaining miss budget. +/// +/// The zero value still permits one lookup: a hit succeeds, while a miss exhausts the entry. +pub(super) struct PollWorkflow { + pub tries: u16, +} + +/// An event with a resolved embedding, awaiting projection. +pub(super) struct Project { + pub embedding: BoxedVecN, +} + +/// One event's fitted-world position or placement failure. +pub(crate) struct Completed(pub Result>); + +/// A placement request retaining its event identity through each phase. +pub(crate) struct PendingEntry { + /// The coalesced event version this request belongs to. + pub event: EventId, + /// The entity awaiting its embedding and position. + pub entity: EntityId, + /// This request's current placement stage. + pub phase: T, +} + +impl PendingEntry { + /// Advances this entry to `phase`, retaining its event and entity identity. + pub(super) fn transition(self, phase: U) -> PendingEntry { + PendingEntry { + event: self.event, + entity: self.entity, + phase, + } + } +} + +/// Placement requests partitioned by their embedding-resolution stage. +#[derive(Default)] +pub(super) struct Pending { + pub poll_database: Vec>, + pub submit_workflow: Vec>, + pub poll_workflow: Vec>, + pub project: Vec>, + pub completed: VecDeque>, +} + +impl Pending { + /// Returns the total count of entries summed across every stage. + pub(crate) fn len(&self) -> usize { + self.poll_database.len() + + self.submit_workflow.len() + + self.poll_workflow.len() + + self.project.len() + + self.completed.len() + } + + /// Sends completed results to `tx` until it fills or closes. + /// + /// When the channel is full, this method restores the popped result to the front and returns + /// [`ControlFlow::Continue`]. Closure discards that result, leaves later results queued and + /// returns [`ControlFlow::Break`]. + pub(crate) fn flush(&mut self, tx: &mpsc::Sender>) -> ControlFlow<()> { + while let Some(next) = self.completed.pop_front() { + match tx.try_send(next) { + Ok(()) => {} + Err(TrySendError::Full(next)) => { + self.completed.push_front(next); + break; + } + Err(TrySendError::Closed(_)) => return ControlFlow::Break(()), + } + } + ControlFlow::Continue(()) + } + + /// Admits a new placement request to await its first database embedding poll. + pub(crate) fn enqueue(&mut self, entry: PendingEntry, tries: u16) { + self.poll_database + .push(entry.transition(PollDatabase { tries })); + } + + /// Consumes matching database embeddings and applies each entry's miss budget. + /// + /// Matches move to projection regardless of the remaining budget. Misses decrement it with + /// saturation, and an unmatched entry at zero moves to workflow submission. + pub(crate) fn transition_database( + &mut self, + embeddings: &mut FastHashMap>, + ) { + self.poll_database.retain_mut(|entry| { + entry.phase.tries = entry.phase.tries.saturating_sub(1); + if let Some(embedding) = embeddings.remove(&ArchivedEntityId::from(entry.entity)) { + self.project.push(PendingEntry { + event: entry.event, + entity: entry.entity, + phase: Project { embedding }, + }); + return false; + } + true + }); + self.submit_workflow.extend( + self.poll_database + .extract_if(.., |entry| entry.phase.tries == 0) + .map(|entry| entry.transition(SubmitWorkflow)), + ); + } + + /// Consumes matching workflow embeddings and applies each entry's miss budget. + /// + /// Matches move to projection regardless of the remaining budget. Misses decrement it with + /// saturation, and an unmatched entry at zero completes with [`PlacementError::Exhaustion`]. + pub(crate) fn transition_workflow( + &mut self, + embeddings: &mut FastHashMap>, + ) { + self.poll_workflow.retain_mut(|entry| { + entry.phase.tries = entry.phase.tries.saturating_sub(1); + if let Some(embedding) = embeddings.remove(&ArchivedEntityId::from(entry.entity)) { + self.project.push(PendingEntry { + event: entry.event, + entity: entry.entity, + phase: Project { embedding }, + }); + return false; + } + true + }); + + self.completed.extend( + self.poll_workflow + .extract_if(.., |entry| entry.phase.tries == 0) + .map(|entry| { + entry.transition(Completed(Err(Report::new(PlacementError::Exhaustion)))) + }), + ); + } + + /// Completes every entry awaiting workflow submission with [`PlacementError::Exhaustion`]. + pub(crate) fn exhaust_submit_workflow(&mut self) { + self.completed + .extend(self.submit_workflow.drain(..).map(|entry| { + entry.transition(Completed(Err(Report::new(PlacementError::Exhaustion)))) + })); + } +} diff --git a/libs/@local/graph/atlas/src/serve/delta/placement/pending/tests.rs b/libs/@local/graph/atlas/src/serve/delta/placement/pending/tests.rs new file mode 100644 index 00000000000..1266b760bfe --- /dev/null +++ b/libs/@local/graph/atlas/src/serve/delta/placement/pending/tests.rs @@ -0,0 +1,460 @@ +use core::{assert_matches, ops::ControlFlow}; + +use error_stack::Report; +use hashql_core::collections::FastHashMap; +use tokio::sync::mpsc; +use type_system::knowledge::entity::EntityId; +use uuid::Uuid; + +use super::{ + Completed, Initial, Pending, PendingEntry, PollDatabase, PollWorkflow, Project, SubmitWorkflow, +}; +use crate::{ + dataset::PROJECTOR_DIMENSIONS, + math::BoxedVecN, + postgres::id::ArchivedEntityId, + serve::delta::{feed::EventId, placement::PlacementError}, +}; + +/// Builds an entity id from a fixed web and `value`, for compact literal test identities. +fn entity(value: u128) -> EntityId { + EntityId::from(ArchivedEntityId { + web_id: Uuid::from_u128(1).into(), + entity_uuid: Uuid::from_u128(value).into(), + }) +} + +/// Builds an embedding with every component set to `fill`. +fn embedding(fill: f32) -> BoxedVecN { + BoxedVecN::from([fill; PROJECTOR_DIMENSIONS]) +} + +/// Builds a completed entry that failed with an exhaustion error. +fn errored(event: u32, entity_value: u128) -> PendingEntry { + PendingEntry { + event: EventId::new(event), + entity: entity(entity_value), + phase: Completed(Err(Report::new(PlacementError::Exhaustion))), + } +} + +/// Enqueues a fresh request with its full database-poll budget. +#[test] +fn enqueue_budget() { + let mut pending = Pending::default(); + let initial = PendingEntry { + event: EventId::new(1), + entity: entity(1), + phase: Initial, + }; + + pending.enqueue(initial, 4); + + assert_eq!(pending.poll_database.len(), 1); + assert_eq!(pending.poll_database[0].event, EventId::new(1)); + assert_eq!(pending.poll_database[0].entity, entity(1)); + assert_eq!(pending.poll_database[0].phase.tries, 4); +} + +/// Moves found database embeddings into projection without cloning them. +/// +/// An embedding found in the batch moves that entry to projection by relocating the map's +/// allocation rather than cloning it, while an entry missing from the batch remains polling. +#[test] +fn database_hit() { + let mut pending = Pending::default(); + let hit_entity = entity(1); + let miss_entity = entity(2); + + pending.poll_database.push(PendingEntry { + event: EventId::new(10), + entity: hit_entity, + phase: PollDatabase { tries: 3 }, + }); + pending.poll_database.push(PendingEntry { + event: EventId::new(11), + entity: miss_entity, + phase: PollDatabase { tries: 3 }, + }); + + let source = embedding(0.25); + let source_ptr = source.as_array().as_ptr(); + let mut embeddings = FastHashMap::default(); + embeddings.extend([(ArchivedEntityId::from(hit_entity), source)]); + + pending.transition_database(&mut embeddings); + + assert_eq!(pending.poll_database.len(), 1); + assert_eq!(pending.poll_database[0].entity, miss_entity); + assert!(pending.submit_workflow.is_empty()); + + assert_eq!(pending.project.len(), 1); + let projected = &pending.project[0]; + assert_eq!(projected.event, EventId::new(10)); + assert_eq!(projected.entity, hit_entity); + assert_eq!( + projected.phase.embedding.as_array().as_ptr(), + source_ptr, + "should move the map's allocation into the project entry, not clone it" + ); + + assert!( + embeddings + .remove(&ArchivedEntityId::from(hit_entity)) + .is_none() + ); +} + +/// Promotes a database hit on the last allowed try. +#[test] +fn database_final_hit() { + let mut pending = Pending::default(); + let hit_entity = entity(1); + + pending.poll_database.push(PendingEntry { + event: EventId::new(20), + entity: hit_entity, + phase: PollDatabase { tries: 1 }, + }); + + let mut embeddings = FastHashMap::default(); + embeddings.extend([(ArchivedEntityId::from(hit_entity), embedding(0.75))]); + + pending.transition_database(&mut embeddings); + + assert!(pending.poll_database.is_empty()); + assert!(pending.submit_workflow.is_empty()); + assert_eq!(pending.project.len(), 1); + assert_eq!(pending.project[0].entity, hit_entity); +} + +/// Promotes a database hit even when no tries remain. +#[test] +fn database_zero_hit() { + let mut pending = Pending::default(); + let hit_entity = entity(1); + + pending.poll_database.push(PendingEntry { + event: EventId::new(21), + entity: hit_entity, + phase: PollDatabase { tries: 0 }, + }); + + let mut embeddings = FastHashMap::default(); + embeddings.extend([(ArchivedEntityId::from(hit_entity), embedding(0.75))]); + + pending.transition_database(&mut embeddings); + + assert!(pending.poll_database.is_empty()); + assert!(pending.submit_workflow.is_empty()); + assert_eq!(pending.project.len(), 1); + assert_eq!(pending.project[0].entity, hit_entity); +} + +/// Moves a database miss to workflow submission only when its budget ends. +/// +/// Misses decrement the budget, and the terminal one promotes the entry exactly once. +#[test] +fn database_miss() { + let mut pending = Pending::default(); + let miss_entity = entity(1); + + pending.poll_database.push(PendingEntry { + event: EventId::new(30), + entity: miss_entity, + phase: PollDatabase { tries: 2 }, + }); + + let mut embeddings = FastHashMap::default(); + + pending.transition_database(&mut embeddings); + + assert_eq!(pending.poll_database.len(), 1); + assert_eq!(pending.poll_database[0].phase.tries, 1); + assert!(pending.submit_workflow.is_empty()); + + pending.transition_database(&mut embeddings); + + assert!(pending.poll_database.is_empty()); + assert_eq!(pending.submit_workflow.len(), 1); + assert_eq!(pending.submit_workflow[0].entity, miss_entity); + assert_eq!(pending.submit_workflow[0].event, EventId::new(30)); + + pending.transition_database(&mut embeddings); + + assert!(pending.poll_database.is_empty()); + assert_eq!( + pending.submit_workflow.len(), + 1, + "should not promote the same entry a second time" + ); +} + +/// Consumes found workflow embeddings exactly once. +/// +/// An embedding found in the batch moves that entry to projection and consumes it from the +/// batch's source map so it is not read twice. +#[test] +fn workflow_hit() { + let mut pending = Pending::default(); + let hit_entity = entity(1); + + pending.poll_workflow.push(PendingEntry { + event: EventId::new(40), + entity: hit_entity, + phase: PollWorkflow { tries: 2 }, + }); + + let mut embeddings = FastHashMap::default(); + let key = ArchivedEntityId::from(hit_entity); + embeddings.extend([(key, embedding(0.5))]); + + pending.transition_workflow(&mut embeddings); + + assert!(pending.poll_workflow.is_empty()); + assert!(pending.completed.is_empty()); + assert_eq!(pending.project.len(), 1); + assert_eq!(pending.project[0].entity, hit_entity); + assert_eq!(pending.project[0].phase.embedding, embedding(0.5)); + assert!(embeddings.remove(&key).is_none()); +} + +/// Promotes a workflow hit on the last allowed try. +#[test] +fn workflow_final_hit() { + let mut pending = Pending::default(); + let hit_entity = entity(1); + + pending.poll_workflow.push(PendingEntry { + event: EventId::new(42), + entity: hit_entity, + phase: PollWorkflow { tries: 1 }, + }); + + let mut embeddings = FastHashMap::default(); + embeddings.extend([(ArchivedEntityId::from(hit_entity), embedding(0.5))]); + + pending.transition_workflow(&mut embeddings); + + assert!(pending.poll_workflow.is_empty()); + assert!(pending.completed.is_empty()); + assert_eq!(pending.project.len(), 1); + assert_eq!(pending.project[0].entity, hit_entity); +} + +/// Promotes a workflow hit even when no tries remain. +#[test] +fn workflow_zero_hit() { + let mut pending = Pending::default(); + let hit_entity = entity(1); + + pending.poll_workflow.push(PendingEntry { + event: EventId::new(43), + entity: hit_entity, + phase: PollWorkflow { tries: 0 }, + }); + + let mut embeddings = FastHashMap::default(); + embeddings.extend([(ArchivedEntityId::from(hit_entity), embedding(0.5))]); + + pending.transition_workflow(&mut embeddings); + + assert!(pending.poll_workflow.is_empty()); + assert!(pending.completed.is_empty()); + assert_eq!(pending.project.len(), 1); + assert_eq!(pending.project[0].entity, hit_entity); +} + +/// Completes a workflow miss only when its budget ends. +/// +/// Misses decrement the budget, and the terminal one emits one exhaustion result. +#[test] +fn workflow_exhaustion() { + let mut pending = Pending::default(); + let miss_entity = entity(1); + + pending.poll_workflow.push(PendingEntry { + event: EventId::new(41), + entity: miss_entity, + phase: PollWorkflow { tries: 2 }, + }); + + let mut embeddings = FastHashMap::default(); + + pending.transition_workflow(&mut embeddings); + + assert_eq!(pending.poll_workflow.len(), 1); + assert_eq!(pending.poll_workflow[0].phase.tries, 1); + assert!(pending.completed.is_empty()); + + pending.transition_workflow(&mut embeddings); + + assert!(pending.poll_workflow.is_empty()); + assert_eq!(pending.completed.len(), 1); + assert_eq!(pending.completed[0].entity, miss_entity); + let Err(report) = &pending.completed[0].phase.0 else { + panic!("should complete with an error after the final workflow miss") + }; + assert_matches!(report.current_context(), PlacementError::Exhaustion); + + pending.transition_workflow(&mut embeddings); + + assert!(pending.poll_workflow.is_empty()); + assert_eq!( + pending.completed.len(), + 1, + "should not exhaust the same entry a second time" + ); +} + +/// Completes queued workflow submissions with exhaustion when submission is unavailable. +/// +/// Completion preserves queue order. +#[test] +fn submit_exhaustion() { + let mut pending = Pending::default(); + let first = entity(1); + let second = entity(2); + + pending.submit_workflow.push(PendingEntry { + event: EventId::new(50), + entity: first, + phase: SubmitWorkflow, + }); + pending.submit_workflow.push(PendingEntry { + event: EventId::new(51), + entity: second, + phase: SubmitWorkflow, + }); + + pending.exhaust_submit_workflow(); + + assert!(pending.submit_workflow.is_empty()); + assert_eq!(pending.completed.len(), 2); + + assert_eq!(pending.completed[0].entity, first); + let Err(first_report) = &pending.completed[0].phase.0 else { + panic!("should complete with an error") + }; + assert_matches!(first_report.current_context(), PlacementError::Exhaustion); + + assert_eq!(pending.completed[1].entity, second); + let Err(second_report) = &pending.completed[1].phase.0 else { + panic!("should complete with an error") + }; + assert_matches!(second_report.current_context(), PlacementError::Exhaustion); +} + +/// Counts entries across every pending phase, including completed results. +#[test] +fn len_phases() { + let mut pending = Pending::default(); + + assert_eq!(pending.len(), 0); + + pending.poll_database.push(PendingEntry { + event: EventId::new(60), + entity: entity(1), + phase: PollDatabase { tries: 3 }, + }); + pending.submit_workflow.push(PendingEntry { + event: EventId::new(61), + entity: entity(2), + phase: SubmitWorkflow, + }); + pending.poll_workflow.push(PendingEntry { + event: EventId::new(62), + entity: entity(3), + phase: PollWorkflow { tries: 2 }, + }); + pending.project.push(PendingEntry { + event: EventId::new(63), + entity: entity(4), + phase: Project { + embedding: embedding(0.1), + }, + }); + pending.completed.push_back(errored(64, 5)); + + assert_eq!(pending.len(), 5); +} + +/// Keeps undelivered completions queued in their original order. +/// +/// Flushing more completions than the channel currently admits sends what fits and leaves the +/// rest queued, oldest first. +#[test] +fn flush_full() { + let (tx, mut rx) = mpsc::channel(1); + let mut pending = Pending::default(); + + let first = entity(1); + let second = entity(2); + let third = entity(3); + + pending.completed.push_back(errored(70, 1)); + pending.completed.push_back(errored(71, 2)); + pending.completed.push_back(errored(72, 3)); + + let result = pending.flush(&tx); + + assert_matches!(result, ControlFlow::Continue(())); + assert_eq!(pending.completed.len(), 2); + assert_eq!(pending.completed[0].entity, second); + assert_eq!(pending.completed[1].entity, third); + + let received = rx + .try_recv() + .expect("should carry the first entry through the one open slot"); + assert_eq!(received.entity, first); +} + +/// Delivers one queued completion per reopened channel slot, in order. +#[test] +fn flush_resume() { + let (tx, mut rx) = mpsc::channel(1); + let mut pending = Pending::default(); + + let first = entity(1); + let second = entity(2); + let third = entity(3); + + pending.completed.push_back(errored(80, 1)); + pending.completed.push_back(errored(81, 2)); + pending.completed.push_back(errored(82, 3)); + + assert_matches!(pending.flush(&tx), ControlFlow::Continue(())); + let received_first = rx + .try_recv() + .expect("should send the first entry before the channel fills"); + assert_eq!(received_first.entity, first); + + assert_matches!(pending.flush(&tx), ControlFlow::Continue(())); + let received_second = rx + .try_recv() + .expect("should send the second entry once capacity frees"); + assert_eq!(received_second.entity, second); + + assert_matches!(pending.flush(&tx), ControlFlow::Continue(())); + let received_third = rx + .try_recv() + .expect("should send the third entry once capacity frees again"); + assert_eq!(received_third.entity, third); + + assert!(pending.completed.is_empty()); +} + +/// Discards the popped completion when its channel has closed. +#[test] +fn flush_closed() { + let (tx, rx) = mpsc::channel::>(1); + drop(rx); + + let mut pending = Pending::default(); + pending.completed.push_back(errored(90, 1)); + + let result = pending.flush(&tx); + + assert_matches!(result, ControlFlow::Break(())); + assert!(pending.completed.is_empty()); +} diff --git a/libs/@local/graph/atlas/src/serve/delta/placement/tests.rs b/libs/@local/graph/atlas/src/serve/delta/placement/tests.rs new file mode 100644 index 00000000000..65949655b67 --- /dev/null +++ b/libs/@local/graph/atlas/src/serve/delta/placement/tests.rs @@ -0,0 +1,466 @@ +use alloc::sync::Arc; +use core::{ + assert_matches, + future::{Future as _, pending, ready}, + ops::ControlFlow, + pin::pin, + task::{Context, Waker}, + time::Duration, +}; + +use error_stack::Report; +use futures::FutureExt as _; +use hash_graph_postgres_store::store::{ + DatabaseConnectionInfo, DatabasePoolConfig, DatabaseType, PostgresStorePool, + PostgresStoreSettings, +}; +use hashql_core::id::Id as _; +use tokio::{sync::mpsc, time::Instant}; +use tokio_postgres::NoTls; +use type_system::{ + knowledge::entity::EntityId, + principal::actor::{ActorId, ActorType}, +}; +use uuid::Uuid; + +use super::{ + Completed, DeltaPlacementError, DeltaPlacementTask, DeltaPlacementTaskOptions, Initial, + PendingEntry, PlacementError, Tick, + pending::{Pending, Project, SubmitWorkflow}, +}; +use crate::{ + dataset::PROJECTOR_DIMENSIONS, + math::{BoxedVecN, nz}, + postgres::id::ArchivedEntityId, + serve::delta::{ + feed::EventId, + projector::{ProjectionError, projector}, + }, +}; + +/// Builds an initial pending entry with `event` as both its event and entity seed. +fn request(event: u32) -> PendingEntry { + PendingEntry { + event: EventId::new(event), + entity: EntityId::from(ArchivedEntityId { + web_id: Uuid::from_u128(1).into(), + entity_uuid: Uuid::from_u128(u128::from(event)).into(), + }), + phase: Initial, + } +} + +/// Builds a completed entry for `event` that failed with an exhaustion error. +fn completion(event: u32) -> PendingEntry { + request(event).transition(Completed(Err(Report::new(PlacementError::Exhaustion)))) +} + +/// Returns slow-ticking placement options with a small retry budget and queue depth 2. +fn options() -> DeltaPlacementTaskOptions { + DeltaPlacementTaskOptions { + tick_rate: Duration::from_hours(24), + tries_workflow: 2, + tries_database: 2, + minimum_projection_interval: 5, + max_pending: nz!(2), + } +} + +/// Constructs a placement task over an unconnected pool, using [`options`]. +/// +/// # Panics +/// +/// Panics if store-pool construction or placement-option validation fails. +async fn task() -> DeltaPlacementTask { + // Pool construction opens no connection. These cases stop before store acquisition. + let pool = PostgresStorePool::new( + &DatabaseConnectionInfo::new( + DatabaseType::Postgres, + "placement-test".to_owned(), + String::new(), + "/no-placement-test-postgres".to_owned(), + 5432, + "placement-test".to_owned(), + ), + &DatabasePoolConfig { + max_connections: nz!(1), + }, + NoTls, + PostgresStoreSettings::default(), + ) + .await + .expect("should construct an unconnected pool"); + DeltaPlacementTask::new(Arc::new(pool), options(), None) + .expect("should accept a non-zero interval") +} + +/// Rejects a zero tick rate as an invalid interval. +#[tokio::test] +async fn interval_zero() { + let task = task().await; + let mut options = options(); + options.tick_rate = Duration::ZERO; + let error = DeltaPlacementTask::new(task.pool, options, None) + .err() + .expect("should refuse a zero interval"); + assert_matches!( + error.current_context(), + DeltaPlacementError::InvalidInterval + ); +} + +/// Spaces nonempty projection batches by the configured minimum interval. +/// +/// Projection is due with queued work when no prior projection exists or the minimum interval +/// since it has elapsed. It is never due with an empty queue. +#[tokio::test] +async fn projection_cadence() { + let mut task = task().await; + assert!(!task.projection_due()); + task.pending.project.push(request(1).transition(Project { + embedding: BoxedVecN::zero(), + })); + assert!(task.projection_due()); + task.last_projection_at = Some(Tick::new(10)); + task.tick = Tick::new(14); + assert!(!task.projection_due()); + task.tick = Tick::new(15); + assert!(task.projection_due()); + task.tick = Tick::new(16); + assert!(task.projection_due()); + task.pending.project.clear(); + assert!(!task.projection_due()); +} + +/// Keeps projection cadence unchanged while delivering completed results. +/// +/// A tick that only delivers already-completed results, with no new admission or projection, +/// leaves the tick counter unchanged and delivers every queued result in order. +#[tokio::test] +async fn delivery_preserves_cadence() { + let mut task = task().await; + task.pending + .completed + .extend([completion(1), completion(2)]); + let (_requests, mut rx) = mpsc::channel(1); + let (tx, mut completed) = mpsc::channel(2); + let mut input_open = true; + let period = Duration::from_hours(24); + let mut interval = tokio::time::interval_at(Instant::now() + period, period); + + assert!( + task.wait_for_tick(&mut rx, &tx, &mut input_open, &mut interval) + .now_or_never() + .is_none() + ); + assert_eq!( + completed + .try_recv() + .expect("should deliver the first result") + .event, + EventId::new(1) + ); + assert_eq!( + completed + .try_recv() + .expect("should deliver the second result") + .event, + EventId::new(2) + ); + assert_eq!(task.pending.len(), 0); + assert_eq!(task.tick, Tick::new(0)); +} + +/// Admits requests only up to the configured pending limit. +/// +/// Admission stops at the configured queue depth, leaving the remainder on the input channel for +/// a later tick, and each admitted entry starts with its full database-try budget. +#[tokio::test] +async fn admission_capacity() { + let mut task = task().await; + let (requests, mut rx) = mpsc::channel(3); + for event in 1..=3 { + requests + .try_send(request(event)) + .expect("should queue each request within input capacity"); + } + let (tx, _completed) = mpsc::channel(1); + let mut input_open = true; + let period = Duration::from_hours(24); + let mut interval = tokio::time::interval_at(Instant::now() + period, period); + + assert!( + task.wait_for_tick(&mut rx, &tx, &mut input_open, &mut interval) + .now_or_never() + .is_none() + ); + assert_eq!(task.pending.len(), 2); + assert_eq!(rx.len(), 1); + assert_eq!(task.pending.poll_database[0].phase.tries, 2); + assert_eq!(task.pending.poll_database[1].phase.tries, 2); +} + +/// Drains every queued result before treating closed placement input as terminal. +#[tokio::test] +async fn input_closed_drain() { + let mut task = task().await; + task.pending + .completed + .extend([completion(1), completion(2)]); + let (requests, mut rx) = mpsc::channel(1); + drop(requests); + let (tx, mut completed) = mpsc::channel(2); + let mut input_open = true; + let period = Duration::from_hours(24); + let mut interval = tokio::time::interval_at(Instant::now() + period, period); + + assert_eq!( + task.wait_for_tick(&mut rx, &tx, &mut input_open, &mut interval) + .now_or_never(), + Some(ControlFlow::Break(())) + ); + assert!(!input_open); + assert_eq!(task.pending.len(), 0); + assert_eq!( + completed + .try_recv() + .expect("should deliver the first result") + .event, + EventId::new(1) + ); + assert_eq!( + completed + .try_recv() + .expect("should deliver the second result") + .event, + EventId::new(2) + ); +} + +/// Stops on completion-channel closure even while placement input remains open. +#[tokio::test] +async fn output_closed() { + let mut task = task().await; + task.pending.completed.push_back(completion(1)); + let (_requests, mut rx) = mpsc::channel(1); + let (tx, completed) = mpsc::channel(1); + drop(completed); + let mut input_open = true; + let period = Duration::from_hours(24); + let mut interval = tokio::time::interval_at(Instant::now() + period, period); + assert_eq!( + task.wait_for_tick(&mut rx, &tx, &mut input_open, &mut interval) + .now_or_never(), + Some(ControlFlow::Break(())) + ); +} + +/// Queues an undeliverable completion and continues when the output channel is full. +#[tokio::test] +async fn ready_tick_full_output() { + let mut task = task().await; + task.pending.completed.push_back(completion(2)); + let (_requests, mut rx) = mpsc::channel(1); + let (tx, _completed) = mpsc::channel(1); + tx.try_send(completion(1)) + .expect("should fill the output channel"); + let mut input_open = true; + let mut interval = tokio::time::interval(Duration::from_hours(24)); + assert_eq!( + task.wait_for_tick(&mut rx, &tx, &mut input_open, &mut interval) + .await, + ControlFlow::Continue(()) + ); + assert_eq!(task.pending.completed.len(), 1); + assert_eq!(task.pending.completed[0].event, EventId::new(2)); +} + +/// Caches actor lookup failure and exhausts every waiting submission. +/// +/// Resolution then breaks. +#[tokio::test] +async fn actor_lookup_failure() { + let mut cached = None; + let mut pending = Pending::default(); + pending + .submit_workflow + .push(request(1).transition(SubmitWorkflow)); + let result = DeltaPlacementTask::resolve_actor( + &mut cached, + &mut pending, + ready(Err::("unavailable")), + ) + .await; + assert_eq!(result, ControlFlow::Break(())); + assert_eq!(cached, Some(ControlFlow::Break(()))); + assert!(pending.submit_workflow.is_empty()); + assert_eq!(pending.completed.len(), 1); + assert_matches!( + pending.completed[0] + .phase + .0 + .as_ref() + .expect_err("should exhaust the submission") + .current_context(), + PlacementError::Exhaustion + ); +} + +/// Reuses a cached unavailable-actor outcome without polling. +/// +/// Queued workflow submissions still exhaust. +#[tokio::test] +async fn actor_unavailable_cached() { + let mut cached = Some(ControlFlow::Break(())); + let mut entries = Pending::default(); + entries + .submit_workflow + .push(request(1).transition(SubmitWorkflow)); + let result = DeltaPlacementTask::resolve_actor( + &mut cached, + &mut entries, + pending::>(), + ) + .now_or_never(); + assert_eq!(result, Some(ControlFlow::Break(()))); + assert!(entries.submit_workflow.is_empty()); + assert_eq!(entries.completed.len(), 1); +} + +/// Reuses a cached actor without polling. +/// +/// Queued workflow submissions remain untouched. +#[tokio::test] +async fn actor_available_cached() { + let actor = ActorId::new(Uuid::from_u128(1), ActorType::Machine); + let mut cached = None; + let mut entries = Pending::default(); + entries + .submit_workflow + .push(request(1).transition(SubmitWorkflow)); + assert_eq!( + DeltaPlacementTask::resolve_actor(&mut cached, &mut entries, ready(Ok::<_, &str>(actor))) + .await, + ControlFlow::Continue(actor) + ); + assert_eq!(cached, Some(ControlFlow::Continue(actor))); + assert_eq!( + DeltaPlacementTask::resolve_actor( + &mut cached, + &mut entries, + pending::>() + ) + .now_or_never(), + Some(ControlFlow::Continue(actor)) + ); + assert_eq!(entries.submit_workflow.len(), 1); + assert!(entries.completed.is_empty()); +} + +/// Returns reusable projector state while isolating per-row failures. +/// +/// Projecting a batch preserves request order in the result and reports a per-row projection +/// error for a non-finite embedding without failing its neighbors. It also returns a projector +/// usable for a later batch. +#[tokio::test] +async fn projection_owned_batch() { + let mut projector = projector(None); + let input = BoxedVecN::from([1.0; PROJECTOR_DIMENSIONS]); + let expected = projector + .project([&input]) + .next() + .expect("should return one result") + .expect("should project a finite input"); + let batch = vec![ + request(17).transition(Project { embedding: input }), + request(3).transition(Project { + embedding: BoxedVecN::from([f32::NAN; PROJECTOR_DIMENSIONS]), + }), + request(9).transition(Project { + embedding: BoxedVecN::from([1.0; PROJECTOR_DIMENSIONS]), + }), + ]; + let (projector, completed) = DeltaPlacementTask::project(projector, batch) + .await + .expect("should complete the offloaded batch"); + let completed: Vec<_> = completed.collect(); + assert_eq!( + completed + .iter() + .map(|entry| entry.event) + .collect::>(), + [EventId::new(17), EventId::new(3), EventId::new(9)] + ); + for entry in &completed { + let entity = request(entry.event.as_u32()).entity; + assert_eq!(entry.entity, entity); + } + assert_eq!( + completed[0] + .phase + .0 + .as_ref() + .expect("should place the first row"), + &expected + ); + assert_eq!( + completed[2] + .phase + .0 + .as_ref() + .expect("should place the last row"), + &expected + ); + let error = completed[1] + .phase + .0 + .as_ref() + .expect_err("should retain the row failure"); + assert_matches!(error.current_context(), PlacementError::Projection); + assert_eq!( + error.downcast_ref::(), + Some(&ProjectionError::NonFiniteProjection) + ); + + let (_, completed) = DeltaPlacementTask::project( + projector, + vec![request(20).transition(Project { + embedding: BoxedVecN::from([1.0; PROJECTOR_DIMENSIONS]), + })], + ) + .await + .expect("should reuse the returned projector"); + let completed: Vec<_> = completed.collect(); + assert_eq!(completed.len(), 1); + assert_eq!( + completed[0] + .phase + .0 + .as_ref() + .expect("should place the next batch"), + &expected + ); +} + +/// Runs projection on the worker pool rather than the calling task. +/// +/// With the pool's only thread occupied, a first poll returns pending, and the batch still +/// completes once the worker runs. +#[tokio::test] +async fn projection_worker_yield() { + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() + .expect("should build the test worker pool"); + let batch = vec![request(1).transition(Project { + embedding: BoxedVecN::zero(), + })]; + let mut job = pin!(DeltaPlacementTask::project(projector(None), batch)); + // The only worker cannot run the queued projection until this first poll returns. + let first = pool.install(|| job.as_mut().poll(&mut Context::from_waker(Waker::noop()))); + assert!(first.is_pending()); + let (_, completed) = job.await.expect("should return after the worker runs"); + let completed: Vec<_> = completed.collect(); + assert_eq!(completed.len(), 1); + assert_eq!(completed[0].event, EventId::new(1)); +} diff --git a/libs/@local/graph/atlas/src/serve/delta/projector/mod.rs b/libs/@local/graph/atlas/src/serve/delta/projector/mod.rs new file mode 100644 index 00000000000..63b6e7bd9db --- /dev/null +++ b/libs/@local/graph/atlas/src/serve/delta/projector/mod.rs @@ -0,0 +1,448 @@ +//! Validated projection of incoming embeddings into fitted-world coordinates. + +use core::{error::Error, fmt, num::NonZero}; +use std::fs::File; + +use error_stack::{Report, ResultExt as _}; +use hashql_core::id::{Id as _, IdSlice, IdVec}; + +use crate::{ + dataset::PROJECTOR_DIMENSIONS, + device::{self, PhysicalDevice}, + file::{generation::Generation, salt::metadata::ProjectorEvidence}, + identity::NodeRowId, + math::{ + AlignedVecN, Bounds2, DNonNegative, DPositive, DVec2, FinitePointField, MatrixN, + NonNegative, Similarity, Vec2, d_positive, + }, + salt::{ + file::{PointFile, VectorFile}, + fit::PlacementOptions, + projector::{ + artifact, + model::{NodeRole, Projector}, + train::{NodeColumns, refresh}, + }, + }, +}; + +#[cfg(test)] +mod tests; + +// The placement tests use the same small projector fixture. +#[cfg(test)] +pub(super) use tests::projector; + +/// A finite position in the fitted world's coordinate frame. +#[derive(Debug, Copy, Clone, PartialEq)] +pub(crate) struct Position(Vec2); + +impl Position { + /// Wraps `point`, or returns `None` if either coordinate is non-finite. + fn new(point: Vec2) -> Option { + point.is_finite().then_some(Self(point)) + } + + /// Returns the wrapped coordinate. + pub(crate) const fn get(self) -> Vec2 { + self.0 + } +} + +hashql_core::id::newtype! { + /// A row in one complete projection request. + #[id(const)] + pub(crate) struct ForwardIndex(u32) +} + +/// A projection failure for one input row. +#[derive(Debug, Copy, Clone, PartialEq)] +pub(crate) enum ProjectionError { + /// An aligned point lies outside the fitted world. + OutOfBounds { global: Vec2 }, + /// The model produced a non-finite point. + NonFiniteProjection, + /// Alignment produced a non-finite point. + NonFiniteAlignment, +} + +impl fmt::Display for ProjectionError { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::OutOfBounds { global } => write!( + fmt, + "projected point at ({}, {}) lies outside the fitted world", + global.x(), + global.y(), + ), + Self::NonFiniteProjection => fmt.write_str("the model produced a non-finite point"), + Self::NonFiniteAlignment => fmt.write_str("alignment produced a non-finite point"), + } + } +} + +impl Error for ProjectionError {} + +/// A projector open or roundtrip-validation failure. +#[derive(Debug)] +pub(crate) enum ProjectorError { + /// A projector checkpoint accompanies a baseline placement configuration. + UnexpectedBaselinePlacement, + /// Opening the checkpoint file failed. + CheckpointNotFound, + /// Decoding the checkpoint for the configured architecture failed. + InvalidCheckpoint, + /// Opening the representation column with the projector's width failed. + CannotLoadRepresentations, + /// Opening the coordinate column as points failed. + CannotLoadCoordinates, + /// The columns are empty or have different row counts. + InvalidSampleCorpus { + representations: usize, + coordinates: usize, + }, + /// A sampled published coordinate is non-finite. + NonFiniteCoordinates { row: NodeRowId }, + /// Projecting the fitted sample failed. + RoundtripSampleForward, + /// The maximum component error meets or exceeds the roundtrip tolerance. + RoundtripSampleToleranceExceeded { + tolerance: DPositive, + error: DNonNegative, + }, +} + +impl fmt::Display for ProjectorError { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::UnexpectedBaselinePlacement => fmt + .write_str("a projector checkpoint accompanies a baseline placement configuration"), + Self::CheckpointNotFound => fmt.write_str("opening the projector checkpoint failed"), + Self::InvalidCheckpoint => fmt.write_str( + "the projector checkpoint does not decode for the configured architecture", + ), + Self::CannotLoadRepresentations => { + fmt.write_str("the representation column does not open with the projector's width") + } + Self::CannotLoadCoordinates => { + fmt.write_str("the coordinate column does not open as points") + } + Self::InvalidSampleCorpus { + representations, + coordinates, + } => write!( + fmt, + "roundtrip columns have {representations} representation rows and {coordinates} \ + coordinate rows" + ), + Self::NonFiniteCoordinates { row } => { + write!(fmt, "the published coordinate at row {row} is non-finite") + } + Self::RoundtripSampleForward => fmt.write_str("projecting the fitted sample failed"), + Self::RoundtripSampleToleranceExceeded { tolerance, error } => write!( + fmt, + "roundtrip error {error} meets or exceeds tolerance {tolerance}" + ), + } + } +} + +impl Error for ProjectorError {} + +/// Reused per-chunk buffers for one forward pass. +struct ProjectorScratch { + forward: IdVec, + inputs: MatrixN, + roles: Vec, +} + +impl ProjectorScratch { + /// Rows processed per model invocation. + const CHUNK_SIZE: usize = 256; + + /// Allocates buffers sized for one chunk. + fn new() -> Self { + Self { + forward: IdVec::with_capacity(Self::CHUNK_SIZE), + inputs: MatrixN::zeroed(Self::CHUNK_SIZE), + roles: vec![NodeRole::KnowledgeEntity; Self::CHUNK_SIZE], + } + } +} + +/// An opened embedding projector, validated against a sample of its fitted coordinates. +pub(crate) struct DeltaProjector { + model: Projector, + device: PhysicalDevice, + condition: NonNegative, + alignment: Option, + world: Bounds2, + forward_rows: NonZero, + + scratch: ProjectorScratch, +} + +impl DeltaProjector { + /// Reopens the projector and checks a sample of its fitted coordinates. + /// + /// Returns `Ok(None)` when the generation has no projector checkpoint. + /// + /// # Errors + /// + /// Returns [`ProjectorError`] if the checkpoint conflicts with the placement configuration, + /// opening or decoding a file fails, or the roundtrip sample fails validation. + #[tracing::instrument(skip_all)] + pub(crate) fn open( + generation: &Generation, + device: PhysicalDevice, + ) -> Result, Report> { + let repository = generation.repository(); + let metadata = &repository.metadata; + + let Some(checkpoint) = &repository.files.projector else { + tracing::info!( + "the generation placed rows by landmark baseline, arrivals stage until a refit" + ); + + return Ok(None); + }; + + let PlacementOptions::Projector(options) = &metadata.reproducibility.config.placement + else { + tracing::warn!( + "the generation stages a projector checkpoint while its configuration echo \ + records a baseline placement" + ); + + return Err(Report::new(ProjectorError::UnexpectedBaselinePlacement)); + }; + + let (condition, alignment) = match &metadata.evidence.projector { + None => { + tracing::warn!( + "no projector evidence available, assuming zero condition placement" + ); + (NonNegative::ZERO, None) + } + Some(ProjectorEvidence { + ladder: Some(ladder), + .. + }) if let Some(step) = ladder.steps.get(ladder.canonical_index) => { + (ladder.canonical, Some(step.alignment)) + } + Some(ProjectorEvidence { + ladder: Some(ladder), + .. + }) => { + tracing::warn!( + canonical_index = ladder.canonical_index, + steps = ladder.steps.len(), + "projector evidence published a canonical out of bounds index" + ); + (NonNegative::ZERO, None) + } + Some(ProjectorEvidence { ladder: None, .. }) => { + tracing::warn!( + "projector evidence available but no ladder, assuming zero condition placement" + ); + (NonNegative::ZERO, None) + } + }; + + let checkpoint = File::open(generation.path_of(&checkpoint.name())) + .change_context(ProjectorError::CheckpointNotFound)?; + let model = artifact::open_model(checkpoint, options.architecture, &device) + .change_context(ProjectorError::InvalidCheckpoint)?; + + let mut this = Self { + model, + device, + condition, + alignment, + world: metadata.evidence.lod.world, + forward_rows: options.forward_rows, + scratch: ProjectorScratch::new(), + }; + + this.try_roundtrip_sample(generation)?; + + Ok(Some(this)) + } + + /// Replays a sample from the fitted representation and coordinate columns. + /// + /// # Errors + /// + /// Returns [`ProjectorError`] for artifact opening or sampled projection failures. + fn try_roundtrip_sample( + &mut self, + generation: &Generation, + ) -> Result<(), Report> { + let files = &generation.repository().files; + + let representations = VectorFile::open(generation.path_of(&files.representations.name())) + .change_context(ProjectorError::CannotLoadRepresentations)?; + let coordinates = PointFile::open(generation.path_of(&files.coordinates.name())) + .change_context(ProjectorError::CannotLoadCoordinates)?; + + self.try_roundtrip_sample_impl(&representations, &coordinates) + } + + /// Checks up to 1,024 regularly spaced fitted rows against their published coordinates. + /// + /// # Errors + /// + /// Returns [`ProjectorError`] for empty or unequal columns, non-finite sampled coordinates, + /// a failed projection, or a maximum component error at or above the roundtrip tolerance. + #[tracing::instrument( + skip_all, + fields(rows = representations.len(), samples = tracing::field::Empty), + )] + #[expect( + clippy::integer_division, + clippy::integer_division_remainder_used, + reason = "the floored stride spreads the sample across the fitted row domain" + )] + fn try_roundtrip_sample_impl( + &mut self, + representations: &IdSlice>, + coordinates: &IdSlice, + ) -> Result<(), Report> { + const SAMPLES: usize = 1024; + const TOLERANCE: DPositive = d_positive!(1e-3); + + if representations.is_empty() || representations.len() != coordinates.len() { + return Err(Report::new(ProjectorError::InvalidSampleCorpus { + representations: representations.len(), + coordinates: coordinates.len(), + })); + } + + let sample_size = SAMPLES.min(representations.len()); + tracing::Span::current().record("samples", sample_size); + + let sampled: IdVec = (0..sample_size) + .map(|index| NodeRowId::from_usize(index * representations.len() / sample_size)) + .collect(); + let published: IdVec = + sampled.iter().map(|&row| coordinates[row]).collect(); + let published = FinitePointField::new(&published).map_err(|error| { + let row = sampled[error.id]; + Report::new(error).change_context(ProjectorError::NonFiniteCoordinates { row }) + })?; + let sampled_rows = sampled.iter().map(|&row| &representations[row]); + + let aligned: Vec<_> = self + .forward(sampled_rows) + .try_collect() + .change_context(ProjectorError::RoundtripSampleForward)?; + + let mut max_error = DNonNegative::ZERO; + for (&published, &reprojected) in published.iter().zip(aligned.iter()) { + let difference = DVec2::from(reprojected.0) - DVec2::from(published); + // Differences of finite f32 coordinates remain finite after widening to f64. + let error = DNonNegative::new_unchecked(difference.x().abs().max(difference.y().abs())); + max_error = max_error.max(error); + } + + if max_error < TOLERANCE { + tracing::info!(%max_error, tolerance = %TOLERANCE, "Verify sampled projector coordinates"); + return Ok(()); + } + + Err(Report::new( + ProjectorError::RoundtripSampleToleranceExceeded { + tolerance: TOLERANCE, + error: max_error, + }, + )) + } + + /// Projects rows into the fitted world in input order without clamping their coordinates. + /// + /// # Errors + /// + /// Each input has one result. A failed row leaves the other results intact. + /// + /// Returns [`ProjectionError`] for non-finite model or alignment output, or a position outside + /// the fitted world. + #[tracing::instrument(skip_all)] + pub(crate) fn project( + &mut self, + rows: impl IntoIterator< + Item: AsRef>, + IntoIter: ExactSizeIterator, + >, + ) -> impl ExactSizeIterator> { + let world = self.world; + + self.forward(rows).map(move |result| { + let position = result?; + + if !world.contains(position.get()) { + return Err(ProjectionError::OutOfBounds { global: position.0 }); + } + + Ok(position) + }) + } + + /// Projects and aligns rows without checking the fitted world bounds. + /// + /// # Errors + /// + /// Returns one result per input, with [`ProjectionError`] for each non-finite model or + /// alignment output. + fn forward( + &mut self, + rows: impl IntoIterator< + Item: AsRef>, + IntoIter: ExactSizeIterator, + >, + ) -> impl ExactSizeIterator> { + let mut rows = rows.into_iter(); + let mut results = Vec::with_capacity(rows.len()); + + loop { + let mut filled = 0_usize; + + for (slot, row) in self.scratch.inputs.rows_mut().iter_mut().zip(&mut rows) { + slot.copy_from(row.as_ref()); + filled += 1; + } + + if filled == 0 { + break; + } + + let columns: NodeColumns<'_, ForwardIndex> = NodeColumns { + representations: IdSlice::from_raw(&self.scratch.inputs.rows()[..filled]), + roles: IdSlice::from_raw(&self.scratch.roles[..filled]), + }; + + refresh::forward_unchecked_in( + &self.model, + columns, + self.condition, + self.forward_rows, + &self.device, + &mut self.scratch.forward, + ); + + results.extend(self.scratch.forward.drain(..).map(|point| { + let position = Position::new(point).ok_or(ProjectionError::NonFiniteProjection)?; + + self.alignment.map_or(Ok(position), |alignment| { + Position::new(alignment.apply(position.get())) + .ok_or(ProjectionError::NonFiniteAlignment) + }) + })); + + if filled < ProjectorScratch::CHUNK_SIZE { + break; + } + } + + results.into_iter() + } +} diff --git a/libs/@local/graph/atlas/src/serve/delta/projector/tests.rs b/libs/@local/graph/atlas/src/serve/delta/projector/tests.rs new file mode 100644 index 00000000000..e7107414f54 --- /dev/null +++ b/libs/@local/graph/atlas/src/serve/delta/projector/tests.rs @@ -0,0 +1,329 @@ +use core::assert_matches; + +use hashql_core::id::IdSlice; +use rand::SeedableRng as _; +use rand_xoshiro::Xoshiro256PlusPlus; + +use super::{ + DeltaProjector, ForwardIndex, Position, ProjectionError, ProjectorError, ProjectorScratch, +}; +use crate::{ + dataset::PROJECTOR_DIMENSIONS, + device::Device, + identity::NodeRowId, + math::{Bounds2, MatrixN, NonNegative, Rotation, Similarity, Vec2, nz, positive}, + salt::projector::{ + model::{Architecture, NodeRole, Projector}, + train::{NodeColumns, refresh}, + }, +}; + +/// Builds a small deterministic projector, optionally composing `alignment` after the model. +/// +/// # Panics +/// +/// Panics if the fixed world bounds are not finite and ordered or the fixed forward bound is zero. +pub(crate) fn projector(alignment: Option) -> DeltaProjector { + let device = Device::Cpu.pin(0).resolve(); + let architecture = Architecture { + width: nz!(8), + residual_blocks: nz!(1), + representation_dimensions: nz!(PROJECTOR_DIMENSIONS), + role_dimensions: nz!(4), + condition_dimensions: nz!(1), + }; + + DeltaProjector { + model: Projector::new(architecture, &device, Xoshiro256PlusPlus::seed_from_u64(7)), + device, + condition: NonNegative::ZERO, + alignment, + world: Bounds2::new(Vec2::splat(-100.0), Vec2::splat(100.0)) + .expect("should have ordered finite bounds"), + forward_rows: nz!(64), + scratch: ProjectorScratch::new(), + } +} + +/// Builds `count` distinct unit-basis representation rows, cycling through each axis in turn. +fn representations(count: usize) -> MatrixN { + let mut inputs = MatrixN::zeroed(count); + for (axis, row) in (0..PROJECTOR_DIMENSIONS).cycle().zip(inputs.rows_mut()) { + row.as_array_mut()[axis] = 1.0; + } + inputs +} + +/// Accepts finite positions and rejects NaN or infinity in either component. +#[test] +fn position_non_finite() { + assert_eq!( + Position::new(Vec2::ZERO).map(Position::get), + Some(Vec2::ZERO) + ); + for value in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] { + assert!(Position::new(Vec2::new(value, 0.0)).is_none()); + assert!(Position::new(Vec2::new(0.0, value)).is_none()); + } +} + +/// Matches the reference projection across full and partial chunks. +/// +/// Forwarding a batch spanning multiple chunks, including a partial final chunk, matches the +/// reference forward pass row for row once the configured alignment is applied. +#[test] +fn forward_chunk_boundaries() { + let inputs = representations(ProjectorScratch::CHUNK_SIZE * 2 + 1); + let alignment = Similarity::new( + positive!(2.0), + Rotation::from_radians(0.5), + Vec2::new(3.0, -4.0), + ) + .expect("should have a valid similarity"); + let mut projector = projector(Some(alignment)); + let roles = vec![NodeRole::KnowledgeEntity; inputs.rows().len()]; + let columns: NodeColumns<'_, ForwardIndex> = NodeColumns { + representations: IdSlice::from_raw(inputs.rows()), + roles: IdSlice::from_raw(&roles), + }; + let expected = refresh::forward( + &projector.model, + columns, + projector.condition, + projector.forward_rows, + &projector.device, + ) + .expect("should project the reference rows"); + let actual: Vec<_> = projector + .forward(inputs.rows()) + .try_collect() + .expect("should project all chunks"); + assert_eq!(actual.len(), inputs.rows().len()); + for (&actual, &expected) in actual.iter().zip(expected.iter()) { + assert!( + actual + .get() + .distance_squared_wide(alignment.apply(expected)) + < 1e-10 + ); + } +} + +/// Handles shorter and empty batches after earlier use. +/// +/// A shorter batch reprojects the same leading row as a longer one, and an empty batch returns no +/// rows. +#[test] +fn forward_empty_and_reused() { + let mut projector = projector(None); + let inputs = representations(ProjectorScratch::CHUNK_SIZE + 1); + let first: Vec<_> = projector + .forward(inputs.rows()) + .try_collect() + .expect("should project the initial batch"); + let short: Vec<_> = projector + .forward(&inputs.rows()[..1]) + .try_collect() + .expect("should project a short batch"); + assert_eq!(short.len(), 1); + assert!(short[0].get().distance_squared_wide(first[0].get()) < 1e-10); + assert_eq!(projector.forward(&inputs.rows()[..0]).len(), 0); +} + +/// Reports non-finite input errors at exactly their rows. +/// +/// Every finite row still matches the all-finite reference pass, and a later finite batch still +/// projects correctly. +#[test] +fn forward_non_finite_rows() { + let alignment = Similarity::new(positive!(2.0), Rotation::IDENTITY, Vec2::new(3.0, -4.0)) + .expect("should have a valid similarity"); + let mut projector = projector(Some(alignment)); + let mut inputs = representations(ProjectorScratch::CHUNK_SIZE * 2 + 1); + let reference: Vec<_> = projector + .forward(inputs.rows()) + .try_collect() + .expect("should project finite inputs"); + let offenders = [1, ProjectorScratch::CHUNK_SIZE + 3]; + for row in offenders { + inputs.rows_mut()[row].as_array_mut()[0] = f32::NAN; + } + let actual: Vec<_> = projector.forward(inputs.rows()).collect(); + assert_eq!(actual.len(), inputs.rows().len()); + for (row, result) in actual.into_iter().enumerate() { + if offenders.contains(&row) { + assert_eq!(result, Err(ProjectionError::NonFiniteProjection)); + } else { + assert!( + result + .expect("should retain finite rows") + .get() + .distance_squared_wide(reference[row].get()) + < 1e-10 + ); + } + } + projector + .forward(&inputs.rows()[..1]) + .next() + .expect("should return one row") + .expect("should project a finite row after a mixed batch"); +} + +/// Reports non-finite alignment when the applied translation overflows. +#[test] +fn forward_alignment_overflow() { + let mut projector = projector(None); + let inputs = representations(1); + let point = projector + .forward(inputs.rows()) + .next() + .expect("should return one row") + .expect("should project a finite point") + .get(); + assert!(point.x().abs().max(point.y().abs()) > 1e-3); + projector.alignment = Some( + Similarity::new( + positive!(1e37), + Rotation::IDENTITY, + Vec2::new(f32::MAX.copysign(point.x()), f32::MAX.copysign(point.y())), + ) + .expect("should accept a finite similarity with a normal reciprocal"), + ); + assert_eq!( + projector.forward(inputs.rows()).next(), + Some(Err(ProjectionError::NonFiniteAlignment)) + ); +} + +/// Enforces fitted-world bounds even when the world is a point. +/// +/// A world shrunk to a single point admits only positions at that exact point and reports every +/// other position as out of bounds, carrying its rejected global coordinate. +#[test] +fn project_world_bounds() { + let mut projector = projector(None); + let inputs = representations(3); + let positions: Vec<_> = projector + .forward(inputs.rows()) + .try_collect() + .expect("should project finite points"); + let point = positions[0].get(); + projector.world = Bounds2::new(point, point).expect("should accept a point-sized world"); + let actual: Vec<_> = projector.project(inputs.rows()).collect(); + assert_eq!(actual.len(), 3); + assert_eq!(actual[0], Ok(positions[0])); + for (&position, result) in positions.iter().zip(actual) { + if position.get() == point { + assert_eq!(result, Ok(position)); + } else { + assert_eq!( + result, + Err(ProjectionError::OutOfBounds { + global: position.get() + }) + ); + } + } + assert!(positions.iter().any(|position| position.get() != point)); +} + +/// Rejects empty or mismatched roundtrip columns and reports both supplied lengths. +#[test] +fn roundtrip_column_lengths() { + let mut projector = projector(None); + for (representations_len, coordinates_len) in [(0, 0), (1, 0), (0, 1), (2, 1), (1, 2)] { + let inputs = representations(representations_len); + let coordinates = vec![Vec2::ZERO; coordinates_len]; + let report = projector + .try_roundtrip_sample_impl( + IdSlice::from_raw(inputs.rows()), + IdSlice::from_raw(&coordinates), + ) + .expect_err("should refuse empty or mismatched columns"); + assert_matches!(report.current_context(), ProjectorError::InvalidSampleCorpus { representations, coordinates } + if *representations == representations_len && *coordinates == coordinates_len); + } +} + +/// Reports the offending row when a published coordinate is non-finite. +#[test] +fn roundtrip_non_finite_coordinate() { + let mut projector = projector(None); + let inputs = representations(2); + let coordinates = [Vec2::ZERO, Vec2::new(f32::NAN, 0.0)]; + let report = projector + .try_roundtrip_sample_impl( + IdSlice::from_raw(inputs.rows()), + IdSlice::from_raw(&coordinates), + ) + .expect_err("should refuse a non-finite published point"); + assert_matches!(report.current_context(), ProjectorError::NonFiniteCoordinates { row } if *row == NodeRowId::new(1)); +} + +/// Accepts coordinates projected under the same alignment. +#[test] +fn roundtrip_aligned_sample() { + let mut projector = projector(Some( + Similarity::new(positive!(2.0), Rotation::IDENTITY, Vec2::new(3.0, -4.0)) + .expect("should have a valid similarity"), + )); + let inputs = representations(ProjectorScratch::CHUNK_SIZE * 2 + 1); + let coordinates: Vec<_> = projector + .forward(inputs.rows()) + .map(|result| result.map(Position::get)) + .try_collect() + .expect("should project sample coordinates"); + projector + .try_roundtrip_sample_impl( + IdSlice::from_raw(inputs.rows()), + IdSlice::from_raw(&coordinates), + ) + .expect("should reproduce aligned sample coordinates"); +} + +/// Rejects coordinates at or beyond the roundtrip tolerance. +/// +/// The reported error is at least as large as the tolerance. +#[test] +fn roundtrip_tolerance_exceeded() { + let mut projector = projector(None); + let inputs = representations(1); + let point = projector + .forward(inputs.rows()) + .next() + .expect("should return one row") + .expect("should project a finite point") + .get(); + let coordinates = [point + Vec2::new(0.01, 0.0)]; + let report = projector + .try_roundtrip_sample_impl( + IdSlice::from_raw(inputs.rows()), + IdSlice::from_raw(&coordinates), + ) + .expect_err("should refuse a changed published point"); + assert_matches!(report.current_context(), ProjectorError::RoundtripSampleToleranceExceeded { tolerance, error } + if error >= tolerance); +} + +/// Preserves the non-finite projection cause beneath the roundtrip forward-error context. +#[test] +fn roundtrip_non_finite_projection() { + let mut projector = projector(None); + let mut inputs = representations(1); + inputs.rows_mut()[0].as_array_mut()[0] = f32::INFINITY; + let report = projector + .try_roundtrip_sample_impl( + IdSlice::from_raw(inputs.rows()), + IdSlice::from_raw(&[Vec2::ZERO]), + ) + .expect_err("should refuse a non-finite sample projection"); + assert_matches!( + report.current_context(), + ProjectorError::RoundtripSampleForward + ); + assert_eq!( + report.downcast_ref::(), + Some(&ProjectionError::NonFiniteProjection) + ); +} diff --git a/libs/@local/graph/atlas/src/serve/delta/task/mod.rs b/libs/@local/graph/atlas/src/serve/delta/task/mod.rs new file mode 100644 index 00000000000..975db91a257 --- /dev/null +++ b/libs/@local/graph/atlas/src/serve/delta/task/mod.rs @@ -0,0 +1,311 @@ +//! Owned feed execution and immutable request publications. +//! +//! [`DeltaTask`] drives feed reads, placement and publication together. Change notifications are +//! wake-ups rather than a publication queue. Changes accumulated before an exchange appear in one +//! publication. [`DeltaReader`] captures the last exchanged value independently of the +//! runner's lifetime. + +use alloc::sync::Arc; +use core::{error::Error, fmt, future::Future}; + +use arc_swap::ArcSwap; +use error_stack::{Report, ReportSink, ResultExt as _}; +use futures::FutureExt as _; +use hash_graph_postgres_store::store::PostgresStorePool; +use hash_graph_temporal_versioning::{Timestamp, TransactionTime}; +use tokio::{sync::mpsc, task::JoinSet}; + +use super::{ + Delta, + epoch::Epoch, + feed::{self, DeltaFeedTask, DeltaFeedTaskOptions, Publication}, + placement::{ + Completed, DeltaPlacementTask, DeltaPlacementTaskOptions, EmbeddingWorkflow, Initial, + PendingEntry, + }, + projector::DeltaProjector, +}; +use crate::device::PhysicalDevice; + +#[cfg(test)] +mod tests; + +/// Polling and placement limits for one generation's feed. +#[derive(Copy, Clone)] +pub(crate) struct DeltaTaskOptions { + /// Event polling and replay settings. + pub feed: DeltaFeedTaskOptions, + /// Embedding lookup, projection and admission settings. + pub placement: DeltaPlacementTaskOptions, +} + +/// A failure opening or running a generation's feed. +#[derive(Debug)] +pub(crate) enum DeltaTaskError { + /// Opening the generation's projector failed. + Projector, + /// Initializing or running event replay failed. + Feed, + /// Initializing or running embedding placement failed. + Placement, + /// Exchanging immutable publications failed. + Publication, + /// Joining a spawned feed, placement or publication task failed. + Panic, +} + +impl fmt::Display for DeltaTaskError { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Projector => fmt.write_str("could not open the generation's projector"), + Self::Feed => fmt.write_str("the entity feed failed"), + Self::Placement => fmt.write_str("the placement task failed"), + Self::Publication => fmt.write_str("the publication task failed"), + Self::Panic => fmt.write_str("the task panicked"), + } + } +} + +impl Error for DeltaTaskError {} + +/// A request-side handle to the latest coherent publication. +/// +/// The handle exposes the supplied delta until replay exchanges a replacement and retains the last +/// exchanged publication after shutdown. Coherence means one immutable delta revision, not that +/// replay is healthy or caught up with the store. +#[derive(Clone)] +pub(crate) struct DeltaReader { + current: Arc>, +} + +impl DeltaReader { + /// Captures one publication for all of a request's component lookups. + /// + /// A later exchange changes what subsequent loads capture without changing the returned + /// [`Epoch`]. + pub(crate) fn load(&self) -> Epoch { + Epoch::from(self.current.load()) + } +} + +impl From for DeltaReader { + fn from(delta: Delta) -> Self { + Self { + current: Arc::new(ArcSwap::from_pointee(delta)), + } + } +} + +/// A generation's placement worker and its request/completion channels. +struct Placement { + task: DeltaPlacementTask, + projector: DeltaProjector, + requests: mpsc::Receiver>, + completed: mpsc::Sender>, +} + +/// One generation's feed, placement worker and publication loop. +/// +/// Opening starts no background tasks. The owner polls [`Self::run`] and observes its result. +/// The reader can outlive the runner without keeping its database or placement work active. +pub(crate) struct DeltaTask { + feed: DeltaFeedTask, + publication: Publication, + placement: Option, + + current: Arc>, + previous: Arc, +} + +impl DeltaTask { + /// Opens temporal replay for a generation and returns its request reader. + /// + /// A generation without temporal axes returns a static reader without opening the projector or + /// validating task options. For a temporal generation, projector opening reads its checkpoint + /// and validates its roundtrip sample synchronously. Without a checkpoint, the feed still + /// processes withdrawals and metadata while new node placements remain pending. + /// + /// # Errors + /// + /// Returns [`DeltaTaskError`] when projector opening or task construction fails. Projector + /// opening precedes construction. With a projector, placement validation precedes feed + /// validation. Without one, only feed validation runs. + /// + /// # Panics + /// + /// Panics if initialization reaches placement-channel construction with + /// [`DeltaPlacementTaskOptions::max_pending`] greater than + /// [`tokio::sync::Semaphore::MAX_PERMITS`]. Generations without temporal axes or a projector do + /// not construct these channels. + /// + /// With temporal axes, feed construction also has [`DeltaFeedTask::new`]'s + /// unit-revision-increment panic condition. + pub(crate) fn open( + delta: Delta, + pool: Arc, + options: DeltaTaskOptions, + device: PhysicalDevice, + workflow: Option>, + ) -> Result<(DeltaReader, Option), Report> { + let generation = delta.world.generation(); + let Some(axes) = generation.repository().metadata.snapshot.axes else { + return Ok((DeltaReader::from(delta), None)); + }; + + let projector = + DeltaProjector::open(generation, device).change_context(DeltaTaskError::Projector)?; + + let task = Self::new( + delta, + pool, + axes.transaction_time, + projector, + workflow, + options, + )?; + + Ok((task.reader(), Some(task))) + } + + /// Wires the feed, an optional placement worker and their publication loop together. + /// + /// Without a `projector`, placement channels are absent and new node placements remain + /// pending on the feed. + /// + /// # Errors + /// + /// Returns [`DeltaTaskError`] when task initialization fails. With a `projector`, placement + /// validation precedes feed validation. Without one, only feed validation runs. + /// + /// # Panics + /// + /// With a projector and a nonzero placement polling interval, panics when + /// [`DeltaPlacementTaskOptions::max_pending`] exceeds [`tokio::sync::Semaphore::MAX_PERMITS`]. + /// + /// Feed construction also has [`DeltaFeedTask::new`]'s unit-revision-increment panic condition. + fn new( + delta: Delta, + pool: Arc, + watermark: Timestamp, + projector: Option, + workflow: Option>, + options: DeltaTaskOptions, + ) -> Result> { + let (placement, channels) = if let Some(projector) = projector { + let capacity = options.placement.max_pending.get(); + let task = DeltaPlacementTask::new(Arc::clone(&pool), options.placement, workflow) + .change_context(DeltaTaskError::Placement)?; + + let (request_tx, request_rx) = mpsc::channel(capacity); + let (complete_tx, complete_rx) = mpsc::channel(capacity); + + ( + Some(Placement { + task, + projector, + requests: request_rx, + completed: complete_tx, + }), + Some(feed::Placement { + requests: request_tx, + completed: complete_rx, + }), + ) + } else { + (None, None) + }; + + let current = Arc::new(ArcSwap::from_pointee(delta.clone())); + let previous = current.load_full(); + + let (feed, publication) = + DeltaFeedTask::new(delta, pool, options.feed, watermark, channels) + .change_context(DeltaTaskError::Feed)?; + + Ok(Self { + feed, + publication, + placement, + current, + previous, + }) + } + + /// Returns a reader sharing this task's publication point. + pub(crate) fn reader(&self) -> DeltaReader { + DeltaReader { + current: Arc::clone(&self.current), + } + } + + /// Runs the task chain until its channels close, joining every child before returning. + /// + /// The generation owner supplies `shutdown`, such as a cancellation token's owned future. The + /// selection order prioritizes shutdown over the next publication: it preserves the value + /// already visible through [`DeltaReader`] but does not force unpublished working changes into + /// that reader. + /// Closing the publication exchange propagates closure through the feed and placement + /// channels. An operation already selected by either task reaches its next channel check before + /// exiting. This method joins every child. + /// + /// Captured epochs and readers remain valid after shutdown. Dropping this future before its + /// first poll starts no child. After the first poll, dropping it drops the Tokio [`JoinSet`] + /// and aborts its asynchronous children without joining them, but cannot cancel projection + /// work already submitted to Rayon, which runs to completion and drops the result its + /// receiver no longer accepts. + /// + /// # Errors + /// + /// Collects feed, publication and placement failures under the corresponding + /// [`DeltaTaskError`] variants, including failures concurrent with shutdown. A child join + /// failure reports [`DeltaTaskError::Panic`]. Requested shutdown does not itself produce an + /// error. + /// + /// # Panics + /// + /// Panics when polled outside a Tokio runtime while `shutdown` remains pending. + #[tracing::instrument(skip_all, err(Debug))] + pub(crate) async fn run( + self, + shutdown: impl Future + Send + 'static, + ) -> Result<(), Report<[DeltaTaskError]>> { + let mut shutdown = Box::pin(shutdown); + if shutdown.as_mut().now_or_never().is_some() { + return Ok(()); + } + + let Self { + feed, + publication, + placement, + current, + previous, + } = self; + + let mut join = JoinSet::new(); + + join.spawn(async move { feed.run().await.change_context(DeltaTaskError::Feed) }); + join.spawn(async move { + publication + .run(current, previous, shutdown) + .await + .change_context(DeltaTaskError::Publication) + }); + + if let Some(placement) = placement { + join.spawn(async move { + placement + .task + .run(placement.projector, placement.requests, placement.completed) + .await + .change_context(DeltaTaskError::Placement) + }); + } + + let mut sink = ReportSink::new(); + while let Some(result) = join.join_next().await { + sink.attempt(result.change_context(DeltaTaskError::Panic).flatten()); + } + sink.finish() + } +} diff --git a/libs/@local/graph/atlas/src/serve/delta/task/tests.rs b/libs/@local/graph/atlas/src/serve/delta/task/tests.rs new file mode 100644 index 00000000000..39909f7f318 --- /dev/null +++ b/libs/@local/graph/atlas/src/serve/delta/task/tests.rs @@ -0,0 +1,272 @@ +use alloc::sync::{Arc, Weak}; +use core::{assert_matches, future, time::Duration}; + +use arc_swap::Guard; +use hash_graph_postgres_store::store::{ + DatabaseConnectionInfo, DatabasePoolConfig, DatabaseType, PostgresStorePool, + PostgresStoreSettings, +}; +use hash_graph_temporal_versioning::Timestamp; +use hashql_core::id::Id as _; +use rand::{SeedableRng as _, rngs::StdRng}; +use tokio::sync::oneshot; +use tokio_postgres::NoTls; + +use super::{DeltaTask, DeltaTaskError, DeltaTaskOptions}; +use crate::{ + device::Device, + identity::NodeRowId, + math::nz, + serve::{ + delta::{ + Delta, DeltaRevision, + epoch::Epoch, + feed::{DeltaFeedTask, DeltaFeedTaskOptions}, + placement::DeltaPlacementTaskOptions, + projector::projector, + }, + tests::fixture::{TamperFixture, secret}, + world::World, + }, +}; + +/// A delta task opened over a synthetic generation, retaining a weak pool handle for leak checks. +struct Fixture { + _generation: TamperFixture, + task: DeltaTask, + pool: Weak, +} + +/// Returns a fast-ticking task configuration with single-entry queues. +fn options() -> DeltaTaskOptions { + DeltaTaskOptions { + feed: DeltaFeedTaskOptions { + tick_rate: Duration::from_secs(5), + safety_lag: Duration::from_secs(60), + }, + placement: DeltaPlacementTaskOptions { + tick_rate: Duration::from_secs(5), + tries_workflow: 1, + tries_database: 1, + minimum_projection_interval: 1, + max_pending: nz!(1), + }, + } +} + +/// Builds a delta task over a synthetic world published under `name`. +/// +/// The task uses an unconnected pool and optionally enables a placement worker. +/// +/// # Panics +/// +/// Panics on failure during generation publication, world opening, store pool construction, or +/// task initialization. +async fn fixture(name: &str, placement: bool) -> Fixture { + let generation = TamperFixture::publish(name); + let world = Arc::new( + World::open(generation.generation().clone(), &secret()).expect("should open the world"), + ); + let delta = Delta::new(world, StdRng::seed_from_u64(17)).expect("should initialize the delta"); + // Failure and shutdown cases use an unconnected pool without a database server. + let pool = Arc::new( + PostgresStorePool::new( + &DatabaseConnectionInfo::new( + DatabaseType::Postgres, + "task-test".to_owned(), + String::new(), + "/no-task-test-postgres".to_owned(), + 5432, + "task-test".to_owned(), + ), + &DatabasePoolConfig { + max_connections: nz!(1), + }, + NoTls, + PostgresStoreSettings::default(), + ) + .await + .expect("should construct an unconnected pool"), + ); + let weak = Arc::downgrade(&pool); + let task = DeltaTask::new( + delta, + pool, + Timestamp::UNIX_EPOCH, + placement.then(|| projector(None)), + None, + options(), + ) + .expect("should construct the runner"); + Fixture { + _generation: generation, + task, + pool: weak, + } +} + +/// Keeps the reader's supplied revision and world valid after its runner drops. +#[tokio::test] +async fn reader_initial() { + let fixture = fixture("delta-task-reader-initial", false).await; + let reader = fixture.task.reader(); + let world = Arc::clone(&fixture.task.previous.world); + let epoch = reader.load(); + assert_eq!(epoch.revision(), DeltaRevision::MIN); + assert!(world.layout.position(&epoch, NodeRowId::MIN).is_some()); + drop(fixture.task); + assert!(fixture.pool.upgrade().is_none()); + assert_eq!(reader.load().reference(), epoch.reference()); + assert!( + world + .layout + .position(&reader.load(), NodeRowId::MIN) + .is_some() + ); +} + +/// Keeps a loaded epoch on its original revision after the reader advances. +#[tokio::test] +async fn reader_captured() { + let fixture = fixture("delta-task-reader-captured", false).await; + let reader = fixture.task.reader(); + let captured = reader.load(); + let mut next = fixture.task.previous.as_ref().clone(); + let world = Arc::clone(&next.world); + let row = NodeRowId::MIN; + let identity = world + .layout + .index + .key_of(&captured, row) + .expect("should resolve the base identity"); + next.revision.increment_by(1); + assert!(next.withdraw(identity)); + fixture.task.current.store(Arc::new(next)); + assert!(world.layout.position(&captured, row).is_some()); + assert_eq!(world.layout.position(&reader.load(), row), None); +} + +/// Releases both workers and their completion receiver when an unpolled runner drops. +#[tokio::test] +async fn run_cancelled() { + let fixture = fixture("delta-task-run-cancelled", true).await; + let reader = fixture.task.reader(); + let completed = fixture + .task + .placement + .as_ref() + .expect("should have a placement worker") + .completed + .clone(); + assert!(!completed.is_closed()); + drop(fixture.task.run(future::pending())); + assert!(completed.is_closed()); + assert!(fixture.pool.upgrade().is_none()); + assert!(reader.load().contains_node(NodeRowId::MIN)); +} + +/// Closes the chain and joins worker siblings after publication failure. +#[tokio::test] +async fn run_publication_closed() { + let mut fixture = fixture("delta-task-publication-closed", true).await; + let reader = fixture.task.reader(); + let previous = fixture.task.previous.as_ref().clone(); + let captured = Epoch::from(Guard::from_inner(Arc::new(previous.clone()))); + let (feed, publication) = DeltaFeedTask::new( + previous, + fixture.pool.upgrade().expect("should retain the pool"), + options().feed, + Timestamp::UNIX_EPOCH, + None, + ) + .expect("should construct the replacement publication channel"); + drop(feed); + fixture.task.publication = publication; + let completed = fixture + .task + .placement + .as_ref() + .expect("should have a placement worker") + .completed + .clone(); + let error = tokio::spawn(fixture.task.run(future::pending())) + .await + .expect("should join the runner") + .expect_err("should report publication failure"); + assert_matches!( + error.current_contexts().collect::>().as_slice(), + [DeltaTaskError::Publication] + ); + assert!(completed.is_closed()); + assert!(fixture.pool.upgrade().is_none()); + assert_eq!(reader.load().revision(), captured.revision()); +} + +/// Starts no workers for pre-requested shutdown and keeps the reader usable. +#[tokio::test] +async fn run_shutdown_ready() { + let fixture = fixture("delta-task-shutdown-ready", true).await; + let reader = fixture.task.reader(); + let completed = fixture + .task + .placement + .as_ref() + .expect("should have a placement worker") + .completed + .clone(); + fixture + .task + .run(future::ready(())) + .await + .expect("should accept an immediate shutdown"); + assert!(completed.is_closed()); + assert!(fixture.pool.upgrade().is_none()); + assert!(reader.load().contains_node(NodeRowId::MIN)); +} + +/// Joins workers on owner shutdown without invalidating an earlier capture. +#[tokio::test] +async fn run_shutdown_active() { + let fixture = fixture("delta-task-shutdown-active", true).await; + let reader = fixture.task.reader(); + let captured = reader.load(); + let completed = fixture + .task + .placement + .as_ref() + .expect("should have a placement worker") + .completed + .clone(); + let (stop, shutdown) = oneshot::channel(); + let running = tokio::spawn(fixture.task.run(async move { + shutdown.await.expect("should receive the shutdown request"); + })); + tokio::task::yield_now().await; + assert!(!running.is_finished()); + stop.send(()).expect("should request shutdown"); + tokio::time::timeout(Duration::from_secs(5), running) + .await + .expect("should drain without waiting for a database poll") + .expect("should join the runner") + .expect("should stop normally"); + assert!(completed.is_closed()); + assert!(fixture.pool.upgrade().is_none()); + assert_eq!(reader.load().revision(), captured.revision()); + assert!(captured.contains_node(NodeRowId::MIN)); +} + +/// Opens a usable static reader without a task when temporal axes are absent. +#[tokio::test] +async fn open_without_axes() { + let fixture = fixture("delta-task-open-without-axes", false).await; + let (reader, task) = DeltaTask::open( + fixture.task.previous.as_ref().clone(), + fixture.pool.upgrade().expect("should retain the pool"), + options(), + Device::Cpu.pin(0).resolve(), + None, + ) + .expect("should skip a generation without temporal axes"); + assert!(task.is_none()); + assert!(reader.load().contains_node(NodeRowId::MIN)); +} diff --git a/libs/@local/graph/atlas/src/serve/mod.rs b/libs/@local/graph/atlas/src/serve/mod.rs index d4dc6a4552a..0bdfeeb846f 100644 --- a/libs/@local/graph/atlas/src/serve/mod.rs +++ b/libs/@local/graph/atlas/src/serve/mod.rs @@ -25,4 +25,5 @@ pub(crate) mod secret; #[cfg(test)] pub(crate) mod tests; pub(crate) mod visibility; +mod walk; mod world; diff --git a/libs/@local/graph/atlas/src/serve/walk/mod.rs b/libs/@local/graph/atlas/src/serve/walk/mod.rs new file mode 100644 index 00000000000..3eabda5775c --- /dev/null +++ b/libs/@local/graph/atlas/src/serve/walk/mod.rs @@ -0,0 +1,114 @@ +//! Delivery from a captured schedule with the rows the supplied epoch refuses removed. +//! +//! A [`Walk`] pairs a [`DeliverySchedule`] with the [`NodeIndex`] its rows address. Each query +//! subtracts the scheduled rows the supplied epoch's identity state does not permit, keeping every +//! bucket run contiguous and its count exact. A refused row is one the epoch has withdrawn or, for +//! an added row, one outside that epoch's allocated domain. + +use super::{ + delta::{epoch::Epoch, overlay::NaiveIdentityProvider}, + schedule::{DeliveredNodes, DeliverySchedule}, + world::NodeIndex, +}; +use crate::{ + identity::NodeRowId, + morton::{MortonCell, Zoom}, +}; + +#[cfg(test)] +mod tests; + +/// Removes each row `withdrawn` accepts from `rows` in place. +/// +/// The caller passes `runs` that partition `rows`, their lengths summing to the row count. Each +/// run's surviving rows remain contiguous, and its count updates. +/// +/// # Panics +/// +/// Panics if a computed run range extends past the end of `rows`. +fn subtract( + DeliveredNodes { rows, runs, .. }: &mut DeliveredNodes, + mut withdrawn: impl FnMut(NodeRowId) -> bool, +) { + let mut kept = 0; + let mut start = 0; + + for run in &mut *runs { + let end = start + *run; + let mut survivors = 0; + + for index in start..end { + let node = rows[index]; + + if !withdrawn(node) { + rows[kept] = node; + kept += 1; + survivors += 1; + } + } + + *run = survivors; + start = end; + } + + debug_assert_eq!( + start, + rows.len(), + "the runs should partition the delivered rows" + ); + rows.truncate(kept); +} + +/// A captured delivery schedule read through the identity state of a supplied epoch. +/// +/// Scoped schedules contain only their admitted placements. Corpus schedules retain recorded base +/// rows through withdrawals. Both modes apply the epoch check without another visibility mask. +#[derive(Copy, Clone)] +pub(crate) struct Walk<'context> { + /// The captured delivery schedule. + pub schedule: DeliverySchedule<'context>, + + /// The node index that answers the withdrawal checks. + pub index: &'context NodeIndex, +} + +impl Walk<'_> { + /// Returns whether `epoch`'s identity state refuses `node`. + /// + /// A refused row is withdrawn at that epoch or, for an added row, outside its allocated + /// domain. + /// + /// # Panics + /// + /// Panics if `index` does not belong to `epoch`'s world. + fn node_withdrawn(&self, epoch: &Epoch, node: NodeRowId) -> bool { + !epoch + .nodes(self.index) + .bind(NaiveIdentityProvider::from_ref(&self.index.identity)) + .permits_row(node, None) + } + + /// Gathers the rows `cell` newly delivers at `zoom`, minus the rows `epoch` refuses. + /// + /// # Panics + /// + /// Panics beyond the schedule's deepest served zoom, or if `index` does not belong to `epoch`'s + /// world. + pub(crate) fn delta(&self, epoch: &Epoch, zoom: Zoom, cell: MortonCell) -> DeliveredNodes { + let mut delivered = self.schedule.delta(zoom, cell); + subtract(&mut delivered, |node| self.node_withdrawn(epoch, node)); + delivered + } + + /// Gathers `cell`'s cumulative rows through `zoom`, minus the rows `epoch` refuses. + /// + /// # Panics + /// + /// Panics beyond the schedule's deepest served zoom, or if `index` does not belong to `epoch`'s + /// world. + pub(crate) fn total(&self, epoch: &Epoch, zoom: Zoom, cell: MortonCell) -> DeliveredNodes { + let mut delivered = self.schedule.total(zoom, cell); + subtract(&mut delivered, |node| self.node_withdrawn(epoch, node)); + delivered + } +} diff --git a/libs/@local/graph/atlas/src/serve/walk/tests.rs b/libs/@local/graph/atlas/src/serve/walk/tests.rs new file mode 100644 index 00000000000..983e1c92a0b --- /dev/null +++ b/libs/@local/graph/atlas/src/serve/walk/tests.rs @@ -0,0 +1,129 @@ +//! Cases covering withdrawal subtraction: rows leave, runs stay contiguous and re-sum. + +use hashql_core::id::Id as _; + +use super::subtract; +use crate::{identity::NodeRowId, morton::Depth, serve::schedule::DeliveredNodes}; + +/// Builds a [`DeliveredNodes`] fixture over `rows` partitioned by `runs`. +/// +/// The rows convert to [`NodeRowId`]s, and the first bucket is the minimum depth. +fn delivered(rows: &[u64], runs: &[usize]) -> DeliveredNodes { + DeliveredNodes { + rows: rows.iter().copied().map(NodeRowId::new).collect(), + first_bucket: Depth::MIN, + runs: runs.to_vec(), + } +} + +/// Converts raw row numbers to [`NodeRowId`]s. +/// +/// The result compares against a [`DeliveredNodes`]'s `rows` after a `subtract` call. +fn nodes(rows: &[u64]) -> Vec { + rows.iter().copied().map(NodeRowId::new).collect() +} + +/// Asserts that `delivered`'s run lengths still sum to its row count. +/// +/// This is the invariant `subtract` must preserve. +/// +/// # Panics +/// +/// Panics when the run lengths do not sum to the number of delivered rows. +#[track_caller] +fn assert_partitioned(delivered: &DeliveredNodes) { + assert_eq!( + delivered.runs.iter().sum::(), + delivered.rows.len(), + "the runs should re-sum to the delivered count" + ); +} + +/// Each withdrawn row leaves the delivery and debits the run that owned it. +#[test] +fn subtract_run_decrement() { + let mut delivery = delivered(&[4, 5, 6, 20, 21], &[2, 2, 1]); + + subtract(&mut delivery, |node| { + node == NodeRowId::new(5) || node == NodeRowId::new(21) + }); + + assert_eq!(delivery.rows, nodes(&[4, 6, 20])); + assert_eq!( + delivery.runs, + [1, 2, 0], + "each withdrawal should debit its owning run" + ); + assert_partitioned(&delivery); +} + +/// Runs that were already empty keep their positional slots through the subtraction. +#[test] +fn subtract_zero_length_runs() { + let mut delivery = delivered(&[0, 4, 5, 6, 7, 8], &[1, 0, 3, 2]); + + subtract(&mut delivery, |node| { + node == NodeRowId::new(0) || node == NodeRowId::new(8) + }); + + assert_eq!(delivery.rows, nodes(&[4, 5, 6, 7])); + assert_eq!( + delivery.runs, + [0, 0, 3, 1], + "zero-length runs should keep their slots" + ); + assert_partitioned(&delivery); +} + +/// Withdrawing the first and last row of a run leaves the interior intact. +#[test] +fn subtract_run_edges() { + let mut delivery = delivered(&[3, 4, 5, 6, 7], &[5]); + + subtract(&mut delivery, |node| { + node == NodeRowId::new(3) || node == NodeRowId::new(7) + }); + + assert_eq!(delivery.rows, nodes(&[4, 5, 6])); + assert_eq!(delivery.runs, [3]); + assert_partitioned(&delivery); +} + +/// A fully withdrawn delivery keeps every run slot at zero length. +#[test] +fn subtract_all() { + let mut delivery = delivered(&[2, 3, 4, 9], &[3, 1]); + + subtract(&mut delivery, |_| true); + + assert!( + delivery.rows.is_empty(), + "no row should survive a full withdrawal" + ); + assert_eq!(delivery.runs, [0, 0], "the run slots should survive"); + assert_partitioned(&delivery); +} + +/// A predicate withdrawing nothing leaves rows and runs untouched. +#[test] +fn subtract_none() { + let mut delivery = delivered(&[1, 2, 3], &[1, 2]); + + subtract(&mut delivery, |_| false); + + assert_eq!(delivery.rows, nodes(&[1, 2, 3])); + assert_eq!(delivery.runs, [1, 2]); + assert_partitioned(&delivery); +} + +/// An empty delivery passes through with its run slots intact. +#[test] +fn subtract_empty() { + let mut delivery = delivered(&[], &[0, 0]); + + subtract(&mut delivery, |_| true); + + assert!(delivery.rows.is_empty(), "no row should appear"); + assert_eq!(delivery.runs, [0, 0]); + assert_partitioned(&delivery); +} From 78d17894fb21e2c8f29d387536263af81ae90ff0 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Mon, 14 Sep 2026 22:59:36 +0200 Subject: [PATCH 2/4] chore: replace reviewer name in atlas fit test fixtures no more Bilal, we are Mustermaxxing now. --- libs/@local/graph/atlas/src/salt/fit/tests.rs | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/libs/@local/graph/atlas/src/salt/fit/tests.rs b/libs/@local/graph/atlas/src/salt/fit/tests.rs index 4c5fb0d99ba..9130847f834 100644 --- a/libs/@local/graph/atlas/src/salt/fit/tests.rs +++ b/libs/@local/graph/atlas/src/salt/fit/tests.rs @@ -982,7 +982,7 @@ async fn supplied_verdicts_publish_verbatim() { r#"{"pair_verdicts":[],"schema":"atlas-reviewed-verdicts/1","#, r#""sources":{"cards.jsonl":"2a9934acae8bf210b6a3428e553b1bcc0e220a4de113940782cd573da1ea4f4b"},"#, r#""type_verdicts":[{"class":"proximal","relation":"hash:https://hash.ai/@h/types/entity-type/delivers/","#, - r#""reviewer":"Bilal Mahmoud","versioned_url":"https://hash.ai/@h/types/entity-type/delivers/v/3"}]}"#, + r#""reviewer":"Max Mustermann","versioned_url":"https://hash.ai/@h/types/entity-type/delivers/v/3"}]}"#, "\n", ); let supplied = SuppliedVerdicts::from_bytes(document.as_bytes()) @@ -1573,7 +1573,7 @@ fn proximal_link_verdicts() -> SuppliedVerdicts { let document = concat!( r#"{"pair_verdicts":[],"schema":"atlas-reviewed-verdicts/1","sources":{},"#, r#""type_verdicts":[{"class":"proximal","relation":"memory:employment-link","#, - r#""reviewer":"Bilal Mahmoud","versioned_url":"memory://2/v/1"}]}"#, + r#""reviewer":"Max Mustermann","versioned_url":"memory://2/v/1"}]}"#, "\n", ); SuppliedVerdicts::from_bytes(document.as_bytes()).expect("the fixture document admits") @@ -2675,11 +2675,11 @@ fn store_identity_verdicts_resolve_by_reviewed_version() { r#"{"pair_verdicts":[],"schema":"atlas-reviewed-verdicts/1","sources":{},"#, r#""type_verdicts":["#, r#"{"class":"overlay","relation":"hash:http://localhost:3000/@linktest/types/entity-type/acquaintance/","#, - r#""reviewer":"Bilal Mahmoud","versioned_url":"http://localhost:3000/@linktest/types/entity-type/acquaintance/v/1"},"#, + r#""reviewer":"Max Mustermann","versioned_url":"http://localhost:3000/@linktest/types/entity-type/acquaintance/v/1"},"#, r#"{"class":"coincident","relation":"hash:https://hash.ai/@h/types/entity-type/arrives-at/","#, - r#""reviewer":"Bilal Mahmoud","versioned_url":"https://hash.ai/@h/types/entity-type/arrives-at/v/1"},"#, + r#""reviewer":"Max Mustermann","versioned_url":"https://hash.ai/@h/types/entity-type/arrives-at/v/1"},"#, r#"{"class":"proximal","relation":"hash:https://hash.ai/@h/types/entity-type/located-at/","#, - r#""reviewer":"Bilal Mahmoud","versioned_url":"https://hash.ai/@h/types/entity-type/located-at/v/1"}"#, + r#""reviewer":"Max Mustermann","versioned_url":"https://hash.ai/@h/types/entity-type/located-at/v/1"}"#, r#"]}"#, "\n", ); @@ -2713,9 +2713,9 @@ fn plain_number_corpus_resolves_the_memory_scheme() { r#"{"pair_verdicts":[],"schema":"atlas-reviewed-verdicts/1","sources":{},"#, r#""type_verdicts":["#, r#"{"class":"proximal","relation":"hash:https://hash.ai/@h/types/entity-type/delivers/","#, - r#""reviewer":"Bilal Mahmoud","versioned_url":"https://hash.ai/@h/types/entity-type/delivers/v/3"},"#, + r#""reviewer":"Max Mustermann","versioned_url":"https://hash.ai/@h/types/entity-type/delivers/v/3"},"#, r#"{"class":"proximal","relation":"memory:employment-link","#, - r#""reviewer":"Bilal Mahmoud","versioned_url":"memory://2/v/1"}"#, + r#""reviewer":"Max Mustermann","versioned_url":"memory://2/v/1"}"#, r#"]}"#, "\n", ); From a50dd2c293f1fff14977fa62930a81aa1ff505cb Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Tue, 15 Sep 2026 10:13:07 +0200 Subject: [PATCH 3/4] chore: add blank lines for code clarity in atlas delta --- .../@local/graph/atlas/src/serve/delta/feed/pending/mod.rs | 7 +++++++ libs/@local/graph/atlas/src/serve/delta/mod.rs | 1 + 2 files changed, 8 insertions(+) diff --git a/libs/@local/graph/atlas/src/serve/delta/feed/pending/mod.rs b/libs/@local/graph/atlas/src/serve/delta/feed/pending/mod.rs index f6a703d4db9..b084ae328b4 100644 --- a/libs/@local/graph/atlas/src/serve/delta/feed/pending/mod.rs +++ b/libs/@local/graph/atlas/src/serve/delta/feed/pending/mod.rs @@ -329,6 +329,7 @@ impl Pending { let Stage::Capture(geometry) = update.stage else { continue; }; + let Some(Some(DisplayParts { label, icon, @@ -337,6 +338,7 @@ impl Pending { else { continue; }; + let Some((representative, registered)) = delta.register_ontology(representative, icon) else { tracing::warn!(entity = ?update.event.entity, "No ontology row remains for the display"); @@ -363,10 +365,12 @@ impl Pending { /// Panics under the revision conditions on [`Delta::update_node`] and [`Delta::update_edge`]. pub(super) fn apply(&mut self, delta: &mut Delta) -> bool { let mut changed = false; + self.updates.retain(|&entity, update| { let Stage::Ready { geometry, legend } = &update.stage else { return true; }; + let outcome = match *geometry { Geometry::Node(position) => delta.update_node(entity, legend.clone(), position), Geometry::Edge(Some([source, target])) => { @@ -375,6 +379,7 @@ impl Pending { else { return true; }; + delta.update_edge(entity, legend.clone(), Some([source, target])) } Geometry::Edge(None) => return true, @@ -384,9 +389,11 @@ impl Pending { tracing::warn!(?entity, "No entity row remains for the update"); return true; }; + changed |= applied; false }); + changed } } diff --git a/libs/@local/graph/atlas/src/serve/delta/mod.rs b/libs/@local/graph/atlas/src/serve/delta/mod.rs index fe884bdec6a..c043f162c57 100644 --- a/libs/@local/graph/atlas/src/serve/delta/mod.rs +++ b/libs/@local/graph/atlas/src/serve/delta/mod.rs @@ -217,6 +217,7 @@ impl Delta { self.revision, entity, ); + changed |= self.edge.withdraw( NaiveIdentityProvider::from_ref(&self.world.topology.identity), self.revision, From 70e7ed62fe4bbfb07dd6bc691accf4d08eea892a Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:31:53 +0200 Subject: [PATCH 4/4] refactor: split pending endpoint and node application --- .../atlas/src/serve/delta/feed/pending/mod.rs | 42 ++++++++++------ .../src/serve/delta/feed/pending/tests.rs | 49 ++++++++++++++++++- 2 files changed, 76 insertions(+), 15 deletions(-) diff --git a/libs/@local/graph/atlas/src/serve/delta/feed/pending/mod.rs b/libs/@local/graph/atlas/src/serve/delta/feed/pending/mod.rs index b084ae328b4..b495eac290d 100644 --- a/libs/@local/graph/atlas/src/serve/delta/feed/pending/mod.rs +++ b/libs/@local/graph/atlas/src/serve/delta/feed/pending/mod.rs @@ -367,26 +367,40 @@ impl Pending { let mut changed = false; self.updates.retain(|&entity, update| { - let Stage::Ready { geometry, legend } = &update.stage else { + let Stage::Ready { + geometry: Geometry::Node(position), + legend, + } = &update.stage + else { return true; }; - let outcome = match *geometry { - Geometry::Node(position) => delta.update_node(entity, legend.clone(), position), - Geometry::Edge(Some([source, target])) => { - let (Some(source), Some(target)) = - (delta.node_row(source), delta.node_row(target)) - else { - return true; - }; + let Some(applied) = delta.update_node(entity, legend.clone(), *position) else { + tracing::warn!(?entity, "no node row remains for the update"); + return true; + }; - delta.update_edge(entity, legend.clone(), Some([source, target])) - } - Geometry::Edge(None) => return true, + changed |= applied; + false + }); + + self.updates.retain(|&entity, update| { + let Stage::Ready { + geometry: Geometry::Edge(Some([source, target])), + legend, + } = &update.stage + else { + return true; + }; + + let (Some(source), Some(target)) = (delta.node_row(*source), delta.node_row(*target)) + else { + return true; }; - let Some(applied) = outcome else { - tracing::warn!(?entity, "No entity row remains for the update"); + let Some(applied) = delta.update_edge(entity, legend.clone(), Some([source, target])) + else { + tracing::warn!(?entity, "no edge row remains for the update"); return true; }; diff --git a/libs/@local/graph/atlas/src/serve/delta/feed/pending/tests.rs b/libs/@local/graph/atlas/src/serve/delta/feed/pending/tests.rs index fd3082b1159..c7edd155122 100644 --- a/libs/@local/graph/atlas/src/serve/delta/feed/pending/tests.rs +++ b/libs/@local/graph/atlas/src/serve/delta/feed/pending/tests.rs @@ -1,5 +1,5 @@ use alloc::sync::Arc; -use core::assert_matches; +use core::{array, assert_matches}; use arc_swap::Guard; use error_stack::Report; @@ -518,6 +518,53 @@ fn receive_wire_coordinates() { ); } +#[test] +fn apply_ready_endpoints() { + let mut fixture = fixture("pending-apply-ready-endpoints"); + let mut pending = Pending::default(); + for seed in 100..103 { + pending.observe(update(seed, 1, false)); + } + + // assign the first visited entry to the edge, ahead of both endpoint updates. + let mut identities = pending.updates.keys().copied(); + let [edge, source, target] = + array::from_fn(|_| identities.next().expect("should retain three updates")); + for (&entity, update) in &mut pending.updates { + update.stage = Stage::Ready { + geometry: if entity == edge { + Geometry::Edge(Some([source, target])) + } else { + Geometry::Node(Vec2::ZERO) + }, + legend: OwnedLegend::new(OntologyRowId::MIN, Label::new("ready")), + }; + } + + assert!(pending.apply(&mut fixture.delta)); + assert!(pending.updates.is_empty()); + let endpoints = [source, target].map(|entity| { + fixture + .delta + .node_row(entity) + .expect("should allocate the endpoint") + }); + let identities = DeltaIdentityProvider::from_parts( + &fixture.delta.edge, + NaiveIdentityProvider::from_ref(&fixture.delta.world.topology.identity), + ); + let row = identities.row_of(edge).expect("should allocate the edge"); + assert_eq!( + fixture + .delta + .world + .topology + .endpoints(&epoch(&fixture.delta), row), + Some(endpoints) + ); + assert!(!pending.apply(&mut fixture.delta)); +} + /// Requires endpoint rows for link updates and never requests an embedding. #[test] fn edge_missing_endpoint() {