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 diff --git a/docs/architecture.md b/docs/architecture.md index 2ffd45ea0..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. **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 (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-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}"); + } } 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-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-vector/src/adaptive_filter.rs b/nodedb-vector/src/adaptive_filter.rs index 767367ee8..5a23acc78 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. @@ -97,17 +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| { - let v = index.get_vector(id)?; - if index.is_deleted(id) { - return None; - } - Some(SearchResult { - id, - distance: distance(query, v, metric), - }) + let distance = index.dist_to_node(&probe, id, metric)?; + Some(SearchResult { id, distance }) }) .collect(); @@ -242,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 a93ee761a..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.nodes[nid].storage.as_bytes().to_vec(); - 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,17 +138,13 @@ fn select_neighbors_heuristic( break; } - let candidate_bytes = index.nodes[candidate.id as usize].storage.as_bytes(); + // 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 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, - 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 }); @@ -168,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}; @@ -232,6 +237,100 @@ mod tests { } } + /// 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 + /// 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 crate::test_support::{graph_only_restore_with_backing, seeded_index, seeded_vectors}; + + const DIM: usize = 8; + let base = seeded_vectors(200, DIM, 0x9E37_79B9_7F4A_7C15); + let extra = seeded_vectors(5, DIM, 0xD1B5_4A32_D192_ED03); + + 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 = graph_only_restore_with_backing(&src, &base); + + 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/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 704626abb..6a6abd2d5 100644 --- a/nodedb-vector/src/hnsw/graph/index/distance_ops.rs +++ b/nodedb-vector/src/hnsw/graph/index/distance_ops.rs @@ -2,13 +2,15 @@ //! 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; +use crate::hnsw::graph::{MAX_LAYER_CAP, NodeStorage, Probe}; impl HnswIndex { /// Assign a random layer using the exponential distribution. @@ -23,74 +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. + /// 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 as a probe, borrowed without copying. /// - /// 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. + /// 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))), + } + } + + /// Distance from `probe` to node `node_id` under `metric`. /// - /// 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 { - 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 + /// `metric` is caller-supplied, so filtered traversals can rank under a + /// metric other than the build metric. + /// + /// 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, + probe: &Probe<'_>, + node_id: u32, + metric: DistanceMetric, + ) -> 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() } - } else { - node_bytes - }; - distance_typed( - self.params.metric, - self.params.dtype, - query_bytes, - node_bytes, - self.dim, - ) - .unwrap_or(f32::INFINITY) + } } - /// Compute distance between a query given as `&[f32]` and a stored node. + /// Node `node_id`'s f32 vector: local storage first, then the backing. /// - /// 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 { - 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) - } + /// 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 2ea98cc76..e4689951b 100644 --- a/nodedb-vector/src/navix/acorn.rs +++ b/nodedb-vector/src/navix/acorn.rs @@ -11,14 +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::distance::distance; 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`). @@ -45,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], @@ -56,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; @@ -63,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, @@ -77,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() @@ -98,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; @@ -114,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; @@ -125,10 +128,10 @@ mod inner { break; } } - vec![Candidate { + Candidate { dist: current_dist, id: current, - }] + } } /// ACORN-1 layer-0 beam search. @@ -137,102 +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 index.is_deleted(id) { - return None; - } - let v = index.get_vector(id)?; - Some(SearchResult { - id, - distance: distance(query, v, metric), - }) + let distance = index.dist_to_node(probe, id, metric)?; + Some(SearchResult { id, distance }) }) .collect(); @@ -252,19 +213,6 @@ mod inner { results } - #[inline] - fn node_dist( - index: &HnswIndex, - query: &[f32], - 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, - } - } - #[cfg(test)] mod tests { use super::*; @@ -358,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 016ad8b9d..057dc490f 100644 --- a/nodedb-vector/src/navix/traversal.rs +++ b/nodedb-vector/src/navix/traversal.rs @@ -9,14 +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::distance::distance; 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`. @@ -61,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 /// @@ -76,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; @@ -83,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, @@ -98,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; } @@ -106,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() @@ -124,22 +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 index.is_deleted(id) { - return None; - } - let v = index.get_vector(id)?; - Some(SearchResult { - id, - distance: distance(query, v, metric), - }) + let distance = index.dist_to_node(probe, id, metric)?; + Some(SearchResult { id, distance }) }) .collect(); @@ -159,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, @@ -172,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. @@ -226,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, @@ -235,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), + ); } } @@ -361,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. @@ -403,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), + ); } } } @@ -424,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, @@ -439,49 +340,20 @@ fn expand_blind(args: ExpandArgs<'_>) { visited.insert(nb1); for &nb2 in index.neighbors_at(nb1, 0) { - if !visited.insert(nb2) { - continue; - } - if index.is_deleted(nb2) { - continue; - } - if !allowed.contains(nb2) { + if !visited.insert(nb2) || !admitted(index, allowed, 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 { - match index.get_vector(node_id) { - Some(v) => distance(query, v, metric), - None => f32::INFINITY, - } -} - // ── Tests ───────────────────────────────────────────────────────────────────── #[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 { @@ -683,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 +} 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" + ); +} 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/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/cluster/calvin/scheduler/driver/core/catch_up.rs b/nodedb/src/control/cluster/calvin/scheduler/driver/core/catch_up.rs index ce1ca86d1..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 @@ -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,212 @@ 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 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); + + 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!( + 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), + 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/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/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..2c37a8549 100644 --- a/nodedb/src/control/metrics/system/render.rs +++ b/nodedb/src/control/metrics/system/render.rs @@ -18,6 +18,8 @@ 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); + 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 719fad66f..68d3390dc 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(); + overloaded(resp) } ApiError::Coded { status, @@ -352,12 +364,14 @@ impl IntoResponse for ApiError { code, cause, } => { + 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, }; - (status, axum::Json(body)).into_response() + let resp = (status, axum::Json(body)).into_response(); + if is_overload { overloaded(resp) } else { resp } } ApiError::Unauthorized(msg) => plain(StatusCode::UNAUTHORIZED, msg), ApiError::Forbidden(msg) => plain(StatusCode::FORBIDDEN, msg), @@ -371,6 +385,26 @@ 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( + 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 +464,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 +475,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 +667,114 @@ 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); + } + + /// 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..deaca5023 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,35 @@ 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 = + 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), + before + 1 + ); + } + #[test] fn response_encode_msgpack() { let resp = NativeResponse::from_query_result( 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/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"); 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/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/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/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/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/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/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 4b73880d2..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,38 +128,398 @@ 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()?; + 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::*; + + /// 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, + } + } - 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)); + /// 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) } - let refs: Vec<&[u8]> = values.iter().map(|v| v.as_slice()).collect(); - Some(def.encoder.encode(&refs)) -} + 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 = 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")] + ); + } + + 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); + } -/// Convert a JSON field value to sortable bytes. -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() + #[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); } } - serde_json::Value::String(s) => s.as_bytes().to_vec(), - _ => Vec::new(), + } + + #[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)); + } + + /// 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 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!( + 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!( + pks(&idx.top_k(3, NOW_MS)), + vec![(1, "ten"), (2, "ten_half"), (3, "eleven")] + ); + + // 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 27103dd69..9d5b66146 100644 --- a/nodedb/src/engine/kv/sorted_index/key.rs +++ b/nodedb/src/engine/kv/sorted_index/key.rs @@ -2,15 +2,17 @@ //! 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)] @@ -19,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 { @@ -31,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 { @@ -53,294 +65,244 @@ 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. + /// `min` and `max` are [`value`](super::value) encodings of the leading + /// column. Two rules of the key layout apply: /// - /// Two consequences of that framing are handled: - /// - /// * 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. + /// The value bytes of column `name` in `sort_key`, as they were passed to + /// [`Self::encode`]. /// - /// 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) - } - - /// Encode a millisecond timestamp as big-endian bytes. - /// - /// 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) + /// 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 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), + SortDirection::Desc => Cow::Owned(value.iter().map(|b| !b).collect()), + }); + } + rest = tail; + } + None } } #[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_timestamp_ms(1000); - let time_late = SortKeyEncoder::encode_timestamp_ms(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, - }]); + 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))); + } - let score_50 = SortKeyEncoder::encode_i64(50); - let score_150 = SortKeyEncoder::encode_i64(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); + } - let key_50 = encoder.encode(&[&score_50]); - let key_150 = encoder.encode(&[&score_150]); + /// 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:?}"), + } + } + } + } - // ASC: lower score sorts first. - assert!(key_50 < key_150); + #[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 = 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 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_a_null_is_the_null_encoding() { + let enc = score_and_ts(SortDirection::Asc); + 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[..])); } - fn encoder(direction: SortDirection) -> SortKeyEncoder { - SortKeyEncoder::new(vec![SortColumn { - name: "score".into(), - direction, - }]) + #[test] + fn column_value_of_an_unknown_column_or_truncated_key_is_none() { + let enc = score_and_ts(SortDirection::Asc); + 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()); } - /// 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" + ); } } @@ -350,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" + ); } } @@ -372,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 2d94a0a5b..7dd6ef1af 100644 --- a/nodedb/src/engine/kv/sorted_index/mod.rs +++ b/nodedb/src/engine/kv/sorted_index/mod.rs @@ -4,7 +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; @@ -12,5 +15,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/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 new file mode 100644 index 000000000..ef443df8e --- /dev/null +++ b/nodedb/src/engine/kv/sorted_index/timestamp.rs @@ -0,0 +1,110 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Read a window timestamp back from the value bytes of one sort column. +//! +//! The value encoding carries a type tag, so the reader needs no guess: +//! +//! - 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::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 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, bool, array, object, and a string that is not +/// ISO 8601. +pub fn decode_column_timestamp(value: &[u8]) -> Option { + 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, i64::MAX] { + assert_eq!( + decode_column_timestamp(&encode_int(ms)), + Some(ColumnTimestamp::Millis(i128::from(ms))) + ); + } + } + + #[test] + fn negative_integer_timestamp_decodes_negative() { + assert_eq!( + 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_714_608_000_000.5, + 1_714_607_999_999.25, + -1_714_608_000_000.5, + ] { + assert_eq!( + decode_column_timestamp(&encode_float(ms)), + Some(ColumnTimestamp::FractionalMillis(ms)) + ); + } + } + + #[test] + 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(&encode_text("2024-05-02T00:00:00Z")), + Some(ColumnTimestamp::Millis(1_714_608_000_000)) + ); + assert_eq!( + decode_column_timestamp(&encode_text("2024-05-02")), + Some(ColumnTimestamp::Millis(1_714_608_000_000)) + ); + } + + #[test] + 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); + } +} 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 3ed5ffb87..ec1f6993a 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,60 @@ 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) => { + 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 +216,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 +257,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 +271,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, i128::from(u64::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(i128::from(u64::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)) } 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; } 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/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 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..52185ef04 --- /dev/null +++ b/nodedb/tests/native/cases/native_sorted_index_range.rs @@ -0,0 +1,126 @@ +// 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; + +fn require_ok(what: &str, response: NativeResponse) -> NativeResponse { + assert_eq!( + response.status, + ResponseStatus::Ok, + "{what} must succeed: {response:?}" + ); + response +} + +/// 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] +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, + ); + + let expected: Vec = expected.into_iter().map(String::from).collect(); + assert_eq!(keys_in(&text, &via_sql), expected, "{text}"); + assert_eq!( + keys_in("KvSortedIndexRange", &via_op), + expected, + "native range [{min:?}, {max:?}] must select what SQL RANGE selects" + ); + } + + server.shutdown().await; +}