From 4fb4b9ecf69faba9191e419a2f23e12e6254aef3 Mon Sep 17 00:00:00 2001 From: Farhan Syah Date: Mon, 12 Oct 2026 11:13:13 +0800 Subject: [PATCH 01/11] fix(bridge): count dispatch capacity refusals and keep them retryable Every dispatcher capacity refusal is now built in one place that increments nodedb_dispatch_capacity_busy_total{scope}. The counter is shared with the node's metrics, so /metrics reads it without taking the dispatcher lock. A capacity refusal keeps its type through the system dispatch door and every protocol surface, so clients see the retryable SERVER_OVERLOAD class (SQLSTATE 57P03) instead of an internal error. Calvin catch-up and deferred dispatch hold a refused transaction and re-send it once capacity frees. --- docs/architecture.md | 2 +- .../src/bridge/dispatch/capacity_metrics.rs | 143 ++++++++++++ nodedb/src/bridge/dispatch/dispatcher.rs | 15 ++ nodedb/src/bridge/dispatch/enqueue.rs | 113 +++++++++- nodedb/src/bridge/dispatch/mod.rs | 4 +- nodedb/src/bridge/dispatch/test_requests.rs | 4 +- .../calvin/scheduler/driver/core/catch_up.rs | 208 +++++++++++++++++- .../calvin/scheduler/driver/core/deferred.rs | 4 +- .../gateway/error_map/capacity_refusal.rs | 80 +++++++ nodedb/src/control/gateway/error_map/mod.rs | 2 + .../error_map/system_dispatch_refusal.rs | 53 +++++ nodedb/src/control/metrics/system/fields.rs | 8 + nodedb/src/control/metrics/system/render.rs | 1 + nodedb/src/control/server/http/auth.rs | 120 +++++++++- .../shared/ddl/sync_dispatch/dispatch.rs | 7 +- .../control/server/shuffle/aggregator_hook.rs | 8 +- .../control/server/shuffle/consumer_hook.rs | 8 +- nodedb/src/control/state/init.rs | 6 +- nodedb/src/control/state/init_prod/open.rs | 3 +- .../cases/event_scheduler_concurrency.rs | 57 +++-- 20 files changed, 784 insertions(+), 62 deletions(-) create mode 100644 nodedb/src/bridge/dispatch/capacity_metrics.rs create mode 100644 nodedb/src/control/gateway/error_map/capacity_refusal.rs diff --git a/docs/architecture.md b/docs/architecture.md index 2ffd45ea0..cba0f607d 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -200,7 +200,7 @@ Admission ordering for requests: 1. Identity & database access (established during authentication) 2. **Tenant** rate/concurrency check → `TENANT_QUOTA_EXCEEDED` 3. **Database** rate/concurrency check → `DATABASE_QUOTA_EXCEEDED` -4. **Global** pressure check → `SERVER_OVERLOAD` +4. **Dispatch** capacity check (tenant in-flight share, database fair share per core, core queue) → `SERVER_OVERLOAD` (code 1403, SQLSTATE `57P03`). Nothing was enqueued, so the client retries with backoff. `nodedb_dispatch_capacity_busy_total{scope}` counts each refusal. Memory reservation flips the order (largest scope first to fail fast): global → database → tenant → engine. Failure at any layer aborts before the next is consulted. diff --git a/nodedb/src/bridge/dispatch/capacity_metrics.rs b/nodedb/src/bridge/dispatch/capacity_metrics.rs new file mode 100644 index 000000000..f4cb334ad --- /dev/null +++ b/nodedb/src/bridge/dispatch/capacity_metrics.rs @@ -0,0 +1,143 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Process-wide count of dispatches the bridge dispatcher refused at a +//! capacity limit, by scope. +//! +//! The dispatcher records each refusal at the one place it builds +//! [`crate::Error::DispatchCapacity`]. The label is the scope only, never a +//! tenant or database id, so the series count stays fixed. Every field is +//! atomic, so `/metrics` reads it without taking the dispatcher lock. + +use std::sync::atomic::{AtomicU64, Ordering}; + +use crate::DispatchCapacityScope; + +/// The `scope` label of each counter, by index. +const SCOPE_LABELS: [&str; 3] = ["tenant_inflight", "database_suspended", "queue_full"]; + +/// The counter index of `scope`. +fn scope_index(scope: &DispatchCapacityScope) -> usize { + match scope { + DispatchCapacityScope::TenantInflight { .. } => 0, + DispatchCapacityScope::DatabaseSuspended { .. } => 1, + DispatchCapacityScope::QueueFull { .. } => 2, + } +} + +/// Capacity refusals since process start, one counter per scope. +#[derive(Debug, Default)] +pub struct DispatchCapacityMetrics { + refusals: [AtomicU64; SCOPE_LABELS.len()], +} + +impl DispatchCapacityMetrics { + pub fn new() -> Self { + Self::default() + } + + /// Count one refusal at `scope`. + pub(super) fn record(&self, scope: &DispatchCapacityScope) { + self.refusals[scope_index(scope)].fetch_add(1, Ordering::Relaxed); + } + + /// Refusals counted at the scope of `scope`, whatever its ids. + pub fn refusals(&self, scope: &DispatchCapacityScope) -> u64 { + self.refusals[scope_index(scope)].load(Ordering::Relaxed) + } + + /// Append the `nodedb_dispatch_capacity_busy_total` counter. + pub fn write_prometheus(&self, out: &mut String) { + use std::fmt::Write as _; + + let _ = out.write_str( + "# HELP nodedb_dispatch_capacity_busy_total Dispatches refused at a capacity \ + limit. Nothing was enqueued, and the client retries.\n\ + # TYPE nodedb_dispatch_capacity_busy_total counter\n", + ); + for (label, count) in SCOPE_LABELS.iter().zip(&self.refusals) { + let _ = writeln!( + out, + r#"nodedb_dispatch_capacity_busy_total{{scope="{label}"}} {}"#, + count.load(Ordering::Relaxed) + ); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::{DatabaseId, TenantId}; + + fn tenant_scope() -> DispatchCapacityScope { + DispatchCapacityScope::TenantInflight { + tenant_id: TenantId::new(7), + inflight: 4, + cap: 4, + } + } + + fn database_scope() -> DispatchCapacityScope { + DispatchCapacityScope::DatabaseSuspended { + database_id: DatabaseId::new(3), + core_id: 1, + } + } + + fn queue_scope() -> DispatchCapacityScope { + DispatchCapacityScope::QueueFull { + core_id: 0, + capacity: 8, + } + } + + #[test] + fn each_scope_counts_apart() { + let metrics = DispatchCapacityMetrics::new(); + metrics.record(&tenant_scope()); + metrics.record(&tenant_scope()); + metrics.record(&queue_scope()); + + assert_eq!(metrics.refusals(&tenant_scope()), 2); + assert_eq!(metrics.refusals(&database_scope()), 0); + assert_eq!(metrics.refusals(&queue_scope()), 1); + } + + #[test] + fn the_label_is_the_scope_not_its_ids() { + let metrics = DispatchCapacityMetrics::new(); + metrics.record(&tenant_scope()); + metrics.record(&DispatchCapacityScope::TenantInflight { + tenant_id: TenantId::new(99), + inflight: 1, + cap: 1, + }); + assert_eq!(metrics.refusals(&tenant_scope()), 2); + } + + #[test] + fn exposition_has_one_series_per_scope() { + let metrics = DispatchCapacityMetrics::new(); + metrics.record(&database_scope()); + let mut out = String::new(); + metrics.write_prometheus(&mut out); + + assert!( + out.contains("# TYPE nodedb_dispatch_capacity_busy_total counter"), + "{out}" + ); + assert!( + out.contains(r#"nodedb_dispatch_capacity_busy_total{scope="tenant_inflight"} 0"#), + "{out}" + ); + assert!( + out.contains(r#"nodedb_dispatch_capacity_busy_total{scope="database_suspended"} 1"#), + "{out}" + ); + assert!( + out.contains(r#"nodedb_dispatch_capacity_busy_total{scope="queue_full"} 0"#), + "{out}" + ); + assert_eq!(out.lines().count(), 2 + SCOPE_LABELS.len(), "{out}"); + } +} diff --git a/nodedb/src/bridge/dispatch/dispatcher.rs b/nodedb/src/bridge/dispatch/dispatcher.rs index 8847beeb8..08f4e6324 100644 --- a/nodedb/src/bridge/dispatch/dispatcher.rs +++ b/nodedb/src/bridge/dispatch/dispatcher.rs @@ -20,6 +20,7 @@ use crate::control::router::vshard::VShardRouter; use crate::data::eventfd::EventFdNotifier; use crate::types::Lsn; +use super::capacity_metrics::DispatchCapacityMetrics; use super::core_channel::{CoreChannel, CoreChannelDataSide}; use super::dispatched_lsns::DispatchedLsns; use super::journal::WriteSetJournal; @@ -131,6 +132,9 @@ pub struct Dispatcher { /// The floor windows of accepted requests that carry a WAL LSN. pub(super) dispatched_lsns: DispatchedLsns, + + /// Capacity refusals by scope, shared with the node's `/metrics`. + pub(super) capacity_metrics: Arc, } impl Dispatcher { @@ -187,6 +191,7 @@ impl Dispatcher { capacity_freed: Arc::new(Notify::new()), outcome_floor: OutcomeFloor::new(), dispatched_lsns: DispatchedLsns::default(), + capacity_metrics: Arc::new(DispatchCapacityMetrics::new()), }, data_sides, ) @@ -202,6 +207,16 @@ impl Dispatcher { Arc::clone(&self.capacity_freed) } + /// Capacity refusals this dispatcher counted, by scope. + pub fn capacity_metrics(&self) -> Arc { + Arc::clone(&self.capacity_metrics) + } + + /// Count capacity refusals into `metrics`, the instance `/metrics` reads. + pub fn adopt_capacity_metrics(&mut self, metrics: Arc) { + self.capacity_metrics = metrics; + } + /// The node's outcome floor, shared with every write that opens a window. pub fn outcome_floor(&self) -> Arc { Arc::clone(&self.outcome_floor) diff --git a/nodedb/src/bridge/dispatch/enqueue.rs b/nodedb/src/bridge/dispatch/enqueue.rs index 5716f6aaf..fd9239c81 100644 --- a/nodedb/src/bridge/dispatch/enqueue.rs +++ b/nodedb/src/bridge/dispatch/enqueue.rs @@ -4,7 +4,8 @@ //! the bridge [`Dispatcher`]. //! //! A capacity limit refuses with [`crate::Error::DispatchCapacity`] and hands -//! the request back. Every other refusal is terminal. +//! the request back. `nodedb_dispatch_capacity_busy_total` counts each such +//! refusal by scope. Every other refusal is terminal. use tracing::warn; @@ -13,6 +14,7 @@ use crate::bridge::admission_chokepoint::{assert_write_admitted, reject_uninject use crate::bridge::envelope; use crate::types::{Lsn, VShardId}; +use super::capacity_metrics::DispatchCapacityMetrics; use super::dispatcher::Dispatcher; use super::journal::JournalGroup; use super::refusal::DispatchRefusal; @@ -64,7 +66,7 @@ impl Dispatcher { cap, }; return Err(DispatchRefusal::boxed( - crate::Error::DispatchCapacity { scope }, + capacity_refusal(&self.capacity_metrics, scope), request, )); } @@ -89,7 +91,7 @@ impl Dispatcher { core_id, }; return Err(DispatchRefusal::boxed( - crate::Error::DispatchCapacity { scope }, + capacity_refusal(&self.capacity_metrics, scope), request, )); } @@ -101,7 +103,7 @@ impl Dispatcher { capacity: self.per_core_capacity, }; return Err(DispatchRefusal::boxed( - crate::Error::DispatchCapacity { scope }, + capacity_refusal(&self.capacity_metrics, scope), request, )); } @@ -151,12 +153,11 @@ impl Dispatcher { channel.wfq.set_priority(database_id, cls); channel.wfq.try_enqueue(database_id, request).map_err(|_| { - crate::Error::DispatchCapacity { - scope: DispatchCapacityScope::QueueFull { - core_id, - capacity: self.per_core_capacity, - }, - } + let scope = DispatchCapacityScope::QueueFull { + core_id, + capacity: self.per_core_capacity, + }; + capacity_refusal(&self.capacity_metrics, scope) })?; self.commit_enqueued(core_id, database_id, tenant_id, req_id, wal_lsn); @@ -214,6 +215,16 @@ impl Dispatcher { } } +/// Build the refusal for `scope` and count it. Every capacity refusal is +/// built here, so the counter sees each one. +fn capacity_refusal( + metrics: &DispatchCapacityMetrics, + scope: DispatchCapacityScope, +) -> crate::Error { + metrics.record(&scope); + crate::Error::DispatchCapacity { scope } +} + #[cfg(test)] mod tests { use super::*; @@ -362,6 +373,88 @@ mod tests { panic!("the weighted-fair queue never filled"); } + /// Every capacity refusal counts once, under its own scope, in the + /// instance the dispatcher adopted. + #[test] + fn each_capacity_refusal_counts_under_its_scope() { + let tenant_scope = DispatchCapacityScope::TenantInflight { + tenant_id: TenantId::new(1), + inflight: 0, + cap: 0, + }; + let database_scope = DispatchCapacityScope::DatabaseSuspended { + database_id: DatabaseId::new(1), + core_id: 0, + }; + let queue_scope = DispatchCapacityScope::QueueFull { + core_id: 0, + capacity: 0, + }; + let counts = |metrics: &DispatchCapacityMetrics| { + [ + metrics.refusals(&tenant_scope), + metrics.refusals(&database_scope), + metrics.refusals(&queue_scope), + ] + }; + + // One tenant holds all four slots. + let (mut dispatcher, _sides) = Dispatcher::new(1, 4); + let metrics = std::sync::Arc::new(DispatchCapacityMetrics::new()); + dispatcher.adopt_capacity_metrics(std::sync::Arc::clone(&metrics)); + for i in 1..=4u64 { + dispatcher + .dispatch(make_request_for_db(0, i, i)) + .expect("within the tenant cap"); + } + dispatcher + .dispatch(make_request_for_db(0, 99, 99)) + .expect_err("the tenant is at its cap"); + assert_eq!(counts(&metrics), [1, 0, 0]); + + // Distinct tenants fill one database's virtual queue. + let (mut dispatcher, _sides) = Dispatcher::new(1, 4); + let metrics = dispatcher.capacity_metrics(); + let refused = (1..=64u64) + .find_map(|i| { + let mut request = make_request_for_db(0, 1, i); + request.tenant_id = TenantId::new(i); + dispatcher.dispatch(request).err() + }) + .expect("the database's virtual queue suspends"); + assert!( + matches!( + refused, + crate::Error::DispatchCapacity { + scope: DispatchCapacityScope::DatabaseSuspended { .. } + } + ), + "expected a database-suspended refusal, got: {refused}" + ); + assert_eq!(counts(&metrics), [0, 1, 0]); + + // A direct core dispatch fills the core's queue. + let (mut dispatcher, _sides) = Dispatcher::new(1, 4); + let metrics = dispatcher.capacity_metrics(); + let refused = (1..=64u64) + .find_map(|i| { + let mut request = make_request_for_db(0, i, i); + request.tenant_id = TenantId::new(i); + dispatcher.dispatch_to_core(0, request).err() + }) + .expect("the core's queue fills"); + assert!( + matches!( + refused, + crate::Error::DispatchCapacity { + scope: DispatchCapacityScope::QueueFull { .. } + } + ), + "expected a queue-full refusal, got: {refused}" + ); + assert_eq!(counts(&metrics), [0, 0, 1]); + } + #[test] fn dispatch_to_core_tracks_request_lifecycle() { let (mut dispatcher, mut data_sides) = Dispatcher::new(2, 64); diff --git a/nodedb/src/bridge/dispatch/mod.rs b/nodedb/src/bridge/dispatch/mod.rs index ce643e371..e75c0eb5f 100644 --- a/nodedb/src/bridge/dispatch/mod.rs +++ b/nodedb/src/bridge/dispatch/mod.rs @@ -1,5 +1,6 @@ // SPDX-License-Identifier: BUSL-1.1 +mod capacity_metrics; mod closed_lsns; mod core_channel; mod dispatched_lsns; @@ -12,8 +13,9 @@ mod refusal; mod response_poll; mod tenant_inflight; #[cfg(test)] -mod test_requests; +pub(crate) mod test_requests; +pub use capacity_metrics::DispatchCapacityMetrics; pub use core_channel::{CoreChannel, CoreChannelDataSide}; pub use dispatcher::{ BridgeRequest, BridgeResponse, DATA_PLANE_QUEUE_CAPACITY, DatabasePriorityResolver, diff --git a/nodedb/src/bridge/dispatch/test_requests.rs b/nodedb/src/bridge/dispatch/test_requests.rs index f2a8592db..04084e9fc 100644 --- a/nodedb/src/bridge/dispatch/test_requests.rs +++ b/nodedb/src/bridge/dispatch/test_requests.rs @@ -11,7 +11,7 @@ use crate::bridge::envelope::*; use crate::types::*; /// A read request with id 1 for tenant 1 in the default database on `vshard`. -pub(super) fn make_request(vshard: u32) -> envelope::Request { +pub(crate) fn make_request(vshard: u32) -> envelope::Request { envelope::Request { request_id: RequestId::new(1), tenant_id: TenantId::new(1), @@ -45,7 +45,7 @@ pub(super) fn make_request(vshard: u32) -> envelope::Request { } /// A read request with id `req_id` for tenant 1 in database `db` on `vshard`. -pub(super) fn make_request_for_db(vshard: u32, db: u64, req_id: u64) -> envelope::Request { +pub(crate) fn make_request_for_db(vshard: u32, db: u64, req_id: u64) -> envelope::Request { envelope::Request { request_id: RequestId::new(req_id), tenant_id: TenantId::new(1), diff --git a/nodedb/src/control/cluster/calvin/scheduler/driver/core/catch_up.rs b/nodedb/src/control/cluster/calvin/scheduler/driver/core/catch_up.rs index ce1ca86d1..af8fa6780 100644 --- a/nodedb/src/control/cluster/calvin/scheduler/driver/core/catch_up.rs +++ b/nodedb/src/control/cluster/calvin/scheduler/driver/core/catch_up.rs @@ -253,12 +253,20 @@ mod tests { use nodedb_types::TenantId; use nodedb_types::id::DatabaseId; + use nodedb_physical::physical_plan::PhysicalPlan; + use nodedb_physical::physical_plan::meta::MetaOp; + + use crate::bridge::dispatch::{BridgeResponse, CoreChannelDataSide}; + use crate::bridge::envelope::{StageVote, Status}; use crate::control::cluster::calvin::scheduler::driver::core::test_proposer::lead_data_group; use crate::control::cluster::calvin::scheduler::driver::core::test_support::{ build_test_scheduler, build_test_scheduler_with_data_side, fill_tenant_inflight, - make_sequenced_txn, make_validate_only_txn, test_coll_vshard, + make_sequenced_txn, make_validate_only_txn, release_filler, staged_response, + test_coll_vshard, }; use crate::control::cluster::calvin::scheduler::lock_manager::{AcquireOutcome, TxnId}; + use crate::control::cluster::calvin::scheduler::metrics::intake_closure_reason; + use crate::types::RequestId; #[tokio::test] async fn drain_catch_up_is_noop_when_no_drop_recorded() { @@ -686,4 +694,202 @@ mod tests { ); assert_eq!(armed(&scheduler), None, "the armed range is fully replayed"); } + + /// One run-loop pass of the rebuild: re-send parked requests, then + /// replay the armed log range while the intake gate is open. + fn rebuild_pass(scheduler: &mut Scheduler) { + if scheduler.resends_deferred() { + scheduler.redispatch_deferred(); + } + if scheduler.refresh_intake_gate() { + scheduler.drain_catch_up(); + } + scheduler.refresh_intake_gate(); + } + + /// The first index armed for catch-up on `scheduler`'s vShard. + fn armed_index(scheduler: &Scheduler) -> Option { + scheduler + .sequencer_state_machine + .lock() + .unwrap_or_else(|p| p.into_inner()) + .peek_catch_up_from(scheduler.vshard_id) + } + + /// Pop the next request off the Data Plane side. It must be a Calvin + /// request that `wanted` accepts. Returns its id. + fn pop_calvin_request( + data_side: &mut CoreChannelDataSide, + wanted: impl Fn(&MetaOp) -> bool, + ) -> RequestId { + let request = data_side + .request_rx + .try_pop() + .expect("a request reached the Data Plane"); + match &request.inner.plan { + PhysicalPlan::Meta(op) if wanted(op) => request.inner.request_id, + other => panic!("unexpected request on the Data Plane: {other:?}"), + } + } + + /// Answer `request_id` `Ok` on the Data Plane side, route the answer + /// back, and hand its completion to the scheduler. + async fn answer( + scheduler: &mut Scheduler, + data_side: &mut CoreChannelDataSide, + request_id: RequestId, + stage_vote: Option, + ) { + let mut response = staged_response(Status::Ok, stage_vote); + response.request_id = request_id; + data_side + .response_tx + .try_push(BridgeResponse { inner: response }) + .expect("response ring has room"); + assert!(scheduler.shared.poll_and_route_responses() >= 1); + let (txn_id, routed, response) = + tokio::time::timeout(Duration::from_secs(5), scheduler.completion_rx.recv()) + .await + .expect("the completion arrives") + .expect("the completion channel is open"); + assert_eq!(routed, request_id); + scheduler.handle_completion(txn_id, routed, response); + } + + /// Drive the read-only staged txn `txn_id` from its stage request to + /// completion: the stage votes commit, the verdict commits, and the + /// slice drops. + async fn complete_staged( + scheduler: &mut Scheduler, + data_side: &mut CoreChannelDataSide, + txn_id: TxnId, + ) { + let stage = pop_calvin_request(data_side, |op| { + matches!(op, MetaOp::CalvinExecuteStatic { epoch, position, .. } + if *epoch == txn_id.epoch && *position == txn_id.position) + }); + answer(scheduler, data_side, stage, Some(StageVote::Commit)).await; + scheduler.resume_on_verdict(txn_id, true); + let drop = pop_calvin_request(data_side, |op| { + matches!(op, MetaOp::CalvinDrop { epoch, position } + if *epoch == txn_id.epoch && *position == txn_id.position) + }); + answer(scheduler, data_side, drop, None).await; + } + + /// The rebuild re-applies two epochs while the txns' tenant holds its + /// whole in-flight cap. The first replayed stage defers: its txn stays + /// pending, nothing halts or completes, and catch-up stays armed behind + /// it. Rebuild passes while saturated refuse once per attempt and close + /// the intake gate once. Once capacity frees, both epochs apply and the + /// scheduler reports caught up. + #[tokio::test] + async fn rebuild_against_a_saturated_tenant_defers_then_catches_up() { + const SATURATED_PASSES: u64 = 32; + let registry = CalvinCompletionRegistry::new_detached(); + let (mut scheduler, _dir, mut data_side) = + build_test_scheduler_with_data_side(test_coll_vshard(), registry); + lead_data_group(&mut scheduler); + scheduler.rebuild_target_epoch = 1; + let vshard = scheduler.vshard_id; + let shared = std::sync::Arc::clone(&scheduler.shared); + let (_idx0, idx1) = arm_two_dropped_stage_batches(&scheduler); + let fillers = fill_tenant_inflight(&shared, &mut data_side, TenantId::new(1)); + let tenant_scope = crate::DispatchCapacityScope::TenantInflight { + tenant_id: TenantId::new(1), + inflight: 0, + cap: 0, + }; + let busy = || { + shared + .system_metrics + .as_ref() + .expect("the test state carries system metrics") + .dispatch_capacity + .refusals(&tenant_scope) + }; + let busy_before = busy(); + let first = TxnId::new(0, 0); + let second = TxnId::new(1, 0); + + rebuild_pass(&mut scheduler); + + assert!( + scheduler.has_deferred_for(first), + "the replayed stage defers" + ); + assert!(scheduler.pending.contains_key(&first)); + assert!( + !scheduler.pending.contains_key(&second) && !scheduler.blocked.contains_key(&second), + "the next epoch waits behind the closed gate" + ); + assert_eq!(armed_index(&scheduler), Some(idx1)); + + for _ in 0..SATURATED_PASSES { + rebuild_pass(&mut scheduler); + } + + let metrics = &scheduler.metrics; + assert!( + !scheduler.is_apply_halted(), + "a capacity refusal never halts" + ); + assert_eq!(scheduler.deferred_dispatch_len(), 1); + assert_eq!(scheduler.pending.len(), 1); + assert!(!scheduler.applied.is_applied(0, 0)); + assert_eq!(armed_index(&scheduler), Some(idx1)); + assert_eq!(metrics.completed_count.load(Ordering::Relaxed), 0); + assert_eq!(metrics.executor_error_count.load(Ordering::Relaxed), 0); + assert_eq!( + metrics.dispatch_deferred_count.load(Ordering::Relaxed), + 1 + SATURATED_PASSES, + "each attempt refuses once" + ); + assert_eq!(busy() - busy_before, 1 + SATURATED_PASSES); + assert_eq!( + metrics.intake_gate_closed_counts[intake_closure_reason::DEFERRED_DISPATCH] + .load(Ordering::Relaxed), + 1, + "the stall closes the intake gate once" + ); + assert!(!scheduler.is_caught_up()); + assert_eq!(shared.calvin.caught_up.lagging(), vec![vshard]); + + for filler in fillers { + release_filler(&shared, &mut data_side, filler); + } + rebuild_pass(&mut scheduler); + + assert!(!scheduler.has_deferred_dispatch(), "the stage is re-sent"); + assert_eq!( + armed_index(&scheduler), + None, + "the rebuild range is replayed" + ); + assert!( + scheduler.blocked.contains_key(&second), + "the next epoch waits for the first one's locks" + ); + + complete_staged(&mut scheduler, &mut data_side, first).await; + complete_staged(&mut scheduler, &mut data_side, second).await; + + assert!(scheduler.applied.is_applied(0, 0)); + assert!(scheduler.applied.is_applied(1, 0)); + assert!(scheduler.pending.is_empty()); + assert!(!scheduler.is_apply_halted()); + assert_eq!( + scheduler + .metrics + .executor_error_count + .load(Ordering::Relaxed), + 0 + ); + assert_eq!(scheduler.metrics.completed_count.load(Ordering::Relaxed), 2); + assert!(scheduler.is_caught_up()); + assert!( + shared.calvin.caught_up.lagging().is_empty(), + "the readiness gate opens" + ); + } } diff --git a/nodedb/src/control/cluster/calvin/scheduler/driver/core/deferred.rs b/nodedb/src/control/cluster/calvin/scheduler/driver/core/deferred.rs index fe7766f9b..5d33e63e3 100644 --- a/nodedb/src/control/cluster/calvin/scheduler/driver/core/deferred.rs +++ b/nodedb/src/control/cluster/calvin/scheduler/driver/core/deferred.rs @@ -252,7 +252,9 @@ impl Scheduler { match error { crate::Error::DispatchCapacity { scope } => { self.metrics.record_dispatch_deferred(); - tracing::debug!( + // Trace per attempt: the intake gate logs once per stall + // window, as it closes for the deferred dispatch and reopens. + tracing::trace!( vshard_id = self.vshard_id, epoch = txn_id.epoch, position = txn_id.position, diff --git a/nodedb/src/control/gateway/error_map/capacity_refusal.rs b/nodedb/src/control/gateway/error_map/capacity_refusal.rs new file mode 100644 index 000000000..1067a48c1 --- /dev/null +++ b/nodedb/src/control/gateway/error_map/capacity_refusal.rs @@ -0,0 +1,80 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! A dispatch capacity refusal answers the retryable `SERVER_OVERLOAD` class +//! on every client surface. +//! +//! The refusal reaches a listener in two forms: raised by this node's +//! dispatcher, or relayed from a core or a peer as a Data-Plane code. Both +//! forms answer the same class, distinct from an internal error. + +use nodedb_types::error::{ErrorCode, sqlstate}; + +use super::gateway_map::GatewayErrorMap; +use crate::Error; + +/// The refusal in each form a listener can receive. +fn capacity_refusals() -> [Error; 2] { + [ + Error::DispatchCapacity { + scope: crate::DispatchCapacityScope::QueueFull { + core_id: 0, + capacity: 4, + }, + }, + Error::DataPlane(crate::bridge::envelope::ErrorCode::DispatchCapacity { + reason: "tenant 1 holds 4/4 in-flight requests".into(), + }), + ] +} + +#[test] +fn pgwire_answers_57p03() { + for err in capacity_refusals() { + let (state, _) = GatewayErrorMap::to_pgwire(&err); + assert_eq!(state, sqlstate::SERVER_OVERLOAD, "{err:?}"); + } +} + +#[test] +fn native_frame_carries_the_server_overload_code_and_57p03() { + for err in capacity_refusals() { + let (code, _) = GatewayErrorMap::to_native(&err); + assert_eq!(code, ErrorCode::SERVER_OVERLOAD, "{err:?}"); + + let frame = crate::control::server::native::dispatch::error_to_native(1, &err); + let payload = frame.error.expect("error frames carry a payload"); + assert_eq!(payload.ndb_code, ErrorCode::SERVER_OVERLOAD.0, "{err:?}"); + assert_eq!(payload.code, sqlstate::SERVER_OVERLOAD, "{err:?}"); + } +} + +#[test] +fn http_answers_503() { + for err in capacity_refusals() { + let (status, _) = GatewayErrorMap::to_http(&err); + assert_eq!(status, 503, "{err:?}"); + } +} + +#[test] +fn resp_answers_busy() { + for err in capacity_refusals() { + let reply = GatewayErrorMap::to_resp(&err); + assert!(reply.starts_with("BUSY "), "{err:?}: {reply}"); + } +} + +#[test] +fn the_class_is_not_an_unclassified_failure() { + for err in capacity_refusals() { + assert!( + !crate::error_classify::is_unclassified_failure(&err), + "{err:?}" + ); + assert_eq!( + crate::error_classify::classify(&err).code(), + ErrorCode::SERVER_OVERLOAD, + "{err:?}" + ); + } +} diff --git a/nodedb/src/control/gateway/error_map/mod.rs b/nodedb/src/control/gateway/error_map/mod.rs index a704ab5c8..ce59a0935 100644 --- a/nodedb/src/control/gateway/error_map/mod.rs +++ b/nodedb/src/control/gateway/error_map/mod.rs @@ -6,6 +6,8 @@ //! One module per protocol surface owns that surface's mapping, so a change //! to its SQLSTATE / HTTP / RESP / native codes is a one-file edit. +#[cfg(test)] +mod capacity_refusal; #[cfg(test)] pub(crate) mod class_parity; mod gateway_map; diff --git a/nodedb/src/control/gateway/error_map/system_dispatch_refusal.rs b/nodedb/src/control/gateway/error_map/system_dispatch_refusal.rs index 578828d75..cd303aadf 100644 --- a/nodedb/src/control/gateway/error_map/system_dispatch_refusal.rs +++ b/nodedb/src/control/gateway/error_map/system_dispatch_refusal.rs @@ -165,6 +165,59 @@ async fn create_vector_index_refusal_keeps_its_sqlstate() { ); } +/// A dispatcher capacity refusal keeps its type through the system door, so +/// the statement answers the retryable `57P03`, never `XX000`. +#[tokio::test] +async fn dispatch_system_keeps_a_capacity_refusal_retryable() { + let (state, mut side, _dir) = fixture(); + // Tenant 1 takes every in-flight slot. Popping each request keeps the + // ring and the weighted-fair queue empty, so only the tenant cap refuses. + { + let mut dispatcher = state + .dispatcher + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let filled = (1_000_000..1_000_128u64) + .take_while(|&id| { + let admitted = dispatcher + .dispatch(crate::bridge::dispatch::test_requests::make_request_for_db( + 0, 1, id, + )) + .is_ok(); + while side.request_rx.try_pop().is_ok() {} + admitted + }) + .count(); + assert!(filled > 0, "the tenant cap admits some requests first"); + assert!(filled < 128, "the tenant reaches its in-flight cap"); + } + + let result = dispatch_system( + &state, + SystemTask::new( + SystemReason::DdlApply, + TenantId::new(1), + nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, COLLECTION), + set_params_plan(), + ), + Duration::from_secs(5), + ) + .await; + + let err = result.expect_err("the tenant is at its in-flight cap"); + assert!( + matches!( + err, + crate::Error::DispatchCapacity { + scope: crate::DispatchCapacityScope::TenantInflight { .. } + } + ), + "expected a tenant-cap refusal, got {err:?}" + ); + let (state_code, _) = crate::control::gateway::GatewayErrorMap::to_pgwire(&err); + assert_eq!(state_code, sqlstate::SERVER_OVERLOAD); +} + /// A refusal with no code has no class of its own, so it is the one case /// that stays internal. #[tokio::test] diff --git a/nodedb/src/control/metrics/system/fields.rs b/nodedb/src/control/metrics/system/fields.rs index 486843dcc..ade5f4159 100644 --- a/nodedb/src/control/metrics/system/fields.rs +++ b/nodedb/src/control/metrics/system/fields.rs @@ -10,6 +10,7 @@ use super::super::histogram::{AtomicHistogram, WAL_FSYNC_BUCKETS_US}; use super::super::purge::PurgeMetrics; use super::core_fail_stop::CoreFailStops; use super::heartbeat::CoreHeartbeats; +use crate::bridge::dispatch::DispatchCapacityMetrics; use crate::data::executor::core_loop::pressure::ThrottleMetrics; use crate::data::io::IoMetrics; @@ -258,6 +259,13 @@ pub struct SystemMetrics { /// The Prometheus handler reads from here. pub io_metrics: Arc, + // ── Bridge dispatch capacity ── + /// Dispatches the bridge dispatcher refused at a capacity limit, by scope. + /// + /// The dispatcher adopts this `Arc` when the shared state opens, so the + /// Prometheus handler reads it without taking the dispatcher lock. + pub dispatch_capacity: Arc, + // ── SPSC intake throttle ── /// Per-level core counts and transition counters for the Data Plane's /// request-intake throttle. diff --git a/nodedb/src/control/metrics/system/render.rs b/nodedb/src/control/metrics/system/render.rs index 8a413e13f..7d75d0e8b 100644 --- a/nodedb/src/control/metrics/system/render.rs +++ b/nodedb/src/control/metrics/system/render.rs @@ -18,6 +18,7 @@ impl SystemMetrics { self.purge.write_prometheus(&mut out); self.io_metrics.write_prometheus(&mut out); self.spsc_throttle.write_prometheus(&mut out); + self.dispatch_capacity.write_prometheus(&mut out); self.core_fail_stops.write_prometheus(&mut out); out } diff --git a/nodedb/src/control/server/http/auth.rs b/nodedb/src/control/server/http/auth.rs index 719fad66f..8cf85c51c 100644 --- a/nodedb/src/control/server/http/auth.rs +++ b/nodedb/src/control/server/http/auth.rs @@ -11,6 +11,9 @@ use axum::extract::FromRequestParts; use axum::http::request::Parts; use axum::http::{HeaderMap, StatusCode}; use axum::response::{IntoResponse, Response}; +use nodedb_types::error::sqlstate; + +use crate::control::server::pgwire::types::error_map::numeric_code_to_sqlstate; use crate::config::auth::AuthMode; use crate::control::security::identity::AuthenticatedIdentity; @@ -312,6 +315,13 @@ pub enum ApiError { message: String, retry_after_secs: u64, }, + /// 503 Service Unavailable for a transient refusal of SQLSTATE `57P03`, + /// such as a dispatch capacity refusal. Nothing ran, so the client + /// retries after `Retry-After`. + Overloaded { + message: String, + code: nodedb_types::error::ErrorCode, + }, /// Arbitrary HTTP status from gateway error mapping. HttpStatus(u16, String), /// An error with its NodeDB numeric code preserved (e.g. a DDL refusal), @@ -340,11 +350,13 @@ impl IntoResponse for ApiError { retry_after_secs, } => { let body = HttpError::new(message); - let mut resp = (StatusCode::TOO_MANY_REQUESTS, axum::Json(body)).into_response(); - if let Ok(val) = retry_after_secs.to_string().parse() { - resp.headers_mut().insert("Retry-After", val); - } - resp + let resp = (StatusCode::TOO_MANY_REQUESTS, axum::Json(body)).into_response(); + with_retry_after(resp, retry_after_secs) + } + ApiError::Overloaded { message, code } => { + let body = HttpError::with_code(message, code.to_string()); + let resp = (StatusCode::SERVICE_UNAVAILABLE, axum::Json(body)).into_response(); + with_retry_after(resp, OVERLOAD_RETRY_AFTER_SECS) } ApiError::Coded { status, @@ -352,12 +364,18 @@ impl IntoResponse for ApiError { code, cause, } => { + let overloaded = numeric_code_to_sqlstate(code) == sqlstate::SERVER_OVERLOAD; let body = HttpError::with_code(message, code.to_string()); let body = match cause { Some(cause) => body.caused_by(&cause), None => body, }; - (status, axum::Json(body)).into_response() + let resp = (status, axum::Json(body)).into_response(); + if overloaded { + with_retry_after(resp, OVERLOAD_RETRY_AFTER_SECS) + } else { + resp + } } ApiError::Unauthorized(msg) => plain(StatusCode::UNAUTHORIZED, msg), ApiError::Forbidden(msg) => plain(StatusCode::FORBIDDEN, msg), @@ -371,6 +389,18 @@ impl IntoResponse for ApiError { } } +/// Seconds a client waits before it retries a `57P03` capacity refusal. +const OVERLOAD_RETRY_AFTER_SECS: u64 = 1; + +/// Set the `Retry-After` header of `resp` to `secs`. +fn with_retry_after(mut resp: Response, secs: u64) -> Response { + resp.headers_mut().insert( + axum::http::header::RETRY_AFTER, + axum::http::HeaderValue::from(secs), + ); + resp +} + /// Axum extractor that resolves and enforces HTTP auth before a handler runs. /// /// Add this as the first parameter to any handler that performs tenant-scoped @@ -430,7 +460,8 @@ impl FromRequestParts for ResolvedAuth { /// The status and message come from the one gateway mapping, /// [`GatewayErrorMap::to_http`](crate::control::gateway::GatewayErrorMap::to_http). -/// A rate refusal also carries its `Retry-After` hint. +/// A rate refusal and a `57P03` capacity refusal also carry a `Retry-After` +/// hint. impl From for ApiError { fn from(e: crate::Error) -> Self { let (status, message) = crate::control::gateway::GatewayErrorMap::to_http(&e); @@ -440,6 +471,14 @@ impl From for ApiError { retry_after_secs: retry_after_ms.div_ceil(1000).max(1), }; } + let (_severity, state, _message) = + crate::control::server::pgwire::types::error_to_sqlstate(&e); + if state == sqlstate::SERVER_OVERLOAD { + return Self::Overloaded { + message, + code: crate::error_classify::classify(&e).code(), + }; + } Self::HttpStatus(status, message) } } @@ -624,6 +663,73 @@ mod tests { ); } + fn retry_after(response: &Response) -> Option<&str> { + response + .headers() + .get(axum::http::header::RETRY_AFTER) + .and_then(|value| value.to_str().ok()) + } + + /// A dispatch capacity refusal answers 503 with `Retry-After` and the + /// `SERVER_OVERLOAD` code, raised locally or relayed from a core. + #[tokio::test] + async fn a_capacity_refusal_is_503_with_retry_after_and_its_code() { + let errors = [ + crate::Error::DispatchCapacity { + scope: crate::DispatchCapacityScope::TenantInflight { + tenant_id: TenantId::new(1), + inflight: 4, + cap: 4, + }, + }, + crate::Error::DataPlane(crate::bridge::envelope::ErrorCode::DispatchCapacity { + reason: "core 0 queue is full at 4 requests".into(), + }), + ]; + for error in errors { + let label = format!("{error:?}"); + let response = ApiError::from(error).into_response(); + assert_eq!( + response.status(), + StatusCode::SERVICE_UNAVAILABLE, + "{label}" + ); + assert_eq!(retry_after(&response), Some("1"), "{label}"); + let body = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .expect("read response body"); + let json: serde_json::Value = serde_json::from_slice(&body).expect("valid JSON body"); + assert_eq!( + json["code"], + nodedb_types::error::ErrorCode::SERVER_OVERLOAD.to_string(), + "{label}" + ); + } + } + + /// A coded `SERVER_OVERLOAD` answer carries `Retry-After`. Another code + /// does not. + #[test] + fn only_a_coded_overload_carries_retry_after() { + let overload = ApiError::Coded { + status: StatusCode::SERVICE_UNAVAILABLE, + message: "busy".into(), + code: nodedb_types::error::ErrorCode::SERVER_OVERLOAD, + cause: None, + } + .into_response(); + assert_eq!(retry_after(&overload), Some("1")); + + let internal = ApiError::Coded { + status: StatusCode::INTERNAL_SERVER_ERROR, + message: "boom".into(), + code: nodedb_types::error::ErrorCode::INTERNAL, + cause: None, + } + .into_response(); + assert_eq!(retry_after(&internal), None); + } + #[test] fn x_on_deny_malformed_value_is_ignored() { let context = auth_context(); diff --git a/nodedb/src/control/server/shared/ddl/sync_dispatch/dispatch.rs b/nodedb/src/control/server/shared/ddl/sync_dispatch/dispatch.rs index a321b3632..f65d83ed0 100644 --- a/nodedb/src/control/server/shared/ddl/sync_dispatch/dispatch.rs +++ b/nodedb/src/control/server/shared/ddl/sync_dispatch/dispatch.rs @@ -226,11 +226,10 @@ async fn dispatch_plan(state: &SharedState, dispatch: PlanDispatch) -> crate::Re Err(p) => p.into_inner().dispatch(request), }; if let Err(error) = dispatched { - // No response will arrive, and no core applied the plan. + // No response will arrive, and no core applied the plan. The refusal + // keeps its type, so a capacity refusal stays retryable. state.tracker.cancel(&request_id); - return Err(crate::Error::Internal { - detail: error.to_string(), - }); + return Err(error); } // Await to the same instant the envelope carries — yields the thread so the // response poller can run. Reaching that instant is the statement running diff --git a/nodedb/src/control/server/shuffle/aggregator_hook.rs b/nodedb/src/control/server/shuffle/aggregator_hook.rs index 6119ef186..42632f42e 100644 --- a/nodedb/src/control/server/shuffle/aggregator_hook.rs +++ b/nodedb/src/control/server/shuffle/aggregator_hook.rs @@ -196,11 +196,9 @@ impl RegistryShuffleAggregator { Ok(mut d) => d.dispatch(request), Err(poisoned) => poisoned.into_inner().dispatch(request), }; - if let Err(e) = dispatch_result { - return Err(TypedClusterError::Internal { - code: 0, - message: format!("shuffle aggregate dispatch failed: {e}"), - }); + if let Err(error) = dispatch_result { + // A capacity refusal crosses as the retryable overload class. + return Err(execution_error_to_typed(error)); } let max_result_bytes = self.state.tuning.network.max_query_result_bytes as usize; diff --git a/nodedb/src/control/server/shuffle/consumer_hook.rs b/nodedb/src/control/server/shuffle/consumer_hook.rs index e4cc0f988..489cbb751 100644 --- a/nodedb/src/control/server/shuffle/consumer_hook.rs +++ b/nodedb/src/control/server/shuffle/consumer_hook.rs @@ -199,11 +199,9 @@ impl RegistryShuffleConsumer { Ok(mut d) => d.dispatch(request), Err(poisoned) => poisoned.into_inner().dispatch(request), }; - if let Err(e) = dispatch_result { - return Err(TypedClusterError::Internal { - code: 0, - message: format!("shuffle consume dispatch failed: {e}"), - }); + if let Err(error) = dispatch_result { + // A capacity refusal crosses as the retryable overload class. + return Err(execution_error_to_typed(error)); } let max_result_bytes = self.state.tuning.network.max_query_result_bytes as usize; diff --git a/nodedb/src/control/state/init.rs b/nodedb/src/control/state/init.rs index a44590b75..e24b1a411 100644 --- a/nodedb/src/control/state/init.rs +++ b/nodedb/src/control/state/init.rs @@ -37,7 +37,7 @@ impl SharedState { /// Construct the test state every constructor in [`super::init_variants`] /// builds on. pub(super) fn new_inner( - dispatcher: Dispatcher, + mut dispatcher: Dispatcher, wal: Arc, ) -> crate::Result> { let shutdown = Arc::new(crate::control::shutdown::ShutdownWatch::new()); @@ -115,6 +115,8 @@ impl SharedState { ), ); let hlc_clock = wal.hlc_clock(); + let system_metrics = Arc::new(crate::control::metrics::SystemMetrics::new()); + dispatcher.adopt_capacity_metrics(Arc::clone(&system_metrics.dispatch_capacity)); let state = Arc::new(Self { outcome_floor: dispatcher.outcome_floor(), dispatcher: Mutex::new(dispatcher), @@ -321,7 +323,7 @@ impl SharedState { raft_propose_leader_change_retries: AtomicU64::new(0), request_id_counter: AtomicU64::new(1), shuffle_id_counter: AtomicU64::new(1), - system_metrics: Some(Arc::new(crate::control::metrics::SystemMetrics::new())), + system_metrics: Some(system_metrics), database_metrics: Arc::new(crate::control::metrics::DatabaseMetricsRegistry::new()), quota_ceiling: Arc::new(std::sync::RwLock::new( crate::control::security::catalog::GlobalQuotaCeiling::default(), diff --git a/nodedb/src/control/state/init_prod/open.rs b/nodedb/src/control/state/init_prod/open.rs index a4ab66d8f..0422ae9f1 100644 --- a/nodedb/src/control/state/init_prod/open.rs +++ b/nodedb/src/control/state/init_prod/open.rs @@ -32,11 +32,12 @@ impl SharedState { governor: Arc, ) -> crate::Result> { let DataPlaneHandles { - dispatcher, + mut dispatcher, quiesce, array_catalog, system_metrics, } = handles; + dispatcher.adopt_capacity_metrics(Arc::clone(&system_metrics.dispatch_capacity)); let super::bootstrap::ProdBootstrap { credentials, producer_registry, diff --git a/nodedb/tests/inproc/cases/event_scheduler_concurrency.rs b/nodedb/tests/inproc/cases/event_scheduler_concurrency.rs index e1934b536..adebbefa6 100644 --- a/nodedb/tests/inproc/cases/event_scheduler_concurrency.rs +++ b/nodedb/tests/inproc/cases/event_scheduler_concurrency.rs @@ -1,41 +1,36 @@ // SPDX-License-Identifier: BUSL-1.1 -//! Regression coverage: scheduler must observe shutdown, cap concurrency, -//! and reject jobs when budget is exceeded. +//! The scheduled-job dispatcher observes shutdown, caps concurrency, and +//! refuses a job over its result budget. //! -//! Today `scheduler_loop` fires `tokio::spawn` per due schedule with: -//! * no shutdown observation inside the spawned job — long-running jobs -//! hold `Arc` and delay teardown; -//! * no concurrency cap — N schedules at `* * * * *` with -//! `allow_overlap = true` spawn N heavy SQLs at every minute-mark, -//! starving request-path queries; -//! * no memory ceiling — `execute_job` may materialise unbounded result -//! sets through `execute_block`. -//! -//! The fix extracts a `JobDispatcher` (or equivalent) that owns a -//! semaphore bound and a shutdown receiver. This file gates on that type -//! and its contract. +//! A refusal at the concurrency cap is `DispatchOutcome::OverBudget`, a typed +//! and transient outcome: nothing ran, and the same job spawns once a running +//! job finishes. The scheduler's next tick retries it. use std::sync::Arc; +use std::time::Duration; use nodedb::event::scheduler::dispatcher::{DispatchOutcome, JobDispatcher, JobDispatcherConfig}; +/// How long the test waits for a finished job to free its slot. +const SLOT_FREE_WAIT: Duration = Duration::from_secs(5); + #[tokio::test] -async fn dispatcher_rejects_when_concurrency_cap_reached() { - // Cap concurrency at 2. Spawn 5 never-returning jobs: the 3rd, 4th, 5th - // must be rejected with `OverBudget`, not queued unbounded and not - // spawned as free-running tasks. +async fn dispatcher_refuses_over_cap_then_admits_once_a_slot_frees() { + // Cap concurrency at 2. Of five jobs that run until released, the first + // two spawn and the other three are refused, neither queued nor spawned. let cfg = JobDispatcherConfig { max_concurrent_jobs: 2, max_result_bytes: 1 << 20, }; let dispatcher = Arc::new(JobDispatcher::new(cfg)); + let (release_tx, release_rx) = tokio::sync::watch::channel(false); let outcomes: Vec = (0..5) .map(|_| { + let mut release = release_rx.clone(); dispatcher.try_spawn(|_shutdown| async move { - // Simulate a job that never returns until aborted. - std::future::pending::<()>().await; + let _ = release.wait_for(|released| *released).await; Ok(()) }) }) @@ -45,13 +40,31 @@ async fn dispatcher_rejects_when_concurrency_cap_reached() { .iter() .filter(|o| matches!(o, DispatchOutcome::Spawned)) .count(); - let rejected = outcomes + let refused = outcomes .iter() .filter(|o| matches!(o, DispatchOutcome::OverBudget)) .count(); assert_eq!(accepted, 2, "only two jobs fit under the concurrency cap"); - assert_eq!(rejected, 3, "excess jobs must be rejected, not queued"); + assert_eq!(refused, 3, "jobs over the cap are refused, not queued"); + + // The refusal is transient: once the running jobs finish, a job spawns. + release_tx + .send(true) + .expect("the running jobs hold the release receiver"); + let freed = tokio::time::timeout(SLOT_FREE_WAIT, async { + while dispatcher.in_flight() > 0 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await; + assert!(freed.is_ok(), "the released jobs free their slots"); + let retried = dispatcher.try_spawn(|_shutdown| async move { Ok(()) }); + assert_eq!( + retried, + DispatchOutcome::Spawned, + "a refused job spawns once a slot frees" + ); dispatcher.shutdown_and_drain().await; } From aa4f3bc9ad88d0d5922efb7af6e313522b2bb4d7 Mon Sep 17 00:00:00 2001 From: Farhan Syah Date: Mon, 12 Oct 2026 12:26:52 +0800 Subject: [PATCH 02/11] fix(vector): read node vectors through the backing on insert and filtered search After a graph-only checkpoint restore, nodes hold no local vector and the vector lives in the attached backing. Neighbor pruning and the diversity heuristic read node storage directly, so every distance against a restored node errored and new nodes gained no incoming edges. Add HnswIndex::node_bytes, which resolves local storage first and then the backing, encoding to the index dtype when needed. Route insert, prune, and the f32 distance helper through it. dist_to_node_f32 now takes the metric, and the adaptive filter, ACORN, and NaviX traversals use it instead of get_vector, so restored and non-F32 indexes rank correctly under a caller-supplied metric. --- nodedb-vector/src/adaptive_filter.rs | 6 +- nodedb-vector/src/hnsw/build.rs | 136 +++++++++++++++++- .../src/hnsw/graph/index/distance_ops.rs | 111 +++++++++----- nodedb-vector/src/navix/acorn.rs | 11 +- nodedb-vector/src/navix/traversal.rs | 13 +- 5 files changed, 213 insertions(+), 64 deletions(-) diff --git a/nodedb-vector/src/adaptive_filter.rs b/nodedb-vector/src/adaptive_filter.rs index 767367ee8..41af942dd 100644 --- a/nodedb-vector/src/adaptive_filter.rs +++ b/nodedb-vector/src/adaptive_filter.rs @@ -9,7 +9,6 @@ use roaring::RoaringBitmap; -use crate::distance::distance; use crate::hnsw::{HnswIndex, SearchResult}; /// Filter strategy thresholds. @@ -100,13 +99,12 @@ pub fn adaptive_search( let mut results: Vec = bitmap .iter() .filter_map(|id| { - let v = index.get_vector(id)?; - if index.is_deleted(id) { + if id as usize >= index.len() || index.is_deleted(id) { return None; } Some(SearchResult { id, - distance: distance(query, v, metric), + distance: index.dist_to_node_f32(query, id, metric), }) }) .collect(); diff --git a/nodedb-vector/src/hnsw/build.rs b/nodedb-vector/src/hnsw/build.rs index a93ee761a..ff1d9d5c9 100644 --- a/nodedb-vector/src/hnsw/build.rs +++ b/nodedb-vector/src/hnsw/build.rs @@ -82,7 +82,7 @@ impl HnswIndex { self.nodes[nid].neighbors[layer].push(new_id); if self.nodes[nid].neighbors[layer].len() > m { - let node_bytes = self.nodes[nid].storage.as_bytes().to_vec(); + let node_bytes = self.node_bytes(neighbor.id).into_owned(); self.prune_neighbors(nid, layer, &node_bytes, m); } } @@ -131,14 +131,13 @@ fn select_neighbors_heuristic( break; } - let candidate_bytes = index.nodes[candidate.id as usize].storage.as_bytes(); + let candidate_bytes = index.node_bytes(candidate.id); let is_diverse = selected.iter().all(|s| { - let selected_bytes = index.nodes[s.id as usize].storage.as_bytes(); let dist_to_selected = crate::distance::dispatch::distance_typed( index.params.metric, index.params.dtype, - candidate_bytes, - selected_bytes, + &candidate_bytes, + &index.node_bytes(s.id), index.dim, ) .unwrap_or(f32::MAX); @@ -232,6 +231,133 @@ mod tests { } } + /// Deterministic vectors in `[0, 1)` from an xorshift64 stream. `seed` must + /// be non-zero. + fn seeded_vectors(count: usize, dim: usize, seed: u64) -> Vec> { + let mut state = seed; + (0..count) + .map(|_| { + (0..dim) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + (state % 10_000) as f32 / 10_000.0 + }) + .collect() + }) + .collect() + } + + /// A graph-only restore with an attached backing must insert exactly like a + /// full restore. Restored nodes hold no local vector, so neighbor pruning + /// and the diversity heuristic read them through the backing. A new node + /// must gain incoming edges, or search can never reach it. + /// + /// `segment_backing` is gated off on wasm32, so the test is too. + #[cfg(not(target_arch = "wasm32"))] + #[test] + fn inserts_after_graph_only_restore_match_a_full_restore() { + use std::sync::Arc; + + use crate::segment_backing::VectorSegmentBacking; + + /// Serves an in-memory copy of the base vectors. + struct VecBacking { + dim: usize, + vectors: Vec>, + } + impl VectorSegmentBacking for VecBacking { + fn len(&self) -> usize { + self.vectors.len() + } + fn dim(&self) -> usize { + self.dim + } + fn get_vector(&self, id: u32) -> Option<&[f32]> { + self.vectors.get(id as usize).map(Vec::as_slice) + } + fn get_surrogate(&self, _id: u32) -> Option { + None + } + } + + const DIM: usize = 8; + let params = HnswParams { + m: 4, + m0: 8, + ef_construction: 32, + metric: DistanceMetric::L2, + dtype: nodedb_types::vector_dtype::VectorStorageDtype::F32, + }; + let base = seeded_vectors(200, DIM, 0x9E37_79B9_7F4A_7C15); + let extra = seeded_vectors(5, DIM, 0xD1B5_4A32_D192_ED03); + + let mut src = HnswIndex::with_seed(DIM, params, 12345); + for v in &base { + src.insert(v.clone()).unwrap(); + } + + let mut full = HnswIndex::from_checkpoint(&src.checkpoint_to_bytes().unwrap()) + .unwrap() + .expect("full checkpoint must be recognized"); + let mut graph_only = HnswIndex::from_checkpoint(&src.graph_checkpoint_to_bytes().unwrap()) + .unwrap() + .expect("graph checkpoint must be recognized"); + graph_only + .with_backing(Arc::new(VecBacking { + dim: DIM, + vectors: base.clone(), + })) + .expect("a complete backing must attach"); + + for v in &extra { + full.insert(v.clone()).unwrap(); + graph_only.insert(v.clone()).unwrap(); + } + + let neighbors = graph_only.export_neighbors(); + assert_eq!( + full.export_neighbors(), + neighbors, + "a graph-only restore must build the same graph as a full restore" + ); + + for (offset, v) in extra.iter().enumerate() { + let new_id = (base.len() + offset) as u32; + let hits = graph_only.search(v, 1, 64).unwrap(); + assert_eq!( + hits.first().map(|h| h.id), + Some(new_id), + "exact query for inserted node {new_id} returns it, got {hits:?}" + ); + assert!( + hits[0].distance < 1e-6, + "exact query for node {new_id} scores ~0, got {}", + hits[0].distance + ); + assert!( + neighbors + .iter() + .any(|layers| layers.first().is_some_and(|l0| l0.contains(&new_id))), + "node {new_id} has no incoming layer-0 edge" + ); + } + + let m = graph_only.params().m; + let m0 = graph_only.params().m0; + for (id, layers) in neighbors.iter().enumerate() { + for (layer, list) in layers.iter().enumerate() { + let cap = if layer == 0 { m0 } else { m }; + assert!( + list.len() <= cap, + "node {id} layer {layer} has {} neighbors, cap {cap}", + list.len() + ); + } + } + } + #[test] fn compact_removes_tombstones() { let mut idx = make_index(); diff --git a/nodedb-vector/src/hnsw/graph/index/distance_ops.rs b/nodedb-vector/src/hnsw/graph/index/distance_ops.rs index 704626abb..c32d637b6 100644 --- a/nodedb-vector/src/hnsw/graph/index/distance_ops.rs +++ b/nodedb-vector/src/hnsw/graph/index/distance_ops.rs @@ -2,11 +2,13 @@ //! Layer assignment and distance computation for [`HnswIndex`]. +use std::borrow::Cow; + use nodedb_types::vector_dtype::VectorStorageDtype; use super::HnswIndex; use crate::distance::dispatch::distance_typed; -use crate::distance::distance; +use crate::distance::{DistanceMetric, distance}; use crate::dtype::cast_from_f32; use crate::hnsw::graph::MAX_LAYER_CAP; @@ -39,54 +41,87 @@ impl HnswIndex { /// has no vector source at all; aborting a search worker (and poisoning the /// locks it held) is strictly worse than ranking one node last. pub(crate) fn dist_to_node(&self, query_bytes: &[u8], node_id: u32) -> f32 { - let node_bytes = self.nodes[node_id as usize].storage.as_bytes(); - #[cfg(not(target_arch = "wasm32"))] - let node_bytes: &[u8] = if node_bytes.is_empty() { - if let Some(ref b) = self.backing { - // Backing stores F32 only; convert slice to bytes for distance_typed. - if let Some(v) = b.get_vector(node_id) { - // SAFETY: &[f32] → &[u8] cast via bytemuck is always safe. - bytemuck::cast_slice(v) - } else { - node_bytes - } - } else { - node_bytes - } - } else { - node_bytes - }; distance_typed( self.params.metric, self.params.dtype, query_bytes, - node_bytes, + &self.node_bytes(node_id), self.dim, ) .unwrap_or(f32::INFINITY) } - /// Compute distance between a query given as `&[f32]` and a stored node. + /// Node `node_id`'s vector bytes, encoded in `self.params.dtype`. + /// + /// Every distance computation on the insert, prune, and search paths reads + /// node vectors through this accessor. After a graph-checkpoint-only + /// restore, local storage is empty and the vector lives in the attached + /// backing. Reading `node.storage` directly there yields an empty slice, + /// and every distance against it errors. + /// + /// The backing holds F32 vectors. An F32 index borrows them as bytes. Any + /// other dtype encodes them into an owned buffer. + /// + /// Returns an empty slice when `node_id` is out of range, or the node has + /// no local vector and no backing serves one. Distance against an empty + /// slice errors, so callers score that node [`f32::INFINITY`]. + pub(crate) fn node_bytes(&self, node_id: u32) -> Cow<'_, [u8]> { + let local = self + .nodes + .get(node_id as usize) + .map_or(&[][..], |n| n.storage.as_bytes()); + #[cfg(not(target_arch = "wasm32"))] + { + if local.is_empty() + && let Some(v) = self.backing.as_ref().and_then(|b| b.get_vector(node_id)) + { + return match self.params.dtype { + VectorStorageDtype::F32 => Cow::Borrowed(bytemuck::cast_slice(v)), + dtype => Cow::Owned(cast_from_f32(v, dtype)), + }; + } + } + Cow::Borrowed(local) + } + + /// Compute distance between an f32 query and a stored node under `metric`. + /// + /// `metric` is caller-supplied, so filtered traversals that rank under a + /// metric other than the build metric use this accessor too. The node + /// vector comes from [`Self::node_bytes`], so restored nodes resolve + /// through the attached backing. + /// + /// F32 indexes compare the query to the node vector without copying. + /// F16/BF16 indexes encode the query on each call, which allocates. /// - /// For F32 indexes this is a direct call to `distance`. For F16/BF16 - /// indexes the query is encoded to the storage dtype on each call, which - /// is an allocation. Prefer pre-encoding the query once via - /// [`crate::dtype::cast_from_f32`] and calling [`Self::dist_to_node`] - /// for hot-path code such as search. - #[allow(dead_code)] - pub(crate) fn dist_to_node_f32(&self, query: &[f32], node_id: u32) -> f32 { + /// A node whose vector cannot be resolved, or a query of the wrong length, + /// scores [`f32::INFINITY`]. + pub(crate) fn dist_to_node_f32( + &self, + query: &[f32], + node_id: u32, + metric: DistanceMetric, + ) -> f32 { + let node_bytes = self.node_bytes(node_id); match self.params.dtype { - VectorStorageDtype::F32 => distance( - query, - self.nodes[node_id as usize] - .storage - .as_f32_slice() - .expect("F32 dtype must have F32 storage"), - self.params.metric, - ), - _ => { - let query_bytes = cast_from_f32(query, self.params.dtype); - self.dist_to_node(&query_bytes, node_id) + VectorStorageDtype::F32 => match bytemuck::try_cast_slice::(&node_bytes) { + Ok(node) if node.len() == query.len() => distance(query, node, metric), + // Empty node or wrong query length. + Ok(_) => f32::INFINITY, + // Bytes not aligned to f32: `distance_typed` copies them out. + Err(_) => distance_typed( + metric, + VectorStorageDtype::F32, + bytemuck::cast_slice(query), + &node_bytes, + self.dim, + ) + .unwrap_or(f32::INFINITY), + }, + dtype => { + let query_bytes = cast_from_f32(query, dtype); + distance_typed(metric, dtype, &query_bytes, &node_bytes, self.dim) + .unwrap_or(f32::INFINITY) } } } diff --git a/nodedb-vector/src/navix/acorn.rs b/nodedb-vector/src/navix/acorn.rs index 2ea98cc76..7cd992044 100644 --- a/nodedb-vector/src/navix/acorn.rs +++ b/nodedb-vector/src/navix/acorn.rs @@ -16,7 +16,6 @@ mod inner { use roaring::RoaringBitmap; - use crate::distance::distance; use crate::error::{VectorError, check_dim}; use crate::hnsw::graph::{Candidate, HnswIndex}; use crate::navix::traversal::SearchResult; @@ -225,13 +224,12 @@ mod inner { let mut results: Vec = allowed .iter() .filter_map(|id| { - if index.is_deleted(id) { + if id as usize >= index.len() || index.is_deleted(id) { return None; } - let v = index.get_vector(id)?; Some(SearchResult { id, - distance: distance(query, v, metric), + distance: index.dist_to_node_f32(query, id, metric), }) }) .collect(); @@ -259,10 +257,7 @@ mod inner { node_id: u32, metric: nodedb_types::vector_distance::DistanceMetric, ) -> f32 { - match index.get_vector(node_id) { - Some(v) => distance(query, v, metric), - None => f32::INFINITY, - } + index.dist_to_node_f32(query, node_id, metric) } #[cfg(test)] diff --git a/nodedb-vector/src/navix/traversal.rs b/nodedb-vector/src/navix/traversal.rs index 016ad8b9d..699c349ab 100644 --- a/nodedb-vector/src/navix/traversal.rs +++ b/nodedb-vector/src/navix/traversal.rs @@ -14,7 +14,6 @@ use std::collections::{BinaryHeap, HashSet}; use roaring::RoaringBitmap; -use crate::distance::distance; use crate::error::{VectorError, check_dim}; use crate::hnsw::graph::{Candidate, HnswIndex}; use crate::navix::selectivity::{NavixHeuristic, local_selectivity_at, pick_heuristic}; @@ -132,13 +131,12 @@ fn brute_force_on_allowed( let mut results: Vec = allowed .iter() .filter_map(|id| { - if index.is_deleted(id) { + if id as usize >= index.len() || index.is_deleted(id) { return None; } - let v = index.get_vector(id)?; Some(SearchResult { id, - distance: distance(query, v, metric), + distance: index.dist_to_node_f32(query, id, metric), }) }) .collect(); @@ -470,10 +468,7 @@ fn dist( node_id: u32, metric: nodedb_types::vector_distance::DistanceMetric, ) -> f32 { - match index.get_vector(node_id) { - Some(v) => distance(query, v, metric), - None => f32::INFINITY, - } + index.dist_to_node_f32(query, node_id, metric) } // ── Tests ───────────────────────────────────────────────────────────────────── @@ -481,7 +476,7 @@ fn dist( #[cfg(test)] mod tests { use super::*; - use crate::distance::DistanceMetric; + use crate::distance::{DistanceMetric, distance}; use crate::hnsw::{HnswIndex, HnswParams}; fn build_index(n: usize) -> HnswIndex { From a5e245901a2ca69507ab705826842179b57e79c5 Mon Sep 17 00:00:00 2001 From: Farhan Syah Date: Mon, 12 Oct 2026 13:26:17 +0800 Subject: [PATCH 03/11] feat(metrics): count overload answers sent to clients by protocol nodedb_dispatch_capacity_busy_total measures dispatch pressure, including refusals the server waits out itself. Add nodedb_server_overload_responses_total{protocol} to measure client impact: each protocol counts at the point it writes the overload answer (pgwire 57P03, native 57P03 frame, HTTP 503 with Retry-After, RESP BUSY). A gap between the two is pressure the server absorbs. --- docs/architecture.md | 2 +- .../calvin/scheduler/driver/core/catch_up.rs | 10 ++ nodedb/src/control/metrics/mod.rs | 2 + .../src/control/metrics/overload_responses.rs | 146 ++++++++++++++++++ nodedb/src/control/metrics/system/render.rs | 1 + nodedb/src/control/server/http/auth.rs | 59 ++++++- nodedb/src/control/server/native/codec.rs | 45 ++++++ .../control/server/pgwire/wire_safe_error.rs | 43 ++++++ nodedb/src/control/server/resp/codec.rs | 32 ++++ 9 files changed, 332 insertions(+), 8 deletions(-) create mode 100644 nodedb/src/control/metrics/overload_responses.rs diff --git a/docs/architecture.md b/docs/architecture.md index cba0f607d..3dbcd2f99 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -200,7 +200,7 @@ Admission ordering for requests: 1. Identity & database access (established during authentication) 2. **Tenant** rate/concurrency check → `TENANT_QUOTA_EXCEEDED` 3. **Database** rate/concurrency check → `DATABASE_QUOTA_EXCEEDED` -4. **Dispatch** capacity check (tenant in-flight share, database fair share per core, core queue) → `SERVER_OVERLOAD` (code 1403, SQLSTATE `57P03`). Nothing was enqueued, so the client retries with backoff. `nodedb_dispatch_capacity_busy_total{scope}` counts each refusal. +4. **Dispatch** capacity check (tenant in-flight share, database fair share per core, core queue) → `SERVER_OVERLOAD` (code 1403, SQLSTATE `57P03`). Nothing was enqueued, so the client retries with backoff. `nodedb_dispatch_capacity_busy_total{scope}` counts each refusal (pressure). `nodedb_server_overload_responses_total{protocol}` counts each overload answer a client receives (impact). A gap between them is pressure the server absorbs. Memory reservation flips the order (largest scope first to fail fast): global → database → tenant → engine. Failure at any layer aborts before the next is consulted. diff --git a/nodedb/src/control/cluster/calvin/scheduler/driver/core/catch_up.rs b/nodedb/src/control/cluster/calvin/scheduler/driver/core/catch_up.rs index af8fa6780..3b5697e78 100644 --- a/nodedb/src/control/cluster/calvin/scheduler/driver/core/catch_up.rs +++ b/nodedb/src/control/cluster/calvin/scheduler/driver/core/catch_up.rs @@ -809,6 +809,11 @@ mod tests { .refusals(&tenant_scope) }; let busy_before = busy(); + let answers = || { + crate::control::metrics::OverloadProtocol::ALL + .map(|protocol| crate::control::metrics::overload_responses().sent(protocol)) + }; + let answers_before = answers(); let first = TxnId::new(0, 0); let second = TxnId::new(1, 0); @@ -846,6 +851,11 @@ mod tests { "each attempt refuses once" ); assert_eq!(busy() - busy_before, 1 + SATURATED_PASSES); + assert_eq!( + answers(), + answers_before, + "a refusal the scheduler absorbs reaches no client" + ); assert_eq!( metrics.intake_gate_closed_counts[intake_closure_reason::DEFERRED_DISPATCH] .load(Ordering::Relaxed), diff --git a/nodedb/src/control/metrics/mod.rs b/nodedb/src/control/metrics/mod.rs index 5bd4a42b4..3711e246c 100644 --- a/nodedb/src/control/metrics/mod.rs +++ b/nodedb/src/control/metrics/mod.rs @@ -2,6 +2,7 @@ pub mod database; pub mod histogram; +pub mod overload_responses; pub mod per_vshard; pub mod prometheus; pub mod purge; @@ -10,6 +11,7 @@ pub mod tenant; pub use database::{DatabaseCounters, DatabaseMetricsRegistry, DatabaseQuotaMetrics}; pub use histogram::AtomicHistogram; +pub use overload_responses::{OverloadProtocol, OverloadResponses, overload_responses}; pub use per_vshard::{PerVShardMetrics, PerVShardMetricsRegistry, VShardStatsSnapshot}; pub use purge::PurgeMetrics; pub use system::{CoreFailStopReport, CoreFailStops, CoreHeartbeats, SystemMetrics}; diff --git a/nodedb/src/control/metrics/overload_responses.rs b/nodedb/src/control/metrics/overload_responses.rs new file mode 100644 index 000000000..616e47841 --- /dev/null +++ b/nodedb/src/control/metrics/overload_responses.rs @@ -0,0 +1,146 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Process-wide count of overload-class answers that reached a client, by +//! protocol. +//! +//! Each protocol counts at the one place it turns an error into wire bytes, +//! so a capacity refusal the server waits out itself is never counted. The +//! counters are a process-wide static because those places hold no shared +//! state: the pgwire error handler, the native frame encoder, the HTTP +//! response builder and the RESP serializer. The label is the protocol only. + +use std::sync::atomic::{AtomicU64, Ordering}; + +/// The protocol that sent an overload answer. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OverloadProtocol { + /// SQLSTATE `57P03` in a pgwire `ErrorResponse`. + Pgwire, + /// SQLSTATE `57P03` in a native error frame. + Native, + /// `503` with `Retry-After` for a `57P03` error. + Http, + /// A `BUSY` error reply. + Resp, +} + +impl OverloadProtocol { + /// Every protocol, in exposition order. + pub const ALL: [Self; 4] = [Self::Pgwire, Self::Native, Self::Http, Self::Resp]; + + /// The Prometheus `protocol` label. + pub fn label(self) -> &'static str { + match self { + Self::Pgwire => "pgwire", + Self::Native => "native", + Self::Http => "http", + Self::Resp => "resp", + } + } + + fn index(self) -> usize { + match self { + Self::Pgwire => 0, + Self::Native => 1, + Self::Http => 2, + Self::Resp => 3, + } + } +} + +/// Overload answers sent since process start, one counter per protocol. +#[derive(Debug)] +pub struct OverloadResponses { + sent: [AtomicU64; OverloadProtocol::ALL.len()], +} + +static OVERLOAD_RESPONSES: OverloadResponses = OverloadResponses::new(); + +/// The process-wide overload answer counters. +pub fn overload_responses() -> &'static OverloadResponses { + &OVERLOAD_RESPONSES +} + +impl OverloadResponses { + pub const fn new() -> Self { + Self { + sent: [const { AtomicU64::new(0) }; OverloadProtocol::ALL.len()], + } + } + + /// Count one overload answer sent over `protocol`. + pub fn record(&self, protocol: OverloadProtocol) { + self.sent[protocol.index()].fetch_add(1, Ordering::Relaxed); + } + + /// Overload answers sent over `protocol`. + pub fn sent(&self, protocol: OverloadProtocol) -> u64 { + self.sent[protocol.index()].load(Ordering::Relaxed) + } + + /// Append the `nodedb_server_overload_responses_total` counter. + pub fn write_prometheus(&self, out: &mut String) { + use std::fmt::Write as _; + + let _ = out.write_str( + "# HELP nodedb_server_overload_responses_total Overload-class answers \ + (SQLSTATE 57P03, HTTP 503, RESP BUSY) that reached a client.\n\ + # TYPE nodedb_server_overload_responses_total counter\n", + ); + for protocol in OverloadProtocol::ALL { + let _ = writeln!( + out, + r#"nodedb_server_overload_responses_total{{protocol="{}"}} {}"#, + protocol.label(), + self.sent(protocol) + ); + } + } +} + +impl Default for OverloadResponses { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn each_protocol_counts_apart() { + let responses = OverloadResponses::new(); + responses.record(OverloadProtocol::Http); + responses.record(OverloadProtocol::Http); + responses.record(OverloadProtocol::Resp); + + assert_eq!(responses.sent(OverloadProtocol::Pgwire), 0); + assert_eq!(responses.sent(OverloadProtocol::Native), 0); + assert_eq!(responses.sent(OverloadProtocol::Http), 2); + assert_eq!(responses.sent(OverloadProtocol::Resp), 1); + } + + #[test] + fn exposition_has_one_series_per_protocol() { + let responses = OverloadResponses::new(); + responses.record(OverloadProtocol::Native); + let mut out = String::new(); + responses.write_prometheus(&mut out); + + assert!( + out.contains("# TYPE nodedb_server_overload_responses_total counter"), + "{out}" + ); + for (label, count) in [("pgwire", 0), ("native", 1), ("http", 0), ("resp", 0)] { + let series = + format!(r#"nodedb_server_overload_responses_total{{protocol="{label}"}} {count}"#); + assert!(out.contains(&series), "{out}"); + } + assert_eq!( + out.lines().count(), + 2 + OverloadProtocol::ALL.len(), + "{out}" + ); + } +} diff --git a/nodedb/src/control/metrics/system/render.rs b/nodedb/src/control/metrics/system/render.rs index 7d75d0e8b..2c37a8549 100644 --- a/nodedb/src/control/metrics/system/render.rs +++ b/nodedb/src/control/metrics/system/render.rs @@ -19,6 +19,7 @@ impl SystemMetrics { self.io_metrics.write_prometheus(&mut out); self.spsc_throttle.write_prometheus(&mut out); self.dispatch_capacity.write_prometheus(&mut out); + super::super::overload_responses::overload_responses().write_prometheus(&mut out); self.core_fail_stops.write_prometheus(&mut out); out } diff --git a/nodedb/src/control/server/http/auth.rs b/nodedb/src/control/server/http/auth.rs index 8cf85c51c..68d3390dc 100644 --- a/nodedb/src/control/server/http/auth.rs +++ b/nodedb/src/control/server/http/auth.rs @@ -356,7 +356,7 @@ impl IntoResponse for ApiError { ApiError::Overloaded { message, code } => { let body = HttpError::with_code(message, code.to_string()); let resp = (StatusCode::SERVICE_UNAVAILABLE, axum::Json(body)).into_response(); - with_retry_after(resp, OVERLOAD_RETRY_AFTER_SECS) + overloaded(resp) } ApiError::Coded { status, @@ -364,18 +364,14 @@ impl IntoResponse for ApiError { code, cause, } => { - let overloaded = numeric_code_to_sqlstate(code) == sqlstate::SERVER_OVERLOAD; + let is_overload = numeric_code_to_sqlstate(code) == sqlstate::SERVER_OVERLOAD; let body = HttpError::with_code(message, code.to_string()); let body = match cause { Some(cause) => body.caused_by(&cause), None => body, }; let resp = (status, axum::Json(body)).into_response(); - if overloaded { - with_retry_after(resp, OVERLOAD_RETRY_AFTER_SECS) - } else { - resp - } + if is_overload { overloaded(resp) } else { resp } } ApiError::Unauthorized(msg) => plain(StatusCode::UNAUTHORIZED, msg), ApiError::Forbidden(msg) => plain(StatusCode::FORBIDDEN, msg), @@ -392,6 +388,14 @@ impl IntoResponse for ApiError { /// Seconds a client waits before it retries a `57P03` capacity refusal. const OVERLOAD_RETRY_AFTER_SECS: u64 = 1; +/// Finish a `57P03` answer: add `Retry-After` and count it under `http`. +/// Every overload answer HTTP sends passes through here. +fn overloaded(resp: Response) -> Response { + crate::control::metrics::overload_responses() + .record(crate::control::metrics::OverloadProtocol::Http); + with_retry_after(resp, OVERLOAD_RETRY_AFTER_SECS) +} + /// Set the `Retry-After` header of `resp` to `secs`. fn with_retry_after(mut resp: Response, secs: u64) -> Response { resp.headers_mut().insert( @@ -730,6 +734,47 @@ mod tests { assert_eq!(retry_after(&internal), None); } + /// Each overload answer counts once under `http`. Another answer does + /// not count. + #[test] + fn an_overload_answer_counts_under_http() { + use crate::control::metrics::{OverloadProtocol, overload_responses}; + let before = overload_responses().sent(OverloadProtocol::Http); + + let _ = ApiError::from(crate::Error::DispatchCapacity { + scope: crate::DispatchCapacityScope::QueueFull { + core_id: 0, + capacity: 4, + }, + }) + .into_response(); + assert_eq!( + overload_responses().sent(OverloadProtocol::Http), + before + 1 + ); + + let _ = ApiError::Coded { + status: StatusCode::SERVICE_UNAVAILABLE, + message: "busy".into(), + code: nodedb_types::error::ErrorCode::SERVER_OVERLOAD, + cause: None, + } + .into_response(); + assert_eq!( + overload_responses().sent(OverloadProtocol::Http), + before + 2 + ); + + let _ = ApiError::from(crate::Error::Internal { + detail: "boom".into(), + }) + .into_response(); + assert_eq!( + overload_responses().sent(OverloadProtocol::Http), + before + 2 + ); + } + #[test] fn x_on_deny_malformed_value_is_ignored() { let context = auth_context(); diff --git a/nodedb/src/control/server/native/codec.rs b/nodedb/src/control/server/native/codec.rs index c44667807..9c0933278 100644 --- a/nodedb/src/control/server/native/codec.rs +++ b/nodedb/src/control/server/native/codec.rs @@ -120,7 +120,24 @@ pub fn decode_request(payload: &[u8], format: FrameFormat) -> crate::Result crate::Result> { + let bytes = encode_frame(resp, format)?; + if resp + .error + .as_ref() + .is_some_and(|error| error.code == nodedb_types::error::sqlstate::SERVER_OVERLOAD) + { + crate::control::metrics::overload_responses() + .record(crate::control::metrics::OverloadProtocol::Native); + } + Ok(bytes) +} + +/// Serialize `resp` in `format`. +fn encode_frame(resp: &NativeResponse, format: FrameFormat) -> crate::Result> { match format { FrameFormat::Json => sonic_rs::to_vec(resp).map_err(|e| crate::Error::Serialization { format: "json".into(), @@ -206,6 +223,34 @@ mod tests { assert!(s.contains("\"seq\":1")); } + /// A `57P03` error frame counts once under `native`. Another error frame + /// does not count. + #[test] + fn an_overload_frame_counts_under_native() { + use crate::control::metrics::{OverloadProtocol, overload_responses}; + let before = overload_responses().sent(OverloadProtocol::Native); + let refusal = crate::Error::DispatchCapacity { + scope: crate::DispatchCapacityScope::QueueFull { + core_id: 0, + capacity: 4, + }, + }; + + let frame = crate::control::server::native::dispatch::error_to_native(1, &refusal); + encode_response(&frame, FrameFormat::MessagePack).expect("frame encodes"); + assert_eq!( + overload_responses().sent(OverloadProtocol::Native), + before + 1 + ); + + let other = NativeResponse::error(2, "XX000", "boom"); + encode_response(&other, FrameFormat::Json).expect("frame encodes"); + assert_eq!( + overload_responses().sent(OverloadProtocol::Native), + before + 1 + ); + } + #[test] fn response_encode_msgpack() { let resp = NativeResponse::from_query_result( diff --git a/nodedb/src/control/server/pgwire/wire_safe_error.rs b/nodedb/src/control/server/pgwire/wire_safe_error.rs index 0c414d3fc..54613c050 100644 --- a/nodedb/src/control/server/pgwire/wire_safe_error.rs +++ b/nodedb/src/control/server/pgwire/wire_safe_error.rs @@ -126,10 +126,19 @@ impl ErrorHandler for WireSafeErrorHandler { // bytes. if let PgWireError::UserError(info) = error { sanitize_error_info(info); + count_overload(info); } } } +/// Count an error the client receives as the overload class (`57P03`). +fn count_overload(info: &ErrorInfo) { + if info.code == nodedb_types::error::sqlstate::SERVER_OVERLOAD { + crate::control::metrics::overload_responses() + .record(crate::control::metrics::OverloadProtocol::Pgwire); + } +} + #[cfg(test)] mod tests { use super::*; @@ -382,4 +391,38 @@ mod tests { "no field value may carry a NUL onto the wire: {fields:?}" ); } + + fn user_error_info(err: &crate::Error) -> ErrorInfo { + match crate::control::server::pgwire::types::error_to_pg(err) { + PgWireError::UserError(info) => *info, + other => panic!("error_to_pg builds a user error, got {other:?}"), + } + } + + /// A capacity refusal sent to the client counts once under `pgwire`. + /// Another class does not count. + #[test] + fn an_overload_answer_counts_under_pgwire() { + use crate::control::metrics::{OverloadProtocol, overload_responses}; + let before = overload_responses().sent(OverloadProtocol::Pgwire); + + count_overload(&user_error_info(&crate::Error::DispatchCapacity { + scope: crate::DispatchCapacityScope::QueueFull { + core_id: 0, + capacity: 4, + }, + })); + assert_eq!( + overload_responses().sent(OverloadProtocol::Pgwire), + before + 1 + ); + + count_overload(&user_error_info(&crate::Error::Internal { + detail: "boom".into(), + })); + assert_eq!( + overload_responses().sent(OverloadProtocol::Pgwire), + before + 1 + ); + } } diff --git a/nodedb/src/control/server/resp/codec.rs b/nodedb/src/control/server/resp/codec.rs index 27e4e8276..c127bb1d9 100644 --- a/nodedb/src/control/server/resp/codec.rs +++ b/nodedb/src/control/server/resp/codec.rs @@ -103,6 +103,12 @@ impl RespValue { buf.extend_from_slice(b"\r\n"); } Self::Error(s) => { + // Every RESP error reply is written here, so a `BUSY` overload + // reply counts here once. + if s.starts_with("BUSY ") { + crate::control::metrics::overload_responses() + .record(crate::control::metrics::OverloadProtocol::Resp); + } buf.push(b'-'); buf.extend_from_slice(s.as_bytes()); buf.extend_from_slice(b"\r\n"); @@ -425,6 +431,32 @@ mod tests { ); } + /// A `BUSY` reply counts once under `resp`, nested replies included. + /// Another error reply does not count. + #[test] + fn a_busy_reply_counts_under_resp() { + use crate::control::metrics::{OverloadProtocol, overload_responses}; + let before = overload_responses().sent(OverloadProtocol::Resp); + let busy = RespValue::from_error(&crate::Error::DispatchCapacity { + scope: crate::DispatchCapacityScope::QueueFull { + core_id: 0, + capacity: 4, + }, + }); + + let _ = busy.to_bytes(); + assert_eq!( + overload_responses().sent(OverloadProtocol::Resp), + before + 1 + ); + + let _ = RespValue::array(vec![busy, RespValue::err("ERR unknown command")]).to_bytes(); + assert_eq!( + overload_responses().sent(OverloadProtocol::Resp), + before + 2 + ); + } + #[test] fn serialize_integer() { assert_eq!(RespValue::integer(42).to_bytes(), b":42\r\n"); From 1e5929b25de573d2b42f995c9ba7c3e931c15d32 Mon Sep 17 00:00:00 2001 From: Farhan Syah Date: Mon, 12 Oct 2026 13:26:18 +0800 Subject: [PATCH 04/11] fix(client): compile the search-hit decoder only with a client feature search_hit serves only the native and remote clients. Gate it on those features so a build with neither feature has no dead code. Add decoder tests for wrong cell types and a per-source score column. --- nodedb-client/src/lib.rs | 6 ++-- nodedb-client/src/row_decode/mod.rs | 15 +++++---- nodedb-client/src/row_decode/search_hit.rs | 39 ++++++++++++++++++++++ 3 files changed, 51 insertions(+), 9 deletions(-) diff --git a/nodedb-client/src/lib.rs b/nodedb-client/src/lib.rs index 68b6cd610..4b7c6ab19 100644 --- a/nodedb-client/src/lib.rs +++ b/nodedb-client/src/lib.rs @@ -14,9 +14,9 @@ pub mod capabilities; pub mod traits; -/// Shared row decoders used by both the trait default impls and the -/// feature-gated clients. Feature-agnostic on purpose — one parser per -/// row shape regardless of which transport delivered the row. +/// Shared row decoders used by the trait default impls and the feature-gated +/// clients. One parser per row shape, whichever transport delivered the row. +/// A decoder used only by the clients compiles under their features. mod row_decode; /// Shared SQL escaping helpers (string-literal and identifier quoting). diff --git a/nodedb-client/src/row_decode/mod.rs b/nodedb-client/src/row_decode/mod.rs index 282d1f17c..fd6bee2d4 100644 --- a/nodedb-client/src/row_decode/mod.rs +++ b/nodedb-client/src/row_decode/mod.rs @@ -1,16 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 -//! Feature-agnostic decoders that turn `QueryResult` rows back into typed -//! domain structs. +//! Decoders that turn `QueryResult` rows back into typed domain structs. //! -//! Lives at the crate root so both the (feature-agnostic) `NodeDb` trait -//! default impls and the feature-gated `remote` / `native` clients share -//! one parser per row shape — there is exactly one place to fix when a -//! system-catalog column layout changes. +//! Each row shape has one parser, shared by the `NodeDb` trait default impls +//! and the feature-gated `remote` / `native` clients. A system-catalog column +//! layout change has one place to fix. +//! +//! A decoder compiles under the features of its callers. `search_hit` serves +//! only the client search methods, so it compiles with `native` or `remote`. pub(crate) mod dropped_collection; +#[cfg(any(feature = "native", feature = "remote"))] pub(crate) mod search_hit; pub(crate) mod value; pub(crate) use dropped_collection::parse_dropped_collection_rows; +#[cfg(any(feature = "native", feature = "remote"))] pub(crate) use search_hit::{HitSource, decode_search_hits}; diff --git a/nodedb-client/src/row_decode/search_hit.rs b/nodedb-client/src/row_decode/search_hit.rs index c14c5e0b1..8e5231771 100644 --- a/nodedb-client/src/row_decode/search_hit.rs +++ b/nodedb-client/src/row_decode/search_hit.rs @@ -193,4 +193,43 @@ mod tests { let null_score = vec![vec![Value::String("a".into()), Value::Null]]; assert!(decode_search_hits(&SOURCE, &columns, &null_score).is_err()); } + + #[test] + fn a_wrong_cell_type_is_an_error_naming_the_source() { + let columns = names(&["id", "distance"]); + let bool_id = vec![vec![Value::Bool(true), Value::Float(1.0)]]; + let err = decode_search_hits(&SOURCE, &columns, &bool_id).expect_err("bool id"); + assert!(err.to_string().contains("vector_search 'c'"), "{err}"); + assert!(err.to_string().contains("row 0"), "{err}"); + let bool_score = vec![ + vec![Value::String("a".into()), Value::Float(1.0)], + vec![Value::String("b".into()), Value::Bool(false)], + ]; + let err = decode_search_hits(&SOURCE, &columns, &bool_score).expect_err("bool score"); + assert!( + err.to_string().contains("row 1 ('b') has distance"), + "{err}" + ); + } + + #[test] + fn the_score_column_comes_from_the_source() { + let text = HitSource { + op: "text_search", + collection: "docs", + score_column: "score", + }; + let columns = names(&["id", "distance", "score"]); + let rows = vec![vec![ + Value::String("a".into()), + Value::Float(9.0), + Value::Float(0.75), + ]]; + let hits = decode_search_hits(&text, &columns, &rows).expect("decode"); + assert_eq!(hits[0].distance, 0.75); + let err = decode_search_hits(&text, &names(&["id", "distance"]), &rows) + .expect_err("no score column"); + assert!(err.to_string().contains("text_search 'docs'"), "{err}"); + assert!(err.to_string().contains("'score'"), "{err}"); + } } From 83a32a8a8d95da88e0e73384a2cac378772d836d Mon Sep 17 00:00:00 2001 From: Farhan Syah Date: Mon, 12 Oct 2026 13:26:21 +0800 Subject: [PATCH 05/11] fix(kv): decode window timestamps from the sorted index sort key Windowed rank, top-k, range and count now read each entry's timestamp from the timestamp column of its sort key instead of a separate lookup. Integer, float and ISO 8601 string columns decode to a timestamp. An entry with no decodable timestamp is hidden. WindowConfig::bounds returns a half-open [start, end) range, so a custom window admits nothing outside its range. Backfill extracts bool, null and composite values with the same bytes as the write path. --- nodedb/src/engine/kv/sorted_index/index.rs | 299 +++++++++++++++++- nodedb/src/engine/kv/sorted_index/key.rs | 83 +++-- nodedb/src/engine/kv/sorted_index/mod.rs | 4 +- .../src/engine/kv/sorted_index/timestamp.rs | 123 +++++++ nodedb/src/engine/kv/sorted_index/window.rs | 142 +++++++-- .../engine/kv/sorted_index/windowed_query.rs | 103 ++---- 6 files changed, 634 insertions(+), 120 deletions(-) create mode 100644 nodedb/src/engine/kv/sorted_index/timestamp.rs diff --git a/nodedb/src/engine/kv/sorted_index/index.rs b/nodedb/src/engine/kv/sorted_index/index.rs index 4b73880d2..87ff788f1 100644 --- a/nodedb/src/engine/kv/sorted_index/index.rs +++ b/nodedb/src/engine/kv/sorted_index/index.rs @@ -150,6 +150,9 @@ pub(super) fn extract_sort_key_from_value( } /// Convert a JSON field value to sortable bytes. +/// +/// The bytes match what the KV write path extracts for `on_put`. A row +/// backfilled here and the same row written later get the same sort key. fn field_value_to_sort_bytes(val: &serde_json::Value) -> Vec { match val { serde_json::Value::Number(n) => { @@ -162,6 +165,300 @@ fn field_value_to_sort_bytes(val: &serde_json::Value) -> Vec { } } serde_json::Value::String(s) => s.as_bytes().to_vec(), - _ => Vec::new(), + serde_json::Value::Bool(b) => vec![u8::from(*b)], + serde_json::Value::Null => Vec::new(), + other @ (serde_json::Value::Array(_) | serde_json::Value::Object(_)) => { + other.to_string().into_bytes() + } + } +} + +#[cfg(test)] +mod tests { + use super::super::key::{SortColumn, SortDirection}; + use super::super::window::WindowConfig; + use super::*; + + /// 2024-05-02 15:36:40 UTC, a Thursday. + const NOW_MS: u64 = 1_714_664_200_000; + /// 2024-05-02 00:00:00 UTC. + const DAY_START_MS: u64 = 1_714_608_000_000; + /// Monday 2024-04-29 00:00:00 UTC. + const WEEK_START_MS: u64 = 1_714_348_800_000; + /// 2024-05-01 00:00:00 UTC. + const MONTH_START_MS: u64 = 1_714_521_600_000; + + fn def(window: WindowConfig, ts_direction: SortDirection) -> SortedIndexDef { + SortedIndexDef { + name: "lb".into(), + collection: "scores".into(), + key_column: "player".into(), + encoder: SortKeyEncoder::new(vec![ + SortColumn { + name: "score".into(), + direction: SortDirection::Desc, + }, + SortColumn { + name: "ts".into(), + direction: ts_direction, + }, + ]), + window, + } + } + + /// A row with `ts` set to `ts`, or with no `ts` field when `None`. + fn row(player: &str, score: i64, ts: Option) -> (Vec, Vec) { + let mut doc = serde_json::json!({ "player": player, "score": score }); + if let Some(ts) = ts { + doc["ts"] = ts; + } + let value = nodedb_types::json_to_msgpack(&doc).expect("encode row"); + (player.as_bytes().to_vec(), value) + } + + fn ms(ts: u64) -> Option { + Some(serde_json::json!(ts)) + } + + fn build(def: SortedIndexDef, rows: Vec<(Vec, Vec)>) -> SortedIndex { + SortedIndex::build(def, rows.into_iter()).0 + } + + fn pks(entries: &[(u32, Vec)]) -> Vec<(u32, &str)> { + entries + .iter() + .map(|(rank, pk)| (*rank, std::str::from_utf8(pk).expect("utf-8 pk"))) + .collect() + } + + /// `old` has the top score but a timestamp before the window start. + /// `fresh` and `late` are inside the window. Every windowed read must + /// hide `old` and rank the other two as if `old` did not exist. + fn assert_hides_old_keeps_fresh(idx: &SortedIndex, now_ms: u64) { + assert_eq!(idx.count(now_ms), 2); + assert_eq!(idx.rank(b"old", now_ms), None); + assert_eq!(idx.rank(b"fresh", now_ms), Some(1)); + assert_eq!(idx.rank(b"late", now_ms), Some(2)); + assert_eq!(pks(&idx.top_k(10, now_ms)), vec![(1, "fresh"), (2, "late")]); + assert_eq!(pks(&idx.top_k(1, now_ms)), vec![(1, "fresh")]); + assert_eq!( + pks(&idx.range(None, None, now_ms)), + vec![(1, "fresh"), (2, "late")] + ); + let score_200 = SortKeyEncoder::encode_i64(200); + let score_300 = SortKeyEncoder::encode_i64(300); + assert_eq!( + pks(&idx.range(Some(&score_200), Some(&score_300), now_ms)), + vec![(1, "fresh")] + ); + } + + fn window_rows(start_ms: u64) -> Vec<(Vec, Vec)> { + vec![ + row("old", 300, ms(start_ms - 1)), + row("fresh", 200, ms(start_ms)), + row("late", 100, ms(NOW_MS)), + ] + } + + #[test] + fn daily_window_hides_entries_before_today_utc_start() { + let idx = build( + def(WindowConfig::daily("ts"), SortDirection::Asc), + window_rows(DAY_START_MS), + ); + assert_hides_old_keeps_fresh(&idx, NOW_MS); + } + + #[test] + fn daily_window_hides_entries_from_days_ago() { + let idx = build( + def(WindowConfig::daily("ts"), SortDirection::Asc), + vec![ + row("old", 300, ms(DAY_START_MS - 30 * 86_400_000)), + row("fresh", 200, ms(DAY_START_MS + 1)), + row("late", 100, ms(NOW_MS)), + ], + ); + assert_hides_old_keeps_fresh(&idx, NOW_MS); + } + + #[test] + fn weekly_window_starts_at_monday_utc() { + let idx = build( + def(WindowConfig::weekly("ts"), SortDirection::Asc), + window_rows(WEEK_START_MS), + ); + assert_hides_old_keeps_fresh(&idx, NOW_MS); + } + + #[test] + fn monthly_window_starts_at_the_first_of_the_month_utc() { + let idx = build( + def(WindowConfig::monthly("ts"), SortDirection::Asc), + window_rows(MONTH_START_MS), + ); + assert_hides_old_keeps_fresh(&idx, NOW_MS); + } + + #[test] + fn desc_timestamp_column_windows_the_same_as_asc() { + let idx = build( + def(WindowConfig::daily("ts"), SortDirection::Desc), + window_rows(DAY_START_MS), + ); + assert_hides_old_keeps_fresh(&idx, NOW_MS); + } + + #[test] + fn float_timestamp_column_windows_by_its_value() { + let start = DAY_START_MS as f64; + let idx = build( + def(WindowConfig::daily("ts"), SortDirection::Asc), + vec![ + row("old", 300, Some(serde_json::json!(start - 0.5))), + row("fresh", 200, Some(serde_json::json!(start + 0.5))), + row("late", 100, Some(serde_json::json!(NOW_MS as f64 + 0.25))), + ], + ); + assert_hides_old_keeps_fresh(&idx, NOW_MS); + } + + #[test] + fn float_timestamp_column_with_desc_direction_windows_by_its_value() { + let start = DAY_START_MS as f64; + let idx = build( + def(WindowConfig::daily("ts"), SortDirection::Desc), + vec![ + row("old", 300, Some(serde_json::json!(start - 0.5))), + row("fresh", 200, Some(serde_json::json!(start + 0.5))), + row("late", 100, Some(serde_json::json!(NOW_MS as f64 + 0.25))), + ], + ); + assert_hides_old_keeps_fresh(&idx, NOW_MS); + } + + #[test] + fn iso_8601_string_timestamp_column_windows_by_its_instant() { + let idx = build( + def(WindowConfig::daily("ts"), SortDirection::Asc), + vec![ + row("old", 300, Some(serde_json::json!("2024-05-01T23:59:59Z"))), + row( + "fresh", + 200, + Some(serde_json::json!("2024-05-02T00:00:00Z")), + ), + row("late", 100, Some(serde_json::json!("2024-05-02T15:00:00Z"))), + ], + ); + assert_hides_old_keeps_fresh(&idx, NOW_MS); + } + + #[test] + fn custom_window_shows_nothing_when_now_is_outside_it() { + let (start, end) = (DAY_START_MS, DAY_START_MS + 3_600_000); + let idx = build( + def(WindowConfig::custom("ts", start, end), SortDirection::Asc), + vec![ + row("a", 300, ms(start)), + row("b", 200, ms(start + 1)), + row("c", 100, ms(end + 1)), + ], + ); + for now_ms in [start - 1, end, end + 86_400_000] { + assert_eq!(idx.count(now_ms), 0); + assert!(idx.top_k(10, now_ms).is_empty()); + assert!(idx.range(None, None, now_ms).is_empty()); + for pk in [&b"a"[..], &b"b"[..], &b"c"[..]] { + assert_eq!(idx.rank(pk, now_ms), None); + } + } + } + + #[test] + fn custom_window_shows_only_entries_inside_its_range() { + let (start, end) = (DAY_START_MS, DAY_START_MS + 3_600_000); + let idx = build( + def(WindowConfig::custom("ts", start, end), SortDirection::Asc), + vec![ + row("before", 400, ms(start - 1)), + row("first", 300, ms(start)), + row("last", 200, ms(end - 1)), + row("after", 100, ms(end)), + ], + ); + let now_ms = start + 60_000; + assert_eq!(idx.count(now_ms), 2); + assert_eq!(idx.rank(b"before", now_ms), None); + assert_eq!(idx.rank(b"first", now_ms), Some(1)); + assert_eq!(idx.rank(b"last", now_ms), Some(2)); + assert_eq!(idx.rank(b"after", now_ms), None); + assert_eq!(pks(&idx.top_k(10, now_ms)), vec![(1, "first"), (2, "last")]); + assert_eq!( + pks(&idx.range(None, None, now_ms)), + vec![(1, "first"), (2, "last")] + ); + } + + #[test] + fn entry_with_a_null_timestamp_is_outside_every_window() { + let idx = build( + def(WindowConfig::daily("ts"), SortDirection::Asc), + vec![ + row("null_ts", 300, Some(serde_json::Value::Null)), + row("fresh", 200, ms(DAY_START_MS)), + ], + ); + assert_eq!(idx.count(NOW_MS), 1); + assert_eq!(idx.rank(b"null_ts", NOW_MS), None); + assert_eq!(idx.rank(b"fresh", NOW_MS), Some(1)); + assert_eq!(pks(&idx.top_k(10, NOW_MS)), vec![(1, "fresh")]); + assert_eq!(pks(&idx.range(None, None, NOW_MS)), vec![(1, "fresh")]); + } + + #[test] + fn entry_without_a_timestamp_field_is_not_indexed() { + let (idx, indexed) = SortedIndex::build( + def(WindowConfig::daily("ts"), SortDirection::Asc), + vec![row("no_ts", 300, None), row("fresh", 200, ms(DAY_START_MS))].into_iter(), + ); + assert_eq!(indexed, 1); + assert_eq!(idx.count(NOW_MS), 1); + assert_eq!(idx.rank(b"no_ts", NOW_MS), None); + assert_eq!(idx.rank(b"fresh", NOW_MS), Some(1)); + } + + #[test] + fn top_zero_of_a_windowed_index_is_empty() { + let idx = build( + def(WindowConfig::daily("ts"), SortDirection::Asc), + window_rows(DAY_START_MS), + ); + assert!(idx.top_k(0, NOW_MS).is_empty()); + } + + #[test] + fn unwindowed_index_shows_every_entry() { + let idx = build( + def(WindowConfig::none(), SortDirection::Asc), + window_rows(DAY_START_MS), + ); + assert_eq!(idx.count(NOW_MS), 3); + assert_eq!(idx.rank(b"old", NOW_MS), Some(1)); + } + + #[test] + fn backfilled_bool_sorts_like_a_written_bool() { + assert_eq!( + field_value_to_sort_bytes(&serde_json::json!(true)), + vec![1u8] + ); + assert_eq!( + field_value_to_sort_bytes(&serde_json::json!(false)), + vec![0u8] + ); + assert!(field_value_to_sort_bytes(&serde_json::Value::Null).is_empty()); } } diff --git a/nodedb/src/engine/kv/sorted_index/key.rs b/nodedb/src/engine/kv/sorted_index/key.rs index 27103dd69..289a1f5f4 100644 --- a/nodedb/src/engine/kv/sorted_index/key.rs +++ b/nodedb/src/engine/kv/sorted_index/key.rs @@ -12,6 +12,8 @@ //! Columns are separated by a length prefix (4 bytes, big-endian) to avoid //! ambiguity between field boundaries. +use std::borrow::Cow; + /// Sort direction for a column in a sorted index. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum SortDirection { @@ -182,16 +184,26 @@ impl SortKeyEncoder { f64::from_bits(bits) } - /// Encode a millisecond timestamp as big-endian bytes. + /// The value bytes of column `name` in `sort_key`, as they were passed to + /// [`Self::encode`]. /// - /// Timestamps are u64, so they naturally sort ascending in big-endian. - pub fn encode_timestamp_ms(ts_ms: u64) -> [u8; 8] { - ts_ms.to_be_bytes() - } - - /// Decode big-endian bytes back to a millisecond timestamp. - pub fn decode_timestamp_ms(bytes: &[u8; 8]) -> u64 { - u64::from_be_bytes(*bytes) + /// Strips the length prefix and undoes the DESC complement. Returns `None` + /// when the encoder has no column `name` or the key is truncated. + pub fn column_value<'k>(&self, sort_key: &'k [u8], name: &str) -> Option> { + let mut rest = sort_key; + for col in &self.columns { + let (len_bytes, tail) = rest.split_first_chunk::<4>()?; + let len = u32::from_be_bytes(*len_bytes) as usize; + let (value, tail) = tail.split_at_checked(len)?; + if col.name == name { + return Some(match col.direction { + SortDirection::Asc => Cow::Borrowed(value), + SortDirection::Desc => Cow::Owned(value.iter().map(|b| !b).collect()), + }); + } + rest = tail; + } + None } } @@ -274,8 +286,8 @@ mod tests { // Higher score should sort first (lower bytes due to DESC complement). let score_100 = SortKeyEncoder::encode_i64(100); let score_200 = SortKeyEncoder::encode_i64(200); - let time_early = SortKeyEncoder::encode_timestamp_ms(1000); - let time_late = SortKeyEncoder::encode_timestamp_ms(2000); + let time_early = SortKeyEncoder::encode_i64(1000); + let time_late = SortKeyEncoder::encode_i64(2000); let key_200_early = encoder.encode(&[&score_200, &time_early]); let key_100_early = encoder.encode(&[&score_100, &time_early]); @@ -304,15 +316,48 @@ mod tests { assert!(key_50 < key_150); } + fn score_and_ts(ts_direction: SortDirection) -> SortKeyEncoder { + SortKeyEncoder::new(vec![ + SortColumn { + name: "score".into(), + direction: SortDirection::Desc, + }, + SortColumn { + name: "ts".into(), + direction: ts_direction, + }, + ]) + } + + #[test] + fn column_value_returns_the_bytes_passed_to_encode() { + for direction in [SortDirection::Asc, SortDirection::Desc] { + let enc = score_and_ts(direction); + let score = SortKeyEncoder::encode_i64(7); + let ts = SortKeyEncoder::encode_i64(1_714_608_000_000); + let key = enc.encode(&[&score, &ts]); + + assert_eq!(enc.column_value(&key, "ts").as_deref(), Some(&ts[..])); + assert_eq!(enc.column_value(&key, "score").as_deref(), Some(&score[..])); + } + } + + #[test] + fn column_value_of_an_empty_value_is_empty() { + let enc = score_and_ts(SortDirection::Asc); + let key = enc.encode(&[&SortKeyEncoder::encode_i64(7), b""]); + assert_eq!(enc.column_value(&key, "ts").as_deref(), Some(&b""[..])); + } + #[test] - fn timestamp_roundtrip_and_ordering() { - let ts1 = 1714600000000u64; - let ts2 = 1714600001000u64; - let e1 = SortKeyEncoder::encode_timestamp_ms(ts1); - let e2 = SortKeyEncoder::encode_timestamp_ms(ts2); - assert_eq!(SortKeyEncoder::decode_timestamp_ms(&e1), ts1); - assert_eq!(SortKeyEncoder::decode_timestamp_ms(&e2), ts2); - assert!(e1 < e2); + fn column_value_of_an_unknown_column_or_truncated_key_is_none() { + let enc = score_and_ts(SortDirection::Asc); + let ts = SortKeyEncoder::encode_i64(1_714_608_000_000); + let key = enc.encode(&[&SortKeyEncoder::encode_i64(7), &ts]); + + assert!(enc.column_value(&key, "missing").is_none()); + assert!(enc.column_value(&key[..key.len() - 1], "ts").is_none()); + assert!(enc.column_value(&key[..2], "score").is_none()); } fn encoder(direction: SortDirection) -> SortKeyEncoder { diff --git a/nodedb/src/engine/kv/sorted_index/mod.rs b/nodedb/src/engine/kv/sorted_index/mod.rs index 2d94a0a5b..c6aaf1268 100644 --- a/nodedb/src/engine/kv/sorted_index/mod.rs +++ b/nodedb/src/engine/kv/sorted_index/mod.rs @@ -4,6 +4,7 @@ pub mod checkpoint; pub mod index; pub mod key; pub mod manager; +pub mod timestamp; pub mod tree; pub mod window; mod windowed_query; @@ -12,5 +13,6 @@ pub use checkpoint::SortedIndexSnapshot; pub use index::SortedIndex; pub use key::{SortDirection, SortKeyEncoder}; pub use manager::SortedIndexManager; +pub use timestamp::ColumnTimestamp; pub use tree::OrderStatTree; -pub use window::{WindowConfig, WindowType}; +pub use window::{WindowBounds, WindowConfig, WindowType}; diff --git a/nodedb/src/engine/kv/sorted_index/timestamp.rs b/nodedb/src/engine/kv/sorted_index/timestamp.rs new file mode 100644 index 000000000..ec574d8b8 --- /dev/null +++ b/nodedb/src/engine/kv/sorted_index/timestamp.rs @@ -0,0 +1,123 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Read a window timestamp back from the value bytes of one sort column. +//! +//! A sort column stores the bytes the KV value extractor wrote for it: +//! +//! - an integer: [`SortKeyEncoder::encode_i64`], 8 bytes; +//! - a float: [`SortKeyEncoder::encode_f64`], 8 bytes; +//! - a string: its UTF-8 bytes; +//! - null: no bytes. +//! +//! The bytes carry no type tag. A non-negative integer and a non-negative +//! float both encode with the top bit set, so the payload under that bit +//! picks the decoder. An integer timestamp of 2^53 ms is past the year +//! 287,000. A non-negative float whose bit pattern is below 2^53 is below +//! 2^-1021 ms. Payloads below 2^53 decode as integers, the rest as floats. +//! The two ranges overlap only on values no timestamp column holds. + +use nodedb_types::NdbDateTime; + +use super::key::SortKeyEncoder; + +/// Top bit of an encoded 8-byte number. Set for every non-negative value. +const SIGN_BIT: u64 = 1 << 63; + +/// Payloads below this bound decode as integer milliseconds. +const INTEGER_PAYLOAD_LIMIT: u64 = 1 << 53; + +/// A timestamp read from a sort column, in milliseconds since the Unix epoch. +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum ColumnTimestamp { + /// Whole milliseconds, from an integer or an ISO 8601 string column. + Millis(i64), + /// Fractional milliseconds, from a float column. + FractionalMillis(f64), +} + +/// Decode the value bytes of a timestamp column. +/// +/// Returns `None` for null (no bytes) and for a string that is not ISO 8601. +pub fn decode_column_timestamp(value: &[u8]) -> Option { + if let Ok(bytes) = <&[u8; 8]>::try_from(value) { + return Some(decode_number(bytes)); + } + let text = std::str::from_utf8(value).ok()?; + NdbDateTime::parse(text).map(|instant| ColumnTimestamp::Millis(instant.unix_millis())) +} + +/// Decode an 8-byte number written by `encode_i64` or `encode_f64`. +fn decode_number(bytes: &[u8; 8]) -> ColumnTimestamp { + let encoded = u64::from_be_bytes(*bytes); + // A clear top bit is a negative value in both encodings. Every window + // starts at or after the epoch, so the integer reading keeps it outside. + if (encoded & SIGN_BIT) == 0 || (encoded ^ SIGN_BIT) < INTEGER_PAYLOAD_LIMIT { + ColumnTimestamp::Millis(SortKeyEncoder::decode_i64(bytes)) + } else { + ColumnTimestamp::FractionalMillis(SortKeyEncoder::decode_f64(bytes)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn integer_timestamp_decodes_to_its_millis() { + for ms in [0i64, 1, 1_714_608_000_000, (1 << 53) - 1] { + let bytes = SortKeyEncoder::encode_i64(ms); + assert_eq!( + decode_column_timestamp(&bytes), + Some(ColumnTimestamp::Millis(ms)) + ); + } + } + + #[test] + fn negative_integer_timestamp_decodes_negative() { + let bytes = SortKeyEncoder::encode_i64(-86_400_000); + assert_eq!( + decode_column_timestamp(&bytes), + Some(ColumnTimestamp::Millis(-86_400_000)) + ); + } + + #[test] + fn float_timestamp_decodes_to_its_fractional_millis() { + for ms in [1.0f64, 1_714_608_000_000.5, 1_714_607_999_999.25] { + let bytes = SortKeyEncoder::encode_f64(ms); + assert_eq!( + decode_column_timestamp(&bytes), + Some(ColumnTimestamp::FractionalMillis(ms)) + ); + } + } + + #[test] + fn negative_float_timestamp_decodes_negative() { + let bytes = SortKeyEncoder::encode_f64(-1_714_608_000_000.5); + match decode_column_timestamp(&bytes) { + Some(ColumnTimestamp::Millis(ms)) => assert!(ms < 0), + Some(ColumnTimestamp::FractionalMillis(ms)) => assert!(ms < 0.0), + None => panic!("a negative float must decode"), + } + } + + #[test] + fn iso_8601_string_timestamp_decodes_to_millis() { + assert_eq!( + decode_column_timestamp(b"2024-05-02T00:00:00Z"), + Some(ColumnTimestamp::Millis(1_714_608_000_000)) + ); + assert_eq!( + decode_column_timestamp(b"2024-05-02"), + Some(ColumnTimestamp::Millis(1_714_608_000_000)) + ); + } + + #[test] + fn null_or_non_date_value_has_no_timestamp() { + assert_eq!(decode_column_timestamp(b""), None); + assert_eq!(decode_column_timestamp(b"yesterday"), None); + } +} diff --git a/nodedb/src/engine/kv/sorted_index/window.rs b/nodedb/src/engine/kv/sorted_index/window.rs index 3ed5ffb87..121f3fcbc 100644 --- a/nodedb/src/engine/kv/sorted_index/window.rs +++ b/nodedb/src/engine/kv/sorted_index/window.rs @@ -12,7 +12,10 @@ //! //! The sorted index stores ALL entries (no data deletion). Window filtering //! is applied at query time by comparing the entry's timestamp against the -//! current window boundary. +//! current [`WindowBounds`]. An entry outside the bounds, or with no +//! timestamp, is hidden from rank, top-k, range and count. + +use super::timestamp::ColumnTimestamp; /// Window type for a time-windowed sorted index. #[derive(Debug, Clone, PartialEq, Eq)] @@ -79,31 +82,61 @@ impl WindowConfig { self.window_type == WindowType::None } - /// Compute the start of the current window (inclusive) given `now_ms`. + /// The timestamp bounds of the window that contains `now_ms`. /// /// Returns `None` for `WindowType::None` (no filtering). - /// Returns `Some(start_ms)` for all other types — entries with timestamp - /// >= start_ms are in the current window. - pub fn window_start(&self, now_ms: u64) -> Option { + /// Daily, weekly and monthly windows admit every timestamp at or after + /// their start. A custom window admits `[start_ms, end_ms)` while `now_ms` + /// is inside it, and nothing otherwise. + pub fn bounds(&self, now_ms: u64) -> Option { + let open_from = |start_ms| WindowBounds { + start_ms, + end_ms: None, + }; match &self.window_type { WindowType::None => None, - WindowType::Daily => Some(start_of_day_utc(now_ms)), - WindowType::Weekly => Some(start_of_week_utc(now_ms)), - WindowType::Monthly => Some(start_of_month_utc(now_ms)), + WindowType::Daily => Some(open_from(start_of_day_utc(now_ms))), + WindowType::Weekly => Some(open_from(start_of_week_utc(now_ms))), + WindowType::Monthly => Some(open_from(start_of_month_utc(now_ms))), WindowType::Custom { start_ms, end_ms } => { - if now_ms >= *start_ms && now_ms < *end_ms { - Some(*start_ms) + let end_ms = if now_ms >= *start_ms && now_ms < *end_ms { + *end_ms } else { - // Outside the custom window — no entries visible. - Some(u64::MAX) - } + // An empty range: no timestamp is both >= and < start_ms. + *start_ms + }; + Some(WindowBounds { + start_ms: *start_ms, + end_ms: Some(end_ms), + }) } } } +} + +/// The half-open timestamp range `[start_ms, end_ms)` of the visible entries. +/// +/// `end_ms: None` has no upper bound. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct WindowBounds { + pub start_ms: u64, + pub end_ms: Option, +} - /// Compute the window start for a specific target timestamp (for historical queries). - pub fn window_start_at(&self, target_ms: u64) -> Option { - self.window_start(target_ms) +impl WindowBounds { + /// `true` when `ts` falls inside the range. + pub fn contains(&self, ts: ColumnTimestamp) -> bool { + match ts { + ColumnTimestamp::Millis(ms) => { + let ms = i128::from(ms); + ms >= i128::from(self.start_ms) + && self.end_ms.is_none_or(|end| ms < i128::from(end)) + } + // NaN fails both comparisons and stays outside. + ColumnTimestamp::FractionalMillis(ms) => { + ms >= self.start_ms as f64 && self.end_ms.is_none_or(|end| ms < end as f64) + } + } } } @@ -184,28 +217,34 @@ mod tests { fn unwindowed_always_visible() { let config = WindowConfig::none(); assert!(config.is_unwindowed()); - assert!(config.window_start(1000).is_none()); + assert!(config.bounds(1000).is_none()); + } + + fn start_of(config: &WindowConfig, now_ms: u64) -> u64 { + let bounds = config.bounds(now_ms).expect("windowed config has bounds"); + assert_eq!(bounds.end_ms, None, "calendar windows have no upper bound"); + bounds.start_ms } #[test] fn daily_window_start() { - // 2024-05-02 15:30:00 UTC = 1714664200000 ms + // 2024-05-02 15:36:40 UTC = 1714664200000 ms let now_ms = 1714664200000u64; let config = WindowConfig::daily("updated_at"); - let start = config.window_start(now_ms).unwrap(); - // Should be 2024-05-02 00:00:00 UTC = 1714608000000 ms - assert_eq!(start, start_of_day_utc(now_ms)); - assert!(start <= now_ms); + let start = start_of(&config, now_ms); + // 2024-05-02 00:00:00 UTC = 1714608000000 ms + assert_eq!(start, 1_714_608_000_000); assert!(now_ms - start < MS_PER_DAY); } #[test] fn weekly_window_start() { - // 2024-05-02 (Thursday) 15:30:00 UTC + // 2024-05-02 (Thursday) 15:36:40 UTC let now_ms = 1714664200000u64; let config = WindowConfig::weekly("updated_at"); - let start = config.window_start(now_ms).unwrap(); - // Monday of that week: 2024-04-29 00:00:00 UTC + let start = start_of(&config, now_ms); + // Monday 2024-04-29 00:00:00 UTC. + assert_eq!(start, 1_714_348_800_000); assert!(start <= now_ms); assert!(now_ms - start < 7 * MS_PER_DAY); // Verify it's a Monday. @@ -219,8 +258,8 @@ mod tests { // 2024-05-15 12:00:00 UTC let now_ms = 1715774400000u64; let config = WindowConfig::monthly("updated_at"); - let start = config.window_start(now_ms).unwrap(); - // Should be 2024-05-01 00:00:00 UTC + let start = start_of(&config, now_ms); + // 2024-05-01 00:00:00 UTC let (y, m, d) = seconds_to_ymd(start / MS_PER_SECOND); assert_eq!(y, 2024); assert_eq!(m, 5); @@ -233,11 +272,52 @@ mod tests { let end = 2000u64; let config = WindowConfig::custom("ts", start, end); - // Inside window. - assert_eq!(config.window_start(1500), Some(1000)); + // Inside the window: the configured range. + let inside = config.bounds(1500).expect("custom window has bounds"); + assert_eq!( + inside, + WindowBounds { + start_ms: 1000, + end_ms: Some(2000) + } + ); + + // Before or after the window: a range that admits nothing. + for now_ms in [500, 2000, 2500] { + let outside = config.bounds(now_ms).expect("custom window has bounds"); + for ts in [0, 999, 1000, 1999, 2000, i64::MAX] { + assert!(!outside.contains(ColumnTimestamp::Millis(ts))); + } + assert!(!outside.contains(ColumnTimestamp::FractionalMillis(f64::INFINITY))); + } + } + + #[test] + fn bounds_contain_only_the_half_open_range() { + let bounds = WindowBounds { + start_ms: 1000, + end_ms: Some(2000), + }; + assert!(!bounds.contains(ColumnTimestamp::Millis(999))); + assert!(bounds.contains(ColumnTimestamp::Millis(1000))); + assert!(bounds.contains(ColumnTimestamp::Millis(1999))); + assert!(!bounds.contains(ColumnTimestamp::Millis(2000))); + assert!(!bounds.contains(ColumnTimestamp::Millis(-1))); + assert!(!bounds.contains(ColumnTimestamp::FractionalMillis(999.5))); + assert!(bounds.contains(ColumnTimestamp::FractionalMillis(1000.5))); + assert!(!bounds.contains(ColumnTimestamp::FractionalMillis(2000.0))); + assert!(!bounds.contains(ColumnTimestamp::FractionalMillis(f64::NAN))); + } - // Outside window. - assert_eq!(config.window_start(2500), Some(u64::MAX)); + #[test] + fn open_bounds_have_no_upper_limit() { + let bounds = WindowBounds { + start_ms: 1000, + end_ms: None, + }; + assert!(bounds.contains(ColumnTimestamp::Millis(i64::MAX))); + assert!(bounds.contains(ColumnTimestamp::FractionalMillis(1e300))); + assert!(!bounds.contains(ColumnTimestamp::Millis(999))); } #[test] diff --git a/nodedb/src/engine/kv/sorted_index/windowed_query.rs b/nodedb/src/engine/kv/sorted_index/windowed_query.rs index d603de52c..df85c5493 100644 --- a/nodedb/src/engine/kv/sorted_index/windowed_query.rs +++ b/nodedb/src/engine/kv/sorted_index/windowed_query.rs @@ -3,7 +3,9 @@ //! Windowed query helpers for sorted indexes. //! //! These functions handle time-windowed queries where only entries within -//! a configurable time window are considered for rank/top_k/range/count. +//! the current window are considered for rank/top_k/range/count. An entry is +//! visible when the timestamp column of its sort key decodes to a timestamp +//! inside the window bounds. An entry with no decodable timestamp is hidden. //! //! Performance: uses `for_each_in_order` with early termination instead of //! fetching all entries. windowed_top_k(10) on a 10M-entry tree is O(N) in @@ -12,9 +14,10 @@ use std::collections::HashMap; -use super::key::SortKeyEncoder; use super::manager::SortedIndexDef; +use super::timestamp::decode_column_timestamp; use super::tree::OrderStatTree; +use super::window::WindowBounds; /// Internal state for a sorted index (definition + tree). pub(super) struct SortedIndexRef<'a> { @@ -22,20 +25,25 @@ pub(super) struct SortedIndexRef<'a> { pub tree: &'a OrderStatTree, } -/// Windowed rank: count entries with lower sort key that are in the current window. +/// Windowed rank: count visible entries up to and including `primary_key`. +/// +/// `None` when `primary_key` is absent or outside the window. pub(super) fn windowed_rank( idx: &SortedIndexRef<'_>, primary_key: &[u8], now_ms: u64, ) -> Option { - let window_start = idx.def.window.window_start(now_ms)?; + let bounds = idx.def.window.bounds(now_ms)?; let target_sort = idx.tree.get_sort_key(primary_key)?; + if !entry_in_window(idx.def, target_sort, &bounds) { + return None; + } let mut rank = 0u32; let mut found = false; idx.tree.for_each_in_order(|sort_key, pk| { - if !entry_in_window(idx, pk, window_start) { + if !entry_in_window(idx.def, sort_key, &bounds) { return true; // Skip, continue. } rank += 1; @@ -53,15 +61,18 @@ pub(super) fn windowed_rank( /// /// Early termination: stops as soon as k windowed entries are found. pub(super) fn windowed_top_k(idx: &SortedIndexRef<'_>, k: u32, now_ms: u64) -> Vec<(u32, Vec)> { - let Some(window_start) = idx.def.window.window_start(now_ms) else { + let Some(bounds) = idx.def.window.bounds(now_ms) else { return Vec::new(); }; + if k == 0 { + return Vec::new(); + } let mut result = Vec::with_capacity(k as usize); let mut rank = 0u32; - idx.tree.for_each_in_order(|_, pk| { - if !entry_in_window(idx, pk, window_start) { + idx.tree.for_each_in_order(|sort_key, pk| { + if !entry_in_window(idx.def, sort_key, &bounds) { return true; // Skip, continue. } rank += 1; @@ -72,13 +83,14 @@ pub(super) fn windowed_top_k(idx: &SortedIndexRef<'_>, k: u32, now_ms: u64) -> V result } -/// Windowed range: filter range results by window. +/// Windowed range: keep the visible entries of `entries`, ranked among all +/// visible entries. pub(super) fn windowed_range( idx: &SortedIndexRef<'_>, entries: &[(&[u8], &[u8])], now_ms: u64, ) -> Vec<(u32, Vec)> { - let Some(window_start) = idx.def.window.window_start(now_ms) else { + let Some(bounds) = idx.def.window.bounds(now_ms) else { return Vec::new(); }; @@ -86,8 +98,8 @@ pub(super) fn windowed_range( let mut windowed_ranks: HashMap, u32> = HashMap::new(); let mut rank = 0u32; - idx.tree.for_each_in_order(|_, pk| { - if entry_in_window(idx, pk, window_start) { + idx.tree.for_each_in_order(|sort_key, pk| { + if entry_in_window(idx.def, sort_key, &bounds) { rank += 1; windowed_ranks.insert(pk.to_vec(), rank); } @@ -105,13 +117,13 @@ pub(super) fn windowed_range( /// Windowed count: count entries in the current window. pub(super) fn windowed_count(idx: &SortedIndexRef<'_>, now_ms: u64) -> u32 { - let Some(window_start) = idx.def.window.window_start(now_ms) else { + let Some(bounds) = idx.def.window.bounds(now_ms) else { return 0; }; let mut count = 0u32; - idx.tree.for_each_in_order(|_, pk| { - if entry_in_window(idx, pk, window_start) { + idx.tree.for_each_in_order(|sort_key, _| { + if entry_in_window(idx.def, sort_key, &bounds) { count += 1; } true @@ -119,58 +131,13 @@ pub(super) fn windowed_count(idx: &SortedIndexRef<'_>, now_ms: u64) -> u32 { count } -/// Check if an entry's timestamp is within the window. +/// `true` when the timestamp column of `sort_key` lies inside `bounds`. /// -/// Reads the timestamp from the sort key if the timestamp column is part of -/// the composite sort key. For entries where the timestamp is NOT in the sort -/// key, assumes the entry is in-window (documented limitation). -fn entry_in_window(idx: &SortedIndexRef<'_>, primary_key: &[u8], window_start: u64) -> bool { - let Some(sort_key_bytes) = idx.tree.get_sort_key(primary_key) else { - return false; - }; - - // Find the timestamp column's position in the sort key. - let ts_col = &idx.def.window.timestamp_column; - let columns = idx.def.encoder.columns(); - let mut offset = 0usize; - - for col in columns { - // Read length prefix. - if offset + 4 > sort_key_bytes.len() { - return true; // Can't parse — assume in-window. - } - let len = u32::from_be_bytes( - sort_key_bytes[offset..offset + 4] - .try_into() - .unwrap_or([0; 4]), - ) as usize; - offset += 4; - - if col.name == *ts_col { - if offset + len > sort_key_bytes.len() || len != 8 { - return true; // Can't parse — assume in-window. - } - let mut ts_bytes = [0u8; 8]; - ts_bytes.copy_from_slice(&sort_key_bytes[offset..offset + len]); - - // If DESC, undo the complement. - let ts_bytes = if col.direction == super::key::SortDirection::Desc { - let mut unflipped = [0u8; 8]; - for (i, &b) in ts_bytes.iter().enumerate() { - unflipped[i] = !b; - } - unflipped - } else { - ts_bytes - }; - - let ts_ms = SortKeyEncoder::decode_timestamp_ms(&ts_bytes); - return ts_ms >= window_start; - } - - offset += len; - } - - // Timestamp column not in sort key — can't filter, assume in-window. - true +/// Registration rejects a windowed index whose timestamp column is not a sort +/// column. A key without a decodable timestamp is outside every window. +fn entry_in_window(def: &SortedIndexDef, sort_key: &[u8], bounds: &WindowBounds) -> bool { + def.encoder + .column_value(sort_key, &def.window.timestamp_column) + .and_then(|value| decode_column_timestamp(&value)) + .is_some_and(|ts| bounds.contains(ts)) } From 81291cd56d2a4c7138ea29f17e9c64e166743724 Mon Sep 17 00:00:00 2001 From: Farhan Syah Date: Mon, 12 Oct 2026 16:41:21 +0800 Subject: [PATCH 06/11] fix(vector): probe-based distance, vectorless nodes and validated checkpoint metric Distance evaluation reads node vectors through a probe, graph nodes may carry no vector, and checkpoint load checks the stored metric and structure before accepting an index. --- nodedb-vector/src/adaptive_filter.rs | 110 ++++- nodedb-vector/src/collection/checkpoint.rs | 13 +- nodedb-vector/src/hnsw/build.rs | 153 +++---- nodedb-vector/src/hnsw/checkpoint.rs | 180 +++++++- nodedb-vector/src/hnsw/graph/index/backing.rs | 4 +- .../src/hnsw/graph/index/distance_ops.rs | 150 +++---- nodedb-vector/src/hnsw/graph/index/state.rs | 32 +- nodedb-vector/src/hnsw/graph/mod.rs | 1 + nodedb-vector/src/hnsw/graph/types.rs | 13 + nodedb-vector/src/hnsw/search.rs | 29 +- nodedb-vector/src/navix/acorn.rs | 196 +++++---- nodedb-vector/src/navix/beam.rs | 100 +++++ nodedb-vector/src/navix/mod.rs | 1 + nodedb-vector/src/navix/traversal.rs | 414 +++++++++--------- nodedb-vector/src/test_support.rs | 94 ++++ 15 files changed, 956 insertions(+), 534 deletions(-) create mode 100644 nodedb-vector/src/navix/beam.rs diff --git a/nodedb-vector/src/adaptive_filter.rs b/nodedb-vector/src/adaptive_filter.rs index 41af942dd..5a23acc78 100644 --- a/nodedb-vector/src/adaptive_filter.rs +++ b/nodedb-vector/src/adaptive_filter.rs @@ -96,16 +96,14 @@ pub fn adaptive_search( } FilterStrategy::BruteForceMatching => { let metric = index.params().metric; + let probe = index.probe(query); + // A node without a vector has no distance and is not a hit. let mut results: Vec = bitmap .iter() + .filter(|&id| !index.is_deleted(id)) .filter_map(|id| { - if id as usize >= index.len() || index.is_deleted(id) { - return None; - } - Some(SearchResult { - id, - distance: index.dist_to_node_f32(query, id, metric), - }) + let distance = index.dist_to_node(&probe, id, metric)?; + Some(SearchResult { id, distance }) }) .collect(); @@ -240,4 +238,102 @@ mod tests { ); } } + + const SEED: u64 = 0x9E37_79B9_7F4A_7C15; + + /// Thresholds that force pre-filter, post-filter, and brute force in turn. + const FORCED: [FilterThresholds; 3] = [ + FilterThresholds { + high_selectivity: 2.0, + extreme_selectivity: 2.0, + }, + FilterThresholds { + high_selectivity: -1.0, + extreme_selectivity: 2.0, + }, + BRUTE_FORCE, + ]; + + const BRUTE_FORCE: FilterThresholds = FilterThresholds { + high_selectivity: -1.0, + extreme_selectivity: -1.0, + }; + + fn every_tenth(n: u32) -> RoaringBitmap { + (0..n).step_by(10).collect() + } + + /// Brute force over a graph-only restore reads node vectors through the + /// backing and ranks exactly like the source index. + #[cfg(not(target_arch = "wasm32"))] + #[test] + fn brute_force_after_graph_only_restore_matches_the_source_index() { + use crate::test_support::{graph_only_restore_with_backing, seeded_index, seeded_vectors}; + use nodedb_types::vector_dtype::VectorStorageDtype; + + let vectors = seeded_vectors(300, 8, SEED); + let src = seeded_index(&vectors, DistanceMetric::L2, VectorStorageDtype::F32); + let restored = graph_only_restore_with_backing(&src, &vectors); + let bitmap = every_tenth(300); + + let expected = adaptive_search(&src, &vectors[7], 5, 64, &bitmap, &BRUTE_FORCE).unwrap(); + let got = adaptive_search(&restored, &vectors[7], 5, 64, &bitmap, &BRUTE_FORCE).unwrap(); + assert_eq!(expected.len(), 5); + let pairs = |r: &[SearchResult]| -> Vec<(u32, f32)> { + r.iter().map(|h| (h.id, h.distance)).collect() + }; + assert_eq!(pairs(got.as_slice()), pairs(expected.as_slice())); + } + + /// A graph-only restore searched before `with_backing` has no vectors. No + /// strategy may return its nodes, at infinite distance or otherwise. + #[test] + fn no_strategy_returns_a_node_without_a_vector() { + use crate::test_support::{graph_only_restore, seeded_index, seeded_vectors}; + use nodedb_types::vector_dtype::VectorStorageDtype; + + let vectors = seeded_vectors(300, 8, SEED); + let src = seeded_index(&vectors, DistanceMetric::L2, VectorStorageDtype::F32); + let idx = graph_only_restore(&src); + let bitmap = every_tenth(300); + for thresholds in &FORCED { + let results = adaptive_search(&idx, &vectors[7], 5, 64, &bitmap, thresholds).unwrap(); + assert!(results.is_empty(), "got {results:?}"); + } + } + + /// Brute force on an F16/BF16 index scores the encoded vectors under the + /// build metric. + #[test] + fn brute_force_on_narrow_dtypes_matches_decoded_distances() { + use crate::distance::distance; + use crate::test_support::{seeded_index, seeded_vectors}; + use nodedb_types::vector_dtype::VectorStorageDtype; + + let vectors = seeded_vectors(300, 8, SEED); + let bitmap = every_tenth(300); + for dtype in [VectorStorageDtype::F16, VectorStorageDtype::BF16] { + let idx = seeded_index(&vectors, DistanceMetric::Cosine, dtype); + let encoded_query = crate::dtype::cast_from_f32(&vectors[7], dtype); + let query = crate::dtype::cast_to_f32(&encoded_query, dtype, 8).unwrap(); + let mut reference: Vec<(u32, f32)> = bitmap + .iter() + .map(|id| { + let node = idx.materialize_vector(id).unwrap(); + (id, distance(&query, &node, DistanceMetric::Cosine)) + }) + .collect(); + reference.sort_by(|a, b| a.1.total_cmp(&b.1)); + + let got = adaptive_search(&idx, &vectors[7], 5, 64, &bitmap, &BRUTE_FORCE).unwrap(); + assert_eq!(got.len(), 5, "dtype={dtype:?}"); + for (hit, (id, d)) in got.iter().zip(&reference) { + assert_eq!(hit.id, *id, "dtype={dtype:?}: got {got:?}"); + assert!( + (hit.distance - d).abs() < 1e-4, + "dtype={dtype:?}: {hit:?} vs {d}" + ); + } + } + } } diff --git a/nodedb-vector/src/collection/checkpoint.rs b/nodedb-vector/src/collection/checkpoint.rs index e39985717..66576f986 100644 --- a/nodedb-vector/src/collection/checkpoint.rs +++ b/nodedb-vector/src/collection/checkpoint.rs @@ -26,7 +26,6 @@ use serde::{Deserialize, Serialize}; use crate::collection::payload_index::PayloadIndexSetSnapshot; use crate::collection::segment::{BuildingSegment, DEFAULT_SEAL_THRESHOLD, SealedSegment}; use crate::collection::tier::StorageTier; -use crate::distance::DistanceMetric; use crate::error::VectorError; use crate::flat::FlatIndex; use crate::hnsw::{HnswIndex, HnswParams}; @@ -296,17 +295,7 @@ impl VectorCollection { }); } }; - let metric = match snap.params_metric { - 0 => DistanceMetric::L2, - 1 => DistanceMetric::Cosine, - 2 => DistanceMetric::InnerProduct, - 3 => DistanceMetric::Manhattan, - 4 => DistanceMetric::Chebyshev, - 5 => DistanceMetric::Hamming, - 6 => DistanceMetric::Jaccard, - 7 => DistanceMetric::Pearson, - _ => DistanceMetric::Cosine, - }; + let metric = crate::hnsw::checkpoint::metric_from_code(snap.params_metric)?; let params = HnswParams { m: snap.params_m, m0: snap.params_m0, diff --git a/nodedb-vector/src/hnsw/build.rs b/nodedb-vector/src/hnsw/build.rs index ff1d9d5c9..866ec47d3 100644 --- a/nodedb-vector/src/hnsw/build.rs +++ b/nodedb-vector/src/hnsw/build.rs @@ -50,16 +50,15 @@ impl HnswIndex { return Ok(()); }; - // Encode query once to the index dtype for all dist_to_node calls below. - let query = vector; - let query_bytes = cast_from_f32(&query, self.params.dtype); + // One probe serves every distance against the new vector below. + let probe = self.probe(&vector); let mut current_ep = ep; // Phase 1: Greedy descent from top layer to new_layer + 1. if self.max_layer > new_layer { for layer in (new_layer + 1..=self.max_layer).rev() { - let results = search_layer(self, &query_bytes, current_ep, 1, layer, None, 0); + let results = search_layer(self, &probe, current_ep, 1, layer, None, 0); if let Some(nearest) = results.first() { current_ep = nearest.id; } @@ -70,7 +69,7 @@ impl HnswIndex { let insert_top = new_layer.min(self.max_layer); for layer in (0..=insert_top).rev() { let ef = self.params.ef_construction; - let candidates = search_layer(self, &query_bytes, current_ep, ef, layer, None, 0); + let candidates = search_layer(self, &probe, current_ep, ef, layer, None, 0); let m = self.max_neighbors(layer); let selected = select_neighbors_heuristic(self, &candidates, m); @@ -82,8 +81,7 @@ impl HnswIndex { self.nodes[nid].neighbors[layer].push(new_id); if self.nodes[nid].neighbors[layer].len() > m { - let node_bytes = self.node_bytes(neighbor.id).into_owned(); - self.prune_neighbors(nid, layer, &node_bytes, m); + self.prune_neighbors(neighbor.id, layer, m); } } @@ -101,20 +99,28 @@ impl HnswIndex { } /// Prune a node's neighbor list using the diversity heuristic. - fn prune_neighbors(&mut self, node_idx: usize, layer: usize, node_bytes: &[u8], m: usize) { - let neighbor_ids: Vec = self.nodes[node_idx].neighbors[layer].clone(); - - let mut candidates: Vec = neighbor_ids - .iter() - .map(|&nid| Candidate { - id: nid, - dist: self.dist_to_node(node_bytes, nid), - }) - .collect(); - candidates.sort_unstable_by(|a, b| a.dist.total_cmp(&b.dist)); - - let selected = select_neighbors_heuristic(self, &candidates, m); - self.nodes[node_idx].neighbors[layer] = selected.iter().map(|c| c.id).collect(); + fn prune_neighbors(&mut self, node_id: u32, layer: usize, m: usize) { + let metric = self.params.metric; + let node_idx = node_id as usize; + let selected: Vec = { + let probe = self.node_probe(node_id); + let mut candidates: Vec = self.nodes[node_idx].neighbors[layer] + .iter() + .map(|&nid| Candidate { + id: nid, + dist: probe + .as_ref() + .and_then(|p| self.dist_to_node(p, nid, metric)) + .unwrap_or(f32::INFINITY), + }) + .collect(); + candidates.sort_unstable_by(|a, b| a.dist.total_cmp(&b.dist)); + select_neighbors_heuristic(self, &candidates, m) + .iter() + .map(|c| c.id) + .collect() + }; + self.nodes[node_idx].neighbors[layer] = selected; } } @@ -124,6 +130,7 @@ fn select_neighbors_heuristic( candidates: &[Candidate], m: usize, ) -> Vec { + let metric = index.params.metric; let mut selected: Vec = Vec::with_capacity(m); for candidate in candidates { @@ -131,16 +138,13 @@ fn select_neighbors_heuristic( break; } - let candidate_bytes = index.node_bytes(candidate.id); + // Borrowed from node storage: the pairwise loop below allocates nothing. + let candidate_probe = index.node_probe(candidate.id); let is_diverse = selected.iter().all(|s| { - let dist_to_selected = crate::distance::dispatch::distance_typed( - index.params.metric, - index.params.dtype, - &candidate_bytes, - &index.node_bytes(s.id), - index.dim, - ) - .unwrap_or(f32::MAX); + let dist_to_selected = candidate_probe + .as_ref() + .and_then(|p| index.dist_to_node(p, s.id, metric)) + .unwrap_or(f32::MAX); candidate.dist <= dist_to_selected }); @@ -167,6 +171,8 @@ fn select_neighbors_heuristic( #[cfg(test)] mod tests { + use nodedb_types::vector_dtype::VectorStorageDtype; + use crate::distance::DistanceMetric; use crate::hnsw::{HnswIndex, HnswParams}; @@ -231,22 +237,30 @@ mod tests { } } - /// Deterministic vectors in `[0, 1)` from an xorshift64 stream. `seed` must - /// be non-zero. - fn seeded_vectors(count: usize, dim: usize, seed: u64) -> Vec> { - let mut state = seed; - (0..count) - .map(|_| { - (0..dim) - .map(|_| { - state ^= state << 13; - state ^= state >> 7; - state ^= state << 17; - (state % 10_000) as f32 / 10_000.0 - }) - .collect() - }) - .collect() + /// F16/BF16 inserts score node pairs through the borrowed encoded bytes. + /// Every node must stay reachable by an exact query. + #[test] + fn narrow_dtype_inserts_reach_every_node() { + use crate::test_support::{seeded_index, seeded_vectors}; + + let vectors = seeded_vectors(120, 8, 0x2545_F491_4F6C_DD1D); + for dtype in [VectorStorageDtype::F16, VectorStorageDtype::BF16] { + for metric in [DistanceMetric::L2, DistanceMetric::Cosine] { + let idx = seeded_index(&vectors, metric, dtype); + for (id, v) in vectors.iter().enumerate() { + // `ef` covers every node, so a miss means a node is unreachable. + let hits = idx.search(v, 1, 128).unwrap(); + assert_eq!( + hits.first().map(|h| h.id), + Some(id as u32), + "dtype={dtype:?} metric={metric:?}: exact query misses node {id}" + ); + } + for node in &idx.nodes { + assert!(node.neighbors[0].len() <= idx.params.m0); + } + } + } } /// A graph-only restore with an attached backing must insert exactly like a @@ -258,58 +272,17 @@ mod tests { #[cfg(not(target_arch = "wasm32"))] #[test] fn inserts_after_graph_only_restore_match_a_full_restore() { - use std::sync::Arc; - - use crate::segment_backing::VectorSegmentBacking; - - /// Serves an in-memory copy of the base vectors. - struct VecBacking { - dim: usize, - vectors: Vec>, - } - impl VectorSegmentBacking for VecBacking { - fn len(&self) -> usize { - self.vectors.len() - } - fn dim(&self) -> usize { - self.dim - } - fn get_vector(&self, id: u32) -> Option<&[f32]> { - self.vectors.get(id as usize).map(Vec::as_slice) - } - fn get_surrogate(&self, _id: u32) -> Option { - None - } - } + use crate::test_support::{graph_only_restore_with_backing, seeded_index, seeded_vectors}; const DIM: usize = 8; - let params = HnswParams { - m: 4, - m0: 8, - ef_construction: 32, - metric: DistanceMetric::L2, - dtype: nodedb_types::vector_dtype::VectorStorageDtype::F32, - }; let base = seeded_vectors(200, DIM, 0x9E37_79B9_7F4A_7C15); let extra = seeded_vectors(5, DIM, 0xD1B5_4A32_D192_ED03); - let mut src = HnswIndex::with_seed(DIM, params, 12345); - for v in &base { - src.insert(v.clone()).unwrap(); - } - + let src = seeded_index(&base, DistanceMetric::L2, VectorStorageDtype::F32); let mut full = HnswIndex::from_checkpoint(&src.checkpoint_to_bytes().unwrap()) .unwrap() .expect("full checkpoint must be recognized"); - let mut graph_only = HnswIndex::from_checkpoint(&src.graph_checkpoint_to_bytes().unwrap()) - .unwrap() - .expect("graph checkpoint must be recognized"); - graph_only - .with_backing(Arc::new(VecBacking { - dim: DIM, - vectors: base.clone(), - })) - .expect("a complete backing must attach"); + let mut graph_only = graph_only_restore_with_backing(&src, &base); for v in &extra { full.insert(v.clone()).unwrap(); diff --git a/nodedb-vector/src/hnsw/checkpoint.rs b/nodedb-vector/src/hnsw/checkpoint.rs index ec6766eb3..7da6342ed 100644 --- a/nodedb-vector/src/hnsw/checkpoint.rs +++ b/nodedb-vector/src/hnsw/checkpoint.rs @@ -6,6 +6,8 @@ use std::cell::RefCell; +use nodedb_types::vector_dtype::VectorStorageDtype; + use crate::distance::DistanceMetric; use crate::error::VectorError; use crate::hnsw::arena::BeamSearchArena; @@ -17,6 +19,71 @@ const HNSW_RKYV_MAGIC: &[u8; 6] = b"RKHNS\0"; /// Current format version for rkyv-serialized HNSW snapshots. pub const HNSW_FORMAT_VERSION: u8 = 1; +/// Storage dtype of every index restored from a checkpoint. +/// +/// The snapshot stores f32 vectors and no dtype. An F16/BF16 index +/// checkpoints its vectors decoded to f32 and restores as an F32 index. +const CHECKPOINT_DTYPE: VectorStorageDtype = VectorStorageDtype::F32; + +/// Decode a checkpoint metric code written as `DistanceMetric as u8`. +/// +/// # Errors +/// +/// [`VectorError::CheckpointDeserializationError`] for a code no metric uses. +pub(crate) fn metric_from_code(code: u8) -> Result { + match code { + 0 => Ok(DistanceMetric::L2), + 1 => Ok(DistanceMetric::Cosine), + 2 => Ok(DistanceMetric::InnerProduct), + 3 => Ok(DistanceMetric::Manhattan), + 4 => Ok(DistanceMetric::Chebyshev), + 5 => Ok(DistanceMetric::Hamming), + 6 => Ok(DistanceMetric::Jaccard), + 7 => Ok(DistanceMetric::Pearson), + _ => Err(VectorError::CheckpointDeserializationError { + detail: format!("unknown distance metric code {code}"), + }), + } +} + +/// Check that a decoded snapshot describes one consistent graph. +fn validate_snapshot(snap: &HnswSnapshotRkyv) -> Result<(), VectorError> { + let count = snap.node_vectors.len(); + if snap.node_neighbors.len() != count || snap.node_deleted.len() != count { + return Err(VectorError::CheckpointDeserializationError { + detail: format!( + "node counts disagree: {count} vectors, {} neighbor lists, {} deleted flags", + snap.node_neighbors.len(), + snap.node_deleted.len() + ), + }); + } + // An empty vector is a graph-only placeholder served by a backing. + if let Some(v) = snap + .node_vectors + .iter() + .find(|v| !v.is_empty() && v.len() != snap.dim) + { + return Err(VectorError::StoredDimensionMismatch { + expected: snap.dim, + got: v.len(), + }); + } + let in_range = |id: u32| (id as usize) < count; + let neighbors_in_range = snap + .node_neighbors + .iter() + .flatten() + .flatten() + .all(|&id| in_range(id)); + if !neighbors_in_range || snap.entry_point.is_some_and(|ep| !in_range(ep)) { + return Err(VectorError::CheckpointDeserializationError { + detail: format!("a neighbor or entry point id is outside {count} nodes"), + }); + } + Ok(()) +} + /// rkyv-serialized HNSW snapshot. SoA layout for better rkyv compatibility /// (flat Vecs instead of Vec). #[derive(rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)] @@ -153,15 +220,14 @@ impl HnswIndex { } /// Reconstruct HnswIndex from deserialized snapshot fields. + /// + /// The snapshot is untrusted input. Inconsistent node counts, a vector of + /// the wrong dimension, or an out-of-range node id fail the restore. fn from_hnsw_snapshot(snap: HnswSnapshotRkyv) -> Result { use nodedb_types::hnsw::HnswParams; - let metric = match snap.metric { - 0 => DistanceMetric::L2, - 1 => DistanceMetric::Cosine, - 2 => DistanceMetric::InnerProduct, - _ => DistanceMetric::Cosine, - }; + let metric = metric_from_code(snap.metric)?; + validate_snapshot(&snap)?; let flat = FlatNeighborStore::from_nested(&snap.node_neighbors); @@ -184,7 +250,7 @@ impl HnswIndex { m0: snap.m0, ef_construction: snap.ef_construction, metric, - dtype: nodedb_types::vector_dtype::VectorStorageDtype::F32, + dtype: CHECKPOINT_DTYPE, }, nodes, entry_point: snap.entry_point, @@ -200,7 +266,9 @@ impl HnswIndex { #[cfg(test)] mod tests { + use super::HnswSnapshotRkyv; use crate::distance::DistanceMetric; + use crate::error::VectorError; use crate::hnsw::{HnswIndex, HnswParams}; fn make_index() -> HnswIndex { @@ -291,4 +359,102 @@ mod tests { Ok(_) => panic!("expected CheckpointDeserializationError, got Ok"), } } + + /// Every metric survives a checkpoint round trip. + #[test] + fn every_metric_round_trips() { + for metric in [ + DistanceMetric::L2, + DistanceMetric::Cosine, + DistanceMetric::InnerProduct, + DistanceMetric::Manhattan, + DistanceMetric::Chebyshev, + DistanceMetric::Hamming, + DistanceMetric::Jaccard, + DistanceMetric::Pearson, + ] { + let mut idx = HnswIndex::with_seed( + 3, + HnswParams { + metric, + ..make_index().params().clone() + }, + 1, + ); + idx.insert(vec![1.0, 2.0, 3.0]).unwrap(); + let restored = HnswIndex::from_checkpoint(&idx.checkpoint_to_bytes().unwrap()) + .unwrap() + .unwrap(); + assert_eq!(restored.metric(), metric); + } + } + + /// A valid snapshot of three 2-dim nodes in a line, entry point 0. + fn snapshot() -> HnswSnapshotRkyv { + HnswSnapshotRkyv { + dim: 2, + m: 4, + m0: 8, + ef_construction: 32, + metric: DistanceMetric::L2 as u8, + entry_point: Some(0), + max_layer: 0, + rng_state: 1, + node_vectors: vec![vec![0.0, 0.0], vec![1.0, 0.0], vec![2.0, 0.0]], + node_neighbors: vec![vec![vec![1]], vec![vec![0, 2]], vec![vec![1]]], + node_deleted: vec![false; 3], + } + } + + fn encode(snap: &HnswSnapshotRkyv) -> Vec { + let mut buf = super::HNSW_RKYV_MAGIC.to_vec(); + buf.push(super::HNSW_FORMAT_VERSION); + buf.extend_from_slice(&rkyv::to_bytes::(snap).unwrap()); + buf + } + + #[test] + fn a_valid_snapshot_restores() { + let idx = HnswIndex::from_checkpoint(&encode(&snapshot())) + .unwrap() + .unwrap(); + assert_eq!(idx.search(&[2.0, 0.0], 1, 8).unwrap()[0].id, 2); + } + + /// A corrupt snapshot fails the restore instead of panicking a later search. + #[test] + fn a_corrupt_snapshot_is_a_typed_error() { + let mut wrong_dim = snapshot(); + wrong_dim.node_vectors[1] = vec![1.0, 0.0, 0.0]; + let mut short_flags = snapshot(); + short_flags.node_deleted.pop(); + let mut bad_neighbor = snapshot(); + bad_neighbor.node_neighbors[2][0].push(9); + let mut bad_entry = snapshot(); + bad_entry.entry_point = Some(3); + let mut bad_metric = snapshot(); + bad_metric.metric = 200; + + assert!(matches!( + HnswIndex::from_checkpoint(&encode(&wrong_dim)), + Err(VectorError::StoredDimensionMismatch { + expected: 2, + got: 3 + }) + )); + for (name, snap) in [ + ("short flags", short_flags), + ("bad neighbor", bad_neighbor), + ("bad entry", bad_entry), + ("bad metric", bad_metric), + ] { + assert!( + matches!( + HnswIndex::from_checkpoint(&encode(&snap)), + Err(VectorError::CheckpointDeserializationError { .. }) + ), + "{name} must fail the restore" + ); + } + } } diff --git a/nodedb-vector/src/hnsw/graph/index/backing.rs b/nodedb-vector/src/hnsw/graph/index/backing.rs index fb5d35cdd..40858f5e3 100644 --- a/nodedb-vector/src/hnsw/graph/index/backing.rs +++ b/nodedb-vector/src/hnsw/graph/index/backing.rs @@ -21,8 +21,8 @@ use super::state::HnswIndex; impl HnswIndex { /// Attach a [`crate::segment_backing::VectorSegmentBacking`] to this index. /// - /// After this succeeds, `dist_to_node` falls back to the backing whenever a - /// node's local vector storage is empty. This is used by Lite's + /// After this succeeds, `dist_to_node` falls back to the backing whenever an + /// F32 node's local vector storage is empty. This is used by Lite's /// graph-checkpoint-only restore path: the graph topology is loaded from the /// B+ tree blob, but vector data lives in a pagedb segment. /// diff --git a/nodedb-vector/src/hnsw/graph/index/distance_ops.rs b/nodedb-vector/src/hnsw/graph/index/distance_ops.rs index c32d637b6..6a6abd2d5 100644 --- a/nodedb-vector/src/hnsw/graph/index/distance_ops.rs +++ b/nodedb-vector/src/hnsw/graph/index/distance_ops.rs @@ -10,7 +10,7 @@ use super::HnswIndex; use crate::distance::dispatch::distance_typed; use crate::distance::{DistanceMetric, distance}; use crate::dtype::cast_from_f32; -use crate::hnsw::graph::MAX_LAYER_CAP; +use crate::hnsw::graph::{MAX_LAYER_CAP, NodeStorage, Probe}; impl HnswIndex { /// Assign a random layer using the exponential distribution. @@ -25,107 +25,85 @@ impl HnswIndex { layer.min(MAX_LAYER_CAP) } - /// Compute distance between a pre-encoded query and a stored node. + /// Prepares `query` for scoring against this index's nodes. /// - /// `query_bytes` must already be encoded in `self.params.dtype`; callers - /// encode once at the top of search/insert and pass the same buffer to - /// every `dist_to_node` call within that operation. - /// - /// When the node's local vector storage is empty (graph-checkpoint-only - /// restore) and a `backing` is attached, the vector bytes are fetched from - /// the backing. This is the Lite cold-load path; Origin never hits it. - /// - /// A node whose vector cannot be resolved scores [`f32::INFINITY`] — search - /// then never selects it — rather than panicking. `with_backing` rejects a - /// backing that cannot serve the index, so reaching that branch means a node - /// has no vector source at all; aborting a search worker (and poisoning the - /// locks it held) is strictly worse than ranking one node last. - pub(crate) fn dist_to_node(&self, query_bytes: &[u8], node_id: u32) -> f32 { - distance_typed( - self.params.metric, - self.params.dtype, - query_bytes, - &self.node_bytes(node_id), - self.dim, - ) - .unwrap_or(f32::INFINITY) + /// An F32 index borrows the query. F16/BF16 indexes encode it once here, + /// so callers build one probe per search or insert. + pub(crate) fn probe<'a>(&self, query: &'a [f32]) -> Probe<'a> { + match self.params.dtype { + VectorStorageDtype::F32 => Probe::F32(query), + dtype => Probe::Encoded(Cow::Owned(cast_from_f32(query, dtype))), + } } - /// Node `node_id`'s vector bytes, encoded in `self.params.dtype`. - /// - /// Every distance computation on the insert, prune, and search paths reads - /// node vectors through this accessor. After a graph-checkpoint-only - /// restore, local storage is empty and the vector lives in the attached - /// backing. Reading `node.storage` directly there yields an empty slice, - /// and every distance against it errors. - /// - /// The backing holds F32 vectors. An F32 index borrows them as bytes. Any - /// other dtype encodes them into an owned buffer. + /// Node `node_id`'s vector as a probe, borrowed without copying. /// - /// Returns an empty slice when `node_id` is out of range, or the node has - /// no local vector and no backing serves one. Distance against an empty - /// slice errors, so callers score that node [`f32::INFINITY`]. - pub(crate) fn node_bytes(&self, node_id: u32) -> Cow<'_, [u8]> { - let local = self - .nodes - .get(node_id as usize) - .map_or(&[][..], |n| n.storage.as_bytes()); - #[cfg(not(target_arch = "wasm32"))] - { - if local.is_empty() - && let Some(v) = self.backing.as_ref().and_then(|b| b.get_vector(node_id)) - { - return match self.params.dtype { - VectorStorageDtype::F32 => Cow::Borrowed(bytemuck::cast_slice(v)), - dtype => Cow::Owned(cast_from_f32(v, dtype)), - }; - } + /// Returns `None` when the node has no vector to compare. + pub(crate) fn node_probe(&self, node_id: u32) -> Option> { + match self.params.dtype { + VectorStorageDtype::F32 => self.node_f32(node_id).map(Probe::F32), + _ => self + .node_encoded(node_id) + .map(|bytes| Probe::Encoded(Cow::Borrowed(bytes))), } - Cow::Borrowed(local) } - /// Compute distance between an f32 query and a stored node under `metric`. - /// - /// `metric` is caller-supplied, so filtered traversals that rank under a - /// metric other than the build metric use this accessor too. The node - /// vector comes from [`Self::node_bytes`], so restored nodes resolve - /// through the attached backing. + /// Distance from `probe` to node `node_id` under `metric`. /// - /// F32 indexes compare the query to the node vector without copying. - /// F16/BF16 indexes encode the query on each call, which allocates. + /// `metric` is caller-supplied, so filtered traversals can rank under a + /// metric other than the build metric. /// - /// A node whose vector cannot be resolved, or a query of the wrong length, - /// scores [`f32::INFINITY`]. - pub(crate) fn dist_to_node_f32( + /// Returns `None` when the node has no vector, for example after a + /// graph-only restore with no backing attached. Such a node is never a hit. + pub(crate) fn dist_to_node( &self, - query: &[f32], + probe: &Probe<'_>, node_id: u32, metric: DistanceMetric, - ) -> f32 { - let node_bytes = self.node_bytes(node_id); - match self.params.dtype { - VectorStorageDtype::F32 => match bytemuck::try_cast_slice::(&node_bytes) { - Ok(node) if node.len() == query.len() => distance(query, node, metric), - // Empty node or wrong query length. - Ok(_) => f32::INFINITY, - // Bytes not aligned to f32: `distance_typed` copies them out. - Err(_) => distance_typed( - metric, - VectorStorageDtype::F32, - bytemuck::cast_slice(query), - &node_bytes, - self.dim, - ) - .unwrap_or(f32::INFINITY), - }, - dtype => { - let query_bytes = cast_from_f32(query, dtype); - distance_typed(metric, dtype, &query_bytes, &node_bytes, self.dim) - .unwrap_or(f32::INFINITY) + ) -> Option { + match probe { + Probe::F32(query) => self + .node_f32(node_id) + .map(|node| distance(query, node, metric)), + Probe::Encoded(query) => { + let node = self.node_encoded(node_id)?; + distance_typed(metric, self.params.dtype, query, node, self.dim).ok() } } } + /// Node `node_id`'s f32 vector: local storage first, then the backing. + /// + /// The backing serves f32 vectors. Only F32 indexes hold backed nodes, + /// because a checkpoint restores an F32 index. + fn node_f32(&self, node_id: u32) -> Option<&[f32]> { + match &self.nodes.get(node_id as usize)?.storage { + NodeStorage::F32(v) if !v.is_empty() => Some(v.as_slice()), + NodeStorage::F32(_) => self.backing_f32(node_id), + NodeStorage::Bytes { .. } => None, + } + } + + /// Node `node_id`'s bytes in the index dtype, for F16/BF16 indexes. + fn node_encoded(&self, node_id: u32) -> Option<&[u8]> { + match &self.nodes.get(node_id as usize)?.storage { + NodeStorage::Bytes { bytes, .. } if !bytes.is_empty() => Some(bytes.as_slice()), + NodeStorage::Bytes { .. } | NodeStorage::F32(_) => None, + } + } + + /// Node `node_id`'s vector in the attached segment backing. + #[cfg(not(target_arch = "wasm32"))] + fn backing_f32(&self, node_id: u32) -> Option<&[f32]> { + self.backing.as_ref()?.get_vector(node_id) + } + + /// WASM targets have no segment backing (it requires mmap). + #[cfg(target_arch = "wasm32")] + fn backing_f32(&self, _node_id: u32) -> Option<&[f32]> { + None + } + /// Max neighbors allowed at a given layer. pub(crate) fn max_neighbors(&self, layer: usize) -> usize { if layer == 0 { diff --git a/nodedb-vector/src/hnsw/graph/index/state.rs b/nodedb-vector/src/hnsw/graph/index/state.rs index 7d13547b2..939392834 100644 --- a/nodedb-vector/src/hnsw/graph/index/state.rs +++ b/nodedb-vector/src/hnsw/graph/index/state.rs @@ -402,25 +402,23 @@ mod tests { assert_eq!(idx.materialize_vector(1).unwrap(), vec![4.0, 5.0, 6.0]); } - /// A node with no vector source must rank last, not abort the search worker. + /// A node with no vector source is not a hit. Search must neither abort the + /// worker nor return such a node at infinite distance. #[test] - fn search_on_a_vectorless_index_does_not_panic() { - let mut src = HnswIndex::with_seed(3, make_params(VectorStorageDtype::F32), 1); - src.insert(vec![1.0, 2.0, 3.0]).unwrap(); - let graph_only = src.graph_checkpoint_to_bytes().unwrap(); - // No backing attached: every node's storage is an empty placeholder. - let idx = HnswIndex::from_checkpoint(&graph_only) - .unwrap() - .expect("graph checkpoint must be recognized"); + fn search_before_with_backing_returns_no_vectorless_hits() { + use crate::test_support::{graph_only_restore, seeded_index, seeded_vectors}; - let results = idx.search(&[1.0, 2.0, 3.0], 5, 16).unwrap(); - for r in &results { - assert!( - r.distance.is_infinite(), - "a node with no vector source must score infinity, got {}", - r.distance - ); - } + let vectors = seeded_vectors(50, 3, 0x9E37_79B9_7F4A_7C15); + let src = seeded_index(&vectors, DistanceMetric::L2, VectorStorageDtype::F32); + // No backing attached: every node's storage is an empty placeholder. + let idx = graph_only_restore(&src); + assert_eq!(idx.len(), 50, "graph topology must survive"); + + let filter: roaring::RoaringBitmap = (0..50u32).collect(); + let unfiltered = idx.search(&vectors[0], 5, 64).unwrap(); + let filtered = idx.search_filtered(&vectors[0], 5, 64, &filter).unwrap(); + assert!(unfiltered.is_empty(), "got {unfiltered:?}"); + assert!(filtered.is_empty(), "got {filtered:?}"); } #[test] diff --git a/nodedb-vector/src/hnsw/graph/mod.rs b/nodedb-vector/src/hnsw/graph/mod.rs index 88e51bfde..61ea293e1 100644 --- a/nodedb-vector/src/hnsw/graph/mod.rs +++ b/nodedb-vector/src/hnsw/graph/mod.rs @@ -14,4 +14,5 @@ pub use limits::MAX_LAYER_CAP; pub use index::HnswIndex; pub use nodedb_types::hnsw::HnswParams; +pub(crate) use types::Probe; pub use types::{Candidate, Node, NodeStorage, SearchResult, Xorshift64}; diff --git a/nodedb-vector/src/hnsw/graph/types.rs b/nodedb-vector/src/hnsw/graph/types.rs index 955875a59..066308885 100644 --- a/nodedb-vector/src/hnsw/graph/types.rs +++ b/nodedb-vector/src/hnsw/graph/types.rs @@ -1,5 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 +use std::borrow::Cow; + use nodedb_types::vector_dtype::VectorStorageDtype; /// Result of a k-NN search. @@ -57,6 +59,17 @@ impl NodeStorage { } } +/// A vector scored against index nodes, in the form the index dtype compares. +/// +/// Build one per operation with `HnswIndex::probe` or `HnswIndex::node_probe`. +/// Scoring it against many nodes then copies nothing. +pub(crate) enum Probe<'a> { + /// f32 values, compared directly against f32 node vectors. + F32(&'a [f32]), + /// Bytes in the index dtype, compared with the typed kernels. + Encoded(Cow<'a, [u8]>), +} + /// A node in the HNSW graph. pub struct Node { /// Vector data in the index's configured storage dtype. diff --git a/nodedb-vector/src/hnsw/search.rs b/nodedb-vector/src/hnsw/search.rs index 3b82e9fbd..13c0fadf3 100644 --- a/nodedb-vector/src/hnsw/search.rs +++ b/nodedb-vector/src/hnsw/search.rs @@ -31,9 +31,8 @@ fn prefetch_t0(ptr: *const u8) { use roaring::RoaringBitmap; -use crate::dtype::cast_from_f32; use crate::error::{VectorError, check_dim}; -use crate::hnsw::graph::{Candidate, HnswIndex, SearchResult}; +use crate::hnsw::graph::{Candidate, HnswIndex, Probe, SearchResult}; /// Maximum beam width to prevent runaway search cost. const MAX_EF: usize = 8192; @@ -129,19 +128,19 @@ impl HnswIndex { return Ok(Vec::new()); }; - let query_bytes = cast_from_f32(query, self.params.dtype); + let probe = self.probe(query); // Phase 1: Greedy descent from top layer to layer 1. let mut current_ep = ep; for layer in (1..=self.max_layer).rev() { - let results = search_layer(self, &query_bytes, current_ep, 1, layer, None, 0); + let results = search_layer(self, &probe, current_ep, 1, layer, None, 0); if let Some(nearest) = results.first() { current_ep = nearest.id; } } // Phase 2: Beam search at layer 0. - let results = search_layer(self, &query_bytes, current_ep, ef, 0, filter, id_offset); + let results = search_layer(self, &probe, current_ep, ef, 0, filter, id_offset); Ok(results .into_iter() @@ -167,19 +166,23 @@ pub fn decode_filter_bitmap(bytes: &[u8]) -> Result /// When `filter` is `Some`, only nodes present in the bitmap enter results, /// but all nodes are still traversed for graph connectivity. /// +/// Distances use the build metric. A node without a vector is traversed at +/// infinite distance and never enters the results. +/// /// Scratch buffers are drawn from `index.arena` (a `RefCell`-guarded /// `BeamSearchArena`). The arena is reset at entry and its `Vec`/`HashSet` /// capacity grows to the high-water mark over successive calls, giving /// amortised zero-allocation steady state. pub(crate) fn search_layer( index: &HnswIndex, - query_bytes: &[u8], + probe: &Probe<'_>, entry_point: u32, ef: usize, layer: usize, filter: Option<&RoaringBitmap>, id_offset: u32, ) -> Vec { + let metric = index.params.metric; let mut arena = index.arena.borrow_mut(); // Reset scratch buffers — retains Vec/HashSet capacity from prior calls. @@ -187,9 +190,9 @@ pub(crate) fn search_layer( arena.visited.insert(entry_point); - let ep_dist = index.dist_to_node(query_bytes, entry_point); + let ep_dist = index.dist_to_node(probe, entry_point, metric); let ep_candidate = Candidate { - dist: ep_dist, + dist: ep_dist.unwrap_or(f32::INFINITY), id: entry_point, }; @@ -213,7 +216,7 @@ pub(crate) fn search_layer( } }; - if passes_filter(entry_point) { + if ep_dist.is_some() && passes_filter(entry_point) { res_vec.push(ep_candidate); } let mut results = BinaryHeap::from(res_vec); @@ -245,20 +248,20 @@ pub(crate) fn search_layer( continue; } - let dist = index.dist_to_node(query_bytes, neighbor_id); + let scored = index.dist_to_node(probe, neighbor_id, metric); let neighbor = Candidate { - dist, + dist: scored.unwrap_or(f32::INFINITY), id: neighbor_id, }; let worst_dist = results.peek().map_or(f32::INFINITY, |w| w.dist); - let should_explore = dist < worst_dist || results.len() < ef; + let should_explore = neighbor.dist < worst_dist || results.len() < ef; if should_explore { candidates.push(Reverse(neighbor)); } - if passes_filter(neighbor_id) { + if scored.is_some() && passes_filter(neighbor_id) { results.push(neighbor); if results.len() > ef { results.pop(); diff --git a/nodedb-vector/src/navix/acorn.rs b/nodedb-vector/src/navix/acorn.rs index 7cd992044..e4689951b 100644 --- a/nodedb-vector/src/navix/acorn.rs +++ b/nodedb-vector/src/navix/acorn.rs @@ -11,13 +11,13 @@ #[cfg(feature = "acorn-baseline")] mod inner { - use std::cmp::Reverse; - use std::collections::{BinaryHeap, HashSet}; + use std::collections::HashSet; use roaring::RoaringBitmap; use crate::error::{VectorError, check_dim}; - use crate::hnsw::graph::{Candidate, HnswIndex}; + use crate::hnsw::graph::{Candidate, HnswIndex, Probe}; + use crate::navix::beam::Beam; use crate::navix::traversal::SearchResult; /// Options for an ACORN-1 search (same shape as `NavixSearchOptions`). @@ -44,7 +44,8 @@ mod inner { /// /// Uses static 2-hop expansion: when a 1-hop neighbor is not in `allowed`, /// expand to its 2-hop neighbors unconditionally. A query without the - /// index dimension fails with [`VectorError::DimensionMismatch`]. + /// index dimension fails with [`VectorError::DimensionMismatch`]. A node + /// without a vector is never returned. pub fn acorn_search( index: &HnswIndex, query: &[f32], @@ -55,6 +56,7 @@ mod inner { if index.is_empty() || options.allowed.is_empty() || options.k == 0 { return Ok(Vec::new()); } + let probe = index.probe(query); let total = index.len(); let global_sel = options.allowed.len() as f64 / total as f64; @@ -62,7 +64,7 @@ mod inner { if global_sel < options.brute_force_threshold { return Ok(brute_force_on_allowed( index, - query, + &probe, options.k, &options.allowed, metric, @@ -76,15 +78,12 @@ mod inner { // Phase 1: greedy descent (unfiltered) to find best layer-0 entry. let mut current_ep = ep; for layer in (1..=index.max_layer()).rev() { - let results = greedy_descent_layer(index, query, current_ep, layer, metric); - if let Some(nearest) = results.first() { - current_ep = nearest.id; - } + current_ep = greedy_descent_layer(index, &probe, current_ep, layer, metric).id; } // Phase 2: ACORN-1 filtered beam search at layer 0. let ef = options.ef_search.max(options.k); - let results = acorn_search_layer_0(index, query, current_ep, ef, &options.allowed, metric); + let results = acorn_search_layer_0(index, &probe, current_ep, ef, &options.allowed, metric); Ok(results .into_iter() @@ -97,15 +96,19 @@ mod inner { } /// Minimal greedy single-layer descent used for Phase-1 layer navigation. + /// + /// A node without a vector is never chosen over the entry point. fn greedy_descent_layer( index: &HnswIndex, - query: &[f32], + probe: &Probe<'_>, entry_point: u32, layer: usize, metric: nodedb_types::vector_distance::DistanceMetric, - ) -> Vec { + ) -> Candidate { let mut current = entry_point; - let mut current_dist = node_dist(index, query, current, metric); + let mut current_dist = index + .dist_to_node(probe, current, metric) + .unwrap_or(f32::INFINITY); loop { let mut improved = false; @@ -113,8 +116,9 @@ mod inner { if index.is_deleted(nb) { continue; } - let d = node_dist(index, query, nb, metric); - if d < current_dist { + if let Some(d) = index.dist_to_node(probe, nb, metric) + && d < current_dist + { current_dist = d; current = nb; improved = true; @@ -124,10 +128,10 @@ mod inner { break; } } - vec![Candidate { + Candidate { dist: current_dist, id: current, - }] + } } /// ACORN-1 layer-0 beam search. @@ -136,101 +140,60 @@ mod inner { /// neighbors and score those (the "static 2-hop" rule). fn acorn_search_layer_0( index: &HnswIndex, - query: &[f32], + probe: &Probe<'_>, entry_point: u32, ef: usize, allowed: &RoaringBitmap, metric: nodedb_types::vector_distance::DistanceMetric, ) -> Vec { + let admitted = |id: u32| !index.is_deleted(id) && allowed.contains(id); let mut visited: HashSet = HashSet::new(); visited.insert(entry_point); - let ep_dist = node_dist(index, query, entry_point, metric); - let ep_cand = Candidate { - dist: ep_dist, - id: entry_point, - }; - - let mut candidates: BinaryHeap> = BinaryHeap::new(); - candidates.push(Reverse(ep_cand)); - - let mut results: BinaryHeap = BinaryHeap::new(); - if !index.is_deleted(entry_point) && allowed.contains(entry_point) { - results.push(ep_cand); - } - - while let Some(Reverse(current)) = candidates.pop() { - if let Some(worst) = results.peek() - && current.dist > worst.dist - && results.len() >= ef - { - break; - } + let mut beam = Beam::new(ef); + beam.offer( + entry_point, + index.dist_to_node(probe, entry_point, metric), + admitted(entry_point), + ); + while let Some(current) = beam.pop_frontier() { for &nb in index.neighbors_at(current.id, 0) { if !visited.insert(nb) { continue; } - - let d = node_dist(index, query, nb, metric); - let nb_cand = Candidate { dist: d, id: nb }; - let worst_dist = results.peek().map_or(f32::INFINITY, |w| w.dist); - - if d < worst_dist || results.len() < ef { - candidates.push(Reverse(nb_cand)); + let admit = admitted(nb); + beam.offer(nb, index.dist_to_node(probe, nb, metric), admit); + if admit { + continue; } - - if !index.is_deleted(nb) && allowed.contains(nb) { - // 1-hop neighbor passes filter — add to results. - results.push(nb_cand); - if results.len() > ef { - results.pop(); - } - } else { - // 1-hop neighbor does not pass filter — static 2-hop expansion. - for &nb2 in index.neighbors_at(nb, 0) { - if !visited.insert(nb2) { - continue; - } - let d2 = node_dist(index, query, nb2, metric); - let nb2_cand = Candidate { dist: d2, id: nb2 }; - let worst_dist2 = results.peek().map_or(f32::INFINITY, |w| w.dist); - if d2 < worst_dist2 || results.len() < ef { - candidates.push(Reverse(nb2_cand)); - } - if !index.is_deleted(nb2) && allowed.contains(nb2) { - results.push(nb2_cand); - if results.len() > ef { - results.pop(); - } - } + // 1-hop neighbor does not pass filter — static 2-hop expansion. + for &nb2 in index.neighbors_at(nb, 0) { + if !visited.insert(nb2) { + continue; } + beam.offer(nb2, index.dist_to_node(probe, nb2, metric), admitted(nb2)); } } } - let mut v: Vec = results.into_vec(); - v.sort_unstable_by(|a, b| a.dist.total_cmp(&b.dist)); - v + beam.into_sorted() } fn brute_force_on_allowed( index: &HnswIndex, - query: &[f32], + probe: &Probe<'_>, k: usize, allowed: &RoaringBitmap, metric: nodedb_types::vector_distance::DistanceMetric, ) -> Vec { + // A node without a vector has no distance and is not a hit. let mut results: Vec = allowed .iter() + .filter(|&id| !index.is_deleted(id)) .filter_map(|id| { - if id as usize >= index.len() || index.is_deleted(id) { - return None; - } - Some(SearchResult { - id, - distance: index.dist_to_node_f32(query, id, metric), - }) + let distance = index.dist_to_node(probe, id, metric)?; + Some(SearchResult { id, distance }) }) .collect(); @@ -250,16 +213,6 @@ mod inner { results } - #[inline] - fn node_dist( - index: &HnswIndex, - query: &[f32], - node_id: u32, - metric: nodedb_types::vector_distance::DistanceMetric, - ) -> f32 { - index.dist_to_node_f32(query, node_id, metric) - } - #[cfg(test)] mod tests { use super::*; @@ -353,6 +306,65 @@ mod inner { assert_eq!(r.id, 7); } } + + const SEED: u64 = 0x9E37_79B9_7F4A_7C15; + + /// Options that brute-force (`true`) or walk the graph (`false`). + fn options(allowed: RoaringBitmap, brute: bool) -> AcornSearchOptions { + AcornSearchOptions { + k: 5, + ef_search: 64, + allowed, + brute_force_threshold: if brute { 2.0 } else { 0.0 }, + } + } + + fn pairs(results: &[SearchResult]) -> Vec<(u32, f32)> { + results.iter().map(|r| (r.id, r.distance)).collect() + } + + /// After a graph-only restore with a backing, both ACORN paths rank + /// exactly like the source index, under a caller metric too. + #[cfg(not(target_arch = "wasm32"))] + #[test] + fn acorn_after_graph_only_restore_matches_the_source_index() { + use crate::test_support::{ + graph_only_restore_with_backing, seeded_index, seeded_vectors, + }; + use nodedb_types::vector_dtype::VectorStorageDtype; + + let vectors = seeded_vectors(300, 8, SEED); + let src = seeded_index(&vectors, DistanceMetric::L2, VectorStorageDtype::F32); + let restored = graph_only_restore_with_backing(&src, &vectors); + let allowed: RoaringBitmap = (0..300u32).step_by(3).collect(); + for metric in [DistanceMetric::L2, DistanceMetric::Cosine] { + for brute in [true, false] { + let opts = options(allowed.clone(), brute); + let expected = acorn_search(&src, &vectors[7], &opts, metric).unwrap(); + let got = acorn_search(&restored, &vectors[7], &opts, metric).unwrap(); + assert_eq!(expected.len(), 5, "{metric:?} brute={brute}"); + assert_eq!(pairs(&got), pairs(&expected), "{metric:?} brute={brute}"); + } + } + } + + /// A graph-only restore searched before `with_backing` has no vectors. + /// Neither ACORN path may return its nodes. + #[test] + fn acorn_without_backing_returns_no_hits() { + use crate::test_support::{graph_only_restore, seeded_index, seeded_vectors}; + use nodedb_types::vector_dtype::VectorStorageDtype; + + let vectors = seeded_vectors(300, 8, SEED); + let src = seeded_index(&vectors, DistanceMetric::L2, VectorStorageDtype::F32); + let idx = graph_only_restore(&src); + let allowed: RoaringBitmap = (0..300u32).step_by(3).collect(); + for brute in [true, false] { + let opts = options(allowed.clone(), brute); + let res = acorn_search(&idx, &vectors[7], &opts, DistanceMetric::L2).unwrap(); + assert!(res.is_empty(), "brute={brute}: got {res:?}"); + } + } } } diff --git a/nodedb-vector/src/navix/beam.rs b/nodedb-vector/src/navix/beam.rs new file mode 100644 index 000000000..4dbf59f48 --- /dev/null +++ b/nodedb-vector/src/navix/beam.rs @@ -0,0 +1,100 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Frontier and result heaps shared by the NaviX and ACORN traversals. + +use std::cmp::Reverse; +use std::collections::BinaryHeap; + +use crate::hnsw::graph::Candidate; + +/// Heaps for one single-layer beam search of width `ef`. +pub(crate) struct Beam { + frontier: BinaryHeap>, + results: BinaryHeap, + ef: usize, +} + +impl Beam { + pub(crate) fn new(ef: usize) -> Self { + Self { + frontier: BinaryHeap::new(), + results: BinaryHeap::new(), + ef, + } + } + + /// Offers node `id` at distance `dist`. + /// + /// The node joins the frontier while it can still improve the results. + /// It joins the results only when `admit` holds and it has a distance. + /// A node without a vector stays traversable at infinite distance. + pub(crate) fn offer(&mut self, id: u32, dist: Option, admit: bool) { + let candidate = Candidate { + dist: dist.unwrap_or(f32::INFINITY), + id, + }; + let worst = self.results.peek().map_or(f32::INFINITY, |w| w.dist); + if candidate.dist < worst || self.results.len() < self.ef { + self.frontier.push(Reverse(candidate)); + } + if admit && dist.is_some() { + self.results.push(candidate); + if self.results.len() > self.ef { + self.results.pop(); + } + } + } + + /// Pops the nearest frontier node. + /// + /// Returns `None` when the frontier is empty, or when `ef` results exist + /// and the nearest frontier node is farther than all of them. + pub(crate) fn pop_frontier(&mut self) -> Option { + let Reverse(current) = self.frontier.pop()?; + if let Some(worst) = self.results.peek() + && current.dist > worst.dist + && self.results.len() >= self.ef + { + return None; + } + Some(current) + } + + /// The results, nearest first. + pub(crate) fn into_sorted(self) -> Vec { + let mut results = self.results.into_vec(); + results.sort_unstable_by(|a, b| a.dist.total_cmp(&b.dist)); + results + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_node_without_a_vector_is_traversed_but_never_a_result() { + let mut beam = Beam::new(4); + beam.offer(7, None, true); + beam.offer(3, Some(1.0), true); + let first = beam.pop_frontier().expect("scored node is on the frontier"); + assert_eq!(first.id, 3); + let second = beam + .pop_frontier() + .expect("vectorless node stays traversable"); + assert_eq!(second.id, 7); + let ids: Vec = beam.into_sorted().iter().map(|c| c.id).collect(); + assert_eq!(ids, vec![3]); + } + + #[test] + fn results_keep_the_ef_nearest_admitted_nodes() { + let mut beam = Beam::new(2); + beam.offer(1, Some(3.0), true); + beam.offer(2, Some(1.0), false); + beam.offer(3, Some(2.0), true); + beam.offer(4, Some(0.5), true); + let ids: Vec = beam.into_sorted().iter().map(|c| c.id).collect(); + assert_eq!(ids, vec![4, 3]); + } +} diff --git a/nodedb-vector/src/navix/mod.rs b/nodedb-vector/src/navix/mod.rs index 853ce75a3..603a253b6 100644 --- a/nodedb-vector/src/navix/mod.rs +++ b/nodedb-vector/src/navix/mod.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 pub mod acorn; +mod beam; pub mod selectivity; pub mod traversal; diff --git a/nodedb-vector/src/navix/traversal.rs b/nodedb-vector/src/navix/traversal.rs index 699c349ab..057dc490f 100644 --- a/nodedb-vector/src/navix/traversal.rs +++ b/nodedb-vector/src/navix/traversal.rs @@ -9,13 +9,13 @@ //! //! See `selectivity.rs` for heuristic boundary definitions. -use std::cmp::Reverse; -use std::collections::{BinaryHeap, HashSet}; +use std::collections::HashSet; use roaring::RoaringBitmap; use crate::error::{VectorError, check_dim}; -use crate::hnsw::graph::{Candidate, HnswIndex}; +use crate::hnsw::graph::{Candidate, HnswIndex, Probe}; +use crate::navix::beam::Beam; use crate::navix::selectivity::{NavixHeuristic, local_selectivity_at, pick_heuristic}; /// A k-NN result returned by `navix_search`. @@ -60,6 +60,7 @@ impl Default for NavixSearchOptions { /// candidate IDs must be present in `options.allowed`. /// /// Returns an empty Vec when the index is empty or `options.allowed` is empty. +/// A node without a vector is never returned. /// /// # Errors /// @@ -75,6 +76,7 @@ pub fn navix_search( if index.is_empty() || options.allowed.is_empty() || options.k == 0 { return Ok(Vec::new()); } + let probe = index.probe(query); let total = index.len(); let global_sel = options.allowed.len() as f64 / total as f64; @@ -82,7 +84,7 @@ pub fn navix_search( if global_sel < options.brute_force_threshold { return Ok(brute_force_on_allowed( index, - query, + &probe, options.k, &options.allowed, metric, @@ -97,7 +99,7 @@ pub fn navix_search( // standard HNSW — we just want the best entry point for layer 0). let mut current_ep = ep; for layer in (1..=index.max_layer()).rev() { - let results = unfiltered_search_layer(index, query, current_ep, 1, layer, metric); + let results = unfiltered_search_layer(index, &probe, current_ep, 1, layer, metric); if let Some(nearest) = results.first() { current_ep = nearest.id; } @@ -105,7 +107,7 @@ pub fn navix_search( // Phase 2: adaptive-local filtered beam search at layer 0. let ef = options.ef_search.max(options.k); - let results = navix_search_layer_0(index, query, current_ep, ef, &options.allowed, metric); + let results = navix_search_layer_0(index, &probe, current_ep, ef, &options.allowed, metric); Ok(results .into_iter() @@ -123,21 +125,18 @@ pub fn navix_search( /// selectivity drops below the configured threshold. fn brute_force_on_allowed( index: &HnswIndex, - query: &[f32], + probe: &Probe<'_>, k: usize, allowed: &RoaringBitmap, metric: nodedb_types::vector_distance::DistanceMetric, ) -> Vec { + // A node without a vector has no distance and is not a hit. let mut results: Vec = allowed .iter() + .filter(|&id| !index.is_deleted(id)) .filter_map(|id| { - if id as usize >= index.len() || index.is_deleted(id) { - return None; - } - Some(SearchResult { - id, - distance: index.dist_to_node_f32(query, id, metric), - }) + let distance = index.dist_to_node(probe, id, metric)?; + Some(SearchResult { id, distance }) }) .collect(); @@ -157,11 +156,16 @@ fn brute_force_on_allowed( results } +/// Whether node `id` may enter the results of a filtered search. +fn admitted(index: &HnswIndex, allowed: &RoaringBitmap, id: u32) -> bool { + !index.is_deleted(id) && allowed.contains(id) +} + /// Standard unfiltered single-layer beam search used for Phase-1 greedy /// descent (layers 1..max_layer). fn unfiltered_search_layer( index: &HnswIndex, - query: &[f32], + probe: &Probe<'_>, entry_point: u32, ef: usize, layer: usize, @@ -170,50 +174,27 @@ fn unfiltered_search_layer( let mut visited: HashSet = HashSet::new(); visited.insert(entry_point); - let ep_dist = dist(index, query, entry_point, metric); - let ep_cand = Candidate { - dist: ep_dist, - id: entry_point, - }; - - let mut candidates: BinaryHeap> = BinaryHeap::new(); - candidates.push(Reverse(ep_cand)); - - let mut results: BinaryHeap = BinaryHeap::new(); - if !index.is_deleted(entry_point) { - results.push(ep_cand); - } - - while let Some(Reverse(current)) = candidates.pop() { - if let Some(worst) = results.peek() - && current.dist > worst.dist - && results.len() >= ef - { - break; - } + let mut beam = Beam::new(ef); + beam.offer( + entry_point, + index.dist_to_node(probe, entry_point, metric), + !index.is_deleted(entry_point), + ); + while let Some(current) = beam.pop_frontier() { for &nb in index.neighbors_at(current.id, layer) { if !visited.insert(nb) { continue; } - let d = dist(index, query, nb, metric); - let nb_cand = Candidate { dist: d, id: nb }; - let worst_dist = results.peek().map_or(f32::INFINITY, |w| w.dist); - if d < worst_dist || results.len() < ef { - candidates.push(Reverse(nb_cand)); - } - if !index.is_deleted(nb) { - results.push(nb_cand); - if results.len() > ef { - results.pop(); - } - } + beam.offer( + nb, + index.dist_to_node(probe, nb, metric), + !index.is_deleted(nb), + ); } } - let mut v: Vec = results.into_vec(); - v.sort_unstable_by(|a, b| a.dist.total_cmp(&b.dist)); - v + beam.into_sorted() } /// NaviX adaptive-local filtered beam search at layer 0. @@ -224,7 +205,7 @@ fn unfiltered_search_layer( /// - **Blind**: skip 1-hop scoring; sample 2-hop of all 1-hop neighbors. fn navix_search_layer_0( index: &HnswIndex, - query: &[f32], + probe: &Probe<'_>, entry_point: u32, ef: usize, allowed: &RoaringBitmap, @@ -233,124 +214,66 @@ fn navix_search_layer_0( let mut visited: HashSet = HashSet::new(); visited.insert(entry_point); - let ep_dist = dist(index, query, entry_point, metric); - let ep_cand = Candidate { - dist: ep_dist, - id: entry_point, - }; - - let mut candidates: BinaryHeap> = BinaryHeap::new(); - candidates.push(Reverse(ep_cand)); - - let mut results: BinaryHeap = BinaryHeap::new(); - // Entry point enters results only if it is allowed. - if !index.is_deleted(entry_point) && allowed.contains(entry_point) { - results.push(ep_cand); - } - - while let Some(Reverse(current)) = candidates.pop() { - if let Some(worst) = results.peek() - && current.dist > worst.dist - && results.len() >= ef - { - break; - } - + let mut beam = Beam::new(ef); + beam.offer( + entry_point, + index.dist_to_node(probe, entry_point, metric), + admitted(index, allowed, entry_point), + ); + + while let Some(current) = beam.pop_frontier() { let neighbors_1hop = index.neighbors_at(current.id, 0); - let local_sel = local_selectivity_at(neighbors_1hop, allowed); - let heuristic = pick_heuristic(local_sel); - - match heuristic { - NavixHeuristic::Standard => { - expand_standard(ExpandArgs { - index, - query, - neighbors_1hop, - allowed, - ef, - metric, - visited: &mut visited, - candidates: &mut candidates, - results: &mut results, - }); - } - NavixHeuristic::Directed => { - expand_directed(ExpandArgs { - index, - query, - neighbors_1hop, - allowed, - ef, - metric, - visited: &mut visited, - candidates: &mut candidates, - results: &mut results, - }); - } - NavixHeuristic::Blind => { - expand_blind(ExpandArgs { - index, - query, - neighbors_1hop, - allowed, - ef, - metric, - visited: &mut visited, - candidates: &mut candidates, - results: &mut results, - }); - } + let args = ExpandArgs { + index, + probe, + neighbors_1hop, + allowed, + metric, + visited: &mut visited, + beam: &mut beam, + }; + match pick_heuristic(local_selectivity_at(neighbors_1hop, allowed)) { + NavixHeuristic::Standard => expand_standard(args), + NavixHeuristic::Directed => expand_directed(args), + NavixHeuristic::Blind => expand_blind(args), } } - let mut v: Vec = results.into_vec(); - v.sort_unstable_by(|a, b| a.dist.total_cmp(&b.dist)); - v + beam.into_sorted() } /// Shared beam-search state bundled for the per-heuristic `expand_*` helpers. struct ExpandArgs<'a> { index: &'a HnswIndex, - query: &'a [f32], + probe: &'a Probe<'a>, neighbors_1hop: &'a [u32], allowed: &'a RoaringBitmap, - ef: usize, metric: nodedb_types::vector_distance::DistanceMetric, visited: &'a mut HashSet, - candidates: &'a mut BinaryHeap>, - results: &'a mut BinaryHeap, + beam: &'a mut Beam, } /// Standard expansion: score every allowed 1-hop neighbor and add to heaps. fn expand_standard(args: ExpandArgs<'_>) { let ExpandArgs { index, - query, + probe, neighbors_1hop, allowed, - ef, metric, visited, - candidates, - results, + beam, } = args; for &nb in neighbors_1hop { if !visited.insert(nb) { continue; } - let d = dist(index, query, nb, metric); - let nb_cand = Candidate { dist: d, id: nb }; - let worst_dist = results.peek().map_or(f32::INFINITY, |w| w.dist); - if d < worst_dist || results.len() < ef { - candidates.push(Reverse(nb_cand)); - } - if !index.is_deleted(nb) && allowed.contains(nb) { - results.push(nb_cand); - if results.len() > ef { - results.pop(); - } - } + beam.offer( + nb, + index.dist_to_node(probe, nb, metric), + admitted(index, allowed, nb), + ); } } @@ -359,40 +282,29 @@ fn expand_standard(args: ExpandArgs<'_>) { fn expand_directed(args: ExpandArgs<'_>) { let ExpandArgs { index, - query, + probe, neighbors_1hop, allowed, - ef, metric, visited, - candidates, - results, + beam, } = args; - // Score 1-hop; track the best allowed neighbor. + // Score 1-hop; track the best allowed neighbor that has a distance. let mut best_allowed: Option<(u32, f32)> = None; for &nb in neighbors_1hop { - let already_visited = !visited.insert(nb); - if already_visited { + if !visited.insert(nb) { continue; } - let d = dist(index, query, nb, metric); - let nb_cand = Candidate { dist: d, id: nb }; - - let worst_dist = results.peek().map_or(f32::INFINITY, |w| w.dist); - if d < worst_dist || results.len() < ef { - candidates.push(Reverse(nb_cand)); - } - - if !index.is_deleted(nb) && allowed.contains(nb) { - if best_allowed.is_none_or(|(_, bd)| d < bd) { - best_allowed = Some((nb, d)); - } - results.push(nb_cand); - if results.len() > ef { - results.pop(); - } + let d = index.dist_to_node(probe, nb, metric); + let admit = admitted(index, allowed, nb); + if admit + && let Some(d) = d + && best_allowed.is_none_or(|(_, bd)| d < bd) + { + best_allowed = Some((nb, d)); } + beam.offer(nb, d, admit); } // Expand 2-hop of the single best allowed neighbor. @@ -401,18 +313,11 @@ fn expand_directed(args: ExpandArgs<'_>) { if !visited.insert(nb2) { continue; } - let d = dist(index, query, nb2, metric); - let nb2_cand = Candidate { dist: d, id: nb2 }; - let worst_dist = results.peek().map_or(f32::INFINITY, |w| w.dist); - if d < worst_dist || results.len() < ef { - candidates.push(Reverse(nb2_cand)); - } - if !index.is_deleted(nb2) && allowed.contains(nb2) { - results.push(nb2_cand); - if results.len() > ef { - results.pop(); - } - } + beam.offer( + nb2, + index.dist_to_node(probe, nb2, metric), + admitted(index, allowed, nb2), + ); } } } @@ -422,14 +327,12 @@ fn expand_directed(args: ExpandArgs<'_>) { fn expand_blind(args: ExpandArgs<'_>) { let ExpandArgs { index, - query, + probe, neighbors_1hop, allowed, - ef, metric, visited, - candidates, - results, + beam, } = args; for &nb1 in neighbors_1hop { // Mark 1-hop as visited so we do not double-score them later, @@ -437,40 +340,14 @@ fn expand_blind(args: ExpandArgs<'_>) { visited.insert(nb1); for &nb2 in index.neighbors_at(nb1, 0) { - if !visited.insert(nb2) { + if !visited.insert(nb2) || !admitted(index, allowed, nb2) { continue; } - if index.is_deleted(nb2) { - continue; - } - if !allowed.contains(nb2) { - continue; - } - let d = dist(index, query, nb2, metric); - let nb2_cand = Candidate { dist: d, id: nb2 }; - let worst_dist = results.peek().map_or(f32::INFINITY, |w| w.dist); - if d < worst_dist || results.len() < ef { - candidates.push(Reverse(nb2_cand)); - } - results.push(nb2_cand); - if results.len() > ef { - results.pop(); - } + beam.offer(nb2, index.dist_to_node(probe, nb2, metric), true); } } } -/// Inline helper: distance from query to a stored node using the given metric. -#[inline] -fn dist( - index: &HnswIndex, - query: &[f32], - node_id: u32, - metric: nodedb_types::vector_distance::DistanceMetric, -) -> f32 { - index.dist_to_node_f32(query, node_id, metric) -} - // ── Tests ───────────────────────────────────────────────────────────────────── #[cfg(test)] @@ -678,4 +555,125 @@ mod tests { "{result:?}" ); } + + const SEED: u64 = 0x9E37_79B9_7F4A_7C15; + /// Global selectivity is always below this, so the search brute-forces. + const BRUTE_FORCE: f64 = 2.0; + /// Global selectivity is never below this, so the search walks the graph. + const GRAPH: f64 = 0.0; + + /// Allowed sets dense and sparse enough to reach every NaviX heuristic. + fn allowed_sets() -> Vec { + [2usize, 3, 50] + .into_iter() + .map(|step| (0..300u32).step_by(step).collect()) + .collect() + } + + fn options(allowed: RoaringBitmap, brute_force_threshold: f64) -> NavixSearchOptions { + NavixSearchOptions { + k: 5, + ef_search: 64, + allowed, + brute_force_threshold, + } + } + + fn pairs(results: &[SearchResult]) -> Vec<(u32, f32)> { + results.iter().map(|r| (r.id, r.distance)).collect() + } + + /// After a graph-only restore with a backing, brute force and graph + /// traversal rank exactly like the source index, under the build metric + /// and under a different caller metric. + #[cfg(not(target_arch = "wasm32"))] + #[test] + fn graph_only_restore_with_backing_matches_the_source_index() { + use crate::test_support::{graph_only_restore_with_backing, seeded_index, seeded_vectors}; + use nodedb_types::vector_dtype::VectorStorageDtype; + + let vectors = seeded_vectors(300, 8, SEED); + let src = seeded_index(&vectors, DistanceMetric::L2, VectorStorageDtype::F32); + let restored = graph_only_restore_with_backing(&src, &vectors); + let mut graph_hits = 0; + for metric in [DistanceMetric::L2, DistanceMetric::Cosine] { + for (threshold, brute) in [(BRUTE_FORCE, true), (GRAPH, false)] { + for allowed in allowed_sets() { + let opts = options(allowed, threshold); + let expected = navix_search(&src, &vectors[7], &opts, metric).unwrap(); + let got = navix_search(&restored, &vectors[7], &opts, metric).unwrap(); + assert!(!brute || expected.len() == 5, "{metric:?}: {expected:?}"); + assert_eq!(pairs(&got), pairs(&expected), "{metric:?} brute={brute}"); + if !brute { + graph_hits += got.len(); + } + } + } + } + assert!(graph_hits > 0, "graph traversal must return hits"); + } + + /// A graph-only restore searched before `with_backing` has no vectors. + /// Neither brute force nor graph traversal may return its nodes. + #[test] + fn graph_only_restore_without_backing_returns_no_hits() { + use crate::test_support::{graph_only_restore, seeded_index, seeded_vectors}; + use nodedb_types::vector_dtype::VectorStorageDtype; + + let vectors = seeded_vectors(300, 8, SEED); + let src = seeded_index(&vectors, DistanceMetric::L2, VectorStorageDtype::F32); + let idx = graph_only_restore(&src); + for threshold in [BRUTE_FORCE, GRAPH] { + for allowed in allowed_sets() { + let opts = options(allowed, threshold); + let res = navix_search(&idx, &vectors[7], &opts, DistanceMetric::L2).unwrap(); + assert!(res.is_empty(), "threshold={threshold}: got {res:?}"); + } + } + } + + /// Every dtype scores hits under the caller metric, not the build metric. + /// Brute force returns the exact top-k under that metric. + #[test] + fn caller_metric_overrides_the_build_metric_for_every_dtype() { + use crate::test_support::{seeded_index, seeded_vectors}; + use nodedb_types::vector_dtype::VectorStorageDtype; + + let vectors = seeded_vectors(300, 8, SEED); + let allowed: RoaringBitmap = (0..300u32).step_by(3).collect(); + for dtype in [ + VectorStorageDtype::F32, + VectorStorageDtype::F16, + VectorStorageDtype::BF16, + ] { + let idx = seeded_index(&vectors, DistanceMetric::L2, dtype); + let encoded = crate::dtype::cast_from_f32(&vectors[7], dtype); + let query = crate::dtype::cast_to_f32(&encoded, dtype, 8).unwrap(); + let cosine_to = |id: u32| { + let node = idx.materialize_vector(id).unwrap(); + distance(&query, &node, DistanceMetric::Cosine) + }; + let mut reference: Vec<(u32, f32)> = + allowed.iter().map(|id| (id, cosine_to(id))).collect(); + reference.sort_by(|a, b| a.1.total_cmp(&b.1)); + + for (threshold, brute) in [(BRUTE_FORCE, true), (GRAPH, false)] { + let opts = options(allowed.clone(), threshold); + let res = navix_search(&idx, &vectors[7], &opts, DistanceMetric::Cosine).unwrap(); + assert_eq!(res.len(), 5, "dtype={dtype:?} brute={brute}"); + for hit in &res { + let want = cosine_to(hit.id); + assert!( + (hit.distance - want).abs() < 1e-4, + "dtype={dtype:?} brute={brute}: {hit:?} vs cosine {want}" + ); + } + if brute { + let ids: Vec = res.iter().map(|r| r.id).collect(); + let want: Vec = reference.iter().take(5).map(|r| r.0).collect(); + assert_eq!(ids, want, "dtype={dtype:?}"); + } + } + } + } } diff --git a/nodedb-vector/src/test_support.rs b/nodedb-vector/src/test_support.rs index 2cc3710ae..8940a7b4d 100644 --- a/nodedb-vector/src/test_support.rs +++ b/nodedb-vector/src/test_support.rs @@ -27,3 +27,97 @@ pub(crate) fn test_memory() -> ScopedMemory { EngineId::Vector, ) } + +/// Deterministic vectors in `[0, 1)` from an xorshift64 stream. +/// +/// `seed` must be non-zero. +pub(crate) fn seeded_vectors(count: usize, dim: usize, seed: u64) -> Vec> { + let mut state = seed; + (0..count) + .map(|_| { + (0..dim) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + (state % 10_000) as f32 / 10_000.0 + }) + .collect() + }) + .collect() +} + +/// An HNSW index over `vectors`, built with a fixed seed. +pub(crate) fn seeded_index( + vectors: &[Vec], + metric: crate::distance::DistanceMetric, + dtype: nodedb_types::vector_dtype::VectorStorageDtype, +) -> crate::hnsw::HnswIndex { + let dim = vectors.first().map_or(0, Vec::len); + let params = crate::hnsw::HnswParams { + m: 4, + m0: 8, + ef_construction: 32, + metric, + dtype, + }; + let mut index = crate::hnsw::HnswIndex::with_seed(dim, params, 12345); + for v in vectors { + index + .insert(v.clone()) + .expect("seeded vector has the index dim"); + } + index +} + +/// `src` restored from its graph-only checkpoint with no backing attached. +/// +/// Every node of the result has no vector. +pub(crate) fn graph_only_restore(src: &crate::hnsw::HnswIndex) -> crate::hnsw::HnswIndex { + let bytes = src + .graph_checkpoint_to_bytes() + .expect("graph checkpoint encodes"); + crate::hnsw::HnswIndex::from_checkpoint(&bytes) + .expect("graph checkpoint decodes") + .expect("graph checkpoint is recognized") +} + +/// Serves an in-memory copy of an index's vectors. +#[cfg(not(target_arch = "wasm32"))] +pub(crate) struct VecBacking { + pub(crate) dim: usize, + pub(crate) vectors: Vec>, +} + +#[cfg(not(target_arch = "wasm32"))] +impl crate::segment_backing::VectorSegmentBacking for VecBacking { + fn len(&self) -> usize { + self.vectors.len() + } + fn dim(&self) -> usize { + self.dim + } + fn get_vector(&self, id: u32) -> Option<&[f32]> { + self.vectors.get(id as usize).map(Vec::as_slice) + } + fn get_surrogate(&self, _id: u32) -> Option { + None + } +} + +/// `src` restored from its graph-only checkpoint, its vectors served by a +/// backing over `vectors`. +#[cfg(not(target_arch = "wasm32"))] +pub(crate) fn graph_only_restore_with_backing( + src: &crate::hnsw::HnswIndex, + vectors: &[Vec], +) -> crate::hnsw::HnswIndex { + let mut restored = graph_only_restore(src); + restored + .with_backing(Arc::new(VecBacking { + dim: src.dim(), + vectors: vectors.to_vec(), + })) + .expect("a complete backing attaches"); + restored +} From 5be6b2f81ee177cbeea4464c8c9baae89998e9d8 Mon Sep 17 00:00:00 2001 From: Farhan Syah Date: Mon, 12 Oct 2026 16:41:21 +0800 Subject: [PATCH 07/11] fix(kv): order-preserving sorted index encoding and typed native range bounds Encode sorted index values and numbers so byte order matches value order, parse WINDOW clauses in index DDL, and carry typed range bounds through the native sorted read path. --- .../src/protocol/text_fields/decode.rs | 30 +- .../src/protocol/text_fields/encode.rs | 4 +- .../protocol/text_fields/types/text_fields.rs | 10 +- nodedb/src/control/server/native/codec.rs | 3 +- .../server/native/dispatch/index_ddl_op.rs | 12 +- .../server/native/dispatch/sorted_read_op.rs | 45 +- .../shared/ddl/neutral/kv_sorted_index/ddl.rs | 21 +- .../shared/ddl/neutral/kv_sorted_index/mod.rs | 4 +- .../ddl/neutral/kv_sorted_index/parse.rs | 138 ++--- .../ddl/neutral/kv_sorted_index/tail.rs | 422 +++++++++++++++ .../handlers/kv/sorted_index_compute.rs | 15 +- nodedb/src/engine/kv/engine_atomic.rs | 8 +- nodedb/src/engine/kv/engine_helpers.rs | 51 +- nodedb/src/engine/kv/engine_sorted.rs | 95 +++- nodedb/src/engine/kv/engine_write.rs | 13 +- .../src/engine/kv/sorted_index/checkpoint.rs | 72 ++- nodedb/src/engine/kv/sorted_index/index.rs | 157 ++++-- nodedb/src/engine/kv/sorted_index/key.rs | 484 ++++++++---------- nodedb/src/engine/kv/sorted_index/manager.rs | 63 +-- nodedb/src/engine/kv/sorted_index/mod.rs | 2 + nodedb/src/engine/kv/sorted_index/number.rs | 331 ++++++++++++ .../src/engine/kv/sorted_index/timestamp.rs | 105 ++-- nodedb/src/engine/kv/sorted_index/value.rs | 294 +++++++++++ nodedb/src/engine/kv/sorted_index/window.rs | 5 +- nodedb/tests/inproc/cases/mod.rs | 1 + .../cases/sorted_index_mixed_numbers.rs | 177 +++++++ nodedb/tests/native/cases/mod.rs | 1 + .../native/cases/native_sorted_index_range.rs | 115 +++++ 28 files changed, 2080 insertions(+), 598 deletions(-) create mode 100644 nodedb/src/control/server/shared/ddl/neutral/kv_sorted_index/tail.rs create mode 100644 nodedb/src/engine/kv/sorted_index/number.rs create mode 100644 nodedb/src/engine/kv/sorted_index/value.rs create mode 100644 nodedb/tests/inproc/cases/sorted_index_mixed_numbers.rs create mode 100644 nodedb/tests/native/cases/native_sorted_index_range.rs diff --git a/nodedb-types/src/protocol/text_fields/decode.rs b/nodedb-types/src/protocol/text_fields/decode.rs index 0bbc2e1dc..99d48d535 100644 --- a/nodedb-types/src/protocol/text_fields/decode.rs +++ b/nodedb-types/src/protocol/text_fields/decode.rs @@ -267,10 +267,14 @@ impl<'a> zerompk::FromMessagePack<'a> for TextFields { out.top_k_count = Some(reader.read_u32()?); } FID_SCORE_MIN => { - out.score_min = Some(reader.read_binary()?.into_owned()); + out.score_min = Some(::read( + reader, + )?); } FID_SCORE_MAX => { - out.score_max = Some(reader.read_binary()?.into_owned()); + out.score_max = Some(::read( + reader, + )?); } FID_UPDATES => { out.updates = Some(Vec::<(String, Vec)>::read(reader)?); @@ -508,8 +512,6 @@ mod tests { cursor: Some(payload.clone()), expected: Some(payload.clone()), new_value: Some(payload.clone()), - score_min: Some(payload.clone()), - score_max: Some(payload.clone()), filters: Some(payload.clone()), ..Default::default() }; @@ -523,11 +525,27 @@ mod tests { assert_eq!(decoded.cursor.as_deref(), Some(payload.as_slice())); assert_eq!(decoded.expected.as_deref(), Some(payload.as_slice())); assert_eq!(decoded.new_value.as_deref(), Some(payload.as_slice())); - assert_eq!(decoded.score_min.as_deref(), Some(payload.as_slice())); - assert_eq!(decoded.score_max.as_deref(), Some(payload.as_slice())); assert_eq!(decoded.filters.as_deref(), Some(payload.as_slice())); } + #[test] + fn score_bounds_roundtrip_as_typed_values() { + use crate::value::Value; + for (min, max) in [ + (Value::Integer(11), Value::Float(10.5)), + (Value::String("gold".into()), Value::Null), + ] { + let tf = TextFields { + score_min: Some(min.clone()), + score_max: Some(max.clone()), + ..Default::default() + }; + let decoded = roundtrip(&tf); + assert_eq!(decoded.score_min, Some(min)); + assert_eq!(decoded.score_max, Some(max)); + } + } + #[test] fn binary_field_encoder_uses_bin_marker_not_fixarray() { // Wire-shape regression guard. The MessagePack `bin8` marker is diff --git a/nodedb-types/src/protocol/text_fields/encode.rs b/nodedb-types/src/protocol/text_fields/encode.rs index 180777a44..b68b69ed3 100644 --- a/nodedb-types/src/protocol/text_fields/encode.rs +++ b/nodedb-types/src/protocol/text_fields/encode.rs @@ -110,8 +110,8 @@ impl zerompk::ToMessagePack for TextFields { write_opt_field!(writer, FID_WINDOW_START_MS, self.window_start_ms); write_opt_field!(writer, FID_WINDOW_END_MS, self.window_end_ms); write_opt_field!(writer, FID_TOP_K_COUNT, self.top_k_count); - write_opt_field!(binary writer, FID_SCORE_MIN, self.score_min); - write_opt_field!(binary writer, FID_SCORE_MAX, self.score_max); + write_opt_field!(writer, FID_SCORE_MIN, self.score_min); + write_opt_field!(writer, FID_SCORE_MAX, self.score_max); write_opt_field!(writer, FID_UPDATES, self.updates); write_opt_field!(binary writer, FID_FILTERS, self.filters); write_opt_field!(writer, FID_VECTOR, self.vector); diff --git a/nodedb-types/src/protocol/text_fields/types/text_fields.rs b/nodedb-types/src/protocol/text_fields/types/text_fields.rs index cc4f143f8..270d52bcc 100644 --- a/nodedb-types/src/protocol/text_fields/types/text_fields.rs +++ b/nodedb-types/src/protocol/text_fields/types/text_fields.rs @@ -213,12 +213,14 @@ pub struct TextFields { /// Top-K value for SortedIndexTopK. #[serde(skip_serializing_if = "Option::is_none")] pub top_k_count: Option, - /// Score min for SortedIndexRange (encoded bytes). + /// Inclusive lower score bound for `KvSortedIndexRange`, as a typed value + /// of the index's leading sort column. `None` or `Null` leaves it open. #[serde(skip_serializing_if = "Option::is_none")] - pub score_min: Option>, - /// Score max for SortedIndexRange (encoded bytes). + pub score_min: Option, + /// Inclusive upper score bound for `KvSortedIndexRange`. Same rules as + /// `score_min`. #[serde(skip_serializing_if = "Option::is_none")] - pub score_max: Option>, + pub score_max: Option, // ── Document advanced ─────────────────────────────────── /// Field-level updates: [(field_name, value_bytes), ...]. diff --git a/nodedb/src/control/server/native/codec.rs b/nodedb/src/control/server/native/codec.rs index 9c0933278..deaca5023 100644 --- a/nodedb/src/control/server/native/codec.rs +++ b/nodedb/src/control/server/native/codec.rs @@ -243,7 +243,8 @@ mod tests { before + 1 ); - let other = NativeResponse::error(2, "XX000", "boom"); + let other = + crate::control::server::native::sqlstate_code::sqlstate_error(2, "XX000", "boom"); encode_response(&other, FrameFormat::Json).expect("frame encodes"); assert_eq!( overload_responses().sent(OverloadProtocol::Native), diff --git a/nodedb/src/control/server/native/dispatch/index_ddl_op.rs b/nodedb/src/control/server/native/dispatch/index_ddl_op.rs index 5b21a87d3..e714172cb 100644 --- a/nodedb/src/control/server/native/dispatch/index_ddl_op.rs +++ b/nodedb/src/control/server/native/dispatch/index_ddl_op.rs @@ -158,11 +158,13 @@ fn create_sorted_index(fields: &TextFields, collection: &str) -> crate::Result { - sql.push_str(&format!( - " WINDOW CUSTOM START {} END {}", - fields.window_start_ms.unwrap_or(0), - fields.window_end_ms.unwrap_or(0) - )); + let start_ms = fields + .window_start_ms + .ok_or_else(|| bad_request("missing 'window_start_ms'"))?; + let end_ms = fields + .window_end_ms + .ok_or_else(|| bad_request("missing 'window_end_ms'"))?; + sql.push_str(&format!(" WINDOW CUSTOM START {start_ms} END {end_ms}")); if let Some(column) = fields.window_timestamp_column.as_deref() { sql.push_str(&format!(" ON {}", ident(column)?)); } diff --git a/nodedb/src/control/server/native/dispatch/sorted_read_op.rs b/nodedb/src/control/server/native/dispatch/sorted_read_op.rs index e38d5e755..29b9859f2 100644 --- a/nodedb/src/control/server/native/dispatch/sorted_read_op.rs +++ b/nodedb/src/control/server/native/dispatch/sorted_read_op.rs @@ -18,7 +18,7 @@ use nodedb_physical::physical_plan::SortedIndexRead; use nodedb_types::protocol::{NativeResponse, OpCode, TextFields}; -use crate::control::server::shared::ddl::neutral::kv_sorted_index::run_read; +use crate::control::server::shared::ddl::neutral::kv_sorted_index::{encode_score_bound, run_read}; use crate::control::server::shared::session::DmlTxnCtx; use super::response::data_plane_response_to_native; @@ -73,8 +73,8 @@ fn sorted_read(op: OpCode, fields: &TextFields) -> crate::Result<(String, Sorted k: fields.top_k_count.unwrap_or(10), }, OpCode::KvSortedIndexRange => SortedIndexRead::Range { - score_min: fields.score_min.clone(), - score_max: fields.score_max.clone(), + score_min: fields.score_min.clone().and_then(encode_score_bound), + score_max: fields.score_max.clone().and_then(encode_score_bound), }, OpCode::KvSortedIndexCount => SortedIndexRead::Count, OpCode::KvSortedIndexScore => SortedIndexRead::Score { @@ -126,4 +126,43 @@ mod tests { assert!(sorted_read(OpCode::KvSortedIndexRank, &no_key).is_err()); assert!(sorted_read(OpCode::KvSortedIndexCount, &no_key).is_ok()); } + + /// Typed bounds encode as stored values of equal JSON form. `Null` is open. + #[test] + fn range_bounds_encode_as_stored_values() { + use crate::engine::kv::sorted_index::value::encode_json; + use nodedb_types::Value; + + let fields = TextFields { + index_name: Some("lb".into()), + score_min: Some(Value::Float(10.5)), + score_max: Some(Value::Integer(11)), + ..TextFields::default() + }; + assert_eq!( + sorted_read(OpCode::KvSortedIndexRange, &fields) + .expect("range") + .1, + SortedIndexRead::Range { + score_min: Some(encode_json(&serde_json::json!(10.5))), + score_max: Some(encode_json(&serde_json::json!(11))), + } + ); + + let open = TextFields { + index_name: Some("lb".into()), + score_min: Some(Value::Null), + score_max: Some(Value::String("gold".into())), + ..TextFields::default() + }; + assert_eq!( + sorted_read(OpCode::KvSortedIndexRange, &open) + .expect("range") + .1, + SortedIndexRead::Range { + score_min: None, + score_max: Some(encode_json(&serde_json::json!("gold"))), + } + ); + } } diff --git a/nodedb/src/control/server/shared/ddl/neutral/kv_sorted_index/ddl.rs b/nodedb/src/control/server/shared/ddl/neutral/kv_sorted_index/ddl.rs index 323c81820..44985d99a 100644 --- a/nodedb/src/control/server/shared/ddl/neutral/kv_sorted_index/ddl.rs +++ b/nodedb/src/control/server/shared/ddl/neutral/kv_sorted_index/ddl.rs @@ -27,7 +27,8 @@ use super::super::super::result::{DdlError, DdlResult}; use super::super::refuse_gate::RefusingReadGate; use super::dispatch::{SortedIndexTarget, drop_in_engine, register_in_engine}; use super::gate::owning_collection; -use super::parse::{ddl_err, parse_key_column, parse_sort_columns, parse_window_clause}; +use super::parse::{ddl_err, parse_sort_columns}; +use super::tail::{IndexTail, parse_index_tail}; /// Building the index reads every row of the collection into the order-stat /// tree, so the caller must be allowed to read it, and a read policy makes the @@ -44,8 +45,6 @@ pub async fn create_sorted_index( database_id: DatabaseId, sql: &str, ) -> Result, DdlError> { - let upper = sql.to_ascii_uppercase(); - // Extract index name. let rest = nodedb_types::strip_prefix_ascii_case_insensitive(sql, "CREATE SORTED INDEX ") .ok_or_else(|| ddl_err("42601", "expected CREATE SORTED INDEX"))?; @@ -76,11 +75,9 @@ pub async fn create_sorted_index( let cols_str = &sql[paren_start + 1..paren_end]; let sort_columns = parse_sort_columns(cols_str)?; - // Extract KEY column. - let key_column = parse_key_column(&upper)?; - - // Extract WINDOW clause (optional). - let (window_type, window_ts_col, window_start, window_end) = parse_window_clause(&upper); + // KEY column and optional WINDOW clause, from the original text so quoted + // identifiers keep their case. + let IndexTail { key_column, window } = parse_index_tail(&sql[paren_end + 1..], &sort_columns)?; let tenant_id = identity.tenant_id; if state @@ -104,10 +101,10 @@ pub async fn create_sorted_index( index_name: index_name.clone(), sort_columns, key_column, - window_type, - window_timestamp_column: window_ts_col, - window_start_ms: window_start, - window_end_ms: window_end, + window_type: window.window_type, + window_timestamp_column: window.timestamp_column, + window_start_ms: window.start_ms, + window_end_ms: window.end_ms, }); // Inside an explicit transaction the records are buffered first and the diff --git a/nodedb/src/control/server/shared/ddl/neutral/kv_sorted_index/mod.rs b/nodedb/src/control/server/shared/ddl/neutral/kv_sorted_index/mod.rs index 729be04d7..6f1ed5a81 100644 --- a/nodedb/src/control/server/shared/ddl/neutral/kv_sorted_index/mod.rs +++ b/nodedb/src/control/server/shared/ddl/neutral/kv_sorted_index/mod.rs @@ -4,7 +4,7 @@ //! (leaderboards). //! //! DDL: -//! CREATE SORTED INDEX name ON collection (score DESC [, tiebreak ASC]) KEY key_col [WINDOW DAILY ON ts_col] +//! CREATE SORTED INDEX name ON collection (score DESC [, ts ASC]) KEY key_col [WINDOW DAILY ON ts] //! DROP SORTED INDEX name //! //! SQL functions (intercepted in DDL router): @@ -18,9 +18,11 @@ pub mod dispatch; pub mod gate; pub mod parse; pub mod query; +mod tail; mod txn_read; pub use ddl::{create_sorted_index, drop_sorted_index}; pub use dispatch::{SortedIndexTarget, drop_in_engine}; +pub(crate) use parse::encode_score_bound; pub(crate) use query::run_read; pub use query::{select_range, select_rank, select_sorted_count, select_topk}; diff --git a/nodedb/src/control/server/shared/ddl/neutral/kv_sorted_index/parse.rs b/nodedb/src/control/server/shared/ddl/neutral/kv_sorted_index/parse.rs index cdd7530d6..4ccbe3d61 100644 --- a/nodedb/src/control/server/shared/ddl/neutral/kv_sorted_index/parse.rs +++ b/nodedb/src/control/server/shared/ddl/neutral/kv_sorted_index/parse.rs @@ -2,7 +2,10 @@ //! Statement and argument parsing for the sorted-index family. +use nodedb_types::Value; + use super::super::super::result::DdlError; +use crate::engine::kv::sorted_index::value::encode_json; pub(super) fn ddl_err(sqlstate: &str, message: impl Into) -> DdlError { DdlError::new(sqlstate, message) @@ -35,73 +38,6 @@ pub(super) fn parse_sort_columns(cols_str: &str) -> Result Ok(columns) } -pub(super) fn parse_key_column(upper: &str) -> Result { - // Look for "KEY " after the closing paren. - let key_pos = upper - .find(") KEY ") - .or_else(|| upper.find(")KEY ")) - .ok_or_else(|| ddl_err("42601", "missing KEY clause"))?; - let after_key = upper[key_pos..].trim_start_matches(')').trim(); - let after_key = after_key - .strip_prefix("KEY ") - .ok_or_else(|| ddl_err("42601", "missing KEY clause"))?; - let key_col = after_key - .split_whitespace() - .next() - .ok_or_else(|| ddl_err("42601", "missing key column name"))? - .trim_end_matches(';') - .to_lowercase(); - Ok(key_col) -} - -pub(super) fn parse_window_clause(upper: &str) -> (String, String, u64, u64) { - // Look for "WINDOW ON " or "WINDOW CUSTOM START '...' END '...'" - let Some(win_pos) = upper.find(" WINDOW ") else { - return ("none".into(), String::new(), 0, 0); - }; - - let after_window = &upper[win_pos + 8..]; - let tokens: Vec<&str> = after_window.split_whitespace().collect(); - - if tokens.is_empty() { - return ("none".into(), String::new(), 0, 0); - } - - let win_type = tokens[0].to_lowercase(); - - match win_type.as_str() { - "daily" | "weekly" | "monthly" => { - // WINDOW DAILY ON ts_col - let ts_col = if tokens.len() >= 3 && tokens[1] == "ON" { - tokens[2].to_lowercase() - } else { - "updated_at".into() - }; - (win_type, ts_col, 0, 0) - } - "custom" => { - // WINDOW CUSTOM START '2026-01-01' END '2026-03-31' - let mut start_ms = 0u64; - let mut end_ms = 0u64; - let mut ts_col = "updated_at".to_string(); - - for i in 1..tokens.len() { - if tokens[i] == "START" && i + 1 < tokens.len() { - start_ms = tokens[i + 1].trim_matches('\'').parse().unwrap_or(0); - } - if tokens[i] == "END" && i + 1 < tokens.len() { - end_ms = tokens[i + 1].trim_matches('\'').parse().unwrap_or(0); - } - if tokens[i] == "ON" && i + 1 < tokens.len() { - ts_col = tokens[i + 1].to_lowercase(); - } - } - ("custom".into(), ts_col, start_ms, end_ms) - } - _ => ("none".into(), String::new(), 0, 0), - } -} - pub(super) fn parse_function_args(sql: &str) -> Result, DdlError> { let start = sql .find('(') @@ -125,27 +61,65 @@ pub(super) fn unquote(s: &str) -> String { } } -/// A `RANGE(...)` score bound, as the raw value bytes of the index's leading -/// sort column. +/// The sorted index value encoding of a `RANGE` score bound. `None` for a +/// null bound, which leaves that end of the range open. +/// +/// SQL `RANGE(...)` and the native `KvSortedIndexRange` both encode here. The +/// bound converts to JSON the way a stored row does, then encodes as a stored +/// field value, so a bound and a row of equal value get equal bytes. +pub(crate) fn encode_score_bound(bound: Value) -> Option> { + match serde_json::Value::from(bound) { + serde_json::Value::Null => None, + json => Some(encode_json(&json)), + } +} + +/// A SQL `RANGE(...)` score argument, as the value encoding of the index's +/// leading sort column. /// -/// The bytes must match what the engine derives from a stored row for that -/// column (`json_value_to_index_bytes` / `field_value_to_sort_bytes`), because -/// the Data Plane frames these into a tree bound with the index's own encoder -/// and compares them against keys built from stored values. So an integer -/// literal encodes as a sign-flipped big-endian i64, a decimal literal as an -/// order-preserving f64, and anything else as its UTF-8 bytes — the same three -/// shapes a stored number / number / string produce. +/// An integer literal is a number, past i64 too. A decimal literal is a +/// float. `NULL`, `NONE` and `*` leave the end open. Anything else is a +/// string. A quoted literal unquotes first, so a bound parameter `'100'` is +/// the number 100. pub(super) fn parse_score_arg(s: &str) -> Option> { let t = unquote(s); if t.eq_ignore_ascii_case("NULL") || t.eq_ignore_ascii_case("NONE") || t == "*" { return None; } - use crate::engine::kv::sorted_index::key::SortKeyEncoder; - if let Ok(v) = t.parse::() { - return Some(SortKeyEncoder::encode_i64(v).to_vec()); + let bound = if let Ok(v) = t.parse::() { + Value::Integer(v) + } else if let Ok(v) = t.parse::() { + Value::Decimal(rust_decimal::Decimal::from(v)) + } else if let Ok(v) = t.parse::() { + Value::Float(v) + } else { + Value::String(t) + }; + encode_score_bound(bound) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::engine::kv::sorted_index::value::{encode_float, encode_int, encode_text}; + + #[test] + fn integers_past_i64_stay_exact() { + let past = u64::MAX; + assert_eq!( + parse_score_arg(&past.to_string()), + Some(encode_json(&serde_json::json!(past))) + ); } - if let Ok(v) = t.parse::() { - return Some(SortKeyEncoder::encode_f64(v).to_vec()); + + #[test] + fn score_args_encode_like_stored_values() { + assert_eq!(parse_score_arg("11"), Some(encode_int(11))); + assert_eq!(parse_score_arg("'11'"), Some(encode_int(11))); + assert_eq!(parse_score_arg("10.5"), Some(encode_float(10.5))); + assert_eq!(parse_score_arg("11.0"), Some(encode_int(11))); + assert_eq!(parse_score_arg("'gold'"), Some(encode_text("gold"))); + assert_eq!(parse_score_arg("NULL"), None); + assert_eq!(parse_score_arg("*"), None); } - Some(t.into_bytes()) } diff --git a/nodedb/src/control/server/shared/ddl/neutral/kv_sorted_index/tail.rs b/nodedb/src/control/server/shared/ddl/neutral/kv_sorted_index/tail.rs new file mode 100644 index 000000000..65363db47 --- /dev/null +++ b/nodedb/src/control/server/shared/ddl/neutral/kv_sorted_index/tail.rs @@ -0,0 +1,422 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! The clauses of `CREATE SORTED INDEX` after the sort column list: +//! +//! ```text +//! KEY key_col +//! [WINDOW {DAILY | WEEKLY | MONTHLY} [ON ts_col]] +//! [WINDOW CUSTOM START bound END bound [ON ts_col]] +//! ``` +//! +//! Identifiers keep their quoting, so a quoted mixed-case column resolves +//! the same way the sort column list resolves it. A bound is integer +//! milliseconds since the Unix epoch or an ISO 8601 instant. + +use nodedb_types::NdbDateTime; + +use super::super::super::result::DdlError; +use super::parse::ddl_err; +use crate::control::server::shared::ddl::sql_parse::parse_ident_token; + +/// The timestamp column a window reads when the clause names none. +const DEFAULT_TIMESTAMP_COLUMN: &str = "updated_at"; + +/// The parsed clauses after the sort column list. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(super) struct IndexTail { + pub key_column: String, + pub window: WindowClause, +} + +/// A parsed `WINDOW` clause, in the shape `RegisterSortedIndex` carries. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(super) struct WindowClause { + /// `none`, `daily`, `weekly`, `monthly` or `custom`. + pub window_type: String, + /// Empty when `window_type` is `none`. + pub timestamp_column: String, + /// Custom windows only: the half-open range `[start_ms, end_ms)`. + pub start_ms: u64, + pub end_ms: u64, +} + +impl WindowClause { + fn none() -> Self { + Self { + window_type: "none".into(), + timestamp_column: String::new(), + start_ms: 0, + end_ms: 0, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +enum Token<'a> { + /// A bare word or a double-quoted identifier, quotes kept. + Word(&'a str), + /// A single-quoted string literal, unescaped. + Literal(String), +} + +impl Token<'_> { + fn is_keyword(&self, keyword: &str) -> bool { + matches!(self, Token::Word(word) if word.eq_ignore_ascii_case(keyword)) + } + + fn describe(&self) -> String { + match self { + Token::Word(word) => (*word).to_string(), + Token::Literal(text) => format!("'{text}'"), + } + } +} + +/// Parse the text after the sort column list's closing parenthesis. +/// +/// `sort_columns` are the parsed sort columns. A window timestamp column must +/// be one of them. +pub(super) fn parse_index_tail( + text: &str, + sort_columns: &[(String, String)], +) -> Result { + let tokens = tokenize(text)?; + let (key_column, rest) = match tokens.as_slice() { + [key, Token::Word(name), tail @ ..] if key.is_keyword("KEY") => { + (parse_ident_token(name)?, tail) + } + _ => { + return Err(ddl_err( + "42601", + "missing KEY clause: expected KEY ", + )); + } + }; + + let window = match rest { + [] => WindowClause::none(), + [window, tail @ ..] if window.is_keyword("WINDOW") => parse_window(tail)?, + [other, ..] => { + return Err(ddl_err( + "42601", + format!("unexpected '{}' after KEY {key_column}", other.describe()), + )); + } + }; + + if window.window_type != "none" + && !sort_columns + .iter() + .any(|(name, _)| *name == window.timestamp_column) + { + return Err(ddl_err( + "42601", + format!( + "WINDOW timestamp column '{}' must be one of the sort columns", + window.timestamp_column + ), + )); + } + + Ok(IndexTail { key_column, window }) +} + +fn parse_window(tokens: &[Token<'_>]) -> Result { + let (kind, rest) = match tokens { + [Token::Word(kind), rest @ ..] => (kind.to_ascii_lowercase(), rest), + [other, ..] => return Err(unknown_window(&other.describe())), + [] => return Err(ddl_err("42601", "WINDOW requires a window type")), + }; + match kind.as_str() { + "daily" | "weekly" | "monthly" => { + let timestamp_column = match rest { + [] => DEFAULT_TIMESTAMP_COLUMN.to_string(), + [on, Token::Word(column)] if on.is_keyword("ON") => parse_ident_token(column)?, + [other, ..] => return Err(unexpected_in_window(other)), + }; + Ok(WindowClause { + window_type: kind, + timestamp_column, + start_ms: 0, + end_ms: 0, + }) + } + "custom" => parse_custom_window(rest), + _ => Err(unknown_window(&kind)), + } +} + +/// `START bound END bound [ON ts_col]`, each clause at most once, any order. +fn parse_custom_window(mut rest: &[Token<'_>]) -> Result { + let mut start_ms = None; + let mut end_ms = None; + let mut timestamp_column = None; + while let [keyword, value, tail @ ..] = rest { + let duplicate = if keyword.is_keyword("START") { + start_ms.replace(parse_bound(value)?).is_some() + } else if keyword.is_keyword("END") { + end_ms.replace(parse_bound(value)?).is_some() + } else if keyword.is_keyword("ON") { + let Token::Word(column) = value else { + return Err(unexpected_in_window(value)); + }; + timestamp_column + .replace(parse_ident_token(column)?) + .is_some() + } else { + return Err(unexpected_in_window(keyword)); + }; + if duplicate { + return Err(ddl_err( + "42601", + format!("WINDOW CUSTOM repeats '{}'", keyword.describe()), + )); + } + rest = tail; + } + if let [dangling] = rest { + return Err(ddl_err( + "42601", + format!("WINDOW CUSTOM '{}' has no value", dangling.describe()), + )); + } + + let (Some(start_ms), Some(end_ms)) = (start_ms, end_ms) else { + return Err(ddl_err("42601", "WINDOW CUSTOM requires START and END")); + }; + if end_ms <= start_ms { + return Err(ddl_err( + "42601", + format!("WINDOW CUSTOM END ({end_ms} ms) must be after START ({start_ms} ms)"), + )); + } + Ok(WindowClause { + window_type: "custom".into(), + timestamp_column: timestamp_column.unwrap_or_else(|| DEFAULT_TIMESTAMP_COLUMN.into()), + start_ms, + end_ms, + }) +} + +/// Integer milliseconds since the Unix epoch, or an ISO 8601 instant at or +/// after the epoch. +fn parse_bound(token: &Token<'_>) -> Result { + let text = match token { + Token::Word(word) => *word, + Token::Literal(text) => text.as_str(), + }; + if let Ok(ms) = text.trim().parse::() { + return Ok(ms); + } + NdbDateTime::parse(text) + .and_then(|instant| u64::try_from(instant.unix_millis()).ok()) + .ok_or_else(|| { + ddl_err( + "42601", + format!( + "invalid WINDOW CUSTOM bound {}: expected milliseconds since the Unix \ + epoch or an ISO 8601 instant at or after 1970-01-01", + token.describe() + ), + ) + }) +} + +fn unknown_window(kind: &str) -> DdlError { + ddl_err( + "42601", + format!("unknown WINDOW type '{kind}', expected DAILY, WEEKLY, MONTHLY or CUSTOM"), + ) +} + +fn unexpected_in_window(token: &Token<'_>) -> DdlError { + ddl_err( + "42601", + format!("unexpected '{}' in WINDOW clause", token.describe()), + ) +} + +/// Split `text` into words, double-quoted identifiers and single-quoted +/// literals. A trailing `;` ends the statement. +fn tokenize(text: &str) -> Result>, DdlError> { + let mut tokens = Vec::new(); + let mut rest = text.trim_start(); + while let Some(first) = rest.chars().next() { + let consumed = match first { + ';' => { + if !rest[1..].trim().is_empty() { + return Err(ddl_err( + "42601", + format!("unexpected text after ';': {}", rest[1..].trim()), + )); + } + break; + } + '\'' => { + let (literal, len) = quoted(rest, '\'')?; + tokens.push(Token::Literal(literal)); + len + } + '"' => { + let (_, len) = quoted(rest, '"')?; + tokens.push(Token::Word(&rest[..len])); + len + } + _ => { + let len = rest + .find(|c: char| c.is_whitespace() || c == ';') + .unwrap_or(rest.len()); + tokens.push(Token::Word(&rest[..len])); + len + } + }; + rest = rest[consumed..].trim_start(); + } + Ok(tokens) +} + +/// The unescaped body of the quoted token at the start of `text`, and the +/// token's byte length. A doubled quote inside is one quote. +fn quoted(text: &str, quote: char) -> Result<(String, usize), DdlError> { + let mut body = String::new(); + let mut chars = text.char_indices().skip(1).peekable(); + while let Some((at, ch)) = chars.next() { + if ch != quote { + body.push(ch); + continue; + } + if chars.peek().is_some_and(|&(_, next)| next == quote) { + body.push(quote); + chars.next(); + continue; + } + return Ok((body, at + ch.len_utf8())); + } + Err(ddl_err( + "42601", + format!("unterminated quoted text: {text}"), + )) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn cols(names: &[&str]) -> Vec<(String, String)> { + names + .iter() + .map(|name| (name.to_string(), "ASC".to_string())) + .collect() + } + + fn sqlstate_and_message(err: DdlError) -> (String, String) { + (err.sqlstate, err.message) + } + + #[test] + fn key_only() { + let tail = parse_index_tail(" KEY player_id;", &cols(&["score"])).expect("parse"); + assert_eq!(tail.key_column, "player_id"); + assert_eq!(tail.window, WindowClause::none()); + } + + #[test] + fn daily_window_on_a_sort_column() { + let tail = parse_index_tail( + "key id window daily on updated_at", + &cols(&["score", "updated_at"]), + ) + .expect("parse"); + assert_eq!(tail.window.window_type, "daily"); + assert_eq!(tail.window.timestamp_column, "updated_at"); + } + + #[test] + fn quoted_mixed_case_timestamp_column_matches_its_sort_column() { + let quoted = parse_ident_token("\"UpdatedAt\"").expect("ident"); + let tail = parse_index_tail( + "KEY \"PlayerId\" WINDOW WEEKLY ON \"UpdatedAt\"", + &cols(&["score", quoted.as_str()]), + ) + .expect("parse"); + assert_eq!( + tail.key_column, + parse_ident_token("\"PlayerId\"").expect("ident") + ); + assert_eq!(tail.window.timestamp_column, quoted); + } + + #[test] + fn custom_window_accepts_millis_and_iso_8601() { + let columns = cols(&["score", "ts"]); + let tail = parse_index_tail( + "KEY id WINDOW CUSTOM START '2026-01-01' END 1767312000000 ON ts", + &columns, + ) + .expect("parse"); + assert_eq!(tail.window.window_type, "custom"); + assert_eq!(tail.window.start_ms, 1_767_225_600_000); + assert_eq!(tail.window.end_ms, 1_767_312_000_000); + assert_eq!(tail.window.timestamp_column, "ts"); + + let tail = parse_index_tail( + "KEY id WINDOW CUSTOM ON ts START '2026-01-01T00:00:00Z' END '2026-03-31 12:00:00'", + &columns, + ) + .expect("parse"); + assert_eq!(tail.window.start_ms, 1_767_225_600_000); + assert_eq!(tail.window.end_ms, 1_774_958_400_000); + } + + #[test] + fn custom_window_refuses_a_bad_bound_and_names_it() { + let err = parse_index_tail( + "KEY id WINDOW CUSTOM START 'last tuesday' END 5 ON ts", + &cols(&["ts"]), + ) + .expect_err("bad bound"); + let (state, message) = sqlstate_and_message(err); + assert_eq!(state, "42601"); + assert!(message.contains("'last tuesday'"), "{message}"); + } + + #[test] + fn custom_window_refuses_an_empty_or_inverted_range() { + for clause in [ + "KEY id WINDOW CUSTOM START 10 END 10 ON ts", + "KEY id WINDOW CUSTOM START '2026-03-31' END '2026-01-01' ON ts", + "KEY id WINDOW CUSTOM START 10 ON ts", + "KEY id WINDOW CUSTOM START 1 START 2 END 3 ON ts", + "KEY id WINDOW CUSTOM START 1 END", + ] { + let err = parse_index_tail(clause, &cols(&["ts"])).expect_err(clause); + assert_eq!(err.sqlstate, "42601", "{clause}"); + } + } + + #[test] + fn window_timestamp_column_must_be_a_sort_column() { + let err = parse_index_tail( + "KEY player_id WINDOW DAILY ON updated_at", + &cols(&["score"]), + ) + .expect_err("not a sort column"); + let (state, message) = sqlstate_and_message(err); + assert_eq!(state, "42601"); + assert!(message.contains("updated_at"), "{message}"); + } + + #[test] + fn unknown_window_type_or_trailing_text_is_refused() { + for clause in [ + "KEY id WINDOW HOURLY ON ts", + "KEY id WINDOW DAILY ON ts EXTRA", + "KEY id EXTRA", + "id", + "KEY id; DROP", + ] { + let err = parse_index_tail(clause, &cols(&["ts"])).expect_err(clause); + assert_eq!(err.sqlstate, "42601", "{clause}"); + } + } +} diff --git a/nodedb/src/data/executor/handlers/kv/sorted_index_compute.rs b/nodedb/src/data/executor/handlers/kv/sorted_index_compute.rs index 6314074fe..c53096ca8 100644 --- a/nodedb/src/data/executor/handlers/kv/sorted_index_compute.rs +++ b/nodedb/src/data/executor/handlers/kv/sorted_index_compute.rs @@ -28,7 +28,7 @@ pub(in crate::data::executor) struct BuildSortedIndexDefParams<'a> { /// Build a [`SortedIndexDef`] from the raw `RegisterSortedIndex` fields, /// validating that a windowed index's timestamp column is included in its -/// sort columns. Returns `Err` on that validation failure; never panics. +/// sort columns. That failure is a permanent client error; never panics. pub(in crate::data::executor) fn build_sorted_index_def( params: BuildSortedIndexDefParams<'_>, ) -> Result { @@ -70,10 +70,9 @@ pub(in crate::data::executor) fn build_sorted_index_def( let ts_col = &window.timestamp_column; let has_ts = encoder.columns().iter().any(|c| c.name == *ts_col); if !has_ts { - return Err(ErrorCode::Internal { - detail: format!( - "WINDOW timestamp column '{}' must be included in sort columns", - ts_col + return Err(ErrorCode::RejectedPrevalidation { + reason: format!( + "WINDOW timestamp column '{ts_col}' must be one of the sort columns" ), }); } @@ -126,10 +125,10 @@ mod tests { }) .expect_err("must reject missing timestamp column"); match err { - ErrorCode::Internal { detail } => { - assert!(detail.contains("updated_at")); + ErrorCode::RejectedPrevalidation { reason } => { + assert!(reason.contains("updated_at")); } - other => panic!("expected Internal error, got {other:?}"), + other => panic!("expected RejectedPrevalidation, got {other:?}"), } } diff --git a/nodedb/src/engine/kv/engine_atomic.rs b/nodedb/src/engine/kv/engine_atomic.rs index 92a26881f..41873fca1 100644 --- a/nodedb/src/engine/kv/engine_atomic.rs +++ b/nodedb/src/engine/kv/engine_atomic.rs @@ -462,7 +462,11 @@ impl KvEngine { // Index maintenance — the same pair `KvEngine::put` performs, over the // one field extraction both kinds read. if has_secondary_indexes || has_sorted_indexes { - let new_fields = super::engine_helpers::extract_all_field_values_from_msgpack(value); + let new_document = super::engine_helpers::decode_document(value); + let new_fields = new_document + .as_ref() + .map(super::engine_helpers::document_field_values) + .unwrap_or_default(); if has_secondary_indexes { let old_refs: Option> = old_fields.as_ref().map(|fields| { @@ -481,7 +485,7 @@ impl KvEngine { } if has_sorted_indexes { - self.sorted_indexes.on_put(tkey, key, &new_fields); + self.sorted_indexes.on_put(tkey, key, new_document.as_ref()); } } } diff --git a/nodedb/src/engine/kv/engine_helpers.rs b/nodedb/src/engine/kv/engine_helpers.rs index 1dddb7f55..ea62c624d 100644 --- a/nodedb/src/engine/kv/engine_helpers.rs +++ b/nodedb/src/engine/kv/engine_helpers.rs @@ -2,6 +2,8 @@ //! Key encoding and field extraction helpers for the KV engine. +use super::sorted_index::number; + /// Compute a numeric key for per-collection HashMap lookups. /// /// Uses FxHash on `(database_id, tenant_id, collection)` to produce a `u64` — @@ -56,25 +58,30 @@ pub(super) fn parse_expiry_key(composite: &[u8]) -> Option<(u64, u64, String, Ve // Field extraction from MessagePack values (for secondary index maintenance) // --------------------------------------------------------------------------- +/// Decode a MessagePack-encoded KV value. `None` when it does not decode. +pub(super) fn decode_document(data: &[u8]) -> Option { + nodedb_types::json_from_msgpack(data).ok() +} + /// Extract all field name → value bytes pairs from a MessagePack-encoded document. /// /// The value is assumed to be a MessagePack map (as produced by SQL INSERT). -/// Returns `(field_name, value_bytes)` where `value_bytes` is the JSON-serialized -/// scalar value suitable for B-Tree comparison. +/// Returns `(field_name, value_bytes)` where `value_bytes` is the secondary +/// index encoding of the scalar value. pub(super) fn extract_all_field_values_from_msgpack(data: &[u8]) -> Vec<(String, Vec)> { - let Ok(value) = nodedb_types::json_from_msgpack(data) else { - return Vec::new(); - }; + decode_document(data) + .map(|doc| document_field_values(&doc)) + .unwrap_or_default() +} - let Some(obj) = value.as_object() else { +/// Every top-level field of `document` with its secondary index encoding. +/// Empty when `document` is not an object. +pub(super) fn document_field_values(document: &serde_json::Value) -> Vec<(String, Vec)> { + let Some(obj) = document.as_object() else { return Vec::new(); }; - obj.iter() - .map(|(k, v)| { - let bytes = json_value_to_index_bytes(v); - (k.clone(), bytes) - }) + .map(|(k, v)| (k.clone(), json_value_to_index_bytes(v))) .collect() } @@ -96,28 +103,12 @@ pub(super) fn extract_field_values_from_msgpack(data: &[u8], field: &str) -> Vec /// Convert a JSON value to bytes suitable for B-Tree ordering. /// -/// Strings are stored as UTF-8 bytes. Numbers are stored as big-endian bytes -/// for correct lexicographic ordering. Booleans as 0/1. Null as empty. +/// Strings are stored as UTF-8 bytes. Numbers use the sorted index number +/// encoding, so 11 and 11.0 are one key. Booleans as 0/1. Null as empty. fn json_value_to_index_bytes(v: &serde_json::Value) -> Vec { match v { serde_json::Value::String(s) => s.as_bytes().to_vec(), - serde_json::Value::Number(n) => { - if let Some(i) = n.as_i64() { - // Shift to unsigned for correct B-Tree ordering (negatives sort before positives). - let sortable = (i as u64) ^ (1u64 << 63); - sortable.to_be_bytes().to_vec() - } else if let Some(f) = n.as_f64() { - let bits = f.to_bits(); - let sortable = if bits >> 63 == 1 { - !bits // Negative: flip all bits for correct ordering. - } else { - bits | (1u64 << 63) // Positive: set sign bit. - }; - sortable.to_be_bytes().to_vec() - } else { - Vec::new() - } - } + serde_json::Value::Number(n) => number::encode_json_number(n).to_vec(), serde_json::Value::Bool(b) => vec![*b as u8], serde_json::Value::Null => Vec::new(), // Arrays and objects: serialize as JSON string for comparison. diff --git a/nodedb/src/engine/kv/engine_sorted.rs b/nodedb/src/engine/kv/engine_sorted.rs index 01dd9abc5..880864742 100644 --- a/nodedb/src/engine/kv/engine_sorted.rs +++ b/nodedb/src/engine/kv/engine_sorted.rs @@ -416,7 +416,7 @@ mod tests { put_scored(&mut e, "players", "p3", 20); e.register_sorted_index(0, 1, "players", leaderboard_def("players", "lb")); - let bound = |v: i64| SortKeyEncoder::encode_i64(v).to_vec(); + let bound = |v: i64| crate::engine::kv::sorted_index::value::encode_int(v); let mid = e.sorted_index_range(SortedIndexRangeParams { database_id: 0, @@ -559,4 +559,97 @@ mod tests { "the sorted index itself must be gone, as the secondary indexes are" ); } + + fn put_doc(e: &mut KvEngine, key: &str, doc: serde_json::Value) { + e.put(KvPutParams { + database_id: 0, + tenant_id: 1, + collection: "players", + key: key.as_bytes(), + value: &nodedb_types::json_to_msgpack(&doc).expect("encode row"), + ttl_ms: 0, + now_ms: now(), + surrogate: row_surrogate(key.as_bytes()), + }) + .expect("a bound row writes"); + } + + /// Integer and float scores written through `put` rank by numeric value. + #[test] + fn written_integer_and_float_scores_rank_numerically() { + let mut e = make_engine(); + let n = now(); + e.register_sorted_index(0, 1, "players", leaderboard_def("players", "lb")); + + put_doc(&mut e, "ten_half", serde_json::json!({ "score": 10.5 })); + put_doc(&mut e, "eleven", serde_json::json!({ "score": 11 })); + put_doc(&mut e, "ten", serde_json::json!({ "score": 10 })); + + assert_eq!( + ranked_keys(e.sorted_index_top_k(0, 1, "lb", 10, n)), + vec!["eleven", "ten_half", "ten"] + ); + } + + /// A write that drops the sort column takes the row out of the index. + #[test] + fn a_write_without_the_sort_column_leaves_the_index() { + let mut e = make_engine(); + let n = now(); + e.register_sorted_index(0, 1, "players", leaderboard_def("players", "lb")); + + put_doc(&mut e, "p1", serde_json::json!({ "score": 10 })); + assert_eq!(e.sorted_index_rank(0, 1, "lb", b"p1", n), Some(1)); + + put_doc(&mut e, "p1", serde_json::json!({ "name": "no score" })); + assert_eq!(e.sorted_index_rank(0, 1, "lb", b"p1", n), None); + assert_eq!(e.sorted_index_count(0, 1, "lb", n), Some(0)); + } + + /// A windowed index maintained by `put` ranks only rows inside the window. + #[test] + fn windowed_index_maintained_by_put_ranks_inside_the_window() { + let mut e = make_engine(); + let n = now(); + let mut def = leaderboard_def("players", "lb"); + def.encoder = SortKeyEncoder::new(vec![ + SortColumn { + name: "score".into(), + direction: SortDirection::Desc, + }, + SortColumn { + name: "ts".into(), + direction: SortDirection::Asc, + }, + ]); + def.window = WindowConfig::custom("ts", n - 1_000, n + 1_000); + e.register_sorted_index(0, 1, "players", def); + + put_doc( + &mut e, + "early", + serde_json::json!({ "score": 99, "ts": n - 1_001 }), + ); + put_doc( + &mut e, + "int_ts", + serde_json::json!({ "score": 11, "ts": n }), + ); + put_doc( + &mut e, + "float_ts", + serde_json::json!({ "score": 10.5, "ts": n as f64 + 0.5 }), + ); + put_doc( + &mut e, + "late", + serde_json::json!({ "score": 50, "ts": n + 1_000 }), + ); + + assert_eq!(e.sorted_index_count(0, 1, "lb", n), Some(2)); + assert_eq!(e.sorted_index_rank(0, 1, "lb", b"int_ts", n), Some(1)); + assert_eq!(e.sorted_index_rank(0, 1, "lb", b"float_ts", n), Some(2)); + assert_eq!(e.sorted_index_rank(0, 1, "lb", b"early", n), None); + assert_eq!(e.sorted_index_rank(0, 1, "lb", b"late", n), None); + } } diff --git a/nodedb/src/engine/kv/engine_write.rs b/nodedb/src/engine/kv/engine_write.rs index 97147dc69..d9ea57f55 100644 --- a/nodedb/src/engine/kv/engine_write.rs +++ b/nodedb/src/engine/kv/engine_write.rs @@ -5,7 +5,10 @@ use nodedb_types::Surrogate; use super::engine::KvEngine; -use super::engine_helpers::{expiry_key, extract_all_field_values_from_msgpack, table_key}; +use super::engine_helpers::{ + decode_document, document_field_values, expiry_key, extract_all_field_values_from_msgpack, + table_key, +}; use super::entry::NO_EXPIRY; /// Parameters for [`KvEngine::put`]. @@ -217,7 +220,11 @@ impl KvEngine { .and_then(|t| t.get(key, now_ms)) .map(|v| v.to_vec()) .unwrap_or_default(); - let new_fields = extract_all_field_values_from_msgpack(&new_value_bytes); + let new_document = decode_document(&new_value_bytes); + let new_fields = new_document + .as_ref() + .map(document_field_values) + .unwrap_or_default(); let old_fields = old .as_ref() .map(|v| extract_all_field_values_from_msgpack(v)); @@ -237,7 +244,7 @@ impl KvEngine { } if has_sorted { - self.sorted_indexes.on_put(tkey, key, &new_fields); + self.sorted_indexes.on_put(tkey, key, new_document.as_ref()); } } diff --git a/nodedb/src/engine/kv/sorted_index/checkpoint.rs b/nodedb/src/engine/kv/sorted_index/checkpoint.rs index 5929b5e79..bbea25597 100644 --- a/nodedb/src/engine/kv/sorted_index/checkpoint.rs +++ b/nodedb/src/engine/kv/sorted_index/checkpoint.rs @@ -9,13 +9,8 @@ //! generation, or a restart silently comes back with the rows and no index. //! //! Content is exported as raw `(sort_key, primary_key)` pairs rather than -//! rebuilt from the restored rows. Rebuilding would have to pick one of the two -//! sort-key extraction paths — `register`'s backfill (via -//! `extract_sort_key_from_value`) or live PUT maintenance (via -//! `build_sort_key_from_fields`) — and those two disagree on at least boolean -//! columns, so whichever the restore chose would silently rewrite the keys of an -//! index built by the other. Reinstalling the exported keys verbatim reproduces -//! the tree that was actually live. +//! rebuilt from the restored rows. Reinstalling the exported keys verbatim +//! reproduces the tree that was live, with no second pass over the rows. use super::super::engine_helpers::table_key; use super::index::SortedIndex; @@ -127,8 +122,7 @@ mod tests { mgr.register(0, 1, make_def("lb", "scores"), std::iter::empty()); let tbl_key = scores_table_key(); for (pk, score) in [(&b"alice"[..], 100i64), (b"bob", 300), (b"carol", 200)] { - let bytes = SortKeyEncoder::encode_i64(score).to_vec(); - mgr.on_put(tbl_key, pk, &[("score".into(), bytes)]); + mgr.on_put(tbl_key, pk, Some(&serde_json::json!({ "score": score }))); } let exported = mgr.export_for_table(tbl_key); @@ -159,12 +153,68 @@ mod tests { let mut mgr = SortedIndexManager::new(); mgr.restore(0, 1, make_def("lb", "scores"), &[]); - let bytes = SortKeyEncoder::encode_i64(50).to_vec(); - mgr.on_put(scores_table_key(), b"dave", &[("score".into(), bytes)]); + mgr.on_put( + scores_table_key(), + b"dave", + Some(&serde_json::json!({ "score": 50 })), + ); assert_eq!( mgr.rank(0, 1, "lb", b"dave", 0), Some(1), "a restored registration must keep tracking later writes" ); } + + /// A restored windowed index decodes timestamps from the exported keys. + #[test] + fn restored_windowed_index_ranks_inside_the_window() { + let def = || { + let mut def = make_def("lb", "scores"); + def.encoder = SortKeyEncoder::new(vec![ + SortColumn { + name: "score".into(), + direction: SortDirection::Desc, + }, + SortColumn { + name: "ts".into(), + direction: SortDirection::Desc, + }, + ]); + def.window = WindowConfig::custom("ts", 1_000, 2_000); + def + }; + let mut mgr = SortedIndexManager::new(); + mgr.register(0, 1, def(), std::iter::empty()); + let rows = [ + ( + &b"before"[..], + serde_json::json!({ "score": 99, "ts": 999 }), + ), + (b"int", serde_json::json!({ "score": 11, "ts": 1_000 })), + ( + b"float", + serde_json::json!({ "score": 10.5, "ts": 1_500.5 }), + ), + ( + b"iso", + serde_json::json!({ "score": 5, "ts": "1970-01-01T00:00:01Z" }), + ), + (b"after", serde_json::json!({ "score": 50, "ts": 2_000 })), + ]; + for (pk, doc) in &rows { + mgr.on_put(scores_table_key(), pk, Some(doc)); + } + let entries = mgr.export_for_table(scores_table_key())[0].entries.clone(); + + let mut restored = SortedIndexManager::new(); + restored.restore(0, 1, def(), &entries); + + let now = 1_500; + assert_eq!(restored.count(0, 1, "lb", now), Some(3)); + assert_eq!(restored.rank(0, 1, "lb", b"int", now), Some(1)); + assert_eq!(restored.rank(0, 1, "lb", b"float", now), Some(2)); + assert_eq!(restored.rank(0, 1, "lb", b"iso", now), Some(3)); + assert_eq!(restored.rank(0, 1, "lb", b"before", now), None); + assert_eq!(restored.rank(0, 1, "lb", b"after", now), None); + } } diff --git a/nodedb/src/engine/kv/sorted_index/index.rs b/nodedb/src/engine/kv/sorted_index/index.rs index 87ff788f1..dcde5bcd3 100644 --- a/nodedb/src/engine/kv/sorted_index/index.rs +++ b/nodedb/src/engine/kv/sorted_index/index.rs @@ -8,7 +8,6 @@ //! folded in. Both answer through the same methods, so a read inside a //! transaction ranks, counts and windows exactly like an autocommit read. -use super::key::SortKeyEncoder; use super::manager::SortedIndexDef; use super::tree::OrderStatTree; use super::windowed_query::{self, SortedIndexRef}; @@ -87,9 +86,8 @@ impl SortedIndex { /// The entries in a score range as `(rank, primary_key)` pairs. /// - /// `score_min` and `score_max` are the raw value bytes of the index's - /// LEADING sort column (as [`extract_sort_key_from_value`] produces them), - /// not encoded tree keys. The caller names a score. Only the index's own + /// `score_min` and `score_max` are [`value`](super::value) encodings of + /// the index's LEADING sort column, not encoded tree keys. The caller names a score. Only the index's own /// encoder knows the framing and direction that turn it into a bound the /// tree compares against. pub fn range( @@ -130,52 +128,21 @@ impl SortedIndex { } } -/// Extract the sort columns from a MessagePack-encoded KV value and build a -/// sort key. `None` when the value is not a map or lacks a sort column. +/// The sort key of a MessagePack-encoded KV value. `None` when the value is +/// not a map or lacks a sort column. pub(super) fn extract_sort_key_from_value( def: &SortedIndexDef, value_bytes: &[u8], ) -> Option> { let doc: serde_json::Value = nodedb_types::json_from_msgpack(value_bytes).ok()?; - let obj = doc.as_object()?; - - let mut values: Vec> = Vec::with_capacity(def.encoder.column_count()); - for col in def.encoder.columns() { - let field_val = obj.get(&col.name)?; - values.push(field_value_to_sort_bytes(field_val)); - } - - let refs: Vec<&[u8]> = values.iter().map(|v| v.as_slice()).collect(); - Some(def.encoder.encode(&refs)) -} - -/// Convert a JSON field value to sortable bytes. -/// -/// The bytes match what the KV write path extracts for `on_put`. A row -/// backfilled here and the same row written later get the same sort key. -fn field_value_to_sort_bytes(val: &serde_json::Value) -> Vec { - match val { - serde_json::Value::Number(n) => { - if let Some(i) = n.as_i64() { - SortKeyEncoder::encode_i64(i).to_vec() - } else if let Some(f) = n.as_f64() { - SortKeyEncoder::encode_f64(f).to_vec() - } else { - Vec::new() - } - } - serde_json::Value::String(s) => s.as_bytes().to_vec(), - serde_json::Value::Bool(b) => vec![u8::from(*b)], - serde_json::Value::Null => Vec::new(), - other @ (serde_json::Value::Array(_) | serde_json::Value::Object(_)) => { - other.to_string().into_bytes() - } - } + def.encoder.key_for_document(&doc) } #[cfg(test)] mod tests { + use super::super::key::SortKeyEncoder; use super::super::key::{SortColumn, SortDirection}; + use super::super::value::{encode_float, encode_int}; use super::super::window::WindowConfig; use super::*; @@ -246,8 +213,8 @@ mod tests { pks(&idx.range(None, None, now_ms)), vec![(1, "fresh"), (2, "late")] ); - let score_200 = SortKeyEncoder::encode_i64(200); - let score_300 = SortKeyEncoder::encode_i64(300); + let score_200 = encode_int(200); + let score_300 = encode_int(300); assert_eq!( pks(&idx.range(Some(&score_200), Some(&score_300), now_ms)), vec![(1, "fresh")] @@ -449,16 +416,110 @@ mod tests { assert_eq!(idx.rank(b"old", NOW_MS), Some(1)); } + /// A row keyed by backfill and the same row keyed by write maintenance + /// must get one sort key, for every value kind. + #[test] + fn backfill_and_write_path_build_identical_keys_for_every_value_kind() { + use super::super::manager::SortedIndexManager; + + let kinds = [ + serde_json::json!(null), + serde_json::json!(true), + serde_json::json!(false), + serde_json::json!(11), + serde_json::json!(-11), + serde_json::json!(u64::MAX), + serde_json::json!(10.5), + serde_json::json!(-0.0), + serde_json::json!(""), + serde_json::json!("a\u{0}b"), + serde_json::json!([1, "two"]), + serde_json::json!({"k": [null]}), + ]; + let rows: Vec<(Vec, Vec)> = kinds + .iter() + .enumerate() + .map(|(i, score)| { + let doc = serde_json::json!({ "score": score, "ts": NOW_MS }); + let value = nodedb_types::json_to_msgpack(&doc).expect("encode row"); + (format!("pk{i}").into_bytes(), value) + }) + .collect(); + + let table = super::super::super::engine_helpers::table_key(0, 1, "scores"); + let unwindowed = || def(WindowConfig::none(), SortDirection::Desc); + + let mut backfilled = SortedIndexManager::new(); + let indexed = backfilled.register(0, 1, unwindowed(), rows.clone().into_iter()); + assert_eq!(indexed as usize, kinds.len(), "every value kind is indexed"); + + let mut written = SortedIndexManager::new(); + written.register(0, 1, unwindowed(), std::iter::empty()); + for (pk, value) in &rows { + let doc = super::super::super::engine_helpers::decode_document(value); + written.on_put(table, pk, doc.as_ref()); + } + + for (pk, _) in &rows { + let backfill_key = backfilled.score(0, 1, "lb", pk); + assert!(backfill_key.is_some()); + assert_eq!(backfill_key, written.score(0, 1, "lb", pk)); + } + } + + /// The reported case: 11 must sort above 10.5 in one column. #[test] - fn backfilled_bool_sorts_like_a_written_bool() { + fn integer_and_float_scores_rank_by_numeric_value() { + let rows = |scores: &[(&str, serde_json::Value)]| -> Vec<(Vec, Vec)> { + scores + .iter() + .map(|(pk, score)| { + let doc = serde_json::json!({ "score": score, "ts": NOW_MS }); + let value = nodedb_types::json_to_msgpack(&doc).expect("encode row"); + (pk.as_bytes().to_vec(), value) + }) + .collect() + }; + let scores = [ + ("eleven", serde_json::json!(11)), + ("ten_half", serde_json::json!(10.5)), + ("ten", serde_json::json!(10)), + ("big_int", serde_json::json!((1i64 << 53) + 1)), + ("big_float", serde_json::json!((1i64 << 53) as f64)), + ]; + + let mut desc = def(WindowConfig::none(), SortDirection::Asc); + desc.encoder = SortKeyEncoder::new(vec![SortColumn { + name: "score".into(), + direction: SortDirection::Desc, + }]); + let idx = build(desc, rows(&scores[..])); assert_eq!( - field_value_to_sort_bytes(&serde_json::json!(true)), - vec![1u8] + pks(&idx.top_k(10, NOW_MS)), + vec![ + (1, "big_int"), + (2, "big_float"), + (3, "eleven"), + (4, "ten_half"), + (5, "ten") + ] ); + + let mut asc = def(WindowConfig::none(), SortDirection::Asc); + asc.encoder = SortKeyEncoder::new(vec![SortColumn { + name: "score".into(), + direction: SortDirection::Asc, + }]); + let idx = build(asc, rows(&scores[..])); assert_eq!( - field_value_to_sort_bytes(&serde_json::json!(false)), - vec![0u8] + pks(&idx.top_k(3, NOW_MS)), + vec![(1, "ten"), (2, "ten_half"), (3, "eleven")] ); - assert!(field_value_to_sort_bytes(&serde_json::Value::Null).is_empty()); + + // An inclusive [10.5, 11] range holds exactly 10.5 and 11. + let entries = idx.range(Some(&encode_float(10.5)), Some(&encode_int(11)), NOW_MS); + let mut in_range = pks(&entries); + in_range.sort(); + assert_eq!(in_range, vec![(2, "ten_half"), (3, "eleven")]); } } diff --git a/nodedb/src/engine/kv/sorted_index/key.rs b/nodedb/src/engine/kv/sorted_index/key.rs index 289a1f5f4..9d5b66146 100644 --- a/nodedb/src/engine/kv/sorted_index/key.rs +++ b/nodedb/src/engine/kv/sorted_index/key.rs @@ -2,18 +2,18 @@ //! Composite sort key encoding for sorted indexes. //! -//! Encodes multi-column sort keys into a single `Vec` such that the -//! natural lexicographic byte ordering matches the desired sort order. +//! A composite key concatenates one encoded value per sort column, in column +//! order. The [`value`](super::value) encoding is order-preserving and no +//! encoded value is a prefix of another, so byte order of whole keys matches +//! the multi-column sort order. //! -//! For ascending columns: value bytes are stored as-is (big-endian). -//! For descending columns: every byte is bitwise-complemented (~byte), -//! which reverses the sort order while preserving lexicographic comparison. -//! -//! Columns are separated by a length prefix (4 bytes, big-endian) to avoid -//! ambiguity between field boundaries. +//! An ascending column stores its value bytes as-is. A descending column +//! stores every byte complemented (~byte), which reverses its order. use std::borrow::Cow; +use super::value::{MAX_LEADING_BYTE, encode_json, encoded_len}; + /// Sort direction for a column in a sorted index. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum SortDirection { @@ -21,6 +21,16 @@ pub enum SortDirection { Desc, } +impl SortDirection { + /// The byte XORed into each value byte of a column with this direction. + fn flip(self) -> u8 { + match self { + SortDirection::Asc => 0x00, + SortDirection::Desc => 0xFF, + } + } +} + /// Definition of one column in a composite sort key. #[derive(Debug, Clone)] pub struct SortColumn { @@ -33,7 +43,7 @@ pub struct SortColumn { /// Encodes composite sort keys for a sorted index definition. /// /// Given a list of `(column_name, direction)` pairs, produces a single -/// byte key from extracted field values such that `BTreeMap` ordering +/// byte key from encoded field values such that `BTreeMap` ordering /// matches the desired multi-column sort. #[derive(Debug, Clone)] pub struct SortKeyEncoder { @@ -55,146 +65,88 @@ impl SortKeyEncoder { &self.columns } - /// Encode field values into a composite sort key. + /// Encode column values into a composite sort key. /// - /// `values` must have the same length as `columns`. - /// Each value is the raw bytes of the field (big-endian for numerics, - /// UTF-8 for strings). + /// `values` holds one [`value`](super::value) encoding per column, in + /// column order. pub fn encode(&self, values: &[&[u8]]) -> Vec { debug_assert_eq!(values.len(), self.columns.len()); - let total_len: usize = values.iter().map(|v| 4 + v.len()).sum(); + let total_len: usize = values.iter().map(|v| v.len()).sum(); let mut key = Vec::with_capacity(total_len); - for (value, col) in values.iter().zip(&self.columns) { - // Length prefix (4 bytes, big-endian) — always ascending so that - // shorter values sort before longer values within the same column. - let len = value.len() as u32; - key.extend_from_slice(&len.to_be_bytes()); - - match col.direction { - SortDirection::Asc => { - key.extend_from_slice(value); - } - SortDirection::Desc => { - // Bitwise complement reverses sort order. - for &b in *value { - key.push(!b); - } - } - } + let flip = col.direction.flip(); + key.extend(value.iter().map(|byte| byte ^ flip)); } - key } + /// The sort key of a document. Backfill and write maintenance both build + /// keys here. + /// + /// `None` when `document` is not an object or lacks a sort column. + pub fn key_for_document(&self, document: &serde_json::Value) -> Option> { + let fields = document.as_object()?; + let mut key = Vec::new(); + for col in &self.columns { + let flip = col.direction.flip(); + key.extend(encode_json(fields.get(&col.name)?).iter().map(|b| b ^ flip)); + } + Some(key) + } + /// Translate an inclusive `[min, max]` range over the FIRST sort column /// into inclusive bounds in the encoded key space. /// - /// A caller (`RANGE(index, lo, hi)`, `ZRANGEBYSCORE`) knows only the raw - /// value bytes of the leading score column, but the tree is ordered by - /// whole keys produced by [`Self::encode`] — length-prefixed, and - /// bit-complemented on a descending column. Comparing raw value bytes - /// against those keys matches nothing, so the bounds have to be lifted into - /// the same space here, where the framing rules live. - /// - /// Two consequences of that framing are handled: + /// `min` and `max` are [`value`](super::value) encodings of the leading + /// column. Two rules of the key layout apply: /// - /// * complementing reverses order, so on a descending leading column the - /// caller's `min` is the GREATER encoded bound and the two swap; - /// * a composite key continues past the leading column with the next - /// column's 4-byte big-endian length, so the upper bound carries a - /// `0xFF` tail. That tail exceeds every representable continuation — - /// admitting every key whose leading column equals the bound — while - /// still sorting below any key with a greater leading value. + /// * a descending leading column is complemented, which reverses order, so + /// the caller's `min` is the GREATER key bound and the two swap; + /// * a composite key continues past the leading column with more column + /// values. The upper bound carries one byte above the first byte of any + /// value, so it admits every key whose leading column equals `max` and + /// stays below any key with a greater leading value. pub fn first_column_range_bounds( &self, min: Option<&[u8]>, max: Option<&[u8]>, ) -> (Option>, Option>) { - let descending = matches!( - self.columns.first().map(|col| col.direction), - Some(SortDirection::Desc) - ); + let direction = self + .columns + .first() + .map_or(SortDirection::Asc, |col| col.direction); + let flip = direction.flip(); let framed = |value: &[u8]| -> Vec { - let mut out = Vec::with_capacity(4 + value.len() + 4); - out.extend_from_slice(&(value.len() as u32).to_be_bytes()); - if descending { - out.extend(value.iter().map(|byte| !byte)); - } else { - out.extend_from_slice(value); - } + let mut out = Vec::with_capacity(value.len() + 1); + out.extend(value.iter().map(|byte| byte ^ flip)); out }; - let (lower_source, upper_source) = if descending { (max, min) } else { (min, max) }; + let (lower_source, upper_source) = match direction { + SortDirection::Asc => (min, max), + SortDirection::Desc => (max, min), + }; let lower = lower_source.map(framed); let upper = upper_source.map(|value| { let mut bound = framed(value); - bound.extend_from_slice(&[0xFF; 4]); + bound.push(MAX_LEADING_BYTE + 1); bound }); (lower, upper) } - /// Encode an i64 score as big-endian bytes suitable for sorting. - /// - /// Maps i64 to u64 by flipping the sign bit, so negative values sort - /// before positive values in unsigned byte comparison. - pub fn encode_i64(value: i64) -> [u8; 8] { - let unsigned = (value as u64) ^ (1u64 << 63); - unsigned.to_be_bytes() - } - - /// Decode big-endian bytes back to i64. - pub fn decode_i64(bytes: &[u8; 8]) -> i64 { - let unsigned = u64::from_be_bytes(*bytes); - (unsigned ^ (1u64 << 63)) as i64 - } - - /// Encode an f64 score as big-endian bytes suitable for sorting. - /// - /// Uses the IEEE 754 total-order encoding trick: - /// - Positive floats: flip the sign bit (0x80...00 XOR). - /// - Negative floats: flip all bits (bitwise NOT). - /// - /// This produces a byte sequence where f64 ordering matches byte ordering. - pub fn encode_f64(value: f64) -> [u8; 8] { - let bits = value.to_bits(); - let encoded = if bits & (1u64 << 63) == 0 { - // Positive (or +0): flip sign bit. - bits ^ (1u64 << 63) - } else { - // Negative (or -0): flip all bits. - !bits - }; - encoded.to_be_bytes() - } - - /// Decode big-endian bytes back to f64. - pub fn decode_f64(bytes: &[u8; 8]) -> f64 { - let encoded = u64::from_be_bytes(*bytes); - let bits = if encoded & (1u64 << 63) != 0 { - // Was positive: flip sign bit back. - encoded ^ (1u64 << 63) - } else { - // Was negative: flip all bits back. - !encoded - }; - f64::from_bits(bits) - } - /// The value bytes of column `name` in `sort_key`, as they were passed to /// [`Self::encode`]. /// - /// Strips the length prefix and undoes the DESC complement. Returns `None` - /// when the encoder has no column `name` or the key is truncated. + /// Undoes the DESC complement. Returns `None` when the encoder has no + /// column `name` or the key does not hold whole values up to it. pub fn column_value<'k>(&self, sort_key: &'k [u8], name: &str) -> Option> { let mut rest = sort_key; for col in &self.columns { - let (len_bytes, tail) = rest.split_first_chunk::<4>()?; - let len = u32::from_be_bytes(*len_bytes) as usize; - let (value, tail) = tail.split_at_checked(len)?; + let flip = col.direction.flip(); + let len = encoded_len(rest, flip)?; + let (value, tail) = rest.split_at_checked(len)?; if col.name == name { return Some(match col.direction { SortDirection::Asc => Cow::Borrowed(value), @@ -209,183 +161,148 @@ impl SortKeyEncoder { #[cfg(test)] mod tests { + use super::super::value::{encode_float, encode_int, encode_text}; use super::*; - #[test] - fn i64_roundtrip() { - for v in [i64::MIN, -1, 0, 1, i64::MAX] { - let encoded = SortKeyEncoder::encode_i64(v); - assert_eq!(SortKeyEncoder::decode_i64(&encoded), v); - } - } - - #[test] - fn i64_ordering() { - let values = [i64::MIN, -1000, -1, 0, 1, 1000, i64::MAX]; - let encoded: Vec<_> = values - .iter() - .map(|v| SortKeyEncoder::encode_i64(*v)) - .collect(); - for i in 0..encoded.len() - 1 { - assert!( - encoded[i] < encoded[i + 1], - "encode_i64({}) should be < encode_i64({})", - values[i], - values[i + 1] - ); - } - } - - #[test] - fn f64_roundtrip() { - for v in [f64::NEG_INFINITY, -1.0, -0.0, 0.0, 1.0, f64::INFINITY] { - let encoded = SortKeyEncoder::encode_f64(v); - let decoded = SortKeyEncoder::decode_f64(&encoded); - assert_eq!(v.to_bits(), decoded.to_bits(), "roundtrip failed for {v}"); - } - } - - #[test] - fn f64_ordering() { - let values = [ - f64::NEG_INFINITY, - -100.0, - -0.001, - 0.0, - 0.001, - 100.0, - f64::INFINITY, - ]; - let encoded: Vec<_> = values - .iter() - .map(|v| SortKeyEncoder::encode_f64(*v)) - .collect(); - for i in 0..encoded.len() - 1 { - assert!( - encoded[i] < encoded[i + 1], - "encode_f64({}) should be < encode_f64({})", - values[i], - values[i + 1] - ); - } + fn encoder(direction: SortDirection) -> SortKeyEncoder { + SortKeyEncoder::new(vec![SortColumn { + name: "score".into(), + direction, + }]) } - #[test] - fn composite_key_desc_score_asc_time() { - let encoder = SortKeyEncoder::new(vec![ + fn score_and_ts(ts_direction: SortDirection) -> SortKeyEncoder { + SortKeyEncoder::new(vec![ SortColumn { name: "score".into(), direction: SortDirection::Desc, }, SortColumn { - name: "updated_at".into(), - direction: SortDirection::Asc, + name: "ts".into(), + direction: ts_direction, }, - ]); - - // Higher score should sort first (lower bytes due to DESC complement). - let score_100 = SortKeyEncoder::encode_i64(100); - let score_200 = SortKeyEncoder::encode_i64(200); - let time_early = SortKeyEncoder::encode_i64(1000); - let time_late = SortKeyEncoder::encode_i64(2000); + ]) + } - let key_200_early = encoder.encode(&[&score_200, &time_early]); - let key_100_early = encoder.encode(&[&score_100, &time_early]); - let key_200_late = encoder.encode(&[&score_200, &time_late]); + #[test] + fn composite_key_desc_score_asc_time() { + let enc = score_and_ts(SortDirection::Asc); + let key = |score: i64, ts: i64| enc.encode(&[&encode_int(score), &encode_int(ts)]); // 200 DESC sorts before 100 DESC. - assert!(key_200_early < key_100_early); + assert!(key(200, 1000) < key(100, 1000)); // Same score, earlier time ASC sorts before later time. - assert!(key_200_early < key_200_late); + assert!(key(200, 1000) < key(200, 2000)); } + /// An integer and a float in one column order by numeric value. #[test] - fn composite_key_asc_score() { - let encoder = SortKeyEncoder::new(vec![SortColumn { - name: "score".into(), - direction: SortDirection::Asc, - }]); - - let score_50 = SortKeyEncoder::encode_i64(50); - let score_150 = SortKeyEncoder::encode_i64(150); - - let key_50 = encoder.encode(&[&score_50]); - let key_150 = encoder.encode(&[&score_150]); + fn mixed_integer_and_float_keys_sort_numerically() { + let asc = encoder(SortDirection::Asc); + let key = |v: &[u8]| asc.encode(&[v]); + assert!(key(&encode_int(10)) < key(&encode_float(10.5))); + assert!(key(&encode_float(10.5)) < key(&encode_int(11))); + + let desc = encoder(SortDirection::Desc); + let key = |v: &[u8]| desc.encode(&[v]); + assert!(key(&encode_int(11)) < key(&encode_float(10.5))); + assert!(key(&encode_float(10.5)) < key(&encode_int(10))); + } - // ASC: lower score sorts first. - assert!(key_50 < key_150); + /// A document's key equals the key built from its encoded field values. + #[test] + fn key_for_document_matches_encode_of_its_fields() { + let enc = score_and_ts(SortDirection::Desc); + let doc = serde_json::json!({ "score": 10.5, "ts": "2024-05-02", "other": 1 }); + assert_eq!( + enc.key_for_document(&doc), + Some(enc.encode(&[&encode_float(10.5), &encode_text("2024-05-02")])) + ); + assert_eq!( + enc.key_for_document(&serde_json::json!({ "score": 1 })), + None + ); + assert_eq!(enc.key_for_document(&serde_json::json!([1, 2])), None); } - fn score_and_ts(ts_direction: SortDirection) -> SortKeyEncoder { - SortKeyEncoder::new(vec![ - SortColumn { - name: "score".into(), - direction: SortDirection::Desc, - }, - SortColumn { - name: "ts".into(), - direction: ts_direction, - }, - ]) + /// Strings order bytewise, not by length, in either direction. + #[test] + fn string_keys_sort_lexicographically() { + for direction in [SortDirection::Asc, SortDirection::Desc] { + let enc = score_and_ts(direction); + let key = |s: &str| enc.encode(&[&encode_int(1), &encode_text(s)]); + let ascending = ["", "a", "a\u{0}", "aa", "b", "zz"]; + for pair in ascending.windows(2) { + let (lo, hi) = (key(pair[0]), key(pair[1])); + match direction { + SortDirection::Asc => assert!(lo < hi, "{pair:?}"), + SortDirection::Desc => assert!(lo > hi, "{pair:?}"), + } + } + } } #[test] fn column_value_returns_the_bytes_passed_to_encode() { for direction in [SortDirection::Asc, SortDirection::Desc] { let enc = score_and_ts(direction); - let score = SortKeyEncoder::encode_i64(7); - let ts = SortKeyEncoder::encode_i64(1_714_608_000_000); - let key = enc.encode(&[&score, &ts]); - - assert_eq!(enc.column_value(&key, "ts").as_deref(), Some(&ts[..])); - assert_eq!(enc.column_value(&key, "score").as_deref(), Some(&score[..])); + let score = encode_float(7.5); + for ts in [encode_int(1_714_608_000_000), encode_text("2024-05-02")] { + let key = enc.encode(&[&score, &ts]); + assert_eq!(enc.column_value(&key, "ts").as_deref(), Some(&ts[..])); + assert_eq!(enc.column_value(&key, "score").as_deref(), Some(&score[..])); + } } } #[test] - fn column_value_of_an_empty_value_is_empty() { + fn column_value_of_a_null_is_the_null_encoding() { let enc = score_and_ts(SortDirection::Asc); - let key = enc.encode(&[&SortKeyEncoder::encode_i64(7), b""]); - assert_eq!(enc.column_value(&key, "ts").as_deref(), Some(&b""[..])); + let null = encode_json(&serde_json::Value::Null); + let key = enc.encode(&[&encode_int(7), &null]); + assert_eq!(enc.column_value(&key, "ts").as_deref(), Some(&null[..])); } #[test] fn column_value_of_an_unknown_column_or_truncated_key_is_none() { let enc = score_and_ts(SortDirection::Asc); - let ts = SortKeyEncoder::encode_i64(1_714_608_000_000); - let key = enc.encode(&[&SortKeyEncoder::encode_i64(7), &ts]); + let key = enc.encode(&[&encode_int(7), &encode_int(1_714_608_000_000)]); assert!(enc.column_value(&key, "missing").is_none()); assert!(enc.column_value(&key[..key.len() - 1], "ts").is_none()); assert!(enc.column_value(&key[..2], "score").is_none()); } - fn encoder(direction: SortDirection) -> SortKeyEncoder { - SortKeyEncoder::new(vec![SortColumn { - name: "score".into(), - direction, - }]) - } - - /// A bound expressed in raw value bytes has to land inside the encoded key - /// space, or it is compared against a length prefix and matches nothing. + /// A bound expressed in value bytes has to land inside the encoded key + /// space, or it is compared against the wrong bytes and matches nothing. #[test] fn ascending_bounds_bracket_the_keys_they_name() { let enc = encoder(SortDirection::Asc); - let key_of = |v: i64| enc.encode(&[&SortKeyEncoder::encode_i64(v)[..]]); - let (lower, upper) = enc.first_column_range_bounds( - Some(&SortKeyEncoder::encode_i64(10)), - Some(&SortKeyEncoder::encode_i64(20)), - ); + let key_of = |v: &[u8]| enc.encode(&[v]); + let (lower, upper) = + enc.first_column_range_bounds(Some(&encode_int(10)), Some(&encode_int(20))); let (lower, upper) = (lower.expect("lower"), upper.expect("upper")); - for inside in [10, 15, 20] { - let key = key_of(inside); - assert!(key >= lower && key <= upper, "{inside} must be in range"); + for inside in [ + encode_int(10), + encode_float(10.5), + encode_int(15), + encode_float(20.0), + ] { + let key = key_of(&inside); + assert!(key >= lower && key <= upper, "{inside:?} must be in range"); } - for outside in [9, 21] { - let key = key_of(outside); - assert!(key < lower || key > upper, "{outside} must be out of range"); + for outside in [ + encode_int(9), + encode_float(9.5), + encode_float(20.5), + encode_int(21), + ] { + let key = key_of(&outside); + assert!( + key < lower || key > upper, + "{outside:?} must be out of range" + ); } } @@ -395,21 +312,22 @@ mod tests { #[test] fn descending_bounds_are_swapped_into_key_order() { let enc = encoder(SortDirection::Desc); - let key_of = |v: i64| enc.encode(&[&SortKeyEncoder::encode_i64(v)[..]]); - let (lower, upper) = enc.first_column_range_bounds( - Some(&SortKeyEncoder::encode_i64(10)), - Some(&SortKeyEncoder::encode_i64(20)), - ); + let key_of = |v: &[u8]| enc.encode(&[v]); + let (lower, upper) = + enc.first_column_range_bounds(Some(&encode_float(10.5)), Some(&encode_int(11))); let (lower, upper) = (lower.expect("lower"), upper.expect("upper")); assert!(lower <= upper, "bounds must not be inverted"); - for inside in [10, 15, 20] { - let key = key_of(inside); - assert!(key >= lower && key <= upper, "{inside} must be in range"); + for inside in [encode_float(10.5), encode_float(10.75), encode_int(11)] { + let key = key_of(&inside); + assert!(key >= lower && key <= upper, "{inside:?} must be in range"); } - for outside in [9, 21] { - let key = key_of(outside); - assert!(key < lower || key > upper, "{outside} must be out of range"); + for outside in [encode_int(10), encode_float(11.25), encode_int(12)] { + let key = key_of(&outside); + assert!( + key < lower || key > upper, + "{outside:?} must be out of range" + ); } } @@ -417,29 +335,45 @@ mod tests { /// including composite keys that carry more columns after it. #[test] fn upper_bound_admits_composite_keys_on_the_boundary() { - let enc = SortKeyEncoder::new(vec![ - SortColumn { - name: "score".into(), - direction: SortDirection::Asc, - }, - SortColumn { - name: "tiebreak".into(), - direction: SortDirection::Asc, - }, - ]); - let key = enc.encode(&[&SortKeyEncoder::encode_i64(20)[..], b"zzz"]); - let (lower, upper) = enc.first_column_range_bounds( - Some(&SortKeyEncoder::encode_i64(10)), - Some(&SortKeyEncoder::encode_i64(20)), - ); - let (lower, upper) = (lower.expect("lower"), upper.expect("upper")); - - assert!(key >= lower && key <= upper); + for (lead, tail) in [ + (SortDirection::Asc, SortDirection::Asc), + (SortDirection::Asc, SortDirection::Desc), + (SortDirection::Desc, SortDirection::Asc), + (SortDirection::Desc, SortDirection::Desc), + ] { + let enc = SortKeyEncoder::new(vec![ + SortColumn { + name: "score".into(), + direction: lead, + }, + SortColumn { + name: "tiebreak".into(), + direction: tail, + }, + ]); + let (lower, upper) = + enc.first_column_range_bounds(Some(&encode_int(10)), Some(&encode_int(20))); + let (lower, upper) = (lower.expect("lower"), upper.expect("upper")); + + for tiebreak in [ + encode_json(&serde_json::Value::Null), + encode_text(""), + encode_text("zzz"), + encode_json(&serde_json::json!({"k": "v"})), + ] { + for bound in [encode_int(10), encode_int(20)] { + let key = enc.encode(&[&bound, &tiebreak]); + assert!(key >= lower && key <= upper, "{lead:?}/{tail:?}"); + } + } - let beyond = enc.encode(&[&SortKeyEncoder::encode_i64(21)[..], b""]); - assert!( - beyond > upper, - "the next score must stay outside the window" - ); + for beyond in [encode_int(9), encode_float(20.5), encode_int(21)] { + let key = enc.encode(&[&beyond, &encode_text("")]); + assert!( + key < lower || key > upper, + "a score outside the window must stay outside: {lead:?}/{tail:?}" + ); + } + } } } diff --git a/nodedb/src/engine/kv/sorted_index/manager.rs b/nodedb/src/engine/kv/sorted_index/manager.rs index 18654e2f2..3e5d8d241 100644 --- a/nodedb/src/engine/kv/sorted_index/manager.rs +++ b/nodedb/src/engine/kv/sorted_index/manager.rs @@ -142,29 +142,32 @@ impl SortedIndexManager { true } - /// Called on every KV PUT. Updates all sorted indexes on this collection. + /// Called on every KV PUT. Re-keys `primary_key` in every sorted index on + /// this collection. /// - /// `field_values` is a map of field_name → field_value_bytes extracted from - /// the MessagePack value. + /// `document` is the stored value decoded from MessagePack, `None` when it + /// does not decode. A value without a sort column leaves the index. pub fn on_put( &mut self, table_key: u64, primary_key: &[u8], - field_values: &[(String, Vec)], + document: Option<&serde_json::Value>, ) { let Some(idx_keys) = self.collection_indexes.get(&table_key) else { return; }; - // Clone keys to avoid borrow conflict with self.indexes. - let idx_keys: Vec = idx_keys.iter().cloned().collect(); - for idx_key in &idx_keys { + for idx_key in idx_keys { let Some(idx) = self.indexes.get_mut(idx_key) else { continue; }; - - if let Some(sort_key) = build_sort_key_from_fields(&idx.def, field_values) { - idx.tree.insert(sort_key, primary_key.to_vec()); + match document.and_then(|doc| idx.def.encoder.key_for_document(doc)) { + Some(sort_key) => { + idx.tree.insert(sort_key, primary_key.to_vec()); + } + None => { + idx.tree.remove(primary_key); + } } } } @@ -303,29 +306,6 @@ pub(super) fn index_key(database_id: u64, tenant_id: u64, index_name: &str) -> S format!("{database_id}:{tenant_id}:{index_name}") } -/// Build a sort key from pre-extracted field name/value pairs. -fn build_sort_key_from_fields( - def: &SortedIndexDef, - field_values: &[(String, Vec)], -) -> Option> { - let mut values: Vec> = Vec::with_capacity(def.encoder.column_count()); - - for col in def.encoder.columns() { - let val_bytes = field_values - .iter() - .find(|(name, _)| name == &col.name) - .map(|(_, v)| v.clone())?; - - // The field value bytes from extract_all_field_values_from_msgpack are - // already encoded as sortable bytes (integers as big-endian u64 with - // sign-bit flip, strings as UTF-8). Use them directly. - values.push(val_bytes); - } - - let refs: Vec<&[u8]> = values.iter().map(|v| v.as_slice()).collect(); - Some(def.encoder.encode(&refs)) -} - #[cfg(test)] mod tests { use super::super::key::{SortColumn, SortDirection}; @@ -354,6 +334,10 @@ mod tests { (pk, value) } + fn scored(score: i64) -> serde_json::Value { + serde_json::json!({ "score": score }) + } + #[test] fn register_and_backfill() { let mut mgr = SortedIndexManager::new(); @@ -411,11 +395,8 @@ mod tests { let tbl_key = super::super::super::engine_helpers::table_key(0, 1, "scores"); // Simulate PUTs. - let score_bytes = SortKeyEncoder::encode_i64(100).to_vec(); - mgr.on_put(tbl_key, b"alice", &[("score".into(), score_bytes)]); - - let score_bytes = SortKeyEncoder::encode_i64(200).to_vec(); - mgr.on_put(tbl_key, b"bob", &[("score".into(), score_bytes)]); + mgr.on_put(tbl_key, b"alice", Some(&scored(100))); + mgr.on_put(tbl_key, b"bob", Some(&scored(200))); assert_eq!(mgr.count(0, 1, "lb", 0), Some(2)); // DESC: bob(200) = rank 1, alice(100) = rank 2 @@ -510,8 +491,7 @@ mod tests { assert_eq!(mgr.rank(0, 1, "lb", b"alice", 0), Some(3)); // And it is still maintained on write — one insert, not two trees. - let bytes = SortKeyEncoder::encode_i64(400).to_vec(); - mgr.on_put(tbl_key, b"dave", &[("score".into(), bytes)]); + mgr.on_put(tbl_key, b"dave", Some(&scored(400))); assert_eq!(mgr.rank(0, 1, "lb", b"dave", 0), Some(1)); assert_eq!(mgr.count(0, 1, "lb", 0), Some(4)); } @@ -565,8 +545,7 @@ mod tests { assert_eq!(mgr.count(0, 1, "lb", 0), Some(1)); // A write to the collection the index no longer covers must not reach it. - let bytes = SortKeyEncoder::encode_i64(999).to_vec(); - mgr.on_put(old_key, b"ghost", &[("score".into(), bytes)]); + mgr.on_put(old_key, b"ghost", Some(&scored(999))); assert_eq!( mgr.rank(0, 1, "lb", b"ghost", 0), None, diff --git a/nodedb/src/engine/kv/sorted_index/mod.rs b/nodedb/src/engine/kv/sorted_index/mod.rs index c6aaf1268..7dd6ef1af 100644 --- a/nodedb/src/engine/kv/sorted_index/mod.rs +++ b/nodedb/src/engine/kv/sorted_index/mod.rs @@ -4,8 +4,10 @@ pub mod checkpoint; pub mod index; pub mod key; pub mod manager; +pub mod number; pub mod timestamp; pub mod tree; +pub mod value; pub mod window; mod windowed_query; diff --git a/nodedb/src/engine/kv/sorted_index/number.rs b/nodedb/src/engine/kv/sorted_index/number.rs new file mode 100644 index 000000000..cfb69533f --- /dev/null +++ b/nodedb/src/engine/kv/sorted_index/number.rs @@ -0,0 +1,331 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Exact order-preserving encoding for the numbers of a sort column. +//! +//! Integers and floats share one encoding, ordered by numeric value. A number +//! encodes as two parts, compared in order: +//! +//! - the greatest f64 at or below the value, as 8 order-preserving bytes; +//! - the integer residual above that f64, as 2 big-endian bytes. +//! +//! A float is its own floor, so its residual is 0. An integer whose magnitude +//! exceeds 2^53 can fall between two f64 values. Its floor is the f64 below it +//! and the residual holds the exact remainder. No cast rounds a value. +//! +//! Two values with different floors order by floor: the value with the lower +//! floor is below the next f64, which is at or below the other value. Two +//! values with one floor order by residual. So byte order is numeric order. +//! +//! Equal numbers encode equal: 11 and 11.0 give the same bytes, and so do +//! -0.0 and +0.0. Every NaN encodes as one value above +infinity. + +/// Bytes in an encoded number. +pub const NUMBER_LEN: usize = 10; + +/// Top bit of an f64 bit pattern. +const SIGN_BIT: u64 = 1 << 63; + +/// 2^64. Every integer this module encodes has a magnitude at or below it. +const WHOLE_LIMIT: f64 = 18_446_744_073_709_551_616.0; + +/// A number read back from its encoding. +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum SortNumber { + /// A whole number with a magnitude at or below 2^64, from an integer or a float. + Whole(i128), + /// Any other float: a fraction, a whole number past 2^64, an infinity, or NaN. + Float(f64), +} + +/// Encode a signed integer. +pub fn encode_i64(value: i64) -> [u8; NUMBER_LEN] { + encode_integer(i128::from(value)) +} + +/// Encode an unsigned integer. +pub fn encode_u64(value: u64) -> [u8; NUMBER_LEN] { + encode_integer(i128::from(value)) +} + +/// Encode a float. +pub fn encode_f64(value: f64) -> [u8; NUMBER_LEN] { + pack(value, 0) +} + +/// Encode a JSON number by its exact value. +pub fn encode_json_number(value: &serde_json::Number) -> [u8; NUMBER_LEN] { + match (value.as_i64(), value.as_u64()) { + (Some(i), _) => encode_i64(i), + (None, Some(u)) => encode_u64(u), + // A number that is neither i64 nor u64 is an f64, so `as_f64` is `Some`. + (None, None) => encode_f64(value.as_f64().unwrap_or(f64::NAN)), + } +} + +/// Decode bytes written by this module. `None` for bytes no value encodes to. +pub fn decode(bytes: &[u8; NUMBER_LEN]) -> Option { + let (floor_bytes, residual_bytes) = bytes.split_first_chunk::<8>()?; + let residual_bytes: &[u8; 2] = residual_bytes.try_into().ok()?; + let floor = f64_from_order_bits(u64::from_be_bytes(*floor_bytes)); + let residual = u16::from_be_bytes(*residual_bytes); + + if floor.fract() == 0.0 && floor.abs() <= WHOLE_LIMIT { + // The guard bounds `floor` to an exact i128, so the cast is exact. + return Some(SortNumber::Whole(floor as i128 + i128::from(residual))); + } + (residual == 0).then_some(SortNumber::Float(floor)) +} + +/// Encode an integer with a magnitude at or below 2^64. +fn encode_integer(value: i128) -> [u8; NUMBER_LEN] { + let floor = floor_f64(value); + // `floor` is whole and at most 2^64 in magnitude, so the cast is exact. + let residual = value - floor as i128; + // The residual is below the spacing of f64 values at `floor`. That spacing + // is at most 2^11 for magnitudes up to 2^64, so it fits in a u16. + pack(floor, residual as u16) +} + +/// The greatest f64 at or below `value`. +fn floor_f64(value: i128) -> f64 { + // The conversion rounds to nearest, so it can land one f64 above `value`. + let nearest = value as f64; + if nearest as i128 > value { + nearest.next_down() + } else { + nearest + } +} + +fn pack(floor: f64, residual: u16) -> [u8; NUMBER_LEN] { + let mut out = [0u8; NUMBER_LEN]; + out[..8].copy_from_slice(&order_bits(canonical(floor)).to_be_bytes()); + out[8..].copy_from_slice(&residual.to_be_bytes()); + out +} + +/// One bit pattern per numeric value: a single NaN, and +0.0 for both zeros. +fn canonical(value: f64) -> f64 { + if value.is_nan() { + f64::NAN + } else if value == 0.0 { + 0.0 + } else { + value + } +} + +/// Map f64 bits to u64 so unsigned order matches numeric order. +fn order_bits(value: f64) -> u64 { + let bits = value.to_bits(); + if bits & SIGN_BIT == 0 { + bits ^ SIGN_BIT + } else { + !bits + } +} + +fn f64_from_order_bits(encoded: u64) -> f64 { + if encoded & SIGN_BIT != 0 { + f64::from_bits(encoded ^ SIGN_BIT) + } else { + f64::from_bits(!encoded) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A number of either kind, compared by exact numeric value. + #[derive(Debug, Clone, Copy)] + enum Num { + I(i64), + U(u64), + F(f64), + } + + impl Num { + fn encode(self) -> [u8; NUMBER_LEN] { + match self { + Num::I(v) => encode_i64(v), + Num::U(v) => encode_u64(v), + Num::F(v) => encode_f64(v), + } + } + } + + /// Exact numeric comparison of an integer with a finite float. + fn cmp_int_float(int: i128, float: f64) -> std::cmp::Ordering { + use std::cmp::Ordering; + if float >= 2f64.powi(100) { + return Ordering::Less; + } + if float <= -2f64.powi(100) { + return Ordering::Greater; + } + let whole = float.floor(); + let whole_int = whole as i128; + match int.cmp(&whole_int) { + Ordering::Equal if float > whole => Ordering::Less, + other => other, + } + } + + fn numeric_cmp(a: Num, b: Num) -> std::cmp::Ordering { + let as_int = |n: Num| match n { + Num::I(v) => Some(i128::from(v)), + Num::U(v) => Some(i128::from(v)), + Num::F(_) => None, + }; + match (a, b) { + (Num::F(x), Num::F(y)) => x.partial_cmp(&y).unwrap_or_else(|| { + // NaN sorts above every number and equals itself. + x.is_nan().cmp(&y.is_nan()) + }), + (Num::F(x), other) => { + let int = as_int(other).unwrap_or_default(); + if x.is_nan() { + std::cmp::Ordering::Greater + } else if x.is_infinite() { + if x > 0.0 { + std::cmp::Ordering::Greater + } else { + std::cmp::Ordering::Less + } + } else { + cmp_int_float(int, x).reverse() + } + } + (other, Num::F(_)) => numeric_cmp(b, other).reverse(), + (x, y) => as_int(x).cmp(&as_int(y)), + } + } + + /// Edge values: around ±2^53, ±2^63, 2^64, fractions between large + /// consecutive integers, zeros, infinities and NaN. + fn edge_values() -> Vec { + let p53 = 1i64 << 53; + let mut out = vec![ + Num::I(i64::MIN), + Num::I(i64::MIN + 1), + Num::I(i64::MAX), + Num::I(i64::MAX - 1), + Num::U(u64::MAX), + Num::U(u64::MAX - 1), + Num::U(1 << 63), + Num::U((1 << 63) + 1), + Num::I(0), + Num::I(1), + Num::I(-1), + Num::I(10), + Num::I(11), + Num::F(10.5), + Num::F(11.0), + Num::F(-10.5), + Num::F(0.0), + Num::F(-0.0), + Num::F(0.5), + Num::F(-0.5), + Num::F(f64::MIN_POSITIVE), + Num::F(-f64::MIN_POSITIVE), + Num::F(f64::MAX), + Num::F(f64::MIN), + Num::F(f64::INFINITY), + Num::F(f64::NEG_INFINITY), + Num::F(f64::NAN), + Num::F(-f64::NAN), + Num::F(9_223_372_036_854_775_808.0), + Num::F(-9_223_372_036_854_775_808.0), + Num::F(18_446_744_073_709_551_616.0), + Num::F(9_223_372_036_854_774_784.0), + ]; + for base in [p53, -p53, 1i64 << 62, -(1i64 << 62)] { + for delta in -3i64..=3 { + out.push(Num::I(base + delta)); + out.push(Num::F((base + delta) as f64)); + } + // A float strictly between two consecutive integers. + out.push(Num::F(base as f64 / 4.0 + 0.5)); + } + for v in [p53 - 1, p53 + 1, p53 + 3] { + out.push(Num::F(v as f64 - 0.5)); + } + out + } + + #[test] + fn byte_order_matches_numeric_order_for_every_pair_of_edge_values() { + let values = edge_values(); + for &a in &values { + for &b in &values { + assert_eq!( + a.encode().cmp(&b.encode()), + numeric_cmp(a, b), + "{a:?} vs {b:?}" + ); + } + } + } + + #[test] + fn eleven_sorts_above_ten_and_a_half() { + assert!(encode_f64(10.5) < encode_i64(11)); + assert!(encode_i64(10) < encode_f64(10.5)); + } + + #[test] + fn equal_numbers_encode_equal() { + assert_eq!(encode_i64(11), encode_f64(11.0)); + assert_eq!(encode_u64(11), encode_i64(11)); + assert_eq!(encode_f64(-0.0), encode_f64(0.0)); + assert_eq!(encode_f64(-0.0), encode_i64(0)); + assert_eq!(encode_f64(f64::NAN), encode_f64(-f64::NAN)); + } + + #[test] + fn integers_past_two_pow_53_stay_distinct() { + let p53 = 1i64 << 53; + assert!(encode_i64(p53) < encode_i64(p53 + 1)); + assert!(encode_i64(p53 + 1) < encode_f64((p53 + 2) as f64)); + assert!(encode_i64(i64::MAX - 1) < encode_i64(i64::MAX)); + assert!(encode_i64(i64::MAX) < encode_f64(9_223_372_036_854_775_808.0)); + assert!(encode_f64(-9_223_372_036_854_775_808.0) == encode_i64(i64::MIN)); + } + + #[test] + fn nan_sorts_above_infinity() { + assert!(encode_f64(f64::INFINITY) < encode_f64(f64::NAN)); + assert!(encode_u64(u64::MAX) < encode_f64(f64::INFINITY)); + } + + #[test] + fn decode_returns_the_encoded_value() { + for v in [i64::MIN, -(1 << 53) - 1, -1, 0, 1, (1 << 53) + 1, i64::MAX] { + assert_eq!( + decode(&encode_i64(v)), + Some(SortNumber::Whole(i128::from(v))) + ); + } + assert_eq!( + decode(&encode_u64(u64::MAX)), + Some(SortNumber::Whole(i128::from(u64::MAX))) + ); + assert_eq!(decode(&encode_f64(11.0)), Some(SortNumber::Whole(11))); + assert_eq!(decode(&encode_f64(-0.0)), Some(SortNumber::Whole(0))); + for v in [10.5, -10.5, f64::INFINITY, f64::NEG_INFINITY, 1e30] { + assert_eq!(decode(&encode_f64(v)), Some(SortNumber::Float(v))); + } + assert!(matches!( + decode(&encode_f64(f64::NAN)), + Some(SortNumber::Float(v)) if v.is_nan() + )); + } + + #[test] + fn a_residual_on_a_fraction_is_not_a_value() { + let mut bytes = encode_f64(10.5); + bytes[NUMBER_LEN - 1] = 1; + assert_eq!(decode(&bytes), None); + } +} diff --git a/nodedb/src/engine/kv/sorted_index/timestamp.rs b/nodedb/src/engine/kv/sorted_index/timestamp.rs index ec574d8b8..ef443df8e 100644 --- a/nodedb/src/engine/kv/sorted_index/timestamp.rs +++ b/nodedb/src/engine/kv/sorted_index/timestamp.rs @@ -2,122 +2,109 @@ //! Read a window timestamp back from the value bytes of one sort column. //! -//! A sort column stores the bytes the KV value extractor wrote for it: +//! The value encoding carries a type tag, so the reader needs no guess: //! -//! - an integer: [`SortKeyEncoder::encode_i64`], 8 bytes; -//! - a float: [`SortKeyEncoder::encode_f64`], 8 bytes; -//! - a string: its UTF-8 bytes; -//! - null: no bytes. -//! -//! The bytes carry no type tag. A non-negative integer and a non-negative -//! float both encode with the top bit set, so the payload under that bit -//! picks the decoder. An integer timestamp of 2^53 ms is past the year -//! 287,000. A non-negative float whose bit pattern is below 2^53 is below -//! 2^-1021 ms. Payloads below 2^53 decode as integers, the rest as floats. -//! The two ranges overlap only on values no timestamp column holds. +//! - a number is milliseconds since the Unix epoch, whole or fractional; +//! - a string is an ISO 8601 instant; +//! - anything else, and a string that is not a date, has no timestamp. use nodedb_types::NdbDateTime; -use super::key::SortKeyEncoder; - -/// Top bit of an encoded 8-byte number. Set for every non-negative value. -const SIGN_BIT: u64 = 1 << 63; - -/// Payloads below this bound decode as integer milliseconds. -const INTEGER_PAYLOAD_LIMIT: u64 = 1 << 53; +use super::number::SortNumber; +use super::value::{SortValue, decode}; /// A timestamp read from a sort column, in milliseconds since the Unix epoch. #[derive(Debug, Clone, Copy, PartialEq)] pub enum ColumnTimestamp { - /// Whole milliseconds, from an integer or an ISO 8601 string column. - Millis(i64), - /// Fractional milliseconds, from a float column. + /// Whole milliseconds, from a whole number or an ISO 8601 string. + Millis(i128), + /// Fractional milliseconds, or a float past the whole-number range. FractionalMillis(f64), } /// Decode the value bytes of a timestamp column. /// -/// Returns `None` for null (no bytes) and for a string that is not ISO 8601. +/// Returns `None` for null, bool, array, object, and a string that is not +/// ISO 8601. pub fn decode_column_timestamp(value: &[u8]) -> Option { - if let Ok(bytes) = <&[u8; 8]>::try_from(value) { - return Some(decode_number(bytes)); - } - let text = std::str::from_utf8(value).ok()?; - NdbDateTime::parse(text).map(|instant| ColumnTimestamp::Millis(instant.unix_millis())) -} - -/// Decode an 8-byte number written by `encode_i64` or `encode_f64`. -fn decode_number(bytes: &[u8; 8]) -> ColumnTimestamp { - let encoded = u64::from_be_bytes(*bytes); - // A clear top bit is a negative value in both encodings. Every window - // starts at or after the epoch, so the integer reading keeps it outside. - if (encoded & SIGN_BIT) == 0 || (encoded ^ SIGN_BIT) < INTEGER_PAYLOAD_LIMIT { - ColumnTimestamp::Millis(SortKeyEncoder::decode_i64(bytes)) - } else { - ColumnTimestamp::FractionalMillis(SortKeyEncoder::decode_f64(bytes)) + match decode(value)? { + SortValue::Number(SortNumber::Whole(ms)) => Some(ColumnTimestamp::Millis(ms)), + SortValue::Number(SortNumber::Float(ms)) => Some(ColumnTimestamp::FractionalMillis(ms)), + SortValue::Text(text) => NdbDateTime::parse(&text) + .map(|instant| ColumnTimestamp::Millis(i128::from(instant.unix_millis()))), + SortValue::Null | SortValue::Bool(_) | SortValue::Json(_) => None, } } #[cfg(test)] mod tests { + use super::super::value::{encode_float, encode_int, encode_json, encode_text}; use super::*; #[test] fn integer_timestamp_decodes_to_its_millis() { - for ms in [0i64, 1, 1_714_608_000_000, (1 << 53) - 1] { - let bytes = SortKeyEncoder::encode_i64(ms); + for ms in [0i64, 1, 1_714_608_000_000, (1 << 53) + 1, i64::MAX] { assert_eq!( - decode_column_timestamp(&bytes), - Some(ColumnTimestamp::Millis(ms)) + decode_column_timestamp(&encode_int(ms)), + Some(ColumnTimestamp::Millis(i128::from(ms))) ); } } #[test] fn negative_integer_timestamp_decodes_negative() { - let bytes = SortKeyEncoder::encode_i64(-86_400_000); assert_eq!( - decode_column_timestamp(&bytes), + decode_column_timestamp(&encode_int(-86_400_000)), Some(ColumnTimestamp::Millis(-86_400_000)) ); } #[test] fn float_timestamp_decodes_to_its_fractional_millis() { - for ms in [1.0f64, 1_714_608_000_000.5, 1_714_607_999_999.25] { - let bytes = SortKeyEncoder::encode_f64(ms); + for ms in [ + 1_714_608_000_000.5, + 1_714_607_999_999.25, + -1_714_608_000_000.5, + ] { assert_eq!( - decode_column_timestamp(&bytes), + decode_column_timestamp(&encode_float(ms)), Some(ColumnTimestamp::FractionalMillis(ms)) ); } } #[test] - fn negative_float_timestamp_decodes_negative() { - let bytes = SortKeyEncoder::encode_f64(-1_714_608_000_000.5); - match decode_column_timestamp(&bytes) { - Some(ColumnTimestamp::Millis(ms)) => assert!(ms < 0), - Some(ColumnTimestamp::FractionalMillis(ms)) => assert!(ms < 0.0), - None => panic!("a negative float must decode"), - } + fn whole_float_timestamp_decodes_to_whole_millis() { + assert_eq!( + decode_column_timestamp(&encode_float(1_714_608_000_000.0)), + Some(ColumnTimestamp::Millis(1_714_608_000_000)) + ); } + /// Ten bytes is the length of an ISO date. The type tag keeps it a string. #[test] fn iso_8601_string_timestamp_decodes_to_millis() { assert_eq!( - decode_column_timestamp(b"2024-05-02T00:00:00Z"), + decode_column_timestamp(&encode_text("2024-05-02T00:00:00Z")), Some(ColumnTimestamp::Millis(1_714_608_000_000)) ); assert_eq!( - decode_column_timestamp(b"2024-05-02"), + decode_column_timestamp(&encode_text("2024-05-02")), Some(ColumnTimestamp::Millis(1_714_608_000_000)) ); } #[test] - fn null_or_non_date_value_has_no_timestamp() { + fn null_bool_or_non_date_value_has_no_timestamp() { + assert_eq!( + decode_column_timestamp(&encode_json(&serde_json::Value::Null)), + None + ); + assert_eq!( + decode_column_timestamp(&encode_json(&serde_json::json!(true))), + None + ); + assert_eq!(decode_column_timestamp(&encode_text("yesterday")), None); assert_eq!(decode_column_timestamp(b""), None); - assert_eq!(decode_column_timestamp(b"yesterday"), None); } } diff --git a/nodedb/src/engine/kv/sorted_index/value.rs b/nodedb/src/engine/kv/sorted_index/value.rs new file mode 100644 index 000000000..314e71645 --- /dev/null +++ b/nodedb/src/engine/kv/sorted_index/value.rs @@ -0,0 +1,294 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Order-preserving encoding of one sort column value. +//! +//! Every value starts with a type tag, so types order as +//! null < bool < number < string < array/object, whatever their payload. +//! +//! - null: the tag alone; +//! - bool: the tag, then 0 or 1; +//! - number: the tag, then [`number`](super::number) bytes, one order for +//! integers and floats; +//! - string, array/object: the tag, then the UTF-8 (array/object as JSON +//! text) with each 0x00 written as 0x00 0xFF, then the terminator 0x00 0x00. +//! +//! No encoded value is a prefix of another. So a composite key can concatenate +//! column values, and complementing one value reverses its order. + +use std::borrow::Cow; + +use super::number::{self, NUMBER_LEN, SortNumber}; + +const TAG_NULL: u8 = 0x01; +const TAG_BOOL: u8 = 0x02; +const TAG_NUMBER: u8 = 0x03; +const TAG_TEXT: u8 = 0x04; +const TAG_JSON: u8 = 0x05; + +/// The escape that follows a literal 0x00 inside text. +const ESCAPE: u8 = 0xFF; + +/// The highest first byte of any encoded value, plain or complemented. +/// +/// A range bound byte above it sorts after every continuation of a key. +pub const MAX_LEADING_BYTE: u8 = !TAG_NULL; + +/// A sort column value read back from its encoding. +#[derive(Debug, Clone, PartialEq)] +pub enum SortValue<'a> { + Null, + Bool(bool), + Number(SortNumber), + Text(Cow<'a, str>), + /// An array or object, as JSON text. + Json(Cow<'a, str>), +} + +/// Encode a document field value. +pub fn encode_json(value: &serde_json::Value) -> Vec { + match value { + serde_json::Value::Null => vec![TAG_NULL], + serde_json::Value::Bool(b) => vec![TAG_BOOL, u8::from(*b)], + serde_json::Value::Number(n) => encode_number(number::encode_json_number(n)), + serde_json::Value::String(s) => encode_text(s), + other @ (serde_json::Value::Array(_) | serde_json::Value::Object(_)) => { + terminated(TAG_JSON, other.to_string().as_bytes()) + } + } +} + +/// Encode an integer. +pub fn encode_int(value: i64) -> Vec { + encode_number(number::encode_i64(value)) +} + +/// Encode a float. +pub fn encode_float(value: f64) -> Vec { + encode_number(number::encode_f64(value)) +} + +/// Encode a string. +pub fn encode_text(value: &str) -> Vec { + terminated(TAG_TEXT, value.as_bytes()) +} + +/// The length of the encoded value at the start of `bytes`. +/// +/// `flip` is XORed into every byte read: 0x00 for a plain value, 0xFF for a +/// complemented one. `None` when the bytes do not start with a whole value. +pub fn encoded_len(bytes: &[u8], flip: u8) -> Option { + let len = match *bytes.first()? ^ flip { + TAG_NULL => 1, + TAG_BOOL => 2, + TAG_NUMBER => 1 + NUMBER_LEN, + TAG_TEXT | TAG_JSON => 1 + terminated_len(bytes.get(1..)?, flip)?, + _ => return None, + }; + (len <= bytes.len()).then_some(len) +} + +/// Decode one whole encoded value. `None` for bytes no value encodes to. +pub fn decode(bytes: &[u8]) -> Option> { + if encoded_len(bytes, 0)? != bytes.len() { + return None; + } + let (&tag, payload) = bytes.split_first()?; + match tag { + TAG_NULL => Some(SortValue::Null), + TAG_BOOL => match payload { + [0] => Some(SortValue::Bool(false)), + [1] => Some(SortValue::Bool(true)), + _ => None, + }, + TAG_NUMBER => number::decode(payload.try_into().ok()?).map(SortValue::Number), + TAG_TEXT => unescape(payload).map(SortValue::Text), + TAG_JSON => unescape(payload).map(SortValue::Json), + _ => None, + } +} + +fn encode_number(bytes: [u8; NUMBER_LEN]) -> Vec { + let mut out = Vec::with_capacity(1 + NUMBER_LEN); + out.push(TAG_NUMBER); + out.extend_from_slice(&bytes); + out +} + +fn terminated(tag: u8, payload: &[u8]) -> Vec { + let mut out = Vec::with_capacity(payload.len() + 3); + out.push(tag); + for &byte in payload { + out.push(byte); + if byte == 0 { + out.push(ESCAPE); + } + } + out.extend_from_slice(&[0, 0]); + out +} + +/// The length of an escaped payload, terminator included. +fn terminated_len(bytes: &[u8], flip: u8) -> Option { + let mut at = 0; + loop { + if *bytes.get(at)? ^ flip != 0 { + at += 1; + continue; + } + match *bytes.get(at + 1)? ^ flip { + 0 => return Some(at + 2), + ESCAPE => at += 2, + _ => return None, + } + } +} + +/// The UTF-8 text of an escaped, terminated payload. +fn unescape(payload: &[u8]) -> Option> { + let body = payload.strip_suffix(&[0u8, 0])?; + if !body.contains(&0) { + return std::str::from_utf8(body).ok().map(Cow::Borrowed); + } + let mut out = Vec::with_capacity(body.len()); + let mut bytes = body.iter(); + while let Some(&byte) = bytes.next() { + out.push(byte); + if byte == 0 && bytes.next() != Some(&ESCAPE) { + return None; + } + } + String::from_utf8(out).ok().map(Cow::Owned) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn complement(bytes: &[u8]) -> Vec { + bytes.iter().map(|b| !b).collect() + } + + /// Values in ascending sort order. + fn ascending() -> Vec { + vec![ + json!(null), + json!(false), + json!(true), + json!(i64::MIN), + json!(-10.5), + json!(-1), + json!(0), + json!(10), + json!(10.5), + json!(11), + json!(11.5), + json!(i64::MAX), + json!(u64::MAX), + json!(1e300), + json!(""), + json!("\u{0}"), + json!("\u{0}a"), + json!("a"), + json!("a\u{0}"), + json!("a\u{0}\u{0}"), + json!("a\u{1}"), + json!("aa"), + json!("b"), + json!([1]), + json!({"k": 1}), + ] + } + + #[test] + fn byte_order_matches_sort_order() { + let encoded: Vec> = ascending().iter().map(encode_json).collect(); + for pair in encoded.windows(2) { + assert!( + pair[0] < pair[1], + "{:?} must sort below {:?}", + pair[0], + pair[1] + ); + } + } + + #[test] + fn complemented_byte_order_is_reversed() { + let encoded: Vec> = ascending().iter().map(encode_json).collect(); + for pair in encoded.windows(2) { + assert!(complement(&pair[0]) > complement(&pair[1])); + } + } + + #[test] + fn no_value_is_a_prefix_of_another() { + let encoded: Vec> = ascending().iter().map(encode_json).collect(); + for a in &encoded { + for b in &encoded { + if a != b { + assert!(!b.starts_with(a), "{a:?} is a prefix of {b:?}"); + } + } + } + } + + #[test] + fn encoded_len_finds_the_end_of_each_value_plain_and_complemented() { + for value in ascending() { + let encoded = encode_json(&value); + let mut followed = encoded.clone(); + followed.extend_from_slice(&encode_text("next")); + assert_eq!(encoded_len(&followed, 0), Some(encoded.len()), "{value}"); + + let mut flipped = complement(&encoded); + flipped.extend_from_slice(&encode_int(7)); + assert_eq!(encoded_len(&flipped, 0xFF), Some(encoded.len()), "{value}"); + } + } + + #[test] + fn leading_bytes_stay_at_or_below_the_bound() { + for value in ascending() { + let encoded = encode_json(&value); + assert!(encoded[0] <= MAX_LEADING_BYTE); + assert!(!encoded[0] <= MAX_LEADING_BYTE); + } + } + + #[test] + fn decode_reads_back_each_type() { + assert_eq!(decode(&encode_json(&json!(null))), Some(SortValue::Null)); + assert_eq!( + decode(&encode_json(&json!(true))), + Some(SortValue::Bool(true)) + ); + assert_eq!( + decode(&encode_json(&json!(11))), + Some(SortValue::Number(SortNumber::Whole(11))) + ); + assert_eq!( + decode(&encode_json(&json!(10.5))), + Some(SortValue::Number(SortNumber::Float(10.5))) + ); + assert_eq!( + decode(&encode_text("a\u{0}b")), + Some(SortValue::Text(Cow::Owned("a\u{0}b".into()))) + ); + assert_eq!( + decode(&encode_json(&json!([1]))), + Some(SortValue::Json(Cow::Borrowed("[1]"))) + ); + } + + #[test] + fn decode_refuses_truncated_or_trailing_bytes() { + let text = encode_text("abc"); + assert_eq!(decode(&text[..text.len() - 1]), None); + let mut trailing = encode_int(1); + trailing.push(0); + assert_eq!(decode(&trailing), None); + assert_eq!(decode(b""), None); + assert_eq!(decode(&[0x7F]), None); + } +} diff --git a/nodedb/src/engine/kv/sorted_index/window.rs b/nodedb/src/engine/kv/sorted_index/window.rs index 121f3fcbc..ec1f6993a 100644 --- a/nodedb/src/engine/kv/sorted_index/window.rs +++ b/nodedb/src/engine/kv/sorted_index/window.rs @@ -128,7 +128,6 @@ impl WindowBounds { pub fn contains(&self, ts: ColumnTimestamp) -> bool { match ts { ColumnTimestamp::Millis(ms) => { - let ms = i128::from(ms); ms >= i128::from(self.start_ms) && self.end_ms.is_none_or(|end| ms < i128::from(end)) } @@ -285,7 +284,7 @@ mod tests { // Before or after the window: a range that admits nothing. for now_ms in [500, 2000, 2500] { let outside = config.bounds(now_ms).expect("custom window has bounds"); - for ts in [0, 999, 1000, 1999, 2000, i64::MAX] { + for ts in [0, 999, 1000, 1999, 2000, i128::from(u64::MAX)] { assert!(!outside.contains(ColumnTimestamp::Millis(ts))); } assert!(!outside.contains(ColumnTimestamp::FractionalMillis(f64::INFINITY))); @@ -315,7 +314,7 @@ mod tests { start_ms: 1000, end_ms: None, }; - assert!(bounds.contains(ColumnTimestamp::Millis(i64::MAX))); + assert!(bounds.contains(ColumnTimestamp::Millis(i128::from(u64::MAX)))); assert!(bounds.contains(ColumnTimestamp::FractionalMillis(1e300))); assert!(!bounds.contains(ColumnTimestamp::Millis(999))); } diff --git a/nodedb/tests/inproc/cases/mod.rs b/nodedb/tests/inproc/cases/mod.rs index 1dbd62460..90fefb950 100644 --- a/nodedb/tests/inproc/cases/mod.rs +++ b/nodedb/tests/inproc/cases/mod.rs @@ -178,6 +178,7 @@ mod siem_export; mod snapshot_round_trip; mod snapshot_round_trip_crdt; mod snapshot_round_trip_stale_state; +mod sorted_index_mixed_numbers; mod sorted_index_rows; mod sql_custom_types; mod sql_ddl_owner_persistence; diff --git a/nodedb/tests/inproc/cases/sorted_index_mixed_numbers.rs b/nodedb/tests/inproc/cases/sorted_index_mixed_numbers.rs new file mode 100644 index 000000000..ad4921993 --- /dev/null +++ b/nodedb/tests/inproc/cases/sorted_index_mixed_numbers.rs @@ -0,0 +1,177 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! A sorted index over numbers orders integers and floats by value. +//! +//! `RANGE` bounds arrive as SQL literals: `11` is an integer, `10.5` a float. +//! A FLOAT column stores `11` as 11.0. The index must treat an integer bound +//! and a float row of equal value as one number, and must rank 11 above 10.5. +//! +//! The `WINDOW` clause tests pin the DDL rules the window relies on. + +use nodedb_test_support::pgwire_harness::TestServer; + +const BOARD: &str = "sidx_mix_board"; +const DESC_INDEX: &str = "sidx_mix_desc"; +const ASC_INDEX: &str = "sidx_mix_asc"; + +async fn exec(server: &TestServer, sql: &str) { + server + .exec(sql) + .await + .unwrap_or_else(|e| panic!("{sql}: {e}")); +} + +/// A FLOAT leaderboard holding 10, 10.5 and 11, with a DESC and an ASC index. +async fn mixed_board() -> TestServer { + let server = TestServer::start().await; + exec( + &server, + &format!( + "CREATE COLLECTION {BOARD} (id STRING PRIMARY KEY, score FLOAT) WITH (engine='kv')" + ), + ) + .await; + for (id, score) in [("ten_half", "10.5"), ("eleven", "11"), ("ten", "10")] { + exec( + &server, + &format!("INSERT INTO {BOARD} (id, score) VALUES ('{id}', {score})"), + ) + .await; + } + exec( + &server, + &format!("CREATE SORTED INDEX {DESC_INDEX} ON {BOARD} (score DESC) KEY id"), + ) + .await; + exec( + &server, + &format!("CREATE SORTED INDEX {ASC_INDEX} ON {BOARD} (score ASC) KEY id"), + ) + .await; + server +} + +/// The keys a table function delivers, in row order. Rows are `(rank, key)`. +async fn keys(server: &TestServer, sql: &str) -> Vec { + server + .query_rows(sql) + .await + .unwrap_or_else(|e| panic!("{sql}: {e}")) + .into_iter() + .map(|row| row.get(1).cloned().unwrap_or_default()) + .collect() +} + +async fn sorted_keys(server: &TestServer, sql: &str) -> Vec { + let mut out = keys(server, sql).await; + out.sort(); + out +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn eleven_ranks_above_ten_and_a_half_in_both_directions() { + let server = mixed_board().await; + + assert_eq!( + keys(&server, &format!("SELECT * FROM TOPK({DESC_INDEX}, 10)")).await, + vec!["eleven", "ten_half", "ten"] + ); + assert_eq!( + keys(&server, &format!("SELECT * FROM TOPK({ASC_INDEX}, 10)")).await, + vec!["ten", "ten_half", "eleven"] + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn range_with_integer_and_float_bounds_selects_by_value() { + let server = mixed_board().await; + + for index in [DESC_INDEX, ASC_INDEX] { + assert_eq!( + sorted_keys(&server, &format!("SELECT * FROM RANGE({index}, 10.5, 11)")).await, + vec!["eleven", "ten_half"], + "{index}: [10.5, 11] holds the float row and the row stored as 11.0" + ); + assert_eq!( + sorted_keys(&server, &format!("SELECT * FROM RANGE({index}, 11, 11)")).await, + vec!["eleven"], + "{index}: an integer bound matches a float row of equal value" + ); + assert_eq!( + sorted_keys( + &server, + &format!("SELECT * FROM RANGE({index}, 10.25, 10.75)") + ) + .await, + vec!["ten_half"], + "{index}" + ); + } +} + +/// A custom window takes ISO 8601 bounds and reads a millisecond column. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn custom_window_with_iso_bounds_hides_rows_outside_it() { + let server = TestServer::start().await; + exec( + &server, + "CREATE COLLECTION sidx_mix_win (id STRING PRIMARY KEY, score FLOAT, ts BIGINT) \ + WITH (engine='kv')", + ) + .await; + exec( + &server, + "INSERT INTO sidx_mix_win (id, score, ts) VALUES ('old', 99, 100)", + ) + .await; + exec( + &server, + "INSERT INTO sidx_mix_win (id, score, ts) VALUES ('new', 1.5, 1714608000000)", + ) + .await; + exec( + &server, + "CREATE SORTED INDEX sidx_mix_win_idx ON sidx_mix_win (score DESC, ts ASC) KEY id \ + WINDOW CUSTOM START '2000-01-01' END '2999-01-01T00:00:00Z' ON ts", + ) + .await; + + assert_eq!( + keys(&server, "SELECT * FROM TOPK(sidx_mix_win_idx, 10)").await, + vec!["new"], + "a row from 1970 is outside a window that starts in 2000" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn invalid_window_clauses_are_refused_with_a_syntax_error() { + let server = TestServer::start().await; + exec( + &server, + "CREATE COLLECTION sidx_mix_bad (id STRING PRIMARY KEY, score FLOAT, ts BIGINT) \ + WITH (engine='kv')", + ) + .await; + + for (sql, names) in [ + ( + "CREATE SORTED INDEX sidx_bad1 ON sidx_mix_bad (score DESC) KEY id \ + WINDOW DAILY ON ts", + "ts", + ), + ( + "CREATE SORTED INDEX sidx_bad2 ON sidx_mix_bad (score DESC, ts ASC) KEY id \ + WINDOW CUSTOM START 'next week' END 5 ON ts", + "next week", + ), + ( + "CREATE SORTED INDEX sidx_bad3 ON sidx_mix_bad (score DESC, ts ASC) KEY id \ + WINDOW CUSTOM START '2026-03-31' END '2026-01-01' ON ts", + "END", + ), + ] { + let err = server.exec(sql).await.expect_err(sql); + assert!(err.contains("42601"), "{sql}: {err}"); + assert!(err.contains(names), "{sql}: {err}"); + } +} diff --git a/nodedb/tests/native/cases/mod.rs b/nodedb/tests/native/cases/mod.rs index 8af4be444..186a483b1 100644 --- a/nodedb/tests/native/cases/mod.rs +++ b/nodedb/tests/native/cases/mod.rs @@ -20,6 +20,7 @@ mod native_protocol; mod native_result_projection; mod native_session_parameters; mod native_show_dispatch; +mod native_sorted_index_range; mod native_sql_authorization; mod native_transactions_savepoint; mod native_transactions_staging; diff --git a/nodedb/tests/native/cases/native_sorted_index_range.rs b/nodedb/tests/native/cases/native_sorted_index_range.rs new file mode 100644 index 000000000..2711c8d94 --- /dev/null +++ b/nodedb/tests/native/cases/native_sorted_index_range.rs @@ -0,0 +1,115 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! The native `KvSortedIndexRange` opcode takes typed score bounds. +//! +//! The server encodes them the same way it encodes SQL `RANGE` arguments. +//! An integer bound and a float row of equal value are one number, so the +//! opcode and the SQL function select the same rows. + +use nodedb_test_support::native_harness::{NativeTestServer, do_handshake, send_request, send_sql}; +use nodedb_types::protocol::opcodes::ResponseStatus; +use nodedb_types::protocol::text_fields::TextFields; +use nodedb_types::protocol::{HelloFrame, NativeResponse, OpCode}; +use nodedb_types::value::Value; + +const IDS: [&str; 4] = ["alpha", "bravo", "charlie", "delta"]; + +fn require_ok(what: &str, response: NativeResponse) -> NativeResponse { + assert_eq!( + response.status, + ResponseStatus::Ok, + "{what} must succeed: {response:?}" + ); + response +} + +/// The ids from `IDS` that a response names, sorted. +fn ids_in(response: &NativeResponse) -> Vec<&'static str> { + let text = format!("{:?}", response.rows); + let mut ids: Vec<&'static str> = IDS.into_iter().filter(|id| text.contains(id)).collect(); + ids.sort_unstable(); + ids +} + +#[tokio::test] +async fn native_range_with_typed_bounds_matches_sql_range() { + let server = NativeTestServer::start().await; + let (mut stream, _ack) = do_handshake(server.addr, &HelloFrame::current()) + .await + .expect("handshake"); + let mut seq = 0u64; + + let setup = [ + "CREATE COLLECTION nat_rng (id STRING PRIMARY KEY, score FLOAT) WITH (engine='kv')" + .to_string(), + "INSERT INTO nat_rng (id, score) VALUES ('alpha', 10)".to_string(), + "INSERT INTO nat_rng (id, score) VALUES ('bravo', 10.5)".to_string(), + "INSERT INTO nat_rng (id, score) VALUES ('charlie', 11)".to_string(), + "INSERT INTO nat_rng (id, score) VALUES ('delta', 12)".to_string(), + "CREATE SORTED INDEX nat_rng_idx ON nat_rng (score DESC) KEY id".to_string(), + ]; + for statement in setup { + seq += 1; + require_ok(&statement, send_sql(&mut stream, seq, &statement).await); + } + + for (min, max, sql_min, sql_max, expected) in [ + ( + Value::Float(10.5), + Value::Integer(11), + "10.5", + "11", + vec!["bravo", "charlie"], + ), + ( + Value::Integer(11), + Value::Integer(11), + "11", + "11", + vec!["charlie"], + ), + ( + Value::Integer(10), + Value::Float(10.75), + "10", + "10.75", + vec!["alpha", "bravo"], + ), + ( + Value::Null, + Value::Float(10.25), + "NULL", + "10.25", + vec!["alpha"], + ), + ] { + let text = format!("SELECT * FROM RANGE(nat_rng_idx, {sql_min}, {sql_max})"); + seq += 1; + let via_sql = require_ok(&text, send_sql(&mut stream, seq, &text).await); + + seq += 1; + let fields = TextFields { + index_name: Some("nat_rng_idx".into()), + score_min: Some(min.clone()), + score_max: Some(max.clone()), + ..TextFields::default() + }; + let via_op = require_ok( + "KvSortedIndexRange", + send_request(&mut stream, seq, OpCode::KvSortedIndexRange, fields).await, + ); + + assert_eq!( + ids_in(&via_sql), + expected, + "SQL RANGE [{sql_min}, {sql_max}]" + ); + assert_eq!( + ids_in(&via_op), + expected, + "native range [{min:?}, {max:?}] must select what SQL RANGE selects" + ); + } + + server.shutdown().await; +} From 80855c77f5d16dd350774442aad545451164832f Mon Sep 17 00:00:00 2001 From: Farhan Syah Date: Mon, 12 Oct 2026 16:41:21 +0800 Subject: [PATCH 08/11] feat(array-sync): classify rejects and report per-op results Add reject reasons and classes, per-op results on batch responses, advisory refusals, and a snapshot assembly lock with a diagnostic when its buffer is poisoned. --- .../src/sync/wire/array/advisory_refused.rs | 105 +++ nodedb-types/src/sync/wire/array/mod.rs | 6 +- .../src/sync/wire/array/op_results.rs | 138 ++++ nodedb-types/src/sync/wire/array/reject.rs | 202 +++++- nodedb-types/src/sync/wire/frame.rs | 13 +- nodedb-types/src/sync/wire/mod.rs | 8 +- .../control/array_sync/catalog_register.rs | 13 +- nodedb/src/control/array_sync/inbound.rs | 105 ++- .../control/array_sync/inbound_advisory.rs | 80 ++- .../src/control/array_sync/inbound_batch.rs | 80 +++ .../src/control/array_sync/inbound_propose.rs | 68 +- nodedb/src/control/array_sync/mod.rs | 1 + nodedb/src/control/array_sync/reject.rs | 192 +++++- .../control/array_sync/snapshot_assembly.rs | 244 +++++-- nodedb/src/control/server/sync/mod.rs | 2 +- nodedb/src/control/server/sync/refusal.rs | 8 +- .../server/sync/session_handler/array.rs | 602 ++++++++++++++++-- nodedb/src/diag/context/array_sync.rs | 64 ++ nodedb/src/diag/context/mod.rs | 2 + nodedb/src/diag/mod.rs | 18 +- nodedb/src/diag/recording/array_sync.rs | 30 + nodedb/src/diag/recording/mod.rs | 2 + nodedb/tests/inproc/cases/sync_compat.rs | 15 +- 23 files changed, 1686 insertions(+), 312 deletions(-) create mode 100644 nodedb-types/src/sync/wire/array/advisory_refused.rs create mode 100644 nodedb-types/src/sync/wire/array/op_results.rs create mode 100644 nodedb/src/control/array_sync/inbound_batch.rs create mode 100644 nodedb/src/diag/context/array_sync.rs create mode 100644 nodedb/src/diag/recording/array_sync.rs diff --git a/nodedb-types/src/sync/wire/array/advisory_refused.rs b/nodedb-types/src/sync/wire/array/advisory_refused.rs new file mode 100644 index 000000000..20c8caf3b --- /dev/null +++ b/nodedb-types/src/sync/wire/array/advisory_refused.rs @@ -0,0 +1,105 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Refusal of an advisory array request: an ack or a catch-up request. + +use serde::{Deserialize, Serialize}; + +use super::reject::{ArrayRejectClass, ArrayRejectMsg, ArrayRejectReason}; + +/// The advisory request Origin refused. +#[derive( + Debug, + Clone, + Copy, + PartialEq, + Eq, + Serialize, + Deserialize, + zerompk::ToMessagePack, + zerompk::FromMessagePack, +)] +#[repr(u8)] +pub enum ArrayAdvisoryKind { + /// An `ArrayAck` (0x95). + Ack = 1, + /// An `ArrayCatchupRequest` (0x97). + CatchupRequest = 2, +} + +/// Advisory refusal (server → client, 0x99). +/// +/// Names no local op, so it never touches the client's push window. +/// `reason.class()` names the client action, as for `ArrayRejectMsg`. +#[derive( + Debug, + Clone, + PartialEq, + Eq, + Serialize, + Deserialize, + zerompk::ToMessagePack, + zerompk::FromMessagePack, +)] +pub struct ArrayAdvisoryRefusedMsg { + /// Name of the array the request named. + pub array: String, + /// The refused request. + pub request: ArrayAdvisoryKind, + /// Structured reason for the refusal. + pub reason: ArrayRejectReason, + /// Human-readable detail for logs. Clients never parse it. + pub detail: String, +} + +impl ArrayAdvisoryRefusedMsg { + /// The refusal of `request` for the reason and detail `reject` carries. + /// The reject's HLC is dropped: an advisory request names no op. + pub fn from_reject(request: ArrayAdvisoryKind, reject: ArrayRejectMsg) -> Self { + Self { + array: reject.array, + request, + reason: reject.reason, + detail: reject.detail, + } + } + + /// The action a client takes for this refusal. + pub fn class(&self) -> ArrayRejectClass { + self.reason.class() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn roundtrip_via_msgpack() { + let msg = ArrayAdvisoryRefusedMsg { + array: "scores".into(), + request: ArrayAdvisoryKind::CatchupRequest, + reason: ArrayRejectReason::PermissionDenied, + detail: "no read permission".into(), + }; + let encoded = zerompk::to_msgpack_vec(&msg).expect("encode"); + let decoded: ArrayAdvisoryRefusedMsg = zerompk::from_msgpack(&encoded).expect("decode"); + assert_eq!(msg, decoded); + } + + #[test] + fn a_reject_becomes_a_refusal_with_its_reason_and_class() { + let refusal = ArrayAdvisoryRefusedMsg::from_reject( + ArrayAdvisoryKind::Ack, + ArrayRejectMsg { + array: "a".into(), + op_hlc_bytes: [4; 18], + reason: ArrayRejectReason::PermissionDenied, + detail: "denied".into(), + }, + ); + assert_eq!(refusal.array, "a"); + assert_eq!(refusal.request, ArrayAdvisoryKind::Ack); + assert_eq!(refusal.reason, ArrayRejectReason::PermissionDenied); + assert_eq!(refusal.class(), ArrayRejectClass::Permanent); + } +} diff --git a/nodedb-types/src/sync/wire/array/mod.rs b/nodedb-types/src/sync/wire/array/mod.rs index 933b97325..12d191b51 100644 --- a/nodedb-types/src/sync/wire/array/mod.rs +++ b/nodedb-types/src/sync/wire/array/mod.rs @@ -9,19 +9,23 @@ //! payloads via `nodedb-array::sync::op_codec::decode_op` etc. pub mod ack; +pub mod advisory_refused; pub mod catchup_request; pub mod delta; pub mod delta_batch; +pub mod op_results; pub mod reject; pub mod schema; pub mod snapshot; pub mod snapshot_chunk; pub use ack::ArrayAckMsg; +pub use advisory_refused::{ArrayAdvisoryKind, ArrayAdvisoryRefusedMsg}; pub use catchup_request::ArrayCatchupRequestMsg; pub use delta::ArrayDeltaMsg; pub use delta_batch::ArrayDeltaBatchMsg; -pub use reject::{ArrayRejectMsg, ArrayRejectReason}; +pub use op_results::{ArrayOpOutcome, ArrayOpResult, ArrayOpResultsMsg}; +pub use reject::{ArrayRejectClass, ArrayRejectMsg, ArrayRejectReason}; pub use schema::ArraySchemaSyncMsg; pub use snapshot::ArraySnapshotMsg; pub use snapshot_chunk::ArraySnapshotChunkMsg; diff --git a/nodedb-types/src/sync/wire/array/op_results.rs b/nodedb-types/src/sync/wire/array/op_results.rs new file mode 100644 index 000000000..396dae97c --- /dev/null +++ b/nodedb-types/src/sync/wire/array/op_results.rs @@ -0,0 +1,138 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Per-op outcomes for an array frame that carries several ops. + +use serde::{Deserialize, Serialize}; + +use super::reject::{ArrayRejectMsg, ArrayRejectReason}; + +/// What Origin did with one op of a multi-op array frame. +#[derive( + Debug, + Clone, + PartialEq, + Eq, + Serialize, + Deserialize, + zerompk::ToMessagePack, + zerompk::FromMessagePack, +)] +pub enum ArrayOpOutcome { + /// The op committed and applied. + Applied, + /// Origin already held the op's HLC. Nothing changed. + Duplicate, + /// The op did not apply. `reason.class()` names the client action. + Rejected { + /// Structured reason for the rejection. + reason: ArrayRejectReason, + /// Human-readable detail for logs. Clients never parse it. + detail: String, + }, +} + +/// The outcome of one op, keyed by its HLC. +#[derive( + Debug, + Clone, + PartialEq, + Eq, + Serialize, + Deserialize, + zerompk::ToMessagePack, + zerompk::FromMessagePack, +)] +pub struct ArrayOpResult { + /// HLC of the op (18-byte layout matching + /// `nodedb_array::sync::Hlc::to_bytes()`). All zero when the payload did + /// not decode. The result's position in the frame still names the op. + pub op_hlc_bytes: [u8; 18], + /// What Origin did with the op. + pub outcome: ArrayOpOutcome, +} + +impl ArrayOpResult { + /// The result for an op Origin refused with `reject`. + pub fn rejected(reject: ArrayRejectMsg) -> Self { + Self { + op_hlc_bytes: reject.op_hlc_bytes, + outcome: ArrayOpOutcome::Rejected { + reason: reject.reason, + detail: reject.detail, + }, + } + } +} + +/// Per-op results (server → client, 0x98). +/// +/// Origin sends one in answer to every `ArrayDeltaBatch`, and one when an +/// inbound snapshot completes. `results` holds one entry per op, in the +/// order the ops arrived. Each op applies on its own: a reject never stops +/// the ops after it. +#[derive( + Debug, + Clone, + PartialEq, + Eq, + Serialize, + Deserialize, + zerompk::ToMessagePack, + zerompk::FromMessagePack, +)] +pub struct ArrayOpResultsMsg { + /// Name of the array the frame targeted. + pub array: String, + /// One result per op, in arrival order. + pub results: Vec, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn roundtrip_via_msgpack() { + let msg = ArrayOpResultsMsg { + array: "scores".into(), + results: vec![ + ArrayOpResult { + op_hlc_bytes: [1; 18], + outcome: ArrayOpOutcome::Applied, + }, + ArrayOpResult { + op_hlc_bytes: [2; 18], + outcome: ArrayOpOutcome::Duplicate, + }, + ArrayOpResult { + op_hlc_bytes: [0; 18], + outcome: ArrayOpOutcome::Rejected { + reason: ArrayRejectReason::ShapeInvalid, + detail: "decode error".into(), + }, + }, + ], + }; + let encoded = zerompk::to_msgpack_vec(&msg).expect("encode"); + let decoded: ArrayOpResultsMsg = zerompk::from_msgpack(&encoded).expect("decode"); + assert_eq!(msg, decoded); + } + + #[test] + fn a_reject_becomes_a_result_with_its_hlc_and_reason() { + let result = ArrayOpResult::rejected(ArrayRejectMsg { + array: "a".into(), + op_hlc_bytes: [9; 18], + reason: ArrayRejectReason::EngineUnavailable, + detail: "queue full".into(), + }); + assert_eq!(result.op_hlc_bytes, [9; 18]); + assert_eq!( + result.outcome, + ArrayOpOutcome::Rejected { + reason: ArrayRejectReason::EngineUnavailable, + detail: "queue full".into(), + } + ); + } +} diff --git a/nodedb-types/src/sync/wire/array/reject.rs b/nodedb-types/src/sync/wire/array/reject.rs index 11dad6b5f..0d1a49177 100644 --- a/nodedb-types/src/sync/wire/array/reject.rs +++ b/nodedb-types/src/sync/wire/array/reject.rs @@ -1,10 +1,46 @@ // SPDX-License-Identifier: Apache-2.0 -//! Array CRDT rejection wire message and reason codes. +//! Array CRDT rejection wire message, reason codes, and retry classes. use serde::{Deserialize, Serialize}; -/// Reason codes for array op rejection (Origin → Lite, 0x96). +/// What the sender of a refused array op does next. +/// +/// Every [`ArrayRejectReason`] maps to exactly one class through +/// [`ArrayRejectReason::class`], so a client decides without reading `detail`. +#[derive( + Debug, + Clone, + Copy, + PartialEq, + Eq, + Serialize, + Deserialize, + zerompk::ToMessagePack, + zerompk::FromMessagePack, +)] +#[repr(u8)] +pub enum ArrayRejectClass { + /// Nothing applied and nothing is wrong with the op. Keep it queued and + /// re-send it unchanged after a short backoff, in this session. + Transient = 1, + /// Origin lacks the op's schema. Send `ArraySchema` for the array, then + /// re-send the op. + SchemaRequired = 2, + /// The session holds no current authenticated identity. Complete the + /// handshake or token refresh, then re-send. + Reauthenticate = 3, + /// The reject names this replica's catch-up cursor, not a local op. Send + /// `ArrayCatchupRequest`. The pending queue is unchanged. + CatchUp = 4, + /// The op never applies. Remove it from the queue and dead-letter it. + Permanent = 5, +} + +/// Reason codes for array op rejection (Origin → Lite). +/// +/// Each reason has one fixed [`ArrayRejectClass`]. Origin derives the reason +/// from the typed error, never from error text. #[derive( Debug, Clone, @@ -17,25 +53,64 @@ use serde::{Deserialize, Serialize}; zerompk::FromMessagePack, )] #[repr(u8)] -#[non_exhaustive] pub enum ArrayRejectReason { - /// The op was encoded with a schema version newer than Origin knows about. + /// The op's schema HLC is newer than Origin's copy of the schema. SchemaTooNew = 1, - /// The target array does not exist on Origin. + /// The target array has no schema on Origin. ArrayUnknown = 2, - /// The op payload shape is invalid or failed deserialization. + /// The op payload failed decoding or shape validation. ShapeInvalid = 3, - /// The engine itself rejected the op (constraint violation, etc.). - EngineRejected = 4, - /// The op's HLC falls below Origin's GC retention floor. + /// The cluster gave the op no verdict: capacity refusal, leader change, + /// lost quorum, deadline, rate gate, or an unknown commit outcome. + /// Array apply dedups by HLC, so a re-send is safe. + EngineUnavailable = 4, + /// This replica's catch-up cursor is below Origin's op-log GC floor. RetentionFloor = 5, - /// The replica has exceeded its storage or rate quota. - QuotaExceeded = 6, - /// Wire protocol version mismatch between Lite and Origin. - WireVersionMismatch = 7, + /// The engine refused the op on its merits: constraint, type, fence, + /// encoding, or another error the same bytes reproduce. + EngineRefused = 6, + /// The session identity lacks the permission the op needs. + PermissionDenied = 7, + /// The frame arrived before the session authenticated. + Unauthenticated = 8, + /// The session's inbound snapshot buffers are full. Other assemblies + /// release them when they complete. + SnapshotCapacity = 9, + /// Origin discarded the session's partial snapshot buffers after a + /// fault. Re-send the whole snapshot, header first. + SnapshotBufferLost = 10, } impl ArrayRejectReason { + /// Every reason, in discriminant order. + pub const ALL: [Self; 10] = [ + Self::SchemaTooNew, + Self::ArrayUnknown, + Self::ShapeInvalid, + Self::EngineUnavailable, + Self::RetentionFloor, + Self::EngineRefused, + Self::PermissionDenied, + Self::Unauthenticated, + Self::SnapshotCapacity, + Self::SnapshotBufferLost, + ]; + + /// The action a client takes for this reason. + pub fn class(self) -> ArrayRejectClass { + match self { + Self::EngineUnavailable | Self::SnapshotCapacity | Self::SnapshotBufferLost => { + ArrayRejectClass::Transient + } + Self::SchemaTooNew | Self::ArrayUnknown => ArrayRejectClass::SchemaRequired, + Self::Unauthenticated => ArrayRejectClass::Reauthenticate, + Self::RetentionFloor => ArrayRejectClass::CatchUp, + Self::ShapeInvalid | Self::EngineRefused | Self::PermissionDenied => { + ArrayRejectClass::Permanent + } + } + } + /// Convert a raw `u8` discriminant to `ArrayRejectReason`. /// /// Returns `None` for unrecognized values. @@ -44,10 +119,13 @@ impl ArrayRejectReason { 1 => Some(Self::SchemaTooNew), 2 => Some(Self::ArrayUnknown), 3 => Some(Self::ShapeInvalid), - 4 => Some(Self::EngineRejected), + 4 => Some(Self::EngineUnavailable), 5 => Some(Self::RetentionFloor), - 6 => Some(Self::QuotaExceeded), - 7 => Some(Self::WireVersionMismatch), + 6 => Some(Self::EngineRefused), + 7 => Some(Self::PermissionDenied), + 8 => Some(Self::Unauthenticated), + 9 => Some(Self::SnapshotCapacity), + 10 => Some(Self::SnapshotBufferLost), _ => None, } } @@ -55,9 +133,9 @@ impl ArrayRejectReason { /// Array reject message (server → client, 0x96). /// -/// Sent by Origin as a compensation hint when it rejects a Lite op. -/// The originating Lite uses `op_hlc_bytes` to locate the offending op -/// in its local log and roll it back. +/// Origin sends one per refused frame. The originating Lite uses +/// `op_hlc_bytes` to find the refused op and `reason.class()` to decide +/// whether to re-send, re-announce the schema, or drop it. #[derive( Debug, Clone, @@ -68,18 +146,26 @@ impl ArrayRejectReason { zerompk::FromMessagePack, )] pub struct ArrayRejectMsg { - /// Name of the array the rejected op targeted. + /// Name of the array the rejected op targeted. Empty when the frame + /// body did not decode or the session is unauthenticated. pub array: String, /// HLC of the rejected op (18-byte layout matching - /// `nodedb_array::sync::Hlc::to_bytes()`). Used by the Lite peer - /// to identify and roll back the offending local op. + /// `nodedb_array::sync::Hlc::to_bytes()`). All zero when the op did not + /// decode, so the reject names no single op. pub op_hlc_bytes: [u8; 18], /// Structured reason for the rejection. pub reason: ArrayRejectReason, - /// Human-readable detail string for debugging or user-facing messages. + /// Human-readable detail for logs. Clients never parse it. pub detail: String, } +impl ArrayRejectMsg { + /// The action a client takes for this reject. + pub fn class(&self) -> ArrayRejectClass { + self.reason.class() + } +} + #[cfg(test)] mod tests { use super::*; @@ -89,7 +175,7 @@ mod tests { let msg = ArrayRejectMsg { array: "leaderboard".to_string(), op_hlc_bytes: [0x11u8; 18], - reason: ArrayRejectReason::EngineRejected, + reason: ArrayRejectReason::EngineRefused, detail: "unique constraint violated".to_string(), }; let encoded = zerompk::to_msgpack_vec(&msg).expect("encode"); @@ -99,16 +185,7 @@ mod tests { #[test] fn reject_reason_round_trips() { - let variants = [ - ArrayRejectReason::SchemaTooNew, - ArrayRejectReason::ArrayUnknown, - ArrayRejectReason::ShapeInvalid, - ArrayRejectReason::EngineRejected, - ArrayRejectReason::RetentionFloor, - ArrayRejectReason::QuotaExceeded, - ArrayRejectReason::WireVersionMismatch, - ]; - for variant in variants { + for variant in ArrayRejectReason::ALL { let disc = variant as u8; let recovered = ArrayRejectReason::from_u8(disc); assert_eq!(recovered, Some(variant), "failed for discriminant {disc}"); @@ -118,7 +195,62 @@ mod tests { #[test] fn reject_reason_unknown_returns_none() { assert_eq!(ArrayRejectReason::from_u8(0), None); - assert_eq!(ArrayRejectReason::from_u8(8), None); + assert_eq!(ArrayRejectReason::from_u8(11), None); assert_eq!(ArrayRejectReason::from_u8(255), None); } + + #[test] + fn every_reason_has_its_documented_class() { + let expected = [ + ( + ArrayRejectReason::SchemaTooNew, + ArrayRejectClass::SchemaRequired, + ), + ( + ArrayRejectReason::ArrayUnknown, + ArrayRejectClass::SchemaRequired, + ), + (ArrayRejectReason::ShapeInvalid, ArrayRejectClass::Permanent), + ( + ArrayRejectReason::EngineUnavailable, + ArrayRejectClass::Transient, + ), + (ArrayRejectReason::RetentionFloor, ArrayRejectClass::CatchUp), + ( + ArrayRejectReason::EngineRefused, + ArrayRejectClass::Permanent, + ), + ( + ArrayRejectReason::PermissionDenied, + ArrayRejectClass::Permanent, + ), + ( + ArrayRejectReason::Unauthenticated, + ArrayRejectClass::Reauthenticate, + ), + ( + ArrayRejectReason::SnapshotCapacity, + ArrayRejectClass::Transient, + ), + ( + ArrayRejectReason::SnapshotBufferLost, + ArrayRejectClass::Transient, + ), + ]; + assert_eq!(expected.len(), ArrayRejectReason::ALL.len()); + for (reason, class) in expected { + assert_eq!(reason.class(), class, "{reason:?}"); + } + } + + #[test] + fn a_reject_reports_the_class_of_its_reason() { + let msg = ArrayRejectMsg { + array: "a".into(), + op_hlc_bytes: [0; 18], + reason: ArrayRejectReason::EngineUnavailable, + detail: String::new(), + }; + assert_eq!(msg.class(), ArrayRejectClass::Transient); + } } diff --git a/nodedb-types/src/sync/wire/frame.rs b/nodedb-types/src/sync/wire/frame.rs index dbb9617f8..90ccf0cf4 100644 --- a/nodedb-types/src/sync/wire/frame.rs +++ b/nodedb-types/src/sync/wire/frame.rs @@ -95,6 +95,12 @@ pub enum SyncMessageType { ArrayReject = 0x96, /// Array catchup request (client → server, 0x97). ArrayCatchupRequest = 0x97, + /// Per-op array results (server → client, 0x98). Answers an + /// `ArrayDeltaBatch` or a completed inbound snapshot. + ArrayOpResults = 0x98, + /// Refusal of an `ArrayAck` or `ArrayCatchupRequest` (server → client, + /// 0x99). Names no local op. + ArrayAdvisoryRefused = 0x99, /// Columnar batch insert (client → server, 0xA0). ColumnarInsert = 0xA0, /// Columnar insert acknowledgment (server → client, 0xA1). @@ -165,6 +171,8 @@ impl SyncMessageType { 0x95 => Some(Self::ArrayAck), 0x96 => Some(Self::ArrayReject), 0x97 => Some(Self::ArrayCatchupRequest), + 0x98 => Some(Self::ArrayOpResults), + 0x99 => Some(Self::ArrayAdvisoryRefused), 0xA0 => Some(Self::ColumnarInsert), 0xA1 => Some(Self::ColumnarInsertAck), 0xA2 => Some(Self::VectorInsert), @@ -393,11 +401,12 @@ mod tests { fn message_type_roundtrip() { for v in [ 0x01, 0x02, 0x10, 0x11, 0x12, 0x14, 0x20, 0x21, 0x22, 0x23, 0x30, 0x40, 0x41, 0x50, - 0x52, 0x60, 0x61, 0x70, 0x80, 0x81, 0x82, 0xFF, + 0x52, 0x60, 0x61, 0x70, 0x80, 0x81, 0x82, 0x90, 0x91, 0x92, 0x93, 0x94, 0x95, 0x96, + 0x97, 0x98, 0x99, 0xFF, ] { let mt = SyncMessageType::from_u8(v).unwrap(); assert_eq!(mt as u8, v); } - assert!(SyncMessageType::from_u8(0x99).is_none()); + assert!(SyncMessageType::from_u8(0x9A).is_none()); } } diff --git a/nodedb-types/src/sync/wire/mod.rs b/nodedb-types/src/sync/wire/mod.rs index a6f56e679..3e222e155 100644 --- a/nodedb-types/src/sync/wire/mod.rs +++ b/nodedb-types/src/sync/wire/mod.rs @@ -37,6 +37,8 @@ //! - `0x95` ArrayAck (client → server) //! - `0x96` ArrayReject (server → client) //! - `0x97` ArrayCatchupRequest (client → server) +//! - `0x98` ArrayOpResults (server → client) +//! - `0x99` ArrayAdvisoryRefused (server → client) //! - `0xA0` ColumnarInsert (client → server) //! - `0xA1` ColumnarInsertAck (server → client) //! - `0xA2` VectorInsert (client → server) @@ -77,8 +79,10 @@ pub mod vector; pub use ack_result::{SyncAckResult, SyncOutcome}; pub use ack_status::AckStatus; pub use array::{ - ArrayAckMsg, ArrayCatchupRequestMsg, ArrayDeltaBatchMsg, ArrayDeltaMsg, ArrayRejectMsg, - ArrayRejectReason, ArraySchemaSyncMsg, ArraySnapshotChunkMsg, ArraySnapshotMsg, + ArrayAckMsg, ArrayAdvisoryKind, ArrayAdvisoryRefusedMsg, ArrayCatchupRequestMsg, + ArrayDeltaBatchMsg, ArrayDeltaMsg, ArrayOpOutcome, ArrayOpResult, ArrayOpResultsMsg, + ArrayRejectClass, ArrayRejectMsg, ArrayRejectReason, ArraySchemaSyncMsg, ArraySnapshotChunkMsg, + ArraySnapshotMsg, }; pub use collection_schema::{CollectionDescriptor, CollectionSchemaSyncMsg}; pub use columnar::{ColumnarInsertAckMsg, ColumnarInsertMsg}; diff --git a/nodedb/src/control/array_sync/catalog_register.rs b/nodedb/src/control/array_sync/catalog_register.rs index 0c912dfd8..7a37fd717 100644 --- a/nodedb/src/control/array_sync/catalog_register.rs +++ b/nodedb/src/control/array_sync/catalog_register.rs @@ -14,11 +14,11 @@ use nodedb_array::sync::SchemaDoc; use nodedb_array::sync::hlc::Hlc as SchemaHlc; use nodedb_array::sync::replica_id::ReplicaId; use nodedb_array::types::ArrayId; -use nodedb_types::sync::wire::array::{ArrayRejectMsg, ArrayRejectReason, ArraySchemaSyncMsg}; +use nodedb_types::sync::wire::array::{ArrayRejectMsg, ArraySchemaSyncMsg}; use tracing::warn; use super::inbound::OriginArrayInbound; -use super::reject::build_reject; +use super::reject::engine_reject; use crate::control::array_catalog::entry::ArrayCatalogEntry; use crate::control::catalog_entry::CatalogEntry; use crate::control::state::SharedState; @@ -32,7 +32,7 @@ impl OriginArrayInbound { &self, msg: &ArraySchemaSyncMsg, remote_hlc: SchemaHlc, - ) -> Result<(), Option> { + ) -> Result<(), ArrayRejectMsg> { register_array_catalog_entry( self.shared(), self.tenant_id(), @@ -44,12 +44,7 @@ impl OriginArrayInbound { .await .map_err(|e| { warn!(array = %msg.array, error = %e, "array_inbound: catalog registration failed"); - Some(build_reject( - &msg.array, - remote_hlc, - ArrayRejectReason::EngineRejected, - format!("catalog registration error: {e}"), - )) + engine_reject(&msg.array, remote_hlc, "catalog registration error", &e) }) } } diff --git a/nodedb/src/control/array_sync/inbound.rs b/nodedb/src/control/array_sync/inbound.rs index 38bdbc208..893f8c32a 100644 --- a/nodedb/src/control/array_sync/inbound.rs +++ b/nodedb/src/control/array_sync/inbound.rs @@ -39,13 +39,12 @@ use std::collections::HashMap; use std::sync::{Arc, Mutex}; -use nodedb_array::sync::apply::ApplyRejection; use nodedb_array::sync::hlc::Hlc; use nodedb_array::sync::op::ArrayOp; use nodedb_array::sync::op_codec; use nodedb_types::sync::wire::SyncProvenance; use nodedb_types::sync::wire::array::{ - ArrayDeltaBatchMsg, ArrayDeltaMsg, ArrayRejectMsg, ArrayRejectReason, ArraySchemaSyncMsg, + ArrayDeltaMsg, ArrayOpResultsMsg, ArrayRejectMsg, ArrayRejectReason, ArraySchemaSyncMsg, }; use tracing::warn; @@ -70,12 +69,11 @@ pub enum InboundOutcome { Applied, /// The op was already present; no state was changed (idempotent replay). Idempotent, - /// The op was rejected; the caller sends `ArrayRejectMsg` back. - Rejected(ApplyRejection), /// A snapshot chunk was buffered; more chunks are expected. SnapshotPartial { received: u32, total: u32 }, - /// A snapshot was fully assembled and all contained ops applied. - SnapshotApplied { ops_applied: u64 }, + /// A snapshot was fully assembled. Every contained op was attempted, and + /// the message holds one result per op, in snapshot order. + SnapshotApplied(ArrayOpResultsMsg), /// A schema CRDT snapshot was imported into the local registry. SchemaImported, /// An ack was recorded into the ack-vector (GC frontier tracking). @@ -86,6 +84,24 @@ pub enum InboundOutcome { // ─── Dispatcher ────────────────────────────────────────────────────────────── +/// What [`OriginArrayInbound::apply_op`] did with an op it did not refuse. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum OpApplied { + /// The op committed and applied. + Applied, + /// The op's HLC was already present. Nothing changed. + Idempotent, +} + +impl From for InboundOutcome { + fn from(applied: OpApplied) -> Self { + match applied { + OpApplied::Applied => Self::Applied, + OpApplied::Idempotent => Self::Idempotent, + } + } +} + /// Dispatcher for inbound array CRDT wire messages from Lite peers. /// /// Constructed once per sync session (or shared across sessions via `Arc`) @@ -144,10 +160,8 @@ impl OriginArrayInbound { array: &str, hlc: Hlc, permission: crate::control::security::identity::Permission, - ) -> Result< - crate::control::server::shared::authorization::AuthorizedCollection, - Option, - > { + ) -> Result + { let emitter = crate::control::security::audit::ArcAuditEmitter(Arc::clone(&self.shared.audit)); crate::control::server::shared::authorization::authorize_collection_capability( @@ -160,12 +174,12 @@ impl OriginArrayInbound { &emitter, ) .map_err(|error| { - Some(build_reject( + build_reject( array, hlc, - ArrayRejectReason::EngineRejected, + ArrayRejectReason::PermissionDenied, error.resource().to_string(), - )) + ) }) } @@ -240,54 +254,31 @@ impl OriginArrayInbound { pub async fn handle_delta( &self, msg: &ArrayDeltaMsg, - ) -> Result> { + ) -> Result { let op = match op_codec::decode_op(&msg.op_payload) { Ok(op) => op, Err(e) => { warn!(array = %msg.array, error = %e, "array_inbound: delta decode failed"); - return Err(Some(build_reject( + return Err(build_reject( &msg.array, Hlc::ZERO, ArrayRejectReason::ShapeInvalid, format!("decode error: {e}"), - ))); + )); } }; let provenance = self.session_provenance(&msg.array, msg.seq); - self.apply_op(op, &msg.op_payload, provenance).await + self.apply_op(op, &msg.op_payload, provenance) + .await + .map(InboundOutcome::from) } - /// Handle a batch of delta messages from a Lite peer. - /// - /// Returns one outcome per op. If decoding fails for an op, that op - /// yields a reject; subsequent ops are still attempted. - pub async fn handle_delta_batch( - &self, - msg: &ArrayDeltaBatchMsg, - ) -> Vec>> { - let mut outcomes = Vec::with_capacity(msg.op_payloads.len()); - for payload in &msg.op_payloads { - let outcome = match op_codec::decode_op(payload) { - Ok(op) => { - // Batch carries no per-op seq; array dedups by HLC, so seq - // is informational here. The epoch fence still applies. - let provenance = self.session_provenance(&msg.array, 0); - self.apply_op(op, payload, provenance).await - } - Err(e) => { - warn!(array = %msg.array, error = %e, "array_inbound: batch decode failed"); - Err(Some(build_reject( - &msg.array, - Hlc::ZERO, - ArrayRejectReason::ShapeInvalid, - format!("batch decode error: {e}"), - ))) - } - }; - outcomes.push(outcome); - } - outcomes + /// Server-authoritative provenance for one op of a delta batch. + pub(super) fn batch_provenance(&self, array: &str) -> Option { + // A batch carries no per-op seq. Array dedups by HLC, so seq is + // informational here. The epoch fence still applies. + self.session_provenance(array, 0) } // ─── Schema ────────────────────────────────────────────────────────────── @@ -300,7 +291,7 @@ impl OriginArrayInbound { pub async fn handle_schema( &self, msg: &ArraySchemaSyncMsg, - ) -> Result> { + ) -> Result { let hlc_arr: [u8; 18] = msg.schema_hlc_bytes; let remote_hlc = Hlc::from_bytes(&hlc_arr); let authorization = self.authorize_array( @@ -344,15 +335,15 @@ impl OriginArrayInbound { op: ArrayOp, raw_op_bytes: &[u8], provenance: Option, - ) -> Result> { + ) -> Result { // 1. Shape validation. if let Err(e) = op.validate_shape() { - return Err(Some(build_reject( + return Err(build_reject( &op.header.array, op.header.hlc, ArrayRejectReason::ShapeInvalid, format!("shape validation: {e}"), - ))); + )); } let authorization = self.authorize_array( @@ -368,15 +359,15 @@ impl OriginArrayInbound { &op.header.array, ) { None => { - return Err(Some(build_reject( + return Err(build_reject( &op.header.array, op.header.hlc, ArrayRejectReason::ArrayUnknown, format!("array '{}' not known to this replica", op.header.array), - ))); + )); } Some(local_schema) if op.header.schema_hlc > local_schema => { - return Err(Some(build_reject( + return Err(build_reject( &op.header.array, op.header.hlc, ArrayRejectReason::SchemaTooNew, @@ -384,7 +375,7 @@ impl OriginArrayInbound { "op schema_hlc {:?} > local {:?}; request schema sync", op.header.schema_hlc, local_schema ), - ))); + )); } Some(_) => {} } @@ -397,7 +388,7 @@ impl OriginArrayInbound { op.header.hlc, ) { let _authorized_scope = authorization.into_scope(); - return Ok(InboundOutcome::Idempotent); + return Ok(OpApplied::Idempotent); } // 4. Propose through Raft. A put's cell surrogate is @@ -434,7 +425,7 @@ impl OriginArrayInbound { if let Some(observer) = &self.apply_observer { observer.on_op_applied(&op); } - Ok(InboundOutcome::Applied) + Ok(OpApplied::Applied) } Err(e) => Err(e), } diff --git a/nodedb/src/control/array_sync/inbound_advisory.rs b/nodedb/src/control/array_sync/inbound_advisory.rs index a2264c137..6dfa6f870 100644 --- a/nodedb/src/control/array_sync/inbound_advisory.rs +++ b/nodedb/src/control/array_sync/inbound_advisory.rs @@ -1,28 +1,39 @@ // SPDX-License-Identifier: BUSL-1.1 //! Advisory acknowledgement and catch-up handlers for array sync. +//! +//! A refused advisory request draws an `ArrayAdvisoryRefused` reply. It +//! names no local op, so it never touches the client's push window. use std::sync::Arc; use nodedb_array::sync::hlc::Hlc; -use nodedb_types::sync::wire::array::{ArrayAckMsg, ArrayCatchupRequestMsg, ArrayRejectMsg}; +use nodedb_types::sync::wire::array::{ + ArrayAckMsg, ArrayAdvisoryKind, ArrayAdvisoryRefusedMsg, ArrayCatchupRequestMsg, + ArrayRejectReason, +}; use tracing::warn; use super::catchup::OriginCatchupServer; use super::inbound::{InboundOutcome, OriginArrayInbound}; +use super::reject::reason_for_engine_error; impl OriginArrayInbound { /// Record a peer ack for GC frontier tracking. /// /// Forwards the ack into the `ArrayAckRegistry` on `SharedState` so the /// GC task can compute the min-ack frontier for each array. - pub fn handle_ack(&self, msg: &ArrayAckMsg) -> Result> { + pub fn handle_ack(&self, msg: &ArrayAckMsg) -> Result { let ack_hlc = Hlc::from_bytes(&msg.ack_hlc_bytes); - let _authorization = self.authorize_array( - &msg.array, - ack_hlc, - crate::control::security::identity::Permission::Write, - )?; + let _authorization = self + .authorize_array( + &msg.array, + ack_hlc, + crate::control::security::identity::Permission::Write, + ) + .map_err(|reject| { + ArrayAdvisoryRefusedMsg::from_reject(ArrayAdvisoryKind::Ack, reject) + })?; let replica_id = nodedb_array::sync::replica_id::ReplicaId::new(msg.replica_id); self.shared.array_ack_registry.record_in_database( self.database_id(), @@ -42,19 +53,40 @@ impl OriginArrayInbound { /// Handle a catch-up request from a Lite peer. /// - /// Delegates to [`OriginCatchupServer`] which validates the array, selects - /// the op-stream or snapshot delivery path, and enqueues outbound frames. + /// Delegates to [`OriginCatchupServer`] which selects the op-stream or + /// snapshot delivery path and enqueues outbound frames. A denial, an + /// unknown array, and a serve error each draw a typed refusal. pub fn handle_catchup_request( &self, msg: &ArrayCatchupRequestMsg, session_id: &str, - ) -> Result> { + ) -> Result { + let refuse = |reason: ArrayRejectReason, detail: String| ArrayAdvisoryRefusedMsg { + array: msg.array.clone(), + request: ArrayAdvisoryKind::CatchupRequest, + reason, + detail, + }; let from_hlc = Hlc::from_bytes(&msg.from_hlc_bytes); - let _authorization = self.authorize_array( - &msg.array, - from_hlc, - crate::control::security::identity::Permission::Read, - )?; + let _authorization = self + .authorize_array( + &msg.array, + from_hlc, + crate::control::security::identity::Permission::Read, + ) + .map_err(|reject| { + ArrayAdvisoryRefusedMsg::from_reject(ArrayAdvisoryKind::CatchupRequest, reject) + })?; + if self + .schemas + .schema_hlc_in_database(self.database_id(), self.tenant_id().as_u64(), &msg.array) + .is_none() + { + return Err(refuse( + ArrayRejectReason::ArrayUnknown, + format!("array '{}' not known to this replica", msg.array), + )); + } let server = OriginCatchupServer::new( crate::control::array_sync::ArrayServerScope::new( self.database_id(), @@ -75,6 +107,10 @@ impl OriginArrayInbound { error = %error, "array_inbound: catchup server error" ); + return Err(refuse( + reason_for_engine_error(&error), + format!("catch-up serve: {error}"), + )); } Ok(InboundOutcome::CatchupRequested) @@ -154,13 +190,14 @@ mod tests { let rejection = inbound .handle_ack(&msg) - .expect_err("read permission must not authorize ack mutation") - .expect("authorization denial must produce an array rejection"); + .expect_err("read permission must not authorize ack mutation"); assert_eq!( rejection.reason, - nodedb_types::sync::wire::array::ArrayRejectReason::EngineRejected + nodedb_types::sync::wire::array::ArrayRejectReason::PermissionDenied ); + assert_eq!(rejection.request, ArrayAdvisoryKind::Ack); + assert_eq!(rejection.array, ARRAY); assert_eq!( shared .array_ack_registry @@ -187,13 +224,14 @@ mod tests { let rejection = inbound .handle_catchup_request(&msg, session_id) - .expect_err("write permission must not authorize catchup reads") - .expect("authorization denial must produce an array rejection"); + .expect_err("write permission must not authorize catchup reads"); assert_eq!( rejection.reason, - nodedb_types::sync::wire::array::ArrayRejectReason::EngineRejected + nodedb_types::sync::wire::array::ArrayRejectReason::PermissionDenied ); + assert_eq!(rejection.request, ArrayAdvisoryKind::CatchupRequest); + assert_eq!(rejection.array, ARRAY); assert!( shared .array_subscriber_cursors diff --git a/nodedb/src/control/array_sync/inbound_batch.rs b/nodedb/src/control/array_sync/inbound_batch.rs new file mode 100644 index 000000000..b6b0503e5 --- /dev/null +++ b/nodedb/src/control/array_sync/inbound_batch.rs @@ -0,0 +1,80 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Multi-op inbound array frames: per-op application and per-op results. +//! +//! Each op applies on its own. A reject stops nothing: array ops are +//! independent CRDT writes ordered by HLC, so a later op never depends on an +//! earlier one in the same frame. The result list names every op's fate. + +use nodedb_array::sync::hlc::Hlc; +use nodedb_array::sync::op::ArrayOp; +use nodedb_array::sync::op_codec; +use nodedb_types::sync::wire::SyncProvenance; +use nodedb_types::sync::wire::array::{ + ArrayDeltaBatchMsg, ArrayOpOutcome, ArrayOpResult, ArrayOpResultsMsg, ArrayRejectMsg, + ArrayRejectReason, +}; +use tracing::warn; + +use super::inbound::{OpApplied, OriginArrayInbound}; +use super::reject::build_reject; + +impl OriginArrayInbound { + /// Handle a batch of delta ops from a Lite peer. + /// + /// Returns one result per entry of `op_payloads`, in the same order. An + /// op that does not decode is rejected as `ShapeInvalid` with an all-zero + /// HLC. The ops after any reject are still attempted. + pub async fn handle_delta_batch(&self, msg: &ArrayDeltaBatchMsg) -> ArrayOpResultsMsg { + let mut results = Vec::with_capacity(msg.op_payloads.len()); + for payload in &msg.op_payloads { + let result = match op_codec::decode_op(payload) { + Ok(op) => { + let provenance = self.batch_provenance(&msg.array); + self.apply_op_for_result(op, payload, provenance).await + } + Err(e) => { + warn!(array = %msg.array, error = %e, "array_inbound: batch decode failed"); + ArrayOpResult::rejected(build_reject( + &msg.array, + Hlc::ZERO, + ArrayRejectReason::ShapeInvalid, + format!("batch decode error: {e}"), + )) + } + }; + results.push(result); + } + ArrayOpResultsMsg { + array: msg.array.clone(), + results, + } + } + + /// Apply one decoded op and report its fate as a result entry. + pub(super) async fn apply_op_for_result( + &self, + op: ArrayOp, + raw_op_bytes: &[u8], + provenance: Option, + ) -> ArrayOpResult { + let op_hlc_bytes = op.header.hlc.to_bytes(); + let outcome = self.apply_op(op, raw_op_bytes, provenance).await; + op_result(op_hlc_bytes, outcome) + } +} + +/// The result entry for an op whose apply returned `outcome`. +fn op_result(op_hlc_bytes: [u8; 18], outcome: Result) -> ArrayOpResult { + match outcome { + Ok(OpApplied::Applied) => ArrayOpResult { + op_hlc_bytes, + outcome: ArrayOpOutcome::Applied, + }, + Ok(OpApplied::Idempotent) => ArrayOpResult { + op_hlc_bytes, + outcome: ArrayOpOutcome::Duplicate, + }, + Err(reject) => ArrayOpResult::rejected(reject), + } +} diff --git a/nodedb/src/control/array_sync/inbound_propose.rs b/nodedb/src/control/array_sync/inbound_propose.rs index c225da50c..42fc263a0 100644 --- a/nodedb/src/control/array_sync/inbound_propose.rs +++ b/nodedb/src/control/array_sync/inbound_propose.rs @@ -16,7 +16,7 @@ use crate::control::wal_replication::ReplicatedEntry; use crate::types::{TraceId, VShardId}; use super::inbound::OriginArrayInbound; -use super::reject::build_reject; +use super::reject::{build_reject, engine_reject}; impl OriginArrayInbound { fn consume_array_write_authorization( @@ -24,19 +24,19 @@ impl OriginArrayInbound { authorization: crate::control::server::shared::authorization::AuthorizedCollection, array: &str, hlc: Hlc, - ) -> Result<(), Option> { + ) -> Result<(), ArrayRejectMsg> { let (tenant_id, database_id, collection, permission) = authorization.into_scope(); if tenant_id != self.tenant_id() || database_id != self.database_id() || collection != array || permission != crate::control::security::identity::Permission::Write { - return Err(Some(build_reject( + return Err(build_reject( array, hlc, - ArrayRejectReason::EngineRejected, + ArrayRejectReason::PermissionDenied, "array authorization scope mismatch".to_string(), - ))); + )); } Ok(()) } @@ -44,24 +44,25 @@ impl OriginArrayInbound { /// Propose a `ReplicatedEntry` to Raft and await its commit + apply. /// /// Returns `Ok(())` when the entry has been committed and executed by the - /// distributed applier. Returns a reject on proposal or timeout failure. + /// distributed applier. A failure maps to its reject reason through the + /// typed error. pub(super) async fn propose_and_await( &self, mut entry: ReplicatedEntry, array: &str, hlc: Hlc, authorization: crate::control::server::shared::authorization::AuthorizedCollection, - ) -> Result<(), Option> { + ) -> Result<(), ArrayRejectMsg> { self.consume_array_write_authorization(authorization, array, hlc)?; if entry.tenant_id != self.tenant_id().as_u64() || entry.database_id != self.database_id().as_u64() { - return Err(Some(build_reject( + return Err(build_reject( array, hlc, - ArrayRejectReason::EngineRejected, + ArrayRejectReason::PermissionDenied, "array replicated entry scope mismatch".to_string(), - ))); + )); } // Both proposers below take prebuilt bytes, so the floor is stamped // here: replicas hold the write until they applied the array's DDL. @@ -74,12 +75,7 @@ impl OriginArrayInbound { let idempotency_key = entry.idempotency_key; let data = entry.encode().map_err(|e| { error!(array = %array, error = %e, "array_inbound: replicated entry encode failed"); - Some(build_reject( - array, - hlc, - ArrayRejectReason::EngineRejected, - format!("replicated entry encode failed: {e}"), - )) + engine_reject(array, hlc, "replicated entry encode failed", &e) })?; // The async proposer forwards to the group leader and waits for the @@ -87,14 +83,7 @@ impl OriginArrayInbound { let async_proposer = self .shared() .async_raft_proposer() - .map_err(|e| { - Some(build_reject( - array, - hlc, - ArrayRejectReason::EngineRejected, - format!("raft proposer not available: {e}"), - )) - })? + .map_err(|e| engine_reject(array, hlc, "raft proposer not available", &e))? .as_ref(); let deadline = crate::control::server::shared::session::StatementDeadline::of_statement(self.shared()) @@ -103,12 +92,7 @@ impl OriginArrayInbound { Ok((_payload, _committed_version)) => Ok(()), Err(e) => { warn!(array = %array, error = %e, "array_inbound: raft propose+apply failed"); - Err(Some(build_reject( - array, - hlc, - ArrayRejectReason::EngineRejected, - format!("raft propose failed: {e}"), - ))) + Err(engine_reject(array, hlc, "raft propose failed", &e)) } } } @@ -119,21 +103,23 @@ impl OriginArrayInbound { pub(super) async fn cell_surrogate( &self, op: &ArrayOp, - ) -> Result, Option> { + ) -> Result, ArrayRejectMsg> { use nodedb_array::sync::op::ArrayOpKind; if !matches!(op.kind, ArrayOpKind::Put) { return Ok(None); } - let reject = |detail: String| { - Some(build_reject( - &op.header.array, - op.header.hlc, - ArrayRejectReason::EngineRejected, - detail, - )) + let reject = |context: &str, error: &crate::Error| { + engine_reject(&op.header.array, op.header.hlc, context, error) }; - let pk = zerompk::to_msgpack_vec(&op.coord) - .map_err(|e| reject(format!("array cell coord encode: {e}")))?; + let pk = zerompk::to_msgpack_vec(&op.coord).map_err(|e| { + reject( + "array cell coord encode", + &crate::Error::Serialization { + format: "msgpack".into(), + detail: e.to_string(), + }, + ) + })?; let surrogate = crate::control::server::surrogate_exchange::assign_surrogate_routed( self.shared(), nodedb_types::CollectionKey::from_bare(self.database_id(), &op.header.array), @@ -142,7 +128,7 @@ impl OriginArrayInbound { TraceId::ZERO, ) .await - .map_err(|e| reject(format!("array cell surrogate assign: {e}")))?; + .map_err(|e| reject("array cell surrogate assign", &e))?; Ok(Some(surrogate)) } diff --git a/nodedb/src/control/array_sync/mod.rs b/nodedb/src/control/array_sync/mod.rs index b428c0fda..90927ca96 100644 --- a/nodedb/src/control/array_sync/mod.rs +++ b/nodedb/src/control/array_sync/mod.rs @@ -7,6 +7,7 @@ pub mod catchup; pub mod gc_task; pub mod inbound; pub mod inbound_advisory; +pub mod inbound_batch; pub mod inbound_propose; pub mod op_log; pub mod outbound; diff --git a/nodedb/src/control/array_sync/reject.rs b/nodedb/src/control/array_sync/reject.rs index d17a1c16f..6b6e02363 100644 --- a/nodedb/src/control/array_sync/reject.rs +++ b/nodedb/src/control/array_sync/reject.rs @@ -1,10 +1,14 @@ // SPDX-License-Identifier: BUSL-1.1 -//! Helper to build outbound [`ArrayRejectMsg`] frames. +//! Build outbound [`ArrayRejectMsg`] frames, and map a typed engine error to +//! its reject reason. use nodedb_array::sync::hlc::Hlc; use nodedb_types::sync::wire::array::{ArrayRejectMsg, ArrayRejectReason}; +use crate::bridge::envelope::ErrorCode; +use crate::control::server::sync::refusal::is_retryable_dispatch_error; + /// Build an [`ArrayRejectMsg`] for a rejected inbound op. /// /// `op_hlc` is the HLC of the offending op. `reason` is a structured code. @@ -22,3 +26,189 @@ pub fn build_reject( detail: detail.into(), } } + +/// Build the reject for an op the engine path failed with `error`. +/// +/// The reason comes from the typed error through +/// [`reason_for_engine_error`]. `context` names the failed step in `detail`. +pub fn engine_reject( + array: &str, + op_hlc: Hlc, + context: &str, + error: &crate::Error, +) -> ArrayRejectMsg { + build_reject( + array, + op_hlc, + reason_for_engine_error(error), + format!("{context}: {error}"), + ) +} + +/// The reject reason for an engine-path error, read from its type alone. +/// +/// - `PermissionDenied`: an authorization denial. +/// - `EngineUnavailable`: the shared sync classifier finds the error +/// retryable, or the commit outcome is unknown. Array apply dedups by +/// HLC, so a re-send of an op that did commit is a no-op. +/// - `EngineRefused`: every other error, which the same bytes reproduce. +pub fn reason_for_engine_error(error: &crate::Error) -> ArrayRejectReason { + if is_authorization_denial(error) { + return ArrayRejectReason::PermissionDenied; + } + if is_retryable_dispatch_error(error) || is_unknown_commit_outcome(error) { + return ArrayRejectReason::EngineUnavailable; + } + ArrayRejectReason::EngineRefused +} + +fn is_authorization_denial(error: &crate::Error) -> bool { + matches!( + error, + crate::Error::RejectedAuthz { .. } + | crate::Error::DataPlane(ErrorCode::RejectedAuthz { .. }) + ) +} + +fn is_unknown_commit_outcome(error: &crate::Error) -> bool { + matches!( + error, + crate::Error::ProposalOutcomeUnknown { .. } + | crate::Error::CommittedResultUnavailable { .. } + ) +} + +#[cfg(test)] +mod tests { + use nodedb_types::sync::wire::array::ArrayRejectClass; + + use super::*; + use crate::types::{RequestId, TenantId, VShardId}; + + fn class_of(error: &crate::Error) -> ArrayRejectClass { + reason_for_engine_error(error).class() + } + + #[test] + fn a_dispatch_capacity_refusal_is_transient() { + let error = crate::Error::DispatchCapacity { + scope: crate::DispatchCapacityScope::TenantInflight { + tenant_id: TenantId::new(1), + inflight: 64, + cap: 64, + }, + }; + assert_eq!( + reason_for_engine_error(&error), + ArrayRejectReason::EngineUnavailable + ); + assert_eq!(class_of(&error), ArrayRejectClass::Transient); + } + + #[test] + fn a_data_plane_capacity_refusal_is_transient() { + let error = crate::Error::DataPlane(ErrorCode::DispatchCapacity { + reason: "core queue full".into(), + }); + assert_eq!(class_of(&error), ArrayRejectClass::Transient); + } + + #[test] + fn leader_change_deadline_and_lost_quorum_are_transient() { + let errors = [ + crate::Error::NotLeader { + vshard_id: VShardId::new(0), + leader_node: 2, + leader_addr: "10.0.0.2:9000".into(), + leader_term: 3, + }, + crate::Error::DeadlineExceeded { + request_id: RequestId::new(1), + }, + crate::Error::GroupQuorumUnavailable { + group_id: 1, + voters: vec![1, 2, 3], + unreachable: vec![2, 3], + }, + crate::Error::MetadataLeaderUnavailable, + crate::Error::Backpressure { + engine: nodedb_mem::EngineId::Timeseries, + }, + ]; + for error in errors { + assert_eq!( + reason_for_engine_error(&error), + ArrayRejectReason::EngineUnavailable, + "{error:?}" + ); + } + } + + #[test] + fn an_unknown_commit_outcome_is_transient_because_hlc_dedups_a_resend() { + let error = crate::Error::ProposalOutcomeUnknown { + group_id: 1, + log_index: 9, + }; + assert_eq!(class_of(&error), ArrayRejectClass::Transient); + } + + #[test] + fn an_authorization_denial_is_permission_denied() { + let control = crate::Error::RejectedAuthz { + tenant_id: TenantId::new(1), + resource: "array 'a'".into(), + }; + let data = crate::Error::DataPlane(ErrorCode::RejectedAuthz { + resource: "array 'a'".into(), + }); + for error in [control, data] { + assert_eq!( + reason_for_engine_error(&error), + ArrayRejectReason::PermissionDenied + ); + assert_eq!(class_of(&error), ArrayRejectClass::Permanent); + } + } + + #[test] + fn a_verdict_on_the_op_is_permanent() { + let errors = [ + crate::Error::DataPlane(ErrorCode::RejectedPrevalidation { + reason: "array put fenced".into(), + }), + crate::Error::DataPlane(ErrorCode::TypeMismatch { + collection: "a".into(), + detail: "attr v".into(), + }), + crate::Error::RejectedConstraint { + collection: "a".into(), + constraint: "c".into(), + detail: "violated".into(), + }, + crate::Error::Internal { + detail: "replicated entry encode".into(), + }, + ]; + for error in errors { + assert_eq!( + reason_for_engine_error(&error), + ArrayRejectReason::EngineRefused, + "{error:?}" + ); + } + } + + #[test] + fn an_engine_reject_carries_the_context_and_the_derived_reason() { + let error = crate::Error::DispatchCapacity { + scope: crate::DispatchCapacityScope::QueueFull { + core_id: 0, + capacity: 8, + }, + }; + let reject = engine_reject("a", Hlc::ZERO, "raft propose", &error); + assert_eq!(reject.reason, ArrayRejectReason::EngineUnavailable); + assert!(reject.detail.starts_with("raft propose: ")); + } +} diff --git a/nodedb/src/control/array_sync/snapshot_assembly.rs b/nodedb/src/control/array_sync/snapshot_assembly.rs index b3e37d81c..d5e236325 100644 --- a/nodedb/src/control/array_sync/snapshot_assembly.rs +++ b/nodedb/src/control/array_sync/snapshot_assembly.rs @@ -2,13 +2,15 @@ //! Snapshot chunk buffering and assembly for inbound array sync. -use std::collections::BTreeMap; +use std::collections::{BTreeMap, HashMap}; +use std::sync::MutexGuard; use nodedb_array::sync::hlc::Hlc; use nodedb_array::sync::op_codec; use nodedb_array::sync::snapshot::{SnapshotChunk, SnapshotHeader, assemble_chunks}; use nodedb_types::sync::wire::array::{ - ArrayRejectMsg, ArrayRejectReason, ArraySnapshotChunkMsg, ArraySnapshotMsg, + ArrayOpResult, ArrayOpResultsMsg, ArrayRejectMsg, ArrayRejectReason, ArraySnapshotChunkMsg, + ArraySnapshotMsg, }; use tracing::{error, warn}; @@ -47,25 +49,62 @@ impl SnapshotAssembly { } } -fn reject(array: &str, hlc: Hlc, detail: impl Into) -> Option { - Some(build_reject( - array, - hlc, - ArrayRejectReason::ShapeInvalid, - detail.into(), - )) +type SnapshotBuffers = HashMap; + +fn reject(array: &str, hlc: Hlc, detail: impl Into) -> ArrayRejectMsg { + build_reject(array, hlc, ArrayRejectReason::ShapeInvalid, detail.into()) +} + +/// The session's snapshot buffers are full. Nothing is wrong with the chunk. +fn capacity_reject(array: &str, hlc: Hlc, detail: &str) -> ArrayRejectMsg { + build_reject(array, hlc, ArrayRejectReason::SnapshotCapacity, detail) } impl OriginArrayInbound { - fn discard_snapshot_assembly(&self, key: &SnapshotKey) { - if let Ok(mut snapshots) = self.snapshots().lock() { - snapshots.remove(key); - } + /// Lock the session's snapshot buffers. + /// + /// A poisoned lock means a thread panicked mid-assembly, so no buffer can + /// be trusted. Every buffer is discarded, the poison is cleared, a capture + /// is filed, and the frame draws a transient `SnapshotBufferLost` reject. + fn lock_snapshots( + &self, + array: &str, + hlc: Hlc, + ) -> Result, ArrayRejectMsg> { + let poisoned = match self.snapshots().lock() { + Ok(guard) => return Ok(guard), + Err(poisoned) => poisoned, + }; + let mut buffers = poisoned.into_inner(); + let discarded = buffers.len(); + buffers.clear(); + drop(buffers); + self.snapshots().clear_poison(); + error!( + array = %array, + discarded, + "array_inbound: snapshot lock poisoned; discarded every partial snapshot" + ); + crate::diag::array_snapshot_buffer_poisoned( + array, + self.tenant_id().as_u64(), + self.database_id().as_u64(), + discarded, + ); + Err(build_reject( + array, + hlc, + ArrayRejectReason::SnapshotBufferLost, + "snapshot buffers discarded after a fault; re-send the whole snapshot", + )) + } + + fn discard_snapshot_assembly(&self, key: &SnapshotKey, hlc: Hlc) -> Result<(), ArrayRejectMsg> { + self.lock_snapshots(&key.0, hlc)?.remove(key); + Ok(()) } - fn buffered_snapshot_bytes( - snapshots: &std::collections::HashMap, - ) -> usize { + fn buffered_snapshot_bytes(snapshots: &SnapshotBuffers) -> usize { snapshots.values().map(|entry| entry.payload_bytes).sum() } @@ -73,7 +112,7 @@ impl OriginArrayInbound { pub fn handle_snapshot_header( &self, msg: &ArraySnapshotMsg, - ) -> Result> { + ) -> Result { let header: SnapshotHeader = match zerompk::from_msgpack(&msg.header_payload) { Ok(h) => h, Err(e) => { @@ -97,7 +136,7 @@ impl OriginArrayInbound { )? .into_scope(); if header.array != msg.array { - self.discard_snapshot_assembly(&key); + self.discard_snapshot_assembly(&key, header.snapshot_hlc)?; return Err(reject( &msg.array, header.snapshot_hlc, @@ -105,7 +144,7 @@ impl OriginArrayInbound { )); } if header.total_chunks == 0 || header.total_chunks > MAX_CHUNKS_PER_SNAPSHOT { - self.discard_snapshot_assembly(&key); + self.discard_snapshot_assembly(&key, header.snapshot_hlc)?; return Err(reject( &msg.array, header.snapshot_hlc, @@ -113,15 +152,9 @@ impl OriginArrayInbound { )); } - let mut snapshots = match self.snapshots().lock() { - Ok(g) => g, - Err(_) => { - error!(array = %msg.array, "array_inbound: snapshot mutex poisoned"); - return Err(None); - } - }; + let mut snapshots = self.lock_snapshots(&msg.array, header.snapshot_hlc)?; if !snapshots.contains_key(&key) && snapshots.len() >= MAX_CONCURRENT_ASSEMBLIES { - return Err(reject( + return Err(capacity_reject( &msg.array, header.snapshot_hlc, "too many concurrent snapshot assemblies", @@ -162,7 +195,7 @@ impl OriginArrayInbound { pub async fn handle_snapshot_chunk( &self, msg: &ArraySnapshotChunkMsg, - ) -> Result> { + ) -> Result { let snapshot_hlc = Hlc::from_bytes(&msg.snapshot_hlc_bytes); let key = (msg.array.clone(), msg.snapshot_hlc_bytes); @@ -176,7 +209,7 @@ impl OriginArrayInbound { )? .into_scope(); if msg.total_chunks == 0 || msg.total_chunks > MAX_CHUNKS_PER_SNAPSHOT { - self.discard_snapshot_assembly(&key); + self.discard_snapshot_assembly(&key, snapshot_hlc)?; return Err(reject( &msg.array, snapshot_hlc, @@ -184,7 +217,7 @@ impl OriginArrayInbound { )); } if msg.chunk_index >= msg.total_chunks { - self.discard_snapshot_assembly(&key); + self.discard_snapshot_assembly(&key, snapshot_hlc)?; return Err(reject( &msg.array, snapshot_hlc, @@ -192,7 +225,7 @@ impl OriginArrayInbound { )); } if msg.payload.len() > MAX_CHUNK_BYTES { - self.discard_snapshot_assembly(&key); + self.discard_snapshot_assembly(&key, snapshot_hlc)?; return Err(reject( &msg.array, snapshot_hlc, @@ -201,15 +234,9 @@ impl OriginArrayInbound { } let assembled: Option<(SnapshotHeader, Vec)> = { - let mut snapshots = match self.snapshots().lock() { - Ok(g) => g, - Err(_) => { - error!(array = %msg.array, "array_inbound: snapshot mutex poisoned (chunk)"); - return Err(None); - } - }; + let mut snapshots = self.lock_snapshots(&msg.array, snapshot_hlc)?; if !snapshots.contains_key(&key) && snapshots.len() >= MAX_CONCURRENT_ASSEMBLIES { - return Err(reject( + return Err(capacity_reject( &msg.array, snapshot_hlc, "too many concurrent snapshot assemblies", @@ -267,7 +294,7 @@ impl OriginArrayInbound { > MAX_BUFFERED_SNAPSHOT_BYTES { snapshots.remove(&key); - return Err(reject( + return Err(capacity_reject( &msg.array, snapshot_hlc, "snapshot buffer byte limit exceeded", @@ -295,10 +322,7 @@ impl OriginArrayInbound { }; let Some((header, mut chunks)) = assembled else { - let snapshots = match self.snapshots().lock() { - Ok(g) => g, - Err(_) => return Err(None), - }; + let snapshots = self.lock_snapshots(&msg.array, snapshot_hlc)?; let received = snapshots .get(&key) .map(|e| e.chunks.len() as u32) @@ -332,25 +356,123 @@ impl OriginArrayInbound { } }; - let mut ops_applied = 0; - for op in &ops { - let raw = match op_codec::encode_op(op) { - Ok(bytes) => bytes, - Err(e) => { - return Err(reject( - &msg.array, - header.snapshot_hlc, - format!("snapshot op encode: {e}"), - )); - } + // Every op is attempted. A reject names its op and stops nothing. + let mut results = Vec::with_capacity(ops.len()); + for op in ops { + let result = match op_codec::encode_op(&op) { + Ok(raw) => self.apply_op_for_result(op, &raw, None).await, + Err(e) => ArrayOpResult::rejected(build_reject( + &msg.array, + op.header.hlc, + ArrayRejectReason::ShapeInvalid, + format!("snapshot op encode: {e}"), + )), }; - match self.apply_op(op.clone(), &raw, None).await { - Ok(InboundOutcome::Applied) => ops_applied += 1, - Ok(_) => {} - Err(reject) => return Err(reject), - } + results.push(result); } - Ok(InboundOutcome::SnapshotApplied { ops_applied }) + Ok(InboundOutcome::SnapshotApplied(ArrayOpResultsMsg { + array: msg.array.clone(), + results, + })) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use nodedb_array::sync::replica_id::ReplicaId; + use nodedb_array::sync::snapshot::CoordRange; + use nodedb_types::sync::wire::array::ArrayRejectClass; + + use super::*; + use crate::bridge::dispatch::Dispatcher; + use crate::control::array_sync::OriginApplyEngine; + use crate::control::security::identity::{AuthMethod, AuthenticatedIdentity, Role}; + use crate::control::state::SharedState; + use crate::types::{DatabaseId, TenantId}; + use crate::wal::WalManager; + + const ARRAY: &str = "poisoned-snapshot-array"; + + fn writer(dir: &std::path::Path) -> OriginArrayInbound { + let wal = Arc::new(WalManager::open_for_testing(&dir.join("snap.wal")).expect("wal")); + let (dispatcher, _data_sides) = Dispatcher::new(1, 64); + let shared = SharedState::new(dispatcher, wal).expect("shared state"); + let identity = AuthenticatedIdentity::new_regular( + 1, + "snapshot-writer", + TenantId::new(7), + AuthMethod::Trust, + vec![Role::ReadWrite], + Some(DatabaseId::DEFAULT), + AuthenticatedIdentity::default_database_set(false), + ); + let engine = Arc::new(OriginApplyEngine::new( + Arc::clone(&shared.array_sync_schemas), + Arc::clone(&shared.array_sync_op_log), + )); + OriginArrayInbound::new( + engine, + Arc::clone(&shared.array_sync_schemas), + shared, + identity, + ) + } + + fn header(ms: u64) -> ArraySnapshotMsg { + let header = SnapshotHeader { + array: ARRAY.into(), + coord_range: CoordRange { + lo: Vec::new(), + hi: Vec::new(), + }, + total_chunks: 2, + snapshot_hlc: Hlc::new(ms, 0, ReplicaId(3)).expect("hlc"), + schema_hlc: Hlc::ZERO, + }; + ArraySnapshotMsg { + array: ARRAY.into(), + header_payload: zerompk::to_msgpack_vec(&header).expect("encode header"), + } + } + + #[test] + fn a_poisoned_lock_discards_every_buffer_and_draws_a_transient_reject() { + let dir = tempfile::tempdir().expect("tmpdir"); + let inbound = writer(dir.path()); + assert!(matches!( + inbound.handle_snapshot_header(&header(1)), + Ok(InboundOutcome::SnapshotPartial { .. }) + )); + + std::thread::scope(|scope| { + let panicked = scope + .spawn(|| { + let _guard = inbound.snapshots().lock().expect("unpoisoned lock"); + panic!("poison the snapshot lock"); + }) + .join(); + assert!(panicked.is_err()); + }); + assert!(inbound.snapshots().is_poisoned()); + + let reject = inbound + .handle_snapshot_header(&header(2)) + .expect_err("a poisoned lock must draw a reject, never silence"); + assert_eq!(reject.reason, ArrayRejectReason::SnapshotBufferLost); + assert_eq!(reject.class(), ArrayRejectClass::Transient); + assert_eq!(reject.array, ARRAY); + + assert!(!inbound.snapshots().is_poisoned(), "the poison is cleared"); + assert!( + inbound.snapshots().lock().expect("lock").is_empty(), + "every untrusted buffer is discarded" + ); + assert!(matches!( + inbound.handle_snapshot_header(&header(2)), + Ok(InboundOutcome::SnapshotPartial { .. }) + )); } } diff --git a/nodedb/src/control/server/sync/mod.rs b/nodedb/src/control/server/sync/mod.rs index 4604d8767..236e014c8 100644 --- a/nodedb/src/control/server/sync/mod.rs +++ b/nodedb/src/control/server/sync/mod.rs @@ -13,7 +13,7 @@ pub mod listener; pub mod presence; pub mod raft_dispatch; pub mod rate_limit; -mod refusal; +pub(crate) mod refusal; pub mod security; pub mod session; mod session_handler; diff --git a/nodedb/src/control/server/sync/refusal.rs b/nodedb/src/control/server/sync/refusal.rs index 593148db5..4f531e07f 100644 --- a/nodedb/src/control/server/sync/refusal.rs +++ b/nodedb/src/control/server/sync/refusal.rs @@ -411,6 +411,12 @@ fn is_indeterminate_code(code: &ErrorCode) -> bool { } } +/// Whether the same frame is expected to land once `error`'s condition +/// clears. Every sync path that reports a failed dispatch asks this. +pub(crate) fn is_retryable_dispatch_error(error: &crate::Error) -> bool { + retryable_refusal_reason(error).is_some() || is_indeterminate(error) +} + /// The [`AckStatus`] an engine ack must carry when its dispatch failed. /// /// A dispatch can fail because the write was refused on its merits (terminal — @@ -424,7 +430,7 @@ fn is_indeterminate_code(code: &ErrorCode) -> bool { /// `next_seq` is the sequence the sender resumes from — its own seq for /// this batch, since nothing applied. pub(super) fn ack_status_for_dispatch_error(error: &crate::Error, next_seq: u64) -> AckStatus { - if retryable_refusal_reason(error).is_some() || is_indeterminate(error) { + if is_retryable_dispatch_error(error) { return AckStatus::Gap { expected: next_seq }; } AckStatus::Rejected { diff --git a/nodedb/src/control/server/sync/session_handler/array.rs b/nodedb/src/control/server/sync/session_handler/array.rs index 94b690968..9e2364a97 100644 --- a/nodedb/src/control/server/sync/session_handler/array.rs +++ b/nodedb/src/control/server/sync/session_handler/array.rs @@ -12,11 +12,13 @@ use std::sync::Arc; use tracing::warn; use nodedb_types::sync::wire::array::{ - ArrayAckMsg, ArrayCatchupRequestMsg, ArrayDeltaBatchMsg, ArrayDeltaMsg, ArrayRejectMsg, - ArrayRejectReason, ArraySchemaSyncMsg, ArraySnapshotChunkMsg, ArraySnapshotMsg, + ArrayAckMsg, ArrayAdvisoryKind, ArrayAdvisoryRefusedMsg, ArrayCatchupRequestMsg, + ArrayDeltaBatchMsg, ArrayDeltaMsg, ArrayRejectMsg, ArrayRejectReason, ArraySchemaSyncMsg, + ArraySnapshotChunkMsg, ArraySnapshotMsg, }; use super::super::wire::{SyncFrame, SyncMessageType}; +use crate::control::array_sync::InboundOutcome; use crate::control::state::SharedState; /// Build the per-session inbound array engine bound to `tenant_id`, or `None` @@ -75,6 +77,8 @@ pub(super) fn is_array_frame(msg_type: SyncMessageType) -> bool { | SyncMessageType::ArrayAck | SyncMessageType::ArrayReject | SyncMessageType::ArrayCatchupRequest + | SyncMessageType::ArrayOpResults + | SyncMessageType::ArrayAdvisoryRefused ) } @@ -92,92 +96,142 @@ pub(super) fn unauthenticated_array_reject( return None; } + // An advisory frame names no local op, so its refusal leaves the push + // window alone. + let advisory = match frame.msg_type { + SyncMessageType::ArrayAck => Some(ArrayAdvisoryKind::Ack), + SyncMessageType::ArrayCatchupRequest => Some(ArrayAdvisoryKind::CatchupRequest), + _ => None, + }; + if let Some(request) = advisory { + return advisory_refusal_frame(&ArrayAdvisoryRefusedMsg { + array: String::new(), + request, + reason: ArrayRejectReason::Unauthenticated, + detail: "unauthenticated array frame".into(), + }); + } + let reject = ArrayRejectMsg { // Do not reflect untrusted frame contents into the response: this // keeps the rejection bounded even for a deliberately oversized body. array: String::new(), op_hlc_bytes: [0; 18], - reason: ArrayRejectReason::EngineRejected, + reason: ArrayRejectReason::Unauthenticated, detail: "unauthenticated array frame".into(), }; SyncFrame::try_encode(SyncMessageType::ArrayReject, &reject) } -/// Route one inbound array frame to the inbound engine, returning a reject -/// frame to send back when the engine rejects the operation. +/// The reject for an array frame whose body does not decode. It names no +/// array and no op, so the client holds every op it has in flight. +fn undecodable_frame_reject(msg_type: SyncMessageType) -> Option { + let reject = ArrayRejectMsg { + array: String::new(), + op_hlc_bytes: [0; 18], + reason: ArrayRejectReason::ShapeInvalid, + detail: format!("array frame 0x{:02x} body did not decode", msg_type as u8), + }; + SyncFrame::try_encode(SyncMessageType::ArrayReject, &reject) +} + +fn reject_frame(reject: &ArrayRejectMsg) -> Option { + SyncFrame::try_encode(SyncMessageType::ArrayReject, reject) +} + +fn advisory_refusal_frame(refusal: &ArrayAdvisoryRefusedMsg) -> Option { + SyncFrame::try_encode(SyncMessageType::ArrayAdvisoryRefused, refusal) +} + +/// The refusal for an advisory frame whose body does not decode. +fn undecodable_advisory_refusal(request: ArrayAdvisoryKind) -> Option { + advisory_refusal_frame(&ArrayAdvisoryRefusedMsg { + array: String::new(), + request, + reason: ArrayRejectReason::ShapeInvalid, + detail: "advisory array frame body did not decode".into(), + }) +} + +/// The response for a completed inbound snapshot: its per-op results. +fn snapshot_outcome_frame(outcome: InboundOutcome) -> Option { + match outcome { + InboundOutcome::SnapshotApplied(results) => { + SyncFrame::try_encode(SyncMessageType::ArrayOpResults, &results) + } + InboundOutcome::Applied + | InboundOutcome::Idempotent + | InboundOutcome::SnapshotPartial { .. } + | InboundOutcome::SchemaImported + | InboundOutcome::AckRecorded + | InboundOutcome::CatchupRequested => None, + } +} + +/// Route one inbound array frame to the inbound engine, returning the frame +/// to send back: a reject, or the per-op results of a multi-op frame. pub(super) async fn dispatch_array_frame( frame: &SyncFrame, inbound: &crate::control::array_sync::OriginArrayInbound, session_id: &str, ) -> Option { match frame.msg_type { - SyncMessageType::ArrayDelta => { - if let Some(msg) = frame.decode_body::() { - match inbound.handle_delta(&msg).await { - Ok(_) => None, - Err(Some(r)) => SyncFrame::try_encode(SyncMessageType::ArrayReject, &r), - Err(None) => None, - } - } else { - None - } - } - SyncMessageType::ArrayDeltaBatch => { - if let Some(msg) = frame.decode_body::() { - let outcomes = inbound.handle_delta_batch(&msg).await; - outcomes.into_iter().find_map(|r| match r { - Err(Some(reject)) => { - SyncFrame::try_encode(SyncMessageType::ArrayReject, &reject) - } - _ => None, - }) - } else { - None - } - } - SyncMessageType::ArraySnapshot => { - if let Some(msg) = frame.decode_body::() { - match inbound.handle_snapshot_header(&msg) { - Ok(_) => None, - Err(Some(r)) => SyncFrame::try_encode(SyncMessageType::ArrayReject, &r), - Err(None) => None, - } - } else { - None - } - } - SyncMessageType::ArraySnapshotChunk => { - if let Some(msg) = frame.decode_body::() { - match inbound.handle_snapshot_chunk(&msg).await { - Ok(_) => None, - Err(Some(r)) => SyncFrame::try_encode(SyncMessageType::ArrayReject, &r), - Err(None) => None, - } - } else { - None + SyncMessageType::ArrayDelta => match frame.decode_body::() { + Some(msg) => match inbound.handle_delta(&msg).await { + Ok(_) => None, + Err(reject) => reject_frame(&reject), + }, + None => undecodable_frame_reject(frame.msg_type), + }, + SyncMessageType::ArrayDeltaBatch => match frame.decode_body::() { + Some(msg) => { + let results = inbound.handle_delta_batch(&msg).await; + SyncFrame::try_encode(SyncMessageType::ArrayOpResults, &results) } - } - SyncMessageType::ArraySchema => { - if let Some(msg) = frame.decode_body::() { - match inbound.handle_schema(&msg).await { + None => undecodable_frame_reject(frame.msg_type), + }, + SyncMessageType::ArraySnapshot => match frame.decode_body::() { + Some(msg) => match inbound.handle_snapshot_header(&msg) { + Ok(_) => None, + Err(reject) => reject_frame(&reject), + }, + None => undecodable_frame_reject(frame.msg_type), + }, + SyncMessageType::ArraySnapshotChunk => match frame.decode_body::() { + Some(msg) => match inbound.handle_snapshot_chunk(&msg).await { + Ok(outcome) => snapshot_outcome_frame(outcome), + Err(reject) => reject_frame(&reject), + }, + None => undecodable_frame_reject(frame.msg_type), + }, + SyncMessageType::ArraySchema => match frame.decode_body::() { + Some(msg) => match inbound.handle_schema(&msg).await { + Ok(_) => None, + Err(reject) => reject_frame(&reject), + }, + None => undecodable_frame_reject(frame.msg_type), + }, + SyncMessageType::ArrayAck => match frame.decode_body::() { + Some(msg) => match inbound.handle_ack(&msg) { + Ok(_) => None, + Err(refusal) => advisory_refusal_frame(&refusal), + }, + None => undecodable_advisory_refusal(ArrayAdvisoryKind::Ack), + }, + SyncMessageType::ArrayCatchupRequest => { + match frame.decode_body::() { + Some(msg) => match inbound.handle_catchup_request(&msg, session_id) { Ok(_) => None, - Err(Some(r)) => SyncFrame::try_encode(SyncMessageType::ArrayReject, &r), - Err(None) => None, - } - } else { - None - } - } - SyncMessageType::ArrayAck => { - if let Some(msg) = frame.decode_body::() { - let _ = inbound.handle_ack(&msg); + Err(refusal) => advisory_refusal_frame(&refusal), + }, + None => undecodable_advisory_refusal(ArrayAdvisoryKind::CatchupRequest), } - None } - SyncMessageType::ArrayCatchupRequest => { - if let Some(msg) = frame.decode_body::() { - let _ = inbound.handle_catchup_request(&msg, session_id); - } + SyncMessageType::ArrayAdvisoryRefused => { + warn!( + session = %session_id, + "sync: received ArrayAdvisoryRefused (outbound-only); ignoring" + ); None } SyncMessageType::ArrayReject => { @@ -191,6 +245,13 @@ pub(super) async fn dispatch_array_frame( } None } + SyncMessageType::ArrayOpResults => { + warn!( + session = %session_id, + "sync: received ArrayOpResults (outbound-only); ignoring" + ); + None + } _ => None, } } @@ -305,7 +366,7 @@ mod tests { .expect("unauthenticated array frame must be rejected before inbound dispatch"); assert_eq!(reject.msg_type, SyncMessageType::ArrayReject); let body: ArrayRejectMsg = reject.decode_body().expect("decode rejection"); - assert_eq!(body.reason, ArrayRejectReason::EngineRejected); + assert_eq!(body.reason, ArrayRejectReason::Unauthenticated); assert_eq!(body.detail, "unauthenticated array frame"); assert!(body.array.is_empty(), "rejection must remain bounded"); } @@ -315,4 +376,407 @@ mod tests { ); } } + + #[tokio::test] + async fn an_undecodable_array_frame_draws_an_unattributed_shape_reject() { + let dir = tempfile::tempdir().expect("tmpdir"); + let (_shared, inbound) = reject_fixture::writer(dir.path(), reject_fixture::refuse_all()); + for msg_type in [ + SyncMessageType::ArrayDelta, + SyncMessageType::ArrayDeltaBatch, + SyncMessageType::ArraySchema, + ] { + // 0xc1 is a reserved MessagePack byte: the frame parses, its body does not. + let body = [0xc1_u8]; + let mut wire = vec![SyncFrame::FORMAT_VERSION, msg_type as u8]; + wire.extend_from_slice(&(body.len() as u32).to_le_bytes()); + wire.extend_from_slice(&crc32c::crc32c(&body).to_le_bytes()); + wire.extend_from_slice(&body); + let frame = SyncFrame::from_bytes(&wire).expect("a well-formed frame"); + let reply = dispatch_array_frame(&frame, &inbound, "s") + .await + .expect("an undecodable frame must be answered"); + assert_eq!(reply.msg_type, SyncMessageType::ArrayReject); + let body: ArrayRejectMsg = reply.decode_body().expect("decode reject"); + assert_eq!(body.reason, ArrayRejectReason::ShapeInvalid); + assert_eq!(body.op_hlc_bytes, [0; 18]); + assert!(body.array.is_empty()); + } + } + + #[tokio::test] + async fn a_capacity_refused_delta_is_rejected_as_transient() { + use nodedb_types::sync::wire::array::ArrayRejectClass; + + let dir = tempfile::tempdir().expect("tmpdir"); + let (_shared, inbound) = reject_fixture::writer(dir.path(), reject_fixture::refuse_all()); + let op = reject_fixture::delete_op(reject_fixture::ARRAY, 1, Hlc::ZERO); + let frame = SyncFrame::try_encode( + SyncMessageType::ArrayDelta, + &ArrayDeltaMsg { + array: reject_fixture::ARRAY.into(), + op_payload: op_codec::encode_op(&op).expect("encode op"), + producer_id: 0, + epoch: 0, + seq: 1, + }, + ) + .expect("encode delta"); + + let reply = dispatch_array_frame(&frame, &inbound, "s") + .await + .expect("a refused delta must be answered"); + assert_eq!(reply.msg_type, SyncMessageType::ArrayReject); + let body: ArrayRejectMsg = reply.decode_body().expect("decode reject"); + assert_eq!(body.reason, ArrayRejectReason::EngineUnavailable); + assert_eq!(body.class(), ArrayRejectClass::Transient); + assert_eq!(body.op_hlc_bytes, op.header.hlc.to_bytes()); + assert_eq!(body.array, reject_fixture::ARRAY); + } + + #[tokio::test] + async fn a_batch_reports_every_op_in_order_and_keeps_applying_after_a_reject() { + use nodedb_types::sync::wire::array::{ArrayOpOutcome, ArrayOpResultsMsg}; + + let dir = tempfile::tempdir().expect("tmpdir"); + let (_shared, inbound) = + reject_fixture::writer(dir.path(), reject_fixture::admit_first_then_refuse()); + let applied = reject_fixture::delete_op(reject_fixture::ARRAY, 2, Hlc::ZERO); + let unknown = reject_fixture::delete_op("no-such-array", 3, Hlc::ZERO); + let refused = reject_fixture::delete_op(reject_fixture::ARRAY, 4, Hlc::ZERO); + let too_new_schema = + Hlc::new(nodedb_array::sync::hlc::MAX_PHYSICAL_MS, 0, ReplicaId(9)).expect("max hlc"); + let too_new = reject_fixture::delete_op(reject_fixture::ARRAY, 5, too_new_schema); + let encode = |op: &ArrayOp| op_codec::encode_op(op).expect("encode op"); + let frame = SyncFrame::try_encode( + SyncMessageType::ArrayDeltaBatch, + &ArrayDeltaBatchMsg { + array: reject_fixture::ARRAY.into(), + op_payloads: vec![ + vec![0xc1], + encode(&applied), + encode(&unknown), + encode(&refused), + encode(&too_new), + ], + }, + ) + .expect("encode batch"); + + let reply = dispatch_array_frame(&frame, &inbound, "s") + .await + .expect("a batch is always answered"); + assert_eq!(reply.msg_type, SyncMessageType::ArrayOpResults); + let body: ArrayOpResultsMsg = reply.decode_body().expect("decode results"); + assert_eq!(body.array, reject_fixture::ARRAY); + + let observed: Vec<([u8; 18], Option)> = body + .results + .iter() + .map(|result| { + let reason = match &result.outcome { + ArrayOpOutcome::Applied | ArrayOpOutcome::Duplicate => None, + ArrayOpOutcome::Rejected { reason, .. } => Some(*reason), + }; + (result.op_hlc_bytes, reason) + }) + .collect(); + assert_eq!( + observed, + vec![ + ([0; 18], Some(ArrayRejectReason::ShapeInvalid)), + (applied.header.hlc.to_bytes(), None), + ( + unknown.header.hlc.to_bytes(), + Some(ArrayRejectReason::ArrayUnknown) + ), + ( + refused.header.hlc.to_bytes(), + Some(ArrayRejectReason::EngineUnavailable) + ), + ( + too_new.header.hlc.to_bytes(), + Some(ArrayRejectReason::SchemaTooNew) + ), + ] + ); + assert_eq!(body.results[1].outcome, ArrayOpOutcome::Applied); + } + + #[tokio::test] + async fn a_denied_ack_and_catch_up_draw_a_typed_advisory_refusal() { + use nodedb_types::sync::wire::AckStatus; + use nodedb_types::sync::wire::array::ArrayRejectClass; + + let dir = tempfile::tempdir().expect("tmpdir"); + let (_shared, inbound) = reject_fixture::no_permission(dir.path()); + let ack = SyncFrame::try_encode( + SyncMessageType::ArrayAck, + &ArrayAckMsg { + array: reject_fixture::ARRAY.into(), + replica_id: 3, + ack_hlc_bytes: [0; 18], + applied_seq: 0, + status: AckStatus::Applied, + }, + ) + .expect("encode ack"); + let catch_up = SyncFrame::try_encode( + SyncMessageType::ArrayCatchupRequest, + &ArrayCatchupRequestMsg { + array: reject_fixture::ARRAY.into(), + from_hlc_bytes: [0; 18], + }, + ) + .expect("encode catch-up"); + + for (frame, request) in [ + (&ack, ArrayAdvisoryKind::Ack), + (&catch_up, ArrayAdvisoryKind::CatchupRequest), + ] { + let reply = dispatch_array_frame(frame, &inbound, "s") + .await + .expect("a denied advisory request must be answered"); + assert_eq!( + reply.msg_type, + SyncMessageType::ArrayAdvisoryRefused, + "an advisory denial never travels as an ArrayReject" + ); + let body: ArrayAdvisoryRefusedMsg = reply.decode_body().expect("decode refusal"); + assert_eq!(body.request, request); + assert_eq!(body.array, reject_fixture::ARRAY); + assert_eq!(body.reason, ArrayRejectReason::PermissionDenied); + assert_eq!(body.class(), ArrayRejectClass::Permanent); + } + } + + #[tokio::test] + async fn a_catch_up_for_an_unknown_array_draws_a_schema_refusal() { + let dir = tempfile::tempdir().expect("tmpdir"); + let (_shared, inbound) = reject_fixture::writer(dir.path(), reject_fixture::refuse_all()); + let frame = SyncFrame::try_encode( + SyncMessageType::ArrayCatchupRequest, + &ArrayCatchupRequestMsg { + array: "no-such-array".into(), + from_hlc_bytes: [0; 18], + }, + ) + .expect("encode catch-up"); + let reply = dispatch_array_frame(&frame, &inbound, "s") + .await + .expect("an unknown-array catch-up must be answered"); + assert_eq!(reply.msg_type, SyncMessageType::ArrayAdvisoryRefused); + let body: ArrayAdvisoryRefusedMsg = reply.decode_body().expect("decode refusal"); + assert_eq!(body.request, ArrayAdvisoryKind::CatchupRequest); + assert_eq!(body.reason, ArrayRejectReason::ArrayUnknown); + } + + #[test] + fn an_unauthenticated_advisory_frame_draws_an_advisory_refusal() { + let ack = SyncFrame::try_encode( + SyncMessageType::ArrayAck, + &ArrayAckMsg { + array: "a".into(), + replica_id: 1, + ack_hlc_bytes: [0; 18], + applied_seq: 0, + status: nodedb_types::sync::wire::AckStatus::Applied, + }, + ) + .expect("encode ack"); + let reply = unauthenticated_array_reject(&ack, false, false).expect("answered"); + assert_eq!(reply.msg_type, SyncMessageType::ArrayAdvisoryRefused); + let body: ArrayAdvisoryRefusedMsg = reply.decode_body().expect("decode refusal"); + assert_eq!(body.reason, ArrayRejectReason::Unauthenticated); + assert_eq!(body.request, ArrayAdvisoryKind::Ack); + } + + #[tokio::test] + async fn a_full_snapshot_buffer_is_rejected_as_transient() { + use nodedb_array::sync::snapshot::{CoordRange, SnapshotHeader}; + use nodedb_types::sync::wire::array::ArrayRejectClass; + + let dir = tempfile::tempdir().expect("tmpdir"); + let (_shared, inbound) = reject_fixture::writer(dir.path(), reject_fixture::refuse_all()); + let header_frame = |ms: u64| { + let header = SnapshotHeader { + array: reject_fixture::ARRAY.into(), + coord_range: CoordRange { + lo: Vec::new(), + hi: Vec::new(), + }, + total_chunks: 1, + snapshot_hlc: Hlc::new(ms, 0, ReplicaId(3)).expect("hlc"), + schema_hlc: Hlc::ZERO, + }; + SyncFrame::try_encode( + SyncMessageType::ArraySnapshot, + &ArraySnapshotMsg { + array: reject_fixture::ARRAY.into(), + header_payload: zerompk::to_msgpack_vec(&header).expect("encode header"), + }, + ) + .expect("encode snapshot frame") + }; + + let mut ms = 1; + let reject = loop { + if let Some(reply) = dispatch_array_frame(&header_frame(ms), &inbound, "s").await { + break reply; + } + ms += 1; + assert!(ms < 10_000, "the assembly cap must refuse a header"); + }; + let body: ArrayRejectMsg = reject.decode_body().expect("decode reject"); + assert_eq!(body.reason, ArrayRejectReason::SnapshotCapacity); + assert_eq!(body.class(), ArrayRejectClass::Transient); + } + + use nodedb_array::sync::hlc::Hlc; + use nodedb_array::sync::op::ArrayOp; + use nodedb_array::sync::op_codec; + use nodedb_array::sync::replica_id::ReplicaId; + + /// An authenticated array writer whose Raft proposer is a test double. + mod reject_fixture { + use std::path::Path; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use nodedb_array::schema::ArraySchemaBuilder; + use nodedb_array::schema::attr_spec::{AttrSpec, AttrType}; + use nodedb_array::schema::dim_spec::{DimSpec, DimType}; + use nodedb_array::sync::SchemaDoc; + use nodedb_array::sync::hlc::{Hlc, HlcGenerator}; + use nodedb_array::sync::op::{ArrayOp, ArrayOpHeader, ArrayOpKind}; + use nodedb_array::sync::replica_id::ReplicaId; + use nodedb_array::types::coord::value::CoordValue; + use nodedb_array::types::domain::{Domain, DomainBound}; + + use crate::bridge::dispatch::Dispatcher; + use crate::control::array_sync::OriginArrayInbound; + use crate::control::security::identity::{AuthMethod, AuthenticatedIdentity, Role}; + use crate::control::state::SharedState; + use crate::control::wal_replication::{AppliedOutput, AsyncRaftProposer}; + use crate::types::{DatabaseId, ReadVersions, TenantId}; + use crate::wal::WalManager; + + pub(super) const ARRAY: &str = "reject-class-array"; + const TENANT: u64 = 7; + + fn capacity_refusal() -> crate::Error { + crate::Error::DispatchCapacity { + scope: crate::DispatchCapacityScope::QueueFull { + core_id: 0, + capacity: 8, + }, + } + } + + /// A proposer that refuses every write at dispatch capacity. + pub(super) fn refuse_all() -> Arc { + Arc::new(|_vshard, _key, _bytes, _deadline| { + Box::pin(async { Err::(capacity_refusal()) }) + }) + } + + /// A proposer that applies the first write and refuses the rest at + /// dispatch capacity. + pub(super) fn admit_first_then_refuse() -> Arc { + let calls = Arc::new(AtomicUsize::new(0)); + Arc::new(move |_vshard, _key, _bytes, _deadline| { + let first = calls.fetch_add(1, Ordering::SeqCst) == 0; + Box::pin(async move { + if first { + Ok((Vec::new(), ReadVersions::new())) + } else { + Err(capacity_refusal()) + } + }) + }) + } + + pub(super) fn delete_op(array: &str, ms: u64, schema_hlc: Hlc) -> ArrayOp { + ArrayOp { + header: ArrayOpHeader { + array: array.into(), + hlc: Hlc::new(ms, 0, ReplicaId::new(3)).expect("hlc"), + schema_hlc, + valid_from_ms: 0, + valid_until_ms: -1, + system_from_ms: ms as i64, + }, + kind: ArrayOpKind::Delete, + coord: vec![CoordValue::Int64(ms as i64)], + attrs: None, + } + } + + fn import_schema(shared: &SharedState) { + let schema = ArraySchemaBuilder::new(ARRAY) + .dim(DimSpec::new( + "x", + DimType::Int64, + Domain::new(DomainBound::Int64(0), DomainBound::Int64(1_023)), + )) + .attr(AttrSpec::new("v", AttrType::Int64, true)) + .tile_extents(vec![16]) + .build() + .expect("schema"); + let hlc_gen = HlcGenerator::new(ReplicaId(7)); + let doc = SchemaDoc::from_schema(ReplicaId(7), &schema, &hlc_gen).expect("doc"); + let snapshot = doc.export_snapshot().expect("snapshot"); + shared + .array_sync_schemas + .import_snapshot_in_database( + DatabaseId::DEFAULT, + TENANT, + ARRAY, + &snapshot, + doc.schema_hlc(), + ) + .expect("import schema"); + } + + /// A read-write session on `ARRAY` whose proposer is `proposer`. + pub(super) fn writer( + dir: &Path, + proposer: Arc, + ) -> (Arc, Arc) { + session(dir, proposer, vec![Role::ReadWrite]) + } + + /// A session on `ARRAY` that holds no role, so every request is denied. + pub(super) fn no_permission(dir: &Path) -> (Arc, Arc) { + session(dir, refuse_all(), Vec::new()) + } + + fn session( + dir: &Path, + proposer: Arc, + roles: Vec, + ) -> (Arc, Arc) { + let wal = Arc::new(WalManager::open_for_testing(&dir.join("reject.wal")).expect("wal")); + let (dispatcher, _data_sides) = Dispatcher::new(1, 64); + let shared = SharedState::new(dispatcher, wal).expect("shared state"); + crate::control::vshard_admission::install_async_raft_proposer( + &shared, + crate::control::vshard_admission::applying_submit(proposer), + ) + .expect("install proposer"); + import_schema(&shared); + let identity = AuthenticatedIdentity::new_regular( + 1, + "array-writer", + TenantId::new(TENANT), + AuthMethod::Trust, + roles, + Some(DatabaseId::DEFAULT), + AuthenticatedIdentity::default_database_set(false), + ); + let inbound = super::super::build_array_inbound(&Some(Arc::clone(&shared)), identity) + .expect("inbound engine"); + (shared, inbound) + } + } } diff --git a/nodedb/src/diag/context/array_sync.rs b/nodedb/src/diag/context/array_sync.rs new file mode 100644 index 000000000..f04c112a7 --- /dev/null +++ b/nodedb/src/diag/context/array_sync.rs @@ -0,0 +1,64 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Forensic payloads for inbound array-sync capture sites. + +use faultbox::DomainContext; +use faultbox::serde_json::{Value, json}; + +/// A session's snapshot-assembly lock was poisoned: a thread panicked while +/// it held the lock, so no buffered assembly can be trusted. +pub(in crate::diag) struct ArraySnapshotBufferPoisoned<'a> { + /// The array the frame that found the poison named. + pub array: &'a str, + pub tenant_id: u64, + pub database_id: u64, + /// Partial assemblies discarded with the lock. + pub discarded: usize, +} + +impl DomainContext for ArraySnapshotBufferPoisoned<'_> { + fn domain_kind(&self) -> &'static str { + "nodedb.array_snapshot_buffer_poisoned" + } + + fn grouping_key(&self) -> String { + // One root cause: a panic under the lock. The array, scope and + // count are the occurrence. + "array_snapshot_buffer_poisoned".to_string() + } + + fn to_json(&self) -> Value { + json!({ + "array": self.array, + "tenant_id": self.tenant_id, + "database_id": self.database_id, + "discarded_assemblies": self.discarded, + "effect": "every partial inbound snapshot of the session was discarded. The \ + frame that found the poison drew a SnapshotBufferLost reject", + "operator_action": "find the panic that poisoned the lock in this report's \ + breadcrumbs and fix it. Clients re-send their snapshots", + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn grouping_ignores_the_occurrence() { + let first = ArraySnapshotBufferPoisoned { + array: "a", + tenant_id: 1, + database_id: 0, + discarded: 2, + }; + let second = ArraySnapshotBufferPoisoned { + array: "b", + tenant_id: 9, + database_id: 4, + discarded: 7, + }; + assert_eq!(first.grouping_key(), second.grouping_key()); + } +} diff --git a/nodedb/src/diag/context/mod.rs b/nodedb/src/diag/context/mod.rs index db4048e4d..6b649851c 100644 --- a/nodedb/src/diag/context/mod.rs +++ b/nodedb/src/diag/context/mod.rs @@ -6,6 +6,7 @@ //! A grouping key excludes per-occurrence values, so a retry loop files one //! report with a rising count instead of one per attempt. +mod array_sync; mod catalog; mod columnar; mod continuous_agg; @@ -28,6 +29,7 @@ mod vector; mod vector_build; mod write_path; +pub(in crate::diag) use array_sync::ArraySnapshotBufferPoisoned; pub(in crate::diag) use catalog::{ CatalogApplyOrphanRow, CollectionPurgeRowMissing, ConsumerGroupOffsetsRetained, MetadataApplyWedged, SynonymGroupNotApplied, diff --git a/nodedb/src/diag/mod.rs b/nodedb/src/diag/mod.rs index 7cb9e4d4c..b844aa59b 100644 --- a/nodedb/src/diag/mod.rs +++ b/nodedb/src/diag/mod.rs @@ -10,15 +10,15 @@ mod recording; pub use context::{DATABASE_SCOPE, IlpFlushOutcome, LostResponseWrite, TENANT_SCOPE}; pub use recording::{ - IndexRebuildTarget, VectorBuildTarget, apply_panicked, batch_insert_without_surrogates, - calvin_apply_halted, calvin_barrier_log_store_failed, calvin_completion_timeout, - calvin_left_state_not_forgotten, catalog_apply_orphan_row, collection_purge_row_missing, - columnar_segment_corrupt, consumer_group_offsets_retained, continuous_aggregate_not_applied, - crdt_dead_letter_not_enqueued, crdt_dead_letter_not_restored, crdt_dead_letter_not_stored, - cut_barrier_not_placed, data_plane_core_fail_stopped, data_plane_response_lost, - data_plane_responses_lost, descriptor_lease_not_renewed, entry_kind, fts_index_update_failed, - held_entry_debt_not_recorded, history_compaction_not_applied, ilp_invalid_utf8_drop, - ilp_line_read_drop, index_rebuild_not_installed, metadata_apply_wedged, + IndexRebuildTarget, VectorBuildTarget, apply_panicked, array_snapshot_buffer_poisoned, + batch_insert_without_surrogates, calvin_apply_halted, calvin_barrier_log_store_failed, + calvin_completion_timeout, calvin_left_state_not_forgotten, catalog_apply_orphan_row, + collection_purge_row_missing, columnar_segment_corrupt, consumer_group_offsets_retained, + continuous_aggregate_not_applied, crdt_dead_letter_not_enqueued, crdt_dead_letter_not_restored, + crdt_dead_letter_not_stored, cut_barrier_not_placed, data_plane_core_fail_stopped, + data_plane_response_lost, data_plane_responses_lost, descriptor_lease_not_renewed, entry_kind, + fts_index_update_failed, held_entry_debt_not_recorded, history_compaction_not_applied, + ilp_invalid_utf8_drop, ilp_line_read_drop, index_rebuild_not_installed, metadata_apply_wedged, orphaned_index_entry_after_delete, quota_row_invalid, quota_row_undecodable, quota_row_write_failed, quota_scope_purge_incomplete, quota_scope_replay_aborted, raft_entries_reapplied, raft_entry_reapplied, redo_abandon_given_up, redo_entry_malformed, diff --git a/nodedb/src/diag/recording/array_sync.rs b/nodedb/src/diag/recording/array_sync.rs new file mode 100644 index 000000000..f7a6166b2 --- /dev/null +++ b/nodedb/src/diag/recording/array_sync.rs @@ -0,0 +1,30 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Capture sites for inbound array sync. + +use faultbox::{Capture, EventKind}; + +use crate::diag::context; + +/// Report a poisoned snapshot-assembly lock and the assemblies discarded +/// with it. Called only from the lock site that detects the poison. +pub fn array_snapshot_buffer_poisoned( + array: &str, + tenant_id: u64, + database_id: u64, + discarded: usize, +) { + let ctx = context::ArraySnapshotBufferPoisoned { + array, + tenant_id, + database_id, + discarded, + }; + let _ = Capture::new( + EventKind::InvariantViolation, + "array sync: a panic poisoned the snapshot-assembly lock, so its buffers were discarded", + ) + .domain(&ctx) + .with_backtrace() + .emit(); +} diff --git a/nodedb/src/diag/recording/mod.rs b/nodedb/src/diag/recording/mod.rs index b3c64193b..212a1da63 100644 --- a/nodedb/src/diag/recording/mod.rs +++ b/nodedb/src/diag/recording/mod.rs @@ -8,6 +8,7 @@ //! and returns `None` when unrecorded, so the result is deliberately //! discarded. +mod array_sync; mod catalog; mod columnar; mod continuous_agg; @@ -29,6 +30,7 @@ mod shared; mod vector; mod vector_build; +pub use array_sync::array_snapshot_buffer_poisoned; pub use catalog::{ catalog_apply_orphan_row, collection_purge_row_missing, consumer_group_offsets_retained, metadata_apply_wedged, synonym_group_not_applied, diff --git a/nodedb/tests/inproc/cases/sync_compat.rs b/nodedb/tests/inproc/cases/sync_compat.rs index 02fe1eee5..791549669 100644 --- a/nodedb/tests/inproc/cases/sync_compat.rs +++ b/nodedb/tests/inproc/cases/sync_compat.rs @@ -339,7 +339,7 @@ fn lite_fts_index_frame_with_no_fields_roundtrips() { } #[test] -fn all_11_message_types_valid() { +fn core_message_types_keep_their_codes() { for (code, expected) in [ (0x01u8, SyncMessageType::Handshake), (0x02, SyncMessageType::HandshakeAck), @@ -357,7 +357,18 @@ fn all_11_message_types_valid() { assert_eq!(parsed, expected); assert_eq!(parsed as u8, code); } - assert!(SyncMessageType::from_u8(0x99).is_none()); +} + +/// Every byte either names a message type that round-trips to the same byte, +/// or names none. `0x00` stays unassigned. +#[test] +fn every_code_round_trips_or_is_unknown() { + for code in 0..=u8::MAX { + if let Some(parsed) = SyncMessageType::from_u8(code) { + assert_eq!(parsed as u8, code, "code {code:#04x}"); + } + } + assert!(SyncMessageType::from_u8(0x00).is_none()); } /// Roundtrip the `VectorPrimaryConfig` through the same MessagePack channel From ebaeebd1c6df59e7c7e69501b213e6e5f243e82f Mon Sep 17 00:00:00 2001 From: Farhan Syah Date: Mon, 12 Oct 2026 16:41:21 +0800 Subject: [PATCH 09/11] ci: run clippy on nodedb-client feature subsets --- .github/workflows/test.yml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 2a8aae67e..5b5a03c8d 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -56,6 +56,12 @@ jobs: run: cargo fmt --all -- --check - name: Run clippy run: cargo clippy --workspace --all-targets --all-features --profile ci -- -D warnings + - name: Run clippy nodedb-client (no features) + run: cargo clippy -p nodedb-client --no-default-features --all-targets --profile ci -- -D warnings + - name: Run clippy nodedb-client (native only) + run: cargo clippy -p nodedb-client --no-default-features --features native --all-targets --profile ci -- -D warnings + - name: Run clippy nodedb-client (remote only) + run: cargo clippy -p nodedb-client --no-default-features --features remote --all-targets --profile ci -- -D warnings - name: Warm-storage ObjectStore gate run: bash scripts/ci/check_warm_storage_gate.sh - name: Format-path .expect() gate From 9f536626464273c838041eef4e1e81c317dc6385 Mon Sep 17 00:00:00 2001 From: Farhan Syah Date: Mon, 12 Oct 2026 17:07:06 +0800 Subject: [PATCH 10/11] test(vector): count heap allocations per HNSW search --- nodedb-vector/tests/search_allocations.rs | 156 ++++++++++++++++++++++ 1 file changed, 156 insertions(+) create mode 100644 nodedb-vector/tests/search_allocations.rs diff --git a/nodedb-vector/tests/search_allocations.rs b/nodedb-vector/tests/search_allocations.rs new file mode 100644 index 000000000..3925f4b81 --- /dev/null +++ b/nodedb-vector/tests/search_allocations.rs @@ -0,0 +1,156 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Heap allocations per HNSW search, counted by a global allocator. +//! +//! This file is its own test binary because it installs a +//! `#[global_allocator]`. It uses only the public API. Run it with +//! `--nocapture` to see the `ALLOC` lines. + +use std::alloc::{GlobalAlloc, Layout, System}; +use std::sync::atomic::{AtomicU64, Ordering}; + +use nodedb_types::vector_dtype::VectorStorageDtype; +use nodedb_vector::navix::{NavixSearchOptions, navix_search}; +use nodedb_vector::{DistanceMetric, FilterThresholds, HnswIndex, HnswParams, adaptive_search}; +use roaring::RoaringBitmap; + +/// Counts `alloc` and `realloc` calls, then delegates to [`System`]. +struct CountingAlloc; + +static ALLOCATIONS: AtomicU64 = AtomicU64::new(0); + +// SAFETY: every method delegates to `System` with the caller's arguments. +unsafe impl GlobalAlloc for CountingAlloc { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + ALLOCATIONS.fetch_add(1, Ordering::Relaxed); + // SAFETY: the caller upholds the `GlobalAlloc::alloc` contract. + unsafe { System.alloc(layout) } + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + // SAFETY: the caller upholds the `GlobalAlloc::dealloc` contract. + unsafe { System.dealloc(ptr, layout) } + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + ALLOCATIONS.fetch_add(1, Ordering::Relaxed); + // SAFETY: the caller upholds the `GlobalAlloc::realloc` contract. + unsafe { System.realloc(ptr, layout, new_size) } + } +} + +#[global_allocator] +static GLOBAL: CountingAlloc = CountingAlloc; + +const COUNT: usize = 2000; +const DIM: usize = 64; +const SEARCHES: u64 = 100; +const K: usize = 10; +const EF: usize = 64; + +/// Deterministic vectors in `[0, 1)` from an xorshift64 stream. +fn seeded_vectors(count: usize, dim: usize, seed: u64) -> Vec> { + let mut state = seed; + (0..count) + .map(|_| { + (0..dim) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + (state % 10_000) as f32 / 10_000.0 + }) + .collect() + }) + .collect() +} + +fn build(vectors: &[Vec], dtype: VectorStorageDtype) -> HnswIndex { + let params = HnswParams { + m: 16, + m0: 32, + ef_construction: 100, + metric: DistanceMetric::L2, + dtype, + }; + let mut index = HnswIndex::with_seed(DIM, params, 42); + for v in vectors { + index.insert(v.clone()).expect("vector has the index dim"); + } + index +} + +/// Mean allocations per call of `search` over `SEARCHES` calls, after one +/// warm-up call. Each call takes the query index `i`. +fn per_search(mut search: impl FnMut(usize)) -> u64 { + search(0); + let before = ALLOCATIONS.load(Ordering::Relaxed); + for i in 0..SEARCHES as usize { + search(i); + } + (ALLOCATIONS.load(Ordering::Relaxed) - before) / SEARCHES +} + +/// The single test: one test in the binary keeps other threads out of the count. +#[test] +fn search_allocations_do_not_scale_with_distance_count() { + let vectors = seeded_vectors(COUNT, DIM, 0x9E37_79B9_7F4A_7C15); + let queries = seeded_vectors(SEARCHES as usize, DIM, 0xD1B5_4A32_D192_ED03); + let f32_index = build(&vectors, VectorStorageDtype::F32); + let f16_index = build(&vectors, VectorStorageDtype::F16); + + // 1% of ids: the adaptive search brute-forces, NaviX walks the graph. + let sparse: RoaringBitmap = (0..COUNT as u32).step_by(100).collect(); + let half: RoaringBitmap = (0..COUNT as u32).step_by(2).collect(); + let thresholds = FilterThresholds::default(); + let navix_options = NavixSearchOptions { + k: K, + ef_search: EF, + allowed: half.clone(), + brute_force_threshold: 0.0, + }; + + for (name, index) in [("f32", &f32_index), ("f16", &f16_index)] { + let search = per_search(|i| { + index.search(&queries[i], K, EF).expect("search"); + }); + let filtered = per_search(|i| { + index + .search_filtered(&queries[i], K, EF, &half) + .expect("filtered search"); + }); + let brute = per_search(|i| { + adaptive_search(index, &queries[i], K, EF, &sparse, &thresholds) + .expect("brute-force search"); + }); + let navix = per_search(|i| { + navix_search(index, &queries[i], &navix_options, DistanceMetric::L2) + .expect("navix search"); + }); + println!("ALLOC {name} search: {search} per search"); + println!("ALLOC {name} filtered search: {filtered} per search"); + println!("ALLOC {name} brute-force search: {brute} per search"); + println!("ALLOC {name} navix search: {navix} per search"); + } + + // Bound: an F32 search allocates a fixed set of buffers that grow by + // doubling, so its count grows at most with log(ef). Distances allocate + // nothing. From ef=16 to ef=256 visited nodes grow several times, but + // allocations stay under ef=256 itself and under 4x the ef=16 count + 16. + let narrow = per_search(|i| { + f32_index.search(&queries[i], K, 16).expect("search"); + }); + let wide = per_search(|i| { + f32_index.search(&queries[i], K, 256).expect("search"); + }); + println!("ALLOC f32 search ef=16: {narrow} per search"); + println!("ALLOC f32 search ef=256: {wide} per search"); + assert!( + wide < 256, + "ef=256 search allocates {wide} times, at least once per beam slot" + ); + assert!( + wide <= narrow * 4 + 16, + "ef=256 search allocates {wide} times vs {narrow} at ef=16: count tracks visited nodes" + ); +} From 6b5a7199dcb2fc9ebc6f018b36f7e5484aa0882e Mon Sep 17 00:00:00 2001 From: Farhan Syah Date: Mon, 12 Oct 2026 17:07:06 +0800 Subject: [PATCH 11/11] test(kv): compare native range replies by their key column --- .../native/cases/native_sorted_index_range.rs | 39 ++++++++++++------- 1 file changed, 25 insertions(+), 14 deletions(-) diff --git a/nodedb/tests/native/cases/native_sorted_index_range.rs b/nodedb/tests/native/cases/native_sorted_index_range.rs index 2711c8d94..52185ef04 100644 --- a/nodedb/tests/native/cases/native_sorted_index_range.rs +++ b/nodedb/tests/native/cases/native_sorted_index_range.rs @@ -12,8 +12,6 @@ use nodedb_types::protocol::text_fields::TextFields; use nodedb_types::protocol::{HelloFrame, NativeResponse, OpCode}; use nodedb_types::value::Value; -const IDS: [&str; 4] = ["alpha", "bravo", "charlie", "delta"]; - fn require_ok(what: &str, response: NativeResponse) -> NativeResponse { assert_eq!( response.status, @@ -23,12 +21,28 @@ fn require_ok(what: &str, response: NativeResponse) -> NativeResponse { response } -/// The ids from `IDS` that a response names, sorted. -fn ids_in(response: &NativeResponse) -> Vec<&'static str> { - let text = format!("{:?}", response.rows); - let mut ids: Vec<&'static str> = IDS.into_iter().filter(|id| text.contains(id)).collect(); - ids.sort_unstable(); - ids +/// The `key` column of a range reply, sorted. +/// +/// SQL `RANGE` and the opcode both answer one row per entry with a `rank` and +/// a `key` column. The key is the row's primary key as a string. +fn keys_in(what: &str, response: &NativeResponse) -> Vec { + let columns = response.columns.as_deref().unwrap_or_default(); + let key_at = columns + .iter() + .position(|column| column == "key") + .unwrap_or_else(|| panic!("{what}: no `key` column in {columns:?}")); + let mut keys: Vec = response + .rows + .as_deref() + .unwrap_or_default() + .iter() + .map(|row| match row.get(key_at) { + Some(Value::String(key)) => key.clone(), + other => panic!("{what}: `key` cell must be a string, got {other:?}"), + }) + .collect(); + keys.sort_unstable(); + keys } #[tokio::test] @@ -99,13 +113,10 @@ async fn native_range_with_typed_bounds_matches_sql_range() { send_request(&mut stream, seq, OpCode::KvSortedIndexRange, fields).await, ); + let expected: Vec = expected.into_iter().map(String::from).collect(); + assert_eq!(keys_in(&text, &via_sql), expected, "{text}"); assert_eq!( - ids_in(&via_sql), - expected, - "SQL RANGE [{sql_min}, {sql_max}]" - ); - assert_eq!( - ids_in(&via_op), + keys_in("KvSortedIndexRange", &via_op), expected, "native range [{min:?}, {max:?}] must select what SQL RANGE selects" );