fix: M0-M10 third-pass remediation — durability, replication, and CLI hardening
Resolves the 142 findings from tidal/docs/reviews/CODE_REVIEW_m0-m10.md across the engine, server, net, and CLI surfaces: - WAL/session-journal durability, checkpoint format, and crash-recovery hardening - Replication shipper/receiver, tenant isolation, and migration paths - Cluster scatter-gather, router, standalone server + health/offload endpoints - tidalctl refactored into command modules with JSON output and WAL-state tooling - Cohort, governance, signal-ledger, and vector-registry correctness fixes - Expanded UAT/integration/durability test coverage across all milestones
This commit is contained in:
parent
69ae6e64a1
commit
b55ad70141
1
Cargo.lock
generated
1
Cargo.lock
generated
@ -3587,6 +3587,7 @@ version = "0.1.0"
|
||||
dependencies = [
|
||||
"axum 0.8.8",
|
||||
"clap",
|
||||
"crossbeam",
|
||||
"reqwest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
||||
@ -22,6 +22,7 @@ fn make_payload(seqno: u64) -> WalSegmentPayload {
|
||||
id: WalSegmentId::new(RegionId::SINGLE, ShardId(0), seqno),
|
||||
bytes: vec![0xAB; SEGMENT_SIZE],
|
||||
event_count: 10,
|
||||
leader_last_seq: seqno,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -13,6 +13,13 @@ message ShipSegmentRequest {
|
||||
WalSegmentId id = 1;
|
||||
bytes payload = 2;
|
||||
uint64 event_count = 3;
|
||||
// The segment's authoritative last WAL sequence number, computed by the
|
||||
// leader from the ORIGINAL (pre-community-overlay-filter) bytes. Lets the
|
||||
// receiver advance its replication-lag leader high-water-mark even for an
|
||||
// all-local segment that filters to an empty payload (obs-REPL-1). A 0 value
|
||||
// (e.g. from an older sender that omits this field) means "unknown" and the
|
||||
// receiver falls back to the per-batch boundaries it decodes.
|
||||
uint64 leader_last_seq = 4;
|
||||
}
|
||||
|
||||
// Response to a segment shipment.
|
||||
|
||||
@ -2,7 +2,14 @@
|
||||
//!
|
||||
//! State machine: Closed → Open → `HalfOpen` → Closed.
|
||||
//! After `threshold` consecutive failures, the breaker opens for `reset_duration`.
|
||||
//! The first call after the reset period transitions to `HalfOpen` and allows one probe.
|
||||
//! The first [`check`](CircuitBreaker::check) after the reset period transitions
|
||||
//! to `HalfOpen` and admits **exactly one** probe: that single `check()` returns
|
||||
//! `Ok(())`, and every subsequent `check()` returns [`CircuitOpenError`] until the
|
||||
//! outstanding probe resolves via [`record_success`](CircuitBreaker::record_success)
|
||||
//! (→ Closed) or [`record_failure`](CircuitBreaker::record_failure) (→ Open again).
|
||||
//! This is enforced by the state machine, not merely advertised: the half-open
|
||||
//! state carries an `in_flight` flag so concurrent callers cannot all probe a peer
|
||||
//! that just failed.
|
||||
//!
|
||||
//! Struct-shaped error: `CircuitOpenError` carries no fields and represents the single "breaker is open" outcome; an enum with one variant would be ceremony with no readability gain.
|
||||
|
||||
@ -25,8 +32,16 @@ pub struct CircuitBreaker {
|
||||
|
||||
#[derive(Debug)]
|
||||
enum CircuitState {
|
||||
Closed { consecutive_failures: u32 },
|
||||
Open { opened_at: Instant },
|
||||
Closed {
|
||||
consecutive_failures: u32,
|
||||
},
|
||||
Open {
|
||||
opened_at: Instant,
|
||||
},
|
||||
/// A single probe has been admitted and is in flight. `check()` returns
|
||||
/// `Ok(())` exactly once on the Open→`HalfOpen` transition; while in this
|
||||
/// state every further `check()` errors until `record_success` (→ Closed)
|
||||
/// or `record_failure` (→ Open) resolves the outstanding probe.
|
||||
HalfOpen,
|
||||
}
|
||||
|
||||
@ -45,22 +60,38 @@ impl CircuitBreaker {
|
||||
|
||||
/// Check if a request is allowed.
|
||||
///
|
||||
/// Returns `Ok(())` if the breaker is closed or transitioning to half-open,
|
||||
/// and `Ok(())` on lock poisoning (fail-open: allow the request through).
|
||||
/// Returns `Ok(())` if the breaker is closed, or on the single Open→`HalfOpen`
|
||||
/// transition that admits one probe, and `Ok(())` on lock poisoning (fail-open:
|
||||
/// allow the request through).
|
||||
///
|
||||
/// # Half-open single-probe invariant
|
||||
///
|
||||
/// The first `check()` after `reset_duration` elapses transitions Open→`HalfOpen`
|
||||
/// and returns `Ok(())` — that is the one admitted probe. While that probe is
|
||||
/// outstanding the breaker stays `HalfOpen`, and **every subsequent `check()`
|
||||
/// returns [`CircuitOpenError`]** until the probe resolves via
|
||||
/// [`record_success`](Self::record_success) (→ Closed) or
|
||||
/// [`record_failure`](Self::record_failure) (→ Open). N concurrent callers can
|
||||
/// therefore never all probe a peer that just failed; exactly one gets through.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns [`CircuitOpenError`] if the breaker is open and the reset period
|
||||
/// has not yet elapsed.
|
||||
/// has not yet elapsed, or if it is `HalfOpen` with a probe already in flight.
|
||||
pub fn check(&self) -> Result<(), CircuitOpenError> {
|
||||
let Ok(mut state) = self.state.lock() else {
|
||||
tracing::warn!("circuit breaker lock poisoned; failing open");
|
||||
return Ok(());
|
||||
};
|
||||
match *state {
|
||||
CircuitState::Closed { .. } | CircuitState::HalfOpen => Ok(()),
|
||||
CircuitState::Closed { .. } => Ok(()),
|
||||
// A probe is already outstanding; admit no more until it resolves.
|
||||
CircuitState::HalfOpen => Err(CircuitOpenError),
|
||||
CircuitState::Open { opened_at } => {
|
||||
if opened_at.elapsed() >= self.reset_duration {
|
||||
// Admit exactly one probe: the transition itself is the
|
||||
// single Ok. The next `check()` will see `HalfOpen` and error
|
||||
// until this probe's record_success/record_failure resolves it.
|
||||
*state = CircuitState::HalfOpen;
|
||||
Ok(())
|
||||
} else {
|
||||
@ -81,6 +112,42 @@ impl CircuitBreaker {
|
||||
};
|
||||
}
|
||||
|
||||
/// Record a follower backpressure reply (`accepted=false`).
|
||||
///
|
||||
/// Backpressure is normal flow control from a healthy-but-busy follower whose
|
||||
/// inbound channel is full — the WAL is durable on the leader and the follower
|
||||
/// catches up once capacity frees. Counting it as a failure would let routine
|
||||
/// congestion open the breaker and convert transient backpressure into a
|
||||
/// self-inflicted `reset_duration` replication stall (see `error.rs`:
|
||||
/// `SegmentRejected` is classified transient). Only genuine transport/gRPC
|
||||
/// errors call [`record_failure`](Self::record_failure).
|
||||
///
|
||||
/// Effect by state:
|
||||
/// - **Closed:** no-op. It neither opens the breaker nor resets a genuine
|
||||
/// failure streak (a follower alternating real errors with backpressure must
|
||||
/// still trip the breaker — see `backpressure_does_not_reset_failure_streak`).
|
||||
/// - **`HalfOpen`:** **resolves the outstanding probe by closing the breaker.**
|
||||
/// A backpressure reply proves the *transport* is healthy (the RPC round-tripped;
|
||||
/// only the follower's queue was full), so the half-open probe succeeded at the
|
||||
/// layer the breaker protects. Resolving it is mandatory: under the single-probe
|
||||
/// invariant a `HalfOpen` breaker admits no further `check()` until the probe is
|
||||
/// resolved, so a no-op here would wedge the breaker permanently the first time a
|
||||
/// probe drew a backpressure reply.
|
||||
/// - **Open:** no-op (no probe to resolve).
|
||||
pub fn record_backpressure(&self) {
|
||||
let Ok(mut state) = self.state.lock() else {
|
||||
tracing::warn!("circuit breaker lock poisoned; ignoring backpressure");
|
||||
return;
|
||||
};
|
||||
// Only a HalfOpen probe is resolved; Closed/Open are left untouched so
|
||||
// backpressure never opens the breaker nor resets a genuine failure streak.
|
||||
if matches!(*state, CircuitState::HalfOpen) {
|
||||
*state = CircuitState::Closed {
|
||||
consecutive_failures: 0,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
/// Record a failed request. Increments the failure count and opens if threshold reached.
|
||||
pub fn record_failure(&self) {
|
||||
let Ok(mut state) = self.state.lock() else {
|
||||
@ -169,6 +236,56 @@ mod tests {
|
||||
assert!(cb.check().is_err()); // re-opened
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn half_open_admits_exactly_one_probe() {
|
||||
// Single-probe invariant: once `check()` transitions Open→HalfOpen and
|
||||
// returns Ok (the one admitted probe), every subsequent `check()` must
|
||||
// error until the probe resolves. Without this, N concurrent callers in
|
||||
// half-open would all probe a peer that just failed.
|
||||
let cb = CircuitBreaker::new(2, Duration::from_millis(1));
|
||||
cb.record_failure();
|
||||
cb.record_failure();
|
||||
assert!(cb.check().is_err(), "breaker should be open");
|
||||
|
||||
std::thread::sleep(Duration::from_millis(5));
|
||||
|
||||
// First check after reset admits the single probe.
|
||||
assert!(
|
||||
cb.check().is_ok(),
|
||||
"first half-open check admits exactly one probe"
|
||||
);
|
||||
// The probe is still outstanding: no success/failure recorded yet.
|
||||
assert!(
|
||||
cb.check().is_err(),
|
||||
"second half-open check must error while the probe is in flight"
|
||||
);
|
||||
assert!(
|
||||
cb.check().is_err(),
|
||||
"further half-open checks keep erroring until the probe resolves"
|
||||
);
|
||||
|
||||
// Resolving the probe with success closes the breaker and re-admits traffic.
|
||||
cb.record_success();
|
||||
assert!(cb.check().is_ok(), "closed after probe succeeds");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn half_open_probe_failure_then_check_stays_closed_to_traffic() {
|
||||
// The complementary resolution path: if the single probe fails, the
|
||||
// breaker re-opens and `check()` errors (now because it is Open, not
|
||||
// because a probe is in flight).
|
||||
let cb = CircuitBreaker::new(2, Duration::from_millis(1));
|
||||
cb.record_failure();
|
||||
cb.record_failure();
|
||||
|
||||
std::thread::sleep(Duration::from_millis(5));
|
||||
assert!(cb.check().is_ok(), "single probe admitted");
|
||||
assert!(cb.check().is_err(), "second check errors: probe in flight");
|
||||
|
||||
cb.record_failure(); // probe failed → re-open
|
||||
assert!(cb.check().is_err(), "re-opened after probe failure");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn half_open_success_closes() {
|
||||
let cb = CircuitBreaker::new(2, Duration::from_millis(1));
|
||||
@ -180,4 +297,99 @@ mod tests {
|
||||
cb.record_success(); // probe succeeded
|
||||
assert!(cb.check().is_ok()); // closed
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn backpressure_never_opens_breaker() {
|
||||
// CRITICAL-7 regression: a healthy-but-busy follower replying
|
||||
// `accepted=false` is the path `send_to` routes to `record_backpressure`.
|
||||
// Far more than `threshold` consecutive backpressure replies must leave
|
||||
// the breaker closed so the leader keeps shipping (no self-inflicted
|
||||
// replication stall).
|
||||
let cb = CircuitBreaker::new(5, Duration::from_secs(30));
|
||||
for _ in 0..10 {
|
||||
cb.record_backpressure();
|
||||
}
|
||||
assert!(
|
||||
cb.check().is_ok(),
|
||||
"backpressure must not open the breaker; a subsequent send must still be permitted"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn half_open_backpressure_resolves_probe_and_closes() {
|
||||
// A backpressure reply to the single half-open probe proves the transport
|
||||
// is healthy (the RPC round-tripped; only the follower's queue was full),
|
||||
// so it must resolve the probe by closing the breaker. Otherwise the
|
||||
// single-probe invariant would wedge the breaker in HalfOpen forever the
|
||||
// first time a probe drew `accepted=false`.
|
||||
let cb = CircuitBreaker::new(2, Duration::from_millis(1));
|
||||
cb.record_failure();
|
||||
cb.record_failure();
|
||||
|
||||
std::thread::sleep(Duration::from_millis(5));
|
||||
assert!(cb.check().is_ok(), "single probe admitted");
|
||||
assert!(cb.check().is_err(), "second check errors: probe in flight");
|
||||
|
||||
cb.record_backpressure(); // probe round-tripped, follower just busy
|
||||
assert!(
|
||||
cb.check().is_ok(),
|
||||
"backpressure resolves the half-open probe and re-admits traffic"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn genuine_failures_still_open_breaker() {
|
||||
// The opposite arm of `send_to` (`Err(status)`) calls `record_failure`,
|
||||
// which must still open the breaker after `threshold` consecutive
|
||||
// transport errors — backpressure handling must not weaken this.
|
||||
let cb = CircuitBreaker::new(5, Duration::from_secs(30));
|
||||
for _ in 0..5 {
|
||||
cb.record_failure();
|
||||
}
|
||||
assert!(
|
||||
cb.check().is_err(),
|
||||
"genuine transport errors must still open the breaker at threshold"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn backpressure_does_not_count_toward_open_among_failures() {
|
||||
// Backpressure must be neutral: interleaving it with genuine failures
|
||||
// neither advances the failure count toward the threshold nor resets it.
|
||||
// Here 4 real failures + many backpressures stays below the 5-failure
|
||||
// threshold (closed), and the 5th real failure opens it.
|
||||
let cb = CircuitBreaker::new(5, Duration::from_secs(30));
|
||||
for _ in 0..4 {
|
||||
cb.record_failure();
|
||||
cb.record_backpressure();
|
||||
}
|
||||
assert!(
|
||||
cb.check().is_ok(),
|
||||
"4 failures interleaved with backpressure must remain below threshold"
|
||||
);
|
||||
cb.record_failure(); // 5th genuine failure
|
||||
assert!(
|
||||
cb.check().is_err(),
|
||||
"the 5th genuine failure must open the breaker"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn backpressure_does_not_reset_failure_streak() {
|
||||
// Backpressure must not reset the consecutive-failure count the way
|
||||
// `record_success` does — otherwise a follower alternating real errors
|
||||
// with backpressure could never trip the breaker. 4 failures, a
|
||||
// backpressure, then a 5th failure must open it.
|
||||
let cb = CircuitBreaker::new(5, Duration::from_secs(30));
|
||||
for _ in 0..4 {
|
||||
cb.record_failure();
|
||||
}
|
||||
cb.record_backpressure();
|
||||
assert!(cb.check().is_ok()); // still 4 failures, below threshold
|
||||
cb.record_failure(); // 5th consecutive genuine failure
|
||||
assert!(
|
||||
cb.check().is_err(),
|
||||
"backpressure must not reset the failure streak; the 5th failure opens the breaker"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@ -127,13 +127,22 @@ impl PeerPool {
|
||||
Ok(response) => {
|
||||
// The follower signals backpressure with `accepted=false` when
|
||||
// its inbound channel is full: the segment was NOT enqueued.
|
||||
// Treat this as a transient failure so the shipper does not
|
||||
// advance its cursor and retries on the next poll.
|
||||
// This is normal flow control from a healthy-but-busy peer, not
|
||||
// a connection failure — so it must NOT count toward opening the
|
||||
// circuit breaker (that would convert transient congestion into
|
||||
// a self-inflicted `circuit_breaker_reset` replication stall,
|
||||
// amplifying follower lag exactly when the leader must keep
|
||||
// shipping). We route it through `record_backpressure`, which
|
||||
// leaves Closed/Open breaker state untouched (it only resolves an
|
||||
// outstanding half-open probe), and return a transient
|
||||
// `SegmentRejected` so the shipper does not advance its cursor
|
||||
// and retries on the next poll. Only the genuine transport error
|
||||
// arm below calls `record_failure`.
|
||||
if response.into_inner().accepted {
|
||||
peer.circuit_breaker.record_success();
|
||||
Ok(())
|
||||
} else {
|
||||
peer.circuit_breaker.record_failure();
|
||||
peer.circuit_breaker.record_backpressure();
|
||||
Err(GrpcTransportError::SegmentRejected(shard))
|
||||
}
|
||||
}
|
||||
|
||||
@ -4,9 +4,29 @@ use std::{collections::HashMap, net::SocketAddr, path::PathBuf, time::Duration};
|
||||
|
||||
use tidaldb::replication::shard::ShardId;
|
||||
|
||||
/// Maximum payload size in bytes (64 MB), matching `InProcessTransport`.
|
||||
/// Maximum payload size in bytes (64 MiB).
|
||||
///
|
||||
/// **Source of truth:** the engine crate's `InProcessTransport::MAX_PAYLOAD_BYTES`
|
||||
/// (`tidaldb`, `replication/in_process.rs`). Both transports must reject the same
|
||||
/// over-size payloads so behavior does not diverge by which transport a deployment
|
||||
/// uses. That engine constant is currently private (module-local in `in_process.rs`),
|
||||
/// so it cannot be `use`d here directly; the value is mirrored and pinned by the
|
||||
/// compile-time cross-check below. If the engine ever exposes it (e.g. as
|
||||
/// `tidaldb::replication::MAX_PAYLOAD_BYTES`), replace this literal with a re-export
|
||||
/// of that constant and delete the cross-check.
|
||||
pub const MAX_PAYLOAD_BYTES: usize = 64 * 1024 * 1024;
|
||||
|
||||
/// Compile-time guard pinning [`MAX_PAYLOAD_BYTES`] to the engine's 64 MiB limit.
|
||||
///
|
||||
/// This catches accidental drift in the mirrored literal at build time (it is the
|
||||
/// best available cross-check while the engine constant stays private). If this
|
||||
/// fails to compile, the two crates' payload limits have diverged: reconcile them,
|
||||
/// preferring to make the engine's constant the single referenced source of truth.
|
||||
const _: () = assert!(
|
||||
MAX_PAYLOAD_BYTES == 64 * 1024 * 1024,
|
||||
"tidal-net MAX_PAYLOAD_BYTES must match the engine's InProcessTransport limit (64 MiB)"
|
||||
);
|
||||
|
||||
/// Configuration for a single [`GrpcTransport`](crate::GrpcTransport) instance.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct GrpcTransportConfig {
|
||||
@ -51,7 +71,9 @@ impl Default for GrpcTransportConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
local_shard: ShardId::SINGLE,
|
||||
listen_addr: "127.0.0.1:59530".parse().expect("valid default addr"),
|
||||
// Default replication listener. Inside tidalDB's reserved dev band
|
||||
// (59520–59529); 59529 is the slot claimed for the gRPC transport.
|
||||
listen_addr: "127.0.0.1:59529".parse().expect("valid default addr"),
|
||||
peers: HashMap::new(),
|
||||
tls: None,
|
||||
insecure: true,
|
||||
|
||||
@ -20,6 +20,7 @@ impl From<WalSegmentPayload> for proto::ShipSegmentRequest {
|
||||
}),
|
||||
payload: p.bytes,
|
||||
event_count: p.event_count,
|
||||
leader_last_seq: p.leader_last_seq,
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -41,6 +42,7 @@ impl TryFrom<proto::ShipSegmentRequest> for WalSegmentPayload {
|
||||
id: WalSegmentId::new(RegionId(region_id), ShardId(shard_id), id.seqno),
|
||||
bytes: req.payload,
|
||||
event_count: req.event_count,
|
||||
leader_last_seq: req.leader_last_seq,
|
||||
})
|
||||
}
|
||||
}
|
||||
@ -56,10 +58,12 @@ mod tests {
|
||||
id: WalSegmentId::new(RegionId(2), ShardId(5), 42),
|
||||
bytes: vec![0xAB; 100],
|
||||
event_count: 7,
|
||||
leader_last_seq: 48,
|
||||
};
|
||||
|
||||
let proto_req: proto::ShipSegmentRequest = original.into();
|
||||
assert_eq!(proto_req.event_count, 7);
|
||||
assert_eq!(proto_req.leader_last_seq, 48);
|
||||
|
||||
let restored = WalSegmentPayload::try_from(proto_req).unwrap();
|
||||
assert_eq!(restored.id.region_id, RegionId(2));
|
||||
@ -67,6 +71,28 @@ mod tests {
|
||||
assert_eq!(restored.id.seqno, 42);
|
||||
assert_eq!(restored.bytes.len(), 100);
|
||||
assert_eq!(restored.event_count, 7);
|
||||
assert_eq!(
|
||||
restored.leader_last_seq, 48,
|
||||
"leader_last_seq must survive the proto roundtrip"
|
||||
);
|
||||
}
|
||||
|
||||
/// An older sender that omits `leader_last_seq` decodes it as 0 ("unknown"),
|
||||
/// which the receiver treats as a fall-back to per-batch boundaries.
|
||||
#[test]
|
||||
fn missing_leader_last_seq_defaults_to_zero() {
|
||||
let req = proto::ShipSegmentRequest {
|
||||
id: Some(proto::WalSegmentId {
|
||||
region_id: 0,
|
||||
shard_id: 0,
|
||||
seqno: 1,
|
||||
}),
|
||||
payload: vec![1, 2, 3],
|
||||
event_count: 1,
|
||||
leader_last_seq: 0,
|
||||
};
|
||||
let restored = WalSegmentPayload::try_from(req).unwrap();
|
||||
assert_eq!(restored.leader_last_seq, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@ -75,6 +101,7 @@ mod tests {
|
||||
id: None,
|
||||
payload: vec![],
|
||||
event_count: 0,
|
||||
leader_last_seq: 0,
|
||||
};
|
||||
assert!(WalSegmentPayload::try_from(req).is_err());
|
||||
}
|
||||
@ -89,6 +116,7 @@ mod tests {
|
||||
}),
|
||||
payload: vec![],
|
||||
event_count: 0,
|
||||
leader_last_seq: 0,
|
||||
};
|
||||
assert!(WalSegmentPayload::try_from(req).is_err());
|
||||
}
|
||||
@ -103,6 +131,7 @@ mod tests {
|
||||
}),
|
||||
payload: vec![],
|
||||
event_count: 0,
|
||||
leader_last_seq: 0,
|
||||
};
|
||||
assert!(WalSegmentPayload::try_from(req).is_err());
|
||||
}
|
||||
|
||||
@ -6,13 +6,19 @@
|
||||
//! callers (WAL shipper and segment receiver) run on `std::thread`, not
|
||||
//! inside a tokio context.
|
||||
|
||||
use std::{collections::HashMap, sync::Mutex};
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
sync::{
|
||||
Arc, Mutex,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
},
|
||||
};
|
||||
|
||||
use tidaldb::replication::{
|
||||
shard::ShardId,
|
||||
transport::{Transport, TransportError, WalSegmentPayload},
|
||||
};
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::{Notify, mpsc};
|
||||
|
||||
use crate::{client::PeerPool, config::GrpcTransportConfig, error::GrpcTransportError, server};
|
||||
|
||||
@ -28,6 +34,19 @@ use crate::{client::PeerPool, config::GrpcTransportConfig, error::GrpcTransportE
|
||||
/// The `Mutex` around the receiver is a safety net, not a concurrency mechanism.
|
||||
/// Concurrent callers would serialize on the mutex while each blocks on
|
||||
/// `runtime.block_on(rx.recv())`, which is safe but wasteful.
|
||||
///
|
||||
/// # Deterministic receiver shutdown
|
||||
///
|
||||
/// A parked [`recv_segment`](Self::recv_segment) blocks on the inbound channel,
|
||||
/// which only returns `None` once every `inbound_tx` sender drops — and those live
|
||||
/// inside the server task. Aborting the server task on [`Drop`] eventually frees
|
||||
/// them, but the receiver thread is typically a detached `std::thread`, so a clean
|
||||
/// (non-process-exit) teardown could leave it parked. To make shutdown deterministic
|
||||
/// the transport carries a latching `shutdown` signal that `Drop` (and the explicit
|
||||
/// [`shutdown_receivers`](Self::shutdown_receivers)) trips **before** tearing the
|
||||
/// runtime down: `recv_segment` checks the latch on entry and selects over the channel
|
||||
/// and a wake `Notify`, so a parked receiver returns `None` immediately on shutdown and
|
||||
/// its thread can exit and be joined.
|
||||
pub struct GrpcTransport {
|
||||
config: GrpcTransportConfig,
|
||||
/// `Some` for the transport's whole life; taken in [`Drop`] so the runtime
|
||||
@ -37,6 +56,47 @@ pub struct GrpcTransport {
|
||||
inbound_rx: Mutex<mpsc::Receiver<WalSegmentPayload>>,
|
||||
pool: PeerPool,
|
||||
server_handle: tokio::task::JoinHandle<Result<(), tonic::transport::Error>>,
|
||||
/// Shutdown signal for the receiver. The [`AtomicBool`] **latches** the request
|
||||
/// so a `recv_segment` that has not yet parked still observes it on entry (no
|
||||
/// missed-wakeup race), while the [`Notify`] wakes a receiver that is *already*
|
||||
/// parked. Tripped by [`shutdown_receivers`](Self::shutdown_receivers) and by
|
||||
/// [`Drop`] before the runtime is torn down, so a parked
|
||||
/// [`recv_segment`](Self::recv_segment) returns `None` deterministically instead
|
||||
/// of waiting for every server-side `inbound_tx` sender to drop.
|
||||
shutdown: Arc<ShutdownSignal>,
|
||||
}
|
||||
|
||||
/// Latching shutdown signal: an [`AtomicBool`] that survives the not-yet-parked race
|
||||
/// plus a [`Notify`] to wake an already-parked waiter. Held by [`GrpcTransport`] and
|
||||
/// tripped on [`Drop`] / [`shutdown_receivers`](GrpcTransport::shutdown_receivers).
|
||||
struct ShutdownSignal {
|
||||
requested: AtomicBool,
|
||||
notify: Notify,
|
||||
}
|
||||
|
||||
impl ShutdownSignal {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
requested: AtomicBool::new(false),
|
||||
notify: Notify::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Latch the shutdown request and wake any already-parked receiver.
|
||||
///
|
||||
/// `Release` so the latch is visible to a receiver that subsequently reads it
|
||||
/// with `Acquire` on `recv_segment` entry; `notify_waiters` covers a receiver
|
||||
/// already parked on `notified()` at the moment this is called.
|
||||
fn request(&self) {
|
||||
self.requested.store(true, Ordering::Release);
|
||||
self.notify.notify_waiters();
|
||||
}
|
||||
|
||||
/// Whether shutdown has been requested. `Acquire` pairs with the `Release` in
|
||||
/// [`request`](Self::request) so a not-yet-parked receiver observes the latch.
|
||||
fn is_requested(&self) -> bool {
|
||||
self.requested.load(Ordering::Acquire)
|
||||
}
|
||||
}
|
||||
|
||||
/// Install a process-wide rustls [`CryptoProvider`] before tonic's TLS
|
||||
@ -91,6 +151,7 @@ impl GrpcTransport {
|
||||
inbound_rx: Mutex::new(inbound_rx),
|
||||
pool,
|
||||
server_handle,
|
||||
shutdown: Arc::new(ShutdownSignal::new()),
|
||||
})
|
||||
}
|
||||
|
||||
@ -101,6 +162,30 @@ impl GrpcTransport {
|
||||
.expect("runtime is present for the transport's whole lifetime")
|
||||
}
|
||||
|
||||
/// Signal any current or future [`recv_segment`](Self::recv_segment) to return `None`.
|
||||
///
|
||||
/// Latches the shutdown request and wakes a receiver already blocked in
|
||||
/// `recv_segment` so it returns `None` — its exit condition per the [`Transport`]
|
||||
/// contract — letting the receiver thread exit and be joined. Because the request
|
||||
/// latches, a receiver that has not yet parked also observes it on its next entry.
|
||||
///
|
||||
/// # Why this exists separately from [`Drop`]
|
||||
///
|
||||
/// The receiver thread typically holds its own `Arc<dyn Transport>` clone, so
|
||||
/// [`GrpcTransport::drop`] cannot fire until that thread has already exited and
|
||||
/// released its `Arc` — a cycle: the thread parks in `recv_segment` and never
|
||||
/// releases the `Arc`, so `Drop` never runs, so the thread is never woken. Break
|
||||
/// the cycle by calling `shutdown_receivers` from the owner's shutdown path (which
|
||||
/// holds a *different* `Arc`) BEFORE joining the receiver thread and dropping its
|
||||
/// `Arc`. The notify-trip in [`Drop`] remains a backstop for the non-cyclic case
|
||||
/// (e.g. a transport owned outright, as in this crate's tests).
|
||||
///
|
||||
/// Idempotent: the request latches, so any current or future `recv_segment`
|
||||
/// returns `None`.
|
||||
pub fn shutdown_receivers(&self) {
|
||||
self.shutdown.request();
|
||||
}
|
||||
|
||||
/// Whether the embedded gRPC serve loop has terminated.
|
||||
///
|
||||
/// The serve task runs for the transport's whole lifetime; a `true` here
|
||||
@ -154,7 +239,37 @@ impl Transport for GrpcTransport {
|
||||
poisoned.into_inner()
|
||||
}
|
||||
};
|
||||
self.runtime().block_on(rx.recv())
|
||||
let shutdown = &self.shutdown;
|
||||
// Fast path: if shutdown was already requested before we ever parked, the
|
||||
// latched flag short-circuits to `None`. This closes the missed-wakeup race
|
||||
// where `shutdown_receivers`/`Drop` ran between a prior `recv_segment`
|
||||
// returning and this call constructing its `notified()` future.
|
||||
if shutdown.is_requested() {
|
||||
return None;
|
||||
}
|
||||
// Otherwise select over an inbound segment and the shutdown notify so a
|
||||
// parked receiver returns `None` the instant shutdown is signaled, rather
|
||||
// than waiting for every server-side `inbound_tx` sender to drop. `None`
|
||||
// from either arm is the receiver's exit condition (see the `Transport`
|
||||
// trait contract), so a tripped shutdown looks exactly like channel close.
|
||||
self.runtime().block_on(async {
|
||||
// Register the `notified()` future BEFORE the final flag re-check so a
|
||||
// `request()` that runs concurrently either (a) is seen by this re-check
|
||||
// or (b) finds us already registered and wakes us — never both-missed.
|
||||
let notified = shutdown.notify.notified();
|
||||
tokio::pin!(notified);
|
||||
if shutdown.is_requested() {
|
||||
return None;
|
||||
}
|
||||
tokio::select! {
|
||||
// Bias toward draining already-queued segments before observing
|
||||
// shutdown, so a shutdown that races with in-flight segments does
|
||||
// not silently discard one already sitting in the channel.
|
||||
biased;
|
||||
segment = rx.recv() => segment,
|
||||
() = &mut notified => None,
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn local_shard(&self) -> ShardId {
|
||||
@ -164,15 +279,27 @@ impl Transport for GrpcTransport {
|
||||
|
||||
impl Drop for GrpcTransport {
|
||||
fn drop(&mut self) {
|
||||
// Abort the server task so its `inbound_tx` is released (unblocking any
|
||||
// parked `recv_segment`), then tear the runtime down WITHOUT blocking.
|
||||
// 1. Trip the shutdown notify FIRST, while the runtime is still alive, so a
|
||||
// parked `recv_segment` wakes on its `shutdown.notified()` arm and returns
|
||||
// `None` deterministically — its receiver thread can then exit and join
|
||||
// instead of leaking until process exit. `notify_waiters` wakes a waiter
|
||||
// that is already parked; `recv_segment` re-creates its future on each call
|
||||
// so there is no missed-wakeup window for a not-yet-parked receiver (a
|
||||
// fresh `recv_segment` after Drop sees a closed channel once the server
|
||||
// task is aborted below). The latch in `request()` also makes any future
|
||||
// `recv_segment` short-circuit to `None` without parking.
|
||||
self.shutdown.request();
|
||||
|
||||
// 2. Abort the server task so its `inbound_tx` is released (so a receiver
|
||||
// that calls `recv_segment` AFTER Drop also observes channel close), then
|
||||
// tear the runtime down WITHOUT blocking.
|
||||
//
|
||||
// A `GrpcTransport` can be dropped from within an async context — e.g.
|
||||
// cluster shutdown unwinding on the axum reactor thread — where the
|
||||
// default blocking `Runtime::drop` panics ("Cannot drop a runtime in a
|
||||
// context where blocking is not allowed"). `shutdown_background` returns
|
||||
// immediately and reclaims the worker threads off-thread, so `Drop` is
|
||||
// safe in any context.
|
||||
// A `GrpcTransport` can be dropped from within an async context — e.g.
|
||||
// cluster shutdown unwinding on the axum reactor thread — where the
|
||||
// default blocking `Runtime::drop` panics ("Cannot drop a runtime in a
|
||||
// context where blocking is not allowed"). `shutdown_background` returns
|
||||
// immediately and reclaims the worker threads off-thread, so `Drop` is
|
||||
// safe in any context.
|
||||
self.server_handle.abort();
|
||||
if let Some(runtime) = self.runtime.take() {
|
||||
runtime.shutdown_background();
|
||||
|
||||
@ -44,17 +44,19 @@ fn make_config(
|
||||
|
||||
#[test]
|
||||
fn ships_payload_larger_than_tonic_default_codec_limit() {
|
||||
// Sanity: the test is only meaningful if the payload exceeds tonic's 4 MiB
|
||||
// default AND fits inside the configured ceiling. The payload-exceeds-default
|
||||
// half is a compile-time invariant; the fits-in-ceiling half is runtime.
|
||||
const FOUR_MIB: usize = 4 * 1024 * 1024;
|
||||
const EIGHT_MIB: usize = 8 * 1024 * 1024;
|
||||
const _: () = assert!(EIGHT_MIB > FOUR_MIB, "payload must exceed tonic default");
|
||||
|
||||
let addr0 = free_addr();
|
||||
let addr1 = free_addr();
|
||||
|
||||
let config0 = make_config(ShardId(0), addr0, HashMap::from([(ShardId(1), addr1)]));
|
||||
let config1 = make_config(ShardId(1), addr1, HashMap::from([(ShardId(0), addr0)]));
|
||||
|
||||
// Sanity: the test is only meaningful if the payload exceeds tonic's 4 MiB
|
||||
// default AND fits inside the configured ceiling.
|
||||
const FOUR_MIB: usize = 4 * 1024 * 1024;
|
||||
const EIGHT_MIB: usize = 8 * 1024 * 1024;
|
||||
assert!(EIGHT_MIB > FOUR_MIB, "payload must exceed tonic default");
|
||||
assert!(
|
||||
EIGHT_MIB <= config0.max_payload_bytes,
|
||||
"payload must fit configured max"
|
||||
@ -73,6 +75,7 @@ fn ships_payload_larger_than_tonic_default_codec_limit() {
|
||||
id: WalSegmentId::new(RegionId::SINGLE, ShardId(0), 7),
|
||||
bytes: body.clone(),
|
||||
event_count: 3,
|
||||
leader_last_seq: 7,
|
||||
};
|
||||
|
||||
t0.send_segment(ShardId(1), payload)
|
||||
|
||||
@ -155,6 +155,7 @@ fn untrusted_client_cert_is_rejected() {
|
||||
id: WalSegmentId::new(RegionId::SINGLE, ShardId(0), seq),
|
||||
bytes: vec![0x11; 32],
|
||||
event_count: 1,
|
||||
leader_last_seq: seq,
|
||||
};
|
||||
if client.send_segment(ShardId(1), payload).is_ok() {
|
||||
last_ok = true;
|
||||
@ -208,6 +209,7 @@ fn absent_client_cert_is_rejected() {
|
||||
id: WalSegmentId::new(RegionId::SINGLE, ShardId(0), seq),
|
||||
bytes: vec![0x22; 32],
|
||||
event_count: 1,
|
||||
leader_last_seq: seq,
|
||||
};
|
||||
if client.send_segment(ShardId(1), payload).is_ok() {
|
||||
last_ok = true;
|
||||
@ -258,6 +260,7 @@ fn mtls_send_and_receive() {
|
||||
id: WalSegmentId::new(RegionId::SINGLE, ShardId(0), 99),
|
||||
bytes: vec![0xEF; 64],
|
||||
event_count: 2,
|
||||
leader_last_seq: 99,
|
||||
};
|
||||
t0.send_segment(ShardId(1), payload).unwrap();
|
||||
|
||||
|
||||
@ -184,6 +184,7 @@ fn write_and_ship(
|
||||
id: WalSegmentId::new(RegionId::SINGLE, ShardId(0), seqno),
|
||||
bytes,
|
||||
event_count: 1,
|
||||
leader_last_seq: seqno,
|
||||
};
|
||||
node.transport
|
||||
.send_segment(target_shard, payload)
|
||||
@ -260,6 +261,7 @@ fn uat_step2_idempotent_replay() {
|
||||
id: WalSegmentId::new(RegionId::SINGLE, ShardId(0), 1),
|
||||
bytes: bytes.clone(),
|
||||
event_count: 1,
|
||||
leader_last_seq: 1,
|
||||
};
|
||||
leader.transport.send_segment(ShardId(1), payload).unwrap();
|
||||
}
|
||||
@ -422,6 +424,7 @@ fn uat_step5_three_node_replication() {
|
||||
id: WalSegmentId::new(RegionId::SINGLE, ShardId(0), seq),
|
||||
bytes: bytes.clone(),
|
||||
event_count: 1,
|
||||
leader_last_seq: seq,
|
||||
};
|
||||
transports[0].send_segment(target, payload).unwrap();
|
||||
}
|
||||
@ -533,6 +536,7 @@ fn partition_heal_convergence() {
|
||||
id: WalSegmentId::new(RegionId::SINGLE, ShardId(0), i),
|
||||
bytes,
|
||||
event_count: 1,
|
||||
leader_last_seq: i,
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
@ -44,6 +44,7 @@ fn make_payload(seqno: u64) -> WalSegmentPayload {
|
||||
id: WalSegmentId::new(RegionId::SINGLE, ShardId(0), seqno),
|
||||
bytes: vec![0xCD; 50],
|
||||
event_count: 1,
|
||||
leader_last_seq: seqno,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -6,7 +6,16 @@
|
||||
//!
|
||||
//! Mirrors the test patterns from `InProcessTransport` but over gRPC on localhost.
|
||||
|
||||
use std::{collections::HashMap, net::SocketAddr, thread, time::Duration};
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
net::SocketAddr,
|
||||
sync::{
|
||||
Arc,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
},
|
||||
thread,
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
use tidal_net::{GrpcTransport, config::GrpcTransportConfig};
|
||||
use tidaldb::replication::{
|
||||
@ -41,6 +50,7 @@ fn make_payload(shard: ShardId, seqno: u64) -> WalSegmentPayload {
|
||||
id: WalSegmentId::new(RegionId::SINGLE, shard, seqno),
|
||||
bytes: vec![0xAB; 100],
|
||||
event_count: 5,
|
||||
leader_last_seq: seqno,
|
||||
}
|
||||
}
|
||||
|
||||
@ -103,6 +113,7 @@ fn payload_too_large_rejected() {
|
||||
id: WalSegmentId::new(RegionId::SINGLE, ShardId(0), 1),
|
||||
bytes: vec![0u8; 64 * 1024 * 1024 + 1],
|
||||
event_count: 0,
|
||||
leader_last_seq: 1,
|
||||
};
|
||||
let result = t0.send_segment(ShardId(1), payload);
|
||||
assert!(result.is_err());
|
||||
@ -120,6 +131,59 @@ fn local_shard_returns_correct_id() {
|
||||
assert_eq!(t.local_shard(), ShardId(7));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shutdown_unparks_recv_segment_and_thread_exits() {
|
||||
// Models the real receiver wiring: a worker thread holds its own
|
||||
// Arc<GrpcTransport> clone and parks in recv_segment. Without a cancellation
|
||||
// path that thread never returns (recv_segment only yields None when every
|
||||
// server-side inbound_tx sender drops), so it would leak until process exit
|
||||
// AND its Arc would keep GrpcTransport::drop from ever running. shutdown_receivers
|
||||
// breaks the cycle: a parked recv_segment returns None deterministically.
|
||||
let addr = free_addr();
|
||||
let config = make_config(ShardId(0), addr, HashMap::new());
|
||||
let transport = Arc::new(GrpcTransport::new(config).expect("transport"));
|
||||
thread::sleep(Duration::from_millis(50));
|
||||
|
||||
let returned_none = Arc::new(AtomicBool::new(false));
|
||||
let worker_transport = Arc::clone(&transport);
|
||||
let worker_flag = Arc::clone(&returned_none);
|
||||
let receiver = thread::spawn(move || {
|
||||
// No segment will ever arrive (no peers); this parks until shutdown.
|
||||
let segment = worker_transport.recv_segment();
|
||||
worker_flag.store(segment.is_none(), Ordering::SeqCst);
|
||||
// Drop the worker's Arc clone on return so only the test's Arc remains.
|
||||
});
|
||||
|
||||
// Let the worker reach the parked recv_segment.
|
||||
thread::sleep(Duration::from_millis(100));
|
||||
assert!(
|
||||
!receiver.is_finished(),
|
||||
"receiver should still be parked in recv_segment before shutdown"
|
||||
);
|
||||
|
||||
// Trip the cancellation path.
|
||||
transport.shutdown_receivers();
|
||||
|
||||
// The parked recv_segment must return None promptly and the thread must exit.
|
||||
let deadline = Instant::now() + Duration::from_secs(5);
|
||||
while !receiver.is_finished() && Instant::now() < deadline {
|
||||
thread::sleep(Duration::from_millis(10));
|
||||
}
|
||||
assert!(
|
||||
receiver.is_finished(),
|
||||
"recv_segment must return after shutdown_receivers so the thread can exit"
|
||||
);
|
||||
receiver.join().expect("receiver thread joins cleanly");
|
||||
assert!(
|
||||
returned_none.load(Ordering::SeqCst),
|
||||
"recv_segment must return None on shutdown (the receiver's exit condition)"
|
||||
);
|
||||
|
||||
// Only the test's Arc remains; dropping it runs GrpcTransport::drop without leak.
|
||||
assert_eq!(Arc::strong_count(&transport), 1);
|
||||
drop(transport);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multiple_segments_fifo_order() {
|
||||
let (t0, t1) = build_pair();
|
||||
|
||||
@ -34,6 +34,7 @@ unwrap_used = "deny"
|
||||
[dependencies]
|
||||
axum = "0.8"
|
||||
clap = { version = "4.5", features = ["derive", "env"] }
|
||||
crossbeam = "0.8"
|
||||
subtle = "2"
|
||||
tower = { version = "0.5", features = ["limit"] }
|
||||
tower-http = { version = "0.6", features = ["timeout", "trace", "request-id"] }
|
||||
|
||||
@ -3,3 +3,7 @@ regions:
|
||||
- name: eu-west
|
||||
- name: ap-south
|
||||
leader: us-east
|
||||
# Optional: number of runtime-free OS worker threads serving cluster write/heal
|
||||
# requests (gRPC segment ship). Bounds write concurrency on the hottest path.
|
||||
# Defaults to available parallelism (clamped 2..=8) when omitted.
|
||||
# write_workers: 4
|
||||
|
||||
@ -15,6 +15,18 @@
|
||||
//!
|
||||
//! All regions still run inside this single process; true multi-process /
|
||||
//! multi-host deployment with process isolation is tracked separately (m8p10).
|
||||
//!
|
||||
//! # Size / cohesion note (M0–M10 review Maintainability-S)
|
||||
//!
|
||||
//! This file currently carries four loosely-coupled concerns — topology YAML
|
||||
//! (`TopologySpec`), gRPC transport wiring (`build_grpc_transports` …),
|
||||
//! [`ClusterState`], and the HTTP route handlers. A clean split into
|
||||
//! `cluster/{topology,transport,state,routes}.rs` is the right end-state, but it
|
||||
//! is deferred (not forced mid-campaign) so it does not destabilize the
|
||||
//! freshly-hardened `cluster_ref`/`cluster_arc` Result propagation and the
|
||||
//! offload wiring. The shared health probes were already lifted out to
|
||||
//! [`crate::health`] and the data DTOs to [`crate::dto`]. Tracked for a
|
||||
//! dedicated follow-up.
|
||||
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
@ -54,6 +66,7 @@ use crate::{
|
||||
feed_items, search_items,
|
||||
},
|
||||
error::{Result, ServerError},
|
||||
offload::{ClusterWritePool, ClusterWritePoolConfig, offload_read},
|
||||
};
|
||||
|
||||
/// How long to wait for each follower's gRPC server to bind before serving.
|
||||
@ -70,6 +83,13 @@ const GRPC_BUILD_ATTEMPTS: u32 = 3;
|
||||
pub struct TopologySpec {
|
||||
pub regions: Vec<RegionSpec>,
|
||||
pub leader: String,
|
||||
/// Number of runtime-free OS worker threads serving cluster write/heal
|
||||
/// requests (gRPC segment ship). Bounds write concurrency on the hottest
|
||||
/// cluster path. When omitted, defaults to
|
||||
/// [`ClusterWritePoolConfig::default`] (derived from available
|
||||
/// parallelism). See [`crate::offload::ClusterWritePool`].
|
||||
#[serde(default)]
|
||||
pub write_workers: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
@ -133,6 +153,12 @@ pub struct ClusterState {
|
||||
name_to_id: HashMap<String, RegionId>,
|
||||
id_to_name: HashMap<RegionId, String>,
|
||||
shutting_down: AtomicBool,
|
||||
/// Fixed-size, runtime-free OS-thread pool for cluster write/heal work
|
||||
/// (gRPC segment ship). Created once at startup and shared by every
|
||||
/// `/signals` and `/cluster/heal` request, replacing the old per-request
|
||||
/// `std::thread` spawn (unbounded growth on the hottest cluster path). See
|
||||
/// [`crate::offload::ClusterWritePool`].
|
||||
write_pool: ClusterWritePool,
|
||||
}
|
||||
|
||||
/// Environment variable that opts in to the experimental cluster mode.
|
||||
@ -393,11 +419,27 @@ impl ClusterState {
|
||||
|
||||
let cluster = SimulatedCluster::build(config);
|
||||
|
||||
// Build the shared cluster-write worker pool once, sized from the
|
||||
// topology's `write_workers` (or the default derived from available
|
||||
// parallelism). Every `/signals` and `/cluster/heal` request shares it,
|
||||
// so write concurrency is bounded and threads are reused.
|
||||
let pool_config = topology.write_workers.map_or_else(
|
||||
ClusterWritePoolConfig::default,
|
||||
ClusterWritePoolConfig::with_workers,
|
||||
);
|
||||
tracing::info!(
|
||||
workers = pool_config.workers,
|
||||
queue_depth = pool_config.queue_depth,
|
||||
"cluster write worker pool started"
|
||||
);
|
||||
let write_pool = ClusterWritePool::new(pool_config);
|
||||
|
||||
Ok(Self {
|
||||
cluster: Some(Arc::new(cluster)),
|
||||
name_to_id,
|
||||
id_to_name,
|
||||
shutting_down: AtomicBool::new(false),
|
||||
write_pool,
|
||||
})
|
||||
}
|
||||
|
||||
@ -431,29 +473,36 @@ impl ClusterState {
|
||||
}
|
||||
}
|
||||
|
||||
/// Access the underlying cluster, panicking if it has been shut down.
|
||||
/// Access the underlying cluster, or a 503 error if it has been shut down.
|
||||
///
|
||||
/// The cluster is present for the entire request-serving lifetime; it is
|
||||
/// only taken during [`shutdown`](Self::shutdown), after axum has stopped
|
||||
/// accepting requests, so no live handler can observe `None`.
|
||||
fn cluster_ref(&self) -> &SimulatedCluster {
|
||||
/// accepting requests, so in practice no live handler observes `None`. That
|
||||
/// ordering invariant lives in `serve_cluster`, not in the type, so rather
|
||||
/// than `expect`-panic (which corrupts a database's reactor thread, per
|
||||
/// `CODING_GUIDELINES` §7) a post-shutdown access returns
|
||||
/// [`ServerError::Unavailable`] → 503, handled by the caller's existing
|
||||
/// error mapping.
|
||||
fn cluster_ref(&self) -> Result<&SimulatedCluster> {
|
||||
self.cluster
|
||||
.as_ref()
|
||||
.expect("cluster accessed after shutdown")
|
||||
.map(AsRef::as_ref)
|
||||
.ok_or_else(|| ServerError::Unavailable("server shutting down".into()))
|
||||
}
|
||||
|
||||
/// Clone the cluster `Arc` for scatter-gather fan-out.
|
||||
/// Clone the cluster `Arc` for scatter-gather fan-out, or a 503 if the
|
||||
/// server is shutting down.
|
||||
///
|
||||
/// Scatter-gather spawns one detached `'static` worker per shard, so each
|
||||
/// worker needs an owned handle that can outlive the request future if the
|
||||
/// shard's blocking query exceeds the deadline. See
|
||||
/// [`crate::scatter_gather`].
|
||||
fn cluster_arc(&self) -> Arc<SimulatedCluster> {
|
||||
Arc::clone(
|
||||
self.cluster
|
||||
.as_ref()
|
||||
.expect("cluster accessed after shutdown"),
|
||||
)
|
||||
/// [`crate::scatter_gather`]. A post-shutdown access degrades to a clean 503
|
||||
/// (see [`cluster_ref`](Self::cluster_ref)) rather than panicking.
|
||||
fn cluster_arc(&self) -> Result<Arc<SimulatedCluster>> {
|
||||
self.cluster
|
||||
.as_ref()
|
||||
.map(Arc::clone)
|
||||
.ok_or_else(|| ServerError::Unavailable("server shutting down".into()))
|
||||
}
|
||||
|
||||
fn resolve_region(&self, name: &str) -> Result<RegionId> {
|
||||
@ -474,24 +523,34 @@ impl ClusterState {
|
||||
}
|
||||
|
||||
/// Default region for reads when no `?region=` is specified: the leader.
|
||||
fn default_read_region(&self) -> RegionId {
|
||||
self.cluster_ref().leader_region()
|
||||
fn default_read_region(&self) -> Result<RegionId> {
|
||||
Ok(self.cluster_ref()?.leader_region())
|
||||
}
|
||||
|
||||
fn read_region(&self, region_name: Option<&str>) -> Result<RegionId> {
|
||||
region_name.map_or_else(
|
||||
|| Ok(self.default_read_region()),
|
||||
|| self.default_read_region(),
|
||||
|name| self.resolve_region(name),
|
||||
)
|
||||
}
|
||||
|
||||
/// All region IDs (shard list for scatter-gather).
|
||||
pub fn shard_ids(&self) -> Vec<RegionId> {
|
||||
self.cluster_ref().regions()
|
||||
/// All region IDs (shard list for scatter-gather), or a 503 if shutting down.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns [`ServerError::Unavailable`] if the cluster fabric has been taken
|
||||
/// during shutdown.
|
||||
pub fn shard_ids(&self) -> Result<Vec<RegionId>> {
|
||||
Ok(self.cluster_ref()?.regions())
|
||||
}
|
||||
|
||||
/// Access the underlying cluster.
|
||||
pub fn cluster(&self) -> &SimulatedCluster {
|
||||
/// Access the underlying cluster, or a 503 if shutting down.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns [`ServerError::Unavailable`] if the cluster fabric has been taken
|
||||
/// during shutdown.
|
||||
pub fn cluster(&self) -> Result<&SimulatedCluster> {
|
||||
self.cluster_ref()
|
||||
}
|
||||
|
||||
@ -519,8 +578,10 @@ impl Drop for ClusterState {
|
||||
pub fn build_cluster_router(state: Arc<ClusterState>, api_key: Option<Arc<str>>) -> Router {
|
||||
let public = Router::new()
|
||||
.route("/health", get(cluster_health))
|
||||
.route("/health/startup", get(health_startup))
|
||||
.route("/health/live", get(health_live))
|
||||
// Shared with the standalone router via [`crate::health`] so the two
|
||||
// modes can never advertise a different probe contract.
|
||||
.route("/health/startup", get(crate::health::health_startup))
|
||||
.route("/health/live", get(crate::health::health_live))
|
||||
.route("/cluster/status", get(cluster_status))
|
||||
.with_state(Arc::clone(&state));
|
||||
|
||||
@ -554,14 +615,10 @@ pub fn build_cluster_router(state: Arc<ClusterState>, api_key: Option<Arc<str>>)
|
||||
}
|
||||
|
||||
// ── Health ──────────────────────────────────────────────────────────────────
|
||||
|
||||
async fn health_startup() -> Json<serde_json::Value> {
|
||||
Json(serde_json::json!({ "ok": true, "service": "tidaldb" }))
|
||||
}
|
||||
|
||||
async fn health_live() -> Json<serde_json::Value> {
|
||||
Json(serde_json::json!({ "ok": true, "service": "tidaldb" }))
|
||||
}
|
||||
//
|
||||
// The unconditional startup/live probes are shared with the standalone router
|
||||
// via [`crate::health`]; only the mode-specific readiness probe (`cluster_health`)
|
||||
// lives here.
|
||||
|
||||
async fn cluster_health(
|
||||
State(state): State<Arc<ClusterState>>,
|
||||
@ -577,8 +634,9 @@ async fn cluster_health(
|
||||
));
|
||||
}
|
||||
|
||||
let leader = state.cluster().leader_region();
|
||||
let items = state.cluster().item_count(leader);
|
||||
let cluster = state.cluster().map_err(ClusterAppError)?;
|
||||
let leader = cluster.leader_region();
|
||||
let items = cluster.item_count(leader);
|
||||
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
@ -609,8 +667,10 @@ struct RegionStatus {
|
||||
partitioned: bool,
|
||||
}
|
||||
|
||||
async fn cluster_status(State(state): State<Arc<ClusterState>>) -> Json<ClusterStatusResponse> {
|
||||
let cluster = state.cluster();
|
||||
async fn cluster_status(
|
||||
State(state): State<Arc<ClusterState>>,
|
||||
) -> std::result::Result<Json<ClusterStatusResponse>, ClusterAppError> {
|
||||
let cluster = state.cluster().map_err(ClusterAppError)?;
|
||||
let leader_region = cluster.leader_region();
|
||||
let relay_log_len = cluster.relay_log_len();
|
||||
|
||||
@ -629,11 +689,11 @@ async fn cluster_status(State(state): State<Arc<ClusterState>>) -> Json<ClusterS
|
||||
})
|
||||
.collect();
|
||||
|
||||
Json(ClusterStatusResponse {
|
||||
Ok(Json(ClusterStatusResponse {
|
||||
leader: state.region_name(leader_region).to_string(),
|
||||
relay_log_len,
|
||||
regions,
|
||||
})
|
||||
}))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
@ -646,7 +706,10 @@ async fn cluster_promote(
|
||||
Json(req): Json<RegionRequest>,
|
||||
) -> std::result::Result<Json<serde_json::Value>, ClusterAppError> {
|
||||
let region_id = state.resolve_region(&req.region).map_err(ClusterAppError)?;
|
||||
state.cluster().promote_leader(region_id);
|
||||
state
|
||||
.cluster()
|
||||
.map_err(ClusterAppError)?
|
||||
.promote_leader(region_id);
|
||||
Ok(Json(serde_json::json!({
|
||||
"ok": true,
|
||||
"leader": req.region,
|
||||
@ -658,7 +721,10 @@ async fn cluster_partition(
|
||||
Json(req): Json<RegionRequest>,
|
||||
) -> std::result::Result<Json<serde_json::Value>, ClusterAppError> {
|
||||
let region_id = state.resolve_region(&req.region).map_err(ClusterAppError)?;
|
||||
state.cluster().partition_region(region_id);
|
||||
state
|
||||
.cluster()
|
||||
.map_err(ClusterAppError)?
|
||||
.partition_region(region_id);
|
||||
Ok(Json(serde_json::json!({
|
||||
"ok": true,
|
||||
"partitioned": req.region,
|
||||
@ -670,13 +736,17 @@ async fn cluster_heal(
|
||||
Json(req): Json<RegionRequest>,
|
||||
) -> std::result::Result<Json<serde_json::Value>, ClusterAppError> {
|
||||
let region_id = state.resolve_region(&req.region).map_err(ClusterAppError)?;
|
||||
// heal_region re-ships missed segments over gRPC (blocking), so offload it.
|
||||
let cluster = state.cluster_arc();
|
||||
offload_cluster_write(move || {
|
||||
cluster.heal_region(region_id);
|
||||
Ok(())
|
||||
})
|
||||
.await?;
|
||||
// heal_region re-ships missed segments over gRPC (blocking), so offload it to
|
||||
// the runtime-free write pool. A saturated pool degrades to 429, not a 500.
|
||||
let cluster = state.cluster_arc().map_err(ClusterAppError)?;
|
||||
state
|
||||
.write_pool
|
||||
.submit(move || {
|
||||
cluster.heal_region(region_id);
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
.map_err(ClusterAppError)?;
|
||||
Ok(Json(serde_json::json!({
|
||||
"ok": true,
|
||||
"healed": req.region,
|
||||
@ -691,6 +761,7 @@ async fn create_item(
|
||||
) -> std::result::Result<StatusCode, ClusterAppError> {
|
||||
state
|
||||
.cluster()
|
||||
.map_err(ClusterAppError)?
|
||||
.write_item_with_metadata(EntityId::new(req.entity_id), &req.metadata)
|
||||
.map_err(|e| ClusterAppError(ServerError::from(e)))?;
|
||||
Ok(StatusCode::CREATED)
|
||||
@ -702,6 +773,7 @@ async fn write_embedding(
|
||||
) -> std::result::Result<StatusCode, ClusterAppError> {
|
||||
state
|
||||
.cluster()
|
||||
.map_err(ClusterAppError)?
|
||||
.write_item_embedding(EntityId::new(req.entity_id), &req.values)
|
||||
.map_err(|e| ClusterAppError(ServerError::from(e)))?;
|
||||
Ok(StatusCode::NO_CONTENT)
|
||||
@ -712,17 +784,22 @@ async fn write_signal(
|
||||
Json(req): Json<SignalRequest>,
|
||||
) -> std::result::Result<StatusCode, ClusterAppError> {
|
||||
// write_signal ships to followers over gRPC (a blocking `runtime.block_on`),
|
||||
// so it must run off the async reactor — see [`offload_cluster_write`].
|
||||
let cluster = state.cluster_arc();
|
||||
// so it must run off the async reactor AND off any thread carrying a runtime
|
||||
// handle — hand it to the runtime-free write pool. A saturated pool yields
|
||||
// 429 (backpressure), not a 500. See [`crate::offload::ClusterWritePool`].
|
||||
let cluster = state.cluster_arc().map_err(ClusterAppError)?;
|
||||
let signal = req.signal;
|
||||
let entity = EntityId::new(req.entity_id);
|
||||
let weight = req.weight;
|
||||
offload_cluster_write(move || {
|
||||
cluster
|
||||
.write_signal(&signal, entity, weight)
|
||||
.map_err(ServerError::from)
|
||||
})
|
||||
.await?;
|
||||
state
|
||||
.write_pool
|
||||
.submit(move || {
|
||||
cluster
|
||||
.write_signal(&signal, entity, weight)
|
||||
.map_err(ServerError::from)
|
||||
})
|
||||
.await
|
||||
.map_err(ClusterAppError)?;
|
||||
Ok(StatusCode::NO_CONTENT)
|
||||
}
|
||||
|
||||
@ -747,7 +824,7 @@ async fn feed(
|
||||
// `retrieve` is a blocking `TidalDb` query; running it directly on the axum
|
||||
// reactor would stall every other in-flight request behind one slow shard.
|
||||
// Offload it to the blocking pool — see [`offload_cluster_read`].
|
||||
let cluster = state.cluster_arc();
|
||||
let cluster = state.cluster_arc().map_err(ClusterAppError)?;
|
||||
let result = offload_cluster_read(move || {
|
||||
cluster
|
||||
.retrieve(region, &retrieve)
|
||||
@ -782,7 +859,7 @@ async fn search(
|
||||
// running them on the axum reactor would stall the whole node. Offload the
|
||||
// pair to the blocking pool so the reactor stays free — see
|
||||
// [`offload_cluster_read`].
|
||||
let cluster = state.cluster_arc();
|
||||
let cluster = state.cluster_arc().map_err(ClusterAppError)?;
|
||||
let result = offload_cluster_read(move || {
|
||||
// Reload text index on the target region before searching.
|
||||
cluster
|
||||
@ -809,9 +886,9 @@ async fn sharded_create_item(
|
||||
State(state): State<Arc<ClusterState>>,
|
||||
Json(req): Json<ItemRequest>,
|
||||
) -> std::result::Result<StatusCode, ClusterAppError> {
|
||||
let shards = state.shard_ids();
|
||||
let shards = state.shard_ids().map_err(ClusterAppError)?;
|
||||
crate::scatter_gather::sharded_write_item(
|
||||
state.cluster(),
|
||||
state.cluster().map_err(ClusterAppError)?,
|
||||
EntityId::new(req.entity_id),
|
||||
&req.metadata,
|
||||
&shards,
|
||||
@ -824,9 +901,9 @@ async fn sharded_write_embedding(
|
||||
State(state): State<Arc<ClusterState>>,
|
||||
Json(req): Json<EmbeddingRequest>,
|
||||
) -> std::result::Result<StatusCode, ClusterAppError> {
|
||||
let shards = state.shard_ids();
|
||||
let shards = state.shard_ids().map_err(ClusterAppError)?;
|
||||
crate::scatter_gather::sharded_write_embedding(
|
||||
state.cluster(),
|
||||
state.cluster().map_err(ClusterAppError)?,
|
||||
EntityId::new(req.entity_id),
|
||||
&req.values,
|
||||
&shards,
|
||||
@ -839,9 +916,9 @@ async fn sharded_write_signal(
|
||||
State(state): State<Arc<ClusterState>>,
|
||||
Json(req): Json<SignalRequest>,
|
||||
) -> std::result::Result<StatusCode, ClusterAppError> {
|
||||
let shards = state.shard_ids();
|
||||
let shards = state.shard_ids().map_err(ClusterAppError)?;
|
||||
crate::scatter_gather::sharded_write_signal(
|
||||
state.cluster(),
|
||||
state.cluster().map_err(ClusterAppError)?,
|
||||
&req.signal,
|
||||
EntityId::new(req.entity_id),
|
||||
req.weight,
|
||||
@ -897,8 +974,8 @@ async fn sharded_feed(
|
||||
// The scatter-gather coordinator merges shard results and reads creator
|
||||
// metadata for coordinator-level diversity — all blocking. Offload the
|
||||
// whole coordination off the reactor so a slow shard never stalls it.
|
||||
let cluster = state.cluster_arc();
|
||||
let shards = state.shard_ids();
|
||||
let cluster = state.cluster_arc().map_err(ClusterAppError)?;
|
||||
let shards = state.shard_ids().map_err(ClusterAppError)?;
|
||||
let region_names = state.id_to_name_map().clone();
|
||||
let deadline_ms = query.deadline_ms;
|
||||
let (result, meta) = offload_cluster_read(move || {
|
||||
@ -957,8 +1034,8 @@ async fn sharded_search(
|
||||
|
||||
// Offload the blocking scatter-gather coordination off the reactor (see
|
||||
// [`sharded_feed`] / [`offload_cluster_read`]).
|
||||
let cluster = state.cluster_arc();
|
||||
let shards = state.shard_ids();
|
||||
let cluster = state.cluster_arc().map_err(ClusterAppError)?;
|
||||
let shards = state.shard_ids().map_err(ClusterAppError)?;
|
||||
let region_names = state.id_to_name_map().clone();
|
||||
let deadline_ms = query.deadline_ms;
|
||||
let (result, meta) = offload_cluster_read(move || {
|
||||
@ -987,63 +1064,24 @@ async fn sharded_search(
|
||||
|
||||
// ── Blocking-work offload ────────────────────────────────────────────────────
|
||||
|
||||
/// Run a blocking READ- only cluster query (RETRIEVE / SEARCH) on the tokio
|
||||
/// Run a blocking READ-only cluster query (RETRIEVE / SEARCH) on the tokio
|
||||
/// blocking pool, off the async reactor, and await its result.
|
||||
///
|
||||
/// Thin adapter over the shared [`offload_read`] that maps its [`ServerError`]
|
||||
/// into a [`ClusterAppError`] for the cluster handlers. Both the standalone and
|
||||
/// cluster read paths route through the same helper so they cannot drift.
|
||||
///
|
||||
/// `TidalDb::retrieve`/`search` (and `reload_text_index`) are synchronous,
|
||||
/// CPU/IO-bound calls. Running them directly inside an axum handler would block
|
||||
/// the reactor thread for the whole query, so one slow shard would stall every
|
||||
/// other in-flight request on that worker. Reads never call
|
||||
/// [`GrpcTransport::send_segment`], so unlike [`offload_cluster_write`] they do
|
||||
/// NOT need a runtime-free OS thread — `spawn_blocking` (the purpose-built
|
||||
/// blocking pool) is the right tool and keeps the reactor free.
|
||||
/// CPU/IO-bound calls; running them on the axum reactor would stall every other
|
||||
/// in-flight request behind one slow shard. Reads never call
|
||||
/// [`GrpcTransport::send_segment`], so unlike cluster writes they do NOT need a
|
||||
/// runtime-free OS thread — `spawn_blocking` is the right tool.
|
||||
async fn offload_cluster_read<F, T>(f: F) -> std::result::Result<T, ClusterAppError>
|
||||
where
|
||||
F: FnOnce() -> Result<T> + Send + 'static,
|
||||
T: Send + 'static,
|
||||
{
|
||||
tokio::task::spawn_blocking(f)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
ClusterAppError(ServerError::Cluster(format!(
|
||||
"cluster read worker panicked: {e}"
|
||||
)))
|
||||
})?
|
||||
.map_err(ClusterAppError)
|
||||
}
|
||||
|
||||
/// Run a cluster operation that may ship WAL segments over gRPC on a dedicated
|
||||
/// OS thread, off the async reactor, and await its result.
|
||||
///
|
||||
/// [`GrpcTransport::send_segment`] blocks on its own tokio runtime and asserts
|
||||
/// it is not invoked from within a runtime (`Handle::try_current().is_err()`),
|
||||
/// so it must not run on an axum worker — nor on a `spawn_blocking` thread,
|
||||
/// which still carries a runtime handle and would trip that assert. A fresh
|
||||
/// `std::thread` has no ambient runtime; the result is bridged back through a
|
||||
/// oneshot the handler awaits, so the reactor is never blocked.
|
||||
async fn offload_cluster_write<F, T>(f: F) -> std::result::Result<T, ClusterAppError>
|
||||
where
|
||||
F: FnOnce() -> Result<T> + Send + 'static,
|
||||
T: Send + 'static,
|
||||
{
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
std::thread::Builder::new()
|
||||
.name("cluster-write".into())
|
||||
.spawn(move || {
|
||||
let _ = tx.send(f());
|
||||
})
|
||||
.map_err(|e| {
|
||||
ClusterAppError(ServerError::Cluster(format!(
|
||||
"spawn cluster write worker: {e}"
|
||||
)))
|
||||
})?;
|
||||
rx.await
|
||||
.map_err(|_| {
|
||||
ClusterAppError(ServerError::Cluster(
|
||||
"cluster write worker dropped without responding".into(),
|
||||
))
|
||||
})?
|
||||
.map_err(ClusterAppError)
|
||||
offload_read(f).await.map_err(ClusterAppError)
|
||||
}
|
||||
|
||||
// ── Error handling ─────────────────────────────────────────────────────────
|
||||
|
||||
@ -812,6 +812,63 @@ signals:
|
||||
assert!(profiles.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_config_path_explicit_wins() {
|
||||
// An explicit path always wins, even with a config dir set.
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let explicit = Path::new("/some/explicit/schema.yaml");
|
||||
let resolved =
|
||||
resolve_config_path(Some(explicit), Some(dir.path()), CONFIG_DIR_SCHEMA_FILE).unwrap();
|
||||
assert_eq!(resolved.as_deref(), Some(explicit));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_config_path_no_dir_returns_none() {
|
||||
// No explicit path and no config dir → fall back to the compiled-in
|
||||
// default (None).
|
||||
let resolved = resolve_config_path(None, None, CONFIG_DIR_SCHEMA_FILE).unwrap();
|
||||
assert!(resolved.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_config_path_dir_with_file_resolves() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let file = dir.path().join(CONFIG_DIR_SCHEMA_FILE);
|
||||
std::fs::write(&file, "signals: []").unwrap();
|
||||
let resolved = resolve_config_path(None, Some(dir.path()), CONFIG_DIR_SCHEMA_FILE).unwrap();
|
||||
assert_eq!(resolved.as_deref(), Some(file.as_path()));
|
||||
}
|
||||
|
||||
/// Error branch: a config dir is set but does NOT contain the expected file.
|
||||
/// A misconfigured container must fail loudly, not silently serve defaults.
|
||||
#[test]
|
||||
fn resolve_config_path_dir_missing_file_errors() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
// Dir exists but the schema file was never written into it.
|
||||
let result = resolve_config_path(None, Some(dir.path()), CONFIG_DIR_SCHEMA_FILE);
|
||||
let err = result.unwrap_err();
|
||||
assert!(matches!(err, ServerError::SchemaConfig(_)), "got {err:?}");
|
||||
let msg = err.to_string();
|
||||
assert!(
|
||||
msg.contains(CONFIG_DIR_SCHEMA_FILE) && msg.contains("does not contain"),
|
||||
"error must name the missing file: {msg}"
|
||||
);
|
||||
}
|
||||
|
||||
/// A path that exists but is a directory (not a file) is also rejected — the
|
||||
/// `is_file()` guard must distinguish a readable file from a directory.
|
||||
#[test]
|
||||
fn resolve_config_path_rejects_directory_in_place_of_file() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
// Create a subdirectory named exactly like the expected schema file.
|
||||
std::fs::create_dir(dir.path().join(CONFIG_DIR_SCHEMA_FILE)).unwrap();
|
||||
let result = resolve_config_path(None, Some(dir.path()), CONFIG_DIR_SCHEMA_FILE);
|
||||
assert!(
|
||||
matches!(result, Err(ServerError::SchemaConfig(_))),
|
||||
"a directory is not a readable config file"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn profile_defaults() {
|
||||
let yaml = r"
|
||||
|
||||
@ -25,6 +25,10 @@ pub enum ServerError {
|
||||
ExperimentalDisabled(String),
|
||||
#[error("cluster error: {0}")]
|
||||
Cluster(String),
|
||||
/// The server is shutting down and the cluster fabric has been taken; map to
|
||||
/// 503 so a post-shutdown access degrades cleanly instead of panicking.
|
||||
#[error("service unavailable: {0}")]
|
||||
Unavailable(String),
|
||||
}
|
||||
|
||||
impl ServerError {
|
||||
|
||||
47
tidal-server/src/health.rs
Normal file
47
tidal-server/src/health.rs
Normal file
@ -0,0 +1,47 @@
|
||||
//! Shared health-probe handlers for the standalone and cluster routers.
|
||||
//!
|
||||
//! The standalone router ([`crate::router`]) and the cluster router
|
||||
//! ([`crate::cluster`]) both expose the same Kubernetes-style probe endpoints
|
||||
//! `GET /health/startup` and `GET /health/live`. Their handlers and the JSON
|
||||
//! body they return were previously byte-identical copies in each module; any
|
||||
//! drift between the two (e.g. a renamed field on one side) would silently
|
||||
//! change one mode's probe contract. This module is the single source of truth
|
||||
//! so the two routers cannot diverge, mirroring the [`crate::dto`] anti-drift
|
||||
//! pattern for the data surface.
|
||||
|
||||
use axum::Json;
|
||||
|
||||
/// Shared probe body: `{"ok": true, "service": "tidaldb"}`.
|
||||
///
|
||||
/// Both `/health/startup` and `/health/live` return this fixed body — they are
|
||||
/// unconditional 200s (the process is up and there are no migrations to gate
|
||||
/// readiness on). The richer mode-specific readiness check (`GET /health`)
|
||||
/// stays in each router because it inspects mode-specific state.
|
||||
fn probe_body() -> serde_json::Value {
|
||||
serde_json::json!({ "ok": true, "service": "tidaldb" })
|
||||
}
|
||||
|
||||
/// Startup probe: always 200 (no migrations to check).
|
||||
pub async fn health_startup() -> Json<serde_json::Value> {
|
||||
Json(probe_body())
|
||||
}
|
||||
|
||||
/// Liveness probe: always 200 (process is alive).
|
||||
pub async fn health_live() -> Json<serde_json::Value> {
|
||||
Json(probe_body())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Both probes return the same fixed body, so the standalone and cluster
|
||||
/// routers cannot advertise a different probe contract.
|
||||
#[tokio::test]
|
||||
async fn probes_return_shared_body() {
|
||||
let expected = serde_json::json!({ "ok": true, "service": "tidaldb" });
|
||||
assert_eq!(health_startup().await.0, expected);
|
||||
assert_eq!(health_live().await.0, expected);
|
||||
assert_eq!(probe_body(), expected);
|
||||
}
|
||||
}
|
||||
@ -17,6 +17,8 @@ pub mod cluster;
|
||||
pub mod config;
|
||||
pub mod dto;
|
||||
pub mod error;
|
||||
pub mod health;
|
||||
pub mod offload;
|
||||
pub mod router;
|
||||
pub mod scatter_gather;
|
||||
pub mod state;
|
||||
|
||||
351
tidal-server/src/offload.rs
Normal file
351
tidal-server/src/offload.rs
Normal file
@ -0,0 +1,351 @@
|
||||
//! Blocking-work offload primitives shared by the standalone and cluster HTTP
|
||||
//! surfaces.
|
||||
//!
|
||||
//! Two distinct hazards, two distinct tools:
|
||||
//!
|
||||
//! * **Reads** (`RETRIEVE` / `SEARCH` / text-index reload) are synchronous,
|
||||
//! CPU/IO-bound `TidalDb` calls. Running them inside an axum handler pins the
|
||||
//! reactor worker for the whole query, so a burst can starve every other
|
||||
//! in-flight request. [`offload_read`] hands them to tokio's purpose-built
|
||||
//! blocking pool via `spawn_blocking`. Reads never touch the gRPC transport,
|
||||
//! so a `spawn_blocking` thread (which carries an ambient runtime handle) is
|
||||
//! fine.
|
||||
//!
|
||||
//! * **Cluster writes** (`/signals`, `/cluster/heal`) ship WAL segments to
|
||||
//! followers through [`tidal_net::GrpcTransport`], whose `send_segment` blocks
|
||||
//! on its own runtime and *asserts* it is not invoked from within another
|
||||
//! runtime (`Handle::try_current().is_err()`). That rules out both the axum
|
||||
//! reactor and `spawn_blocking` (its threads carry a runtime handle). The
|
||||
//! previous design spawned a fresh OS thread per request — unbounded growth on
|
||||
//! the hottest cluster path. [`ClusterWritePool`] replaces that with a fixed
|
||||
//! set of runtime-free OS threads draining a bounded queue, so write
|
||||
//! concurrency is capped and threads are reused. A full queue degrades to
|
||||
//! [`TidalError::Backpressure`] (→ 429), matching the engine's own
|
||||
//! flow-control semantics, instead of a hard 500.
|
||||
//!
|
||||
//! Both surfaces route through this module so the standalone and cluster paths
|
||||
//! cannot drift in how they treat blocking work.
|
||||
|
||||
use crossbeam::channel::{Receiver, Sender, TrySendError};
|
||||
use tidaldb::TidalError;
|
||||
|
||||
use crate::error::{Result, ServerError};
|
||||
|
||||
/// Run a blocking READ-only query (RETRIEVE / SEARCH / text-index reload) on
|
||||
/// tokio's blocking pool, off the async reactor, and await its result.
|
||||
///
|
||||
/// `TidalDb::retrieve`/`search`/`reload_text_index` are synchronous and
|
||||
/// CPU/IO-bound; running them directly in an axum handler blocks the reactor
|
||||
/// worker for the whole query, so a burst of feed/search requests can pin every
|
||||
/// worker and stall unrelated requests. Reads never call
|
||||
/// [`tidal_net::GrpcTransport::send_segment`], so unlike cluster writes they do
|
||||
/// NOT need a runtime-free OS thread — `spawn_blocking` is the right tool.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns the closure's own [`ServerError`], or [`ServerError::Cluster`] if the
|
||||
/// blocking task panicked or was cancelled (mapped to a 500 by
|
||||
/// [`crate::router::status_from_error`]).
|
||||
pub async fn offload_read<F, T>(f: F) -> Result<T>
|
||||
where
|
||||
F: FnOnce() -> Result<T> + Send + 'static,
|
||||
T: Send + 'static,
|
||||
{
|
||||
tokio::task::spawn_blocking(f)
|
||||
.await
|
||||
// A JoinError means the blocking task panicked or was cancelled — the
|
||||
// query produced no answer, so surface it as a server error (500).
|
||||
.map_err(|e| ServerError::Cluster(format!("blocking read worker failed: {e}")))?
|
||||
}
|
||||
|
||||
/// Configuration for the cluster write worker pool.
|
||||
///
|
||||
/// Defaults are derived once at startup: `workers` from
|
||||
/// [`std::thread::available_parallelism`] (clamped to a sane floor/ceiling),
|
||||
/// `queue_depth` to a small multiple of `workers` so brief bursts queue rather
|
||||
/// than 429 while a sustained overload still sheds load promptly.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct ClusterWritePoolConfig {
|
||||
/// Number of runtime-free OS worker threads draining the queue.
|
||||
pub workers: usize,
|
||||
/// Maximum number of queued (not-yet-started) write closures before new
|
||||
/// submissions are rejected with backpressure.
|
||||
pub queue_depth: usize,
|
||||
}
|
||||
|
||||
/// Floor on worker count: a single-core host still gets parallel shipping.
|
||||
const MIN_WRITE_WORKERS: usize = 2;
|
||||
/// Ceiling on worker count: cluster writes are gRPC-ship bound, not CPU bound,
|
||||
/// so a large core count does not warrant an unbounded thread set.
|
||||
const MAX_WRITE_WORKERS: usize = 8;
|
||||
/// Queued closures permitted per worker before backpressure trips.
|
||||
const QUEUE_DEPTH_PER_WORKER: usize = 8;
|
||||
|
||||
impl ClusterWritePoolConfig {
|
||||
/// Build a config for an explicit worker count, deriving `queue_depth`
|
||||
/// proportionally so a custom size still buffers brief bursts before 429.
|
||||
#[must_use]
|
||||
pub fn with_workers(workers: usize) -> Self {
|
||||
let workers = workers.max(1);
|
||||
Self {
|
||||
workers,
|
||||
queue_depth: workers.saturating_mul(QUEUE_DEPTH_PER_WORKER),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ClusterWritePoolConfig {
|
||||
fn default() -> Self {
|
||||
let workers = std::thread::available_parallelism()
|
||||
.map_or(MIN_WRITE_WORKERS, std::num::NonZeroUsize::get)
|
||||
.clamp(MIN_WRITE_WORKERS, MAX_WRITE_WORKERS);
|
||||
Self::with_workers(workers)
|
||||
}
|
||||
}
|
||||
|
||||
/// A unit of work for the cluster write pool: a boxed closure plus the oneshot
|
||||
/// used to bridge its result back to the awaiting async handler.
|
||||
type WriteJob = Box<dyn FnOnce() + Send + 'static>;
|
||||
|
||||
/// A fixed-size, runtime-free OS-thread pool for cluster write work.
|
||||
///
|
||||
/// Created once at cluster startup and shared via the cluster state. Each worker
|
||||
/// is a plain `std::thread` with NO ambient tokio runtime, so closures may call
|
||||
/// [`tidal_net::GrpcTransport::send_segment`] (which asserts it is outside a
|
||||
/// runtime). Submissions exceeding [`ClusterWritePoolConfig::queue_depth`] are
|
||||
/// rejected with [`TidalError::Backpressure`] rather than spawning unbounded
|
||||
/// threads.
|
||||
///
|
||||
/// Workers drain the queue until the [`Sender`] is dropped (on
|
||||
/// [`ClusterWritePool::Drop`]), at which point `recv` returns `Err` and each
|
||||
/// worker exits; the threads are then joined so no work is abandoned.
|
||||
pub struct ClusterWritePool {
|
||||
sender: Option<Sender<WriteJob>>,
|
||||
workers: Vec<std::thread::JoinHandle<()>>,
|
||||
}
|
||||
|
||||
impl ClusterWritePool {
|
||||
/// Build the pool and start its worker threads.
|
||||
///
|
||||
/// `workers` and `queue_depth` are clamped to at least 1 so the pool always
|
||||
/// has a live consumer and a bounded queue.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Panics if a worker OS thread cannot be spawned. This runs once at server
|
||||
/// startup (before any request is served), so a thread-exhausted host fails
|
||||
/// loudly at boot rather than degrading every later request — the opposite
|
||||
/// of the old per-request spawn that turned the same failure into a hot-path
|
||||
/// 500.
|
||||
#[must_use]
|
||||
pub fn new(config: ClusterWritePoolConfig) -> Self {
|
||||
let workers = config.workers.max(1);
|
||||
let queue_depth = config.queue_depth.max(1);
|
||||
|
||||
// Bounded so a sustained burst sheds load (429) instead of growing the
|
||||
// queue without limit. crossbeam's MPMC channel lets every worker pull
|
||||
// from the same queue without a shared Mutex<Receiver>.
|
||||
let (sender, receiver) = crossbeam::channel::bounded::<WriteJob>(queue_depth);
|
||||
|
||||
let handles = (0..workers)
|
||||
.map(|i| {
|
||||
let rx: Receiver<WriteJob> = receiver.clone();
|
||||
std::thread::Builder::new()
|
||||
.name(format!("cluster-write-{i}"))
|
||||
.spawn(move || worker_loop(&rx))
|
||||
// A thread that cannot start at construction time is a hard
|
||||
// startup failure, not a per-request degrade. Surfacing it
|
||||
// as a panic here (at process start, before serving) is
|
||||
// acceptable and far better than the old per-request spawn.
|
||||
.expect("spawn cluster write worker thread")
|
||||
})
|
||||
.collect();
|
||||
|
||||
Self {
|
||||
sender: Some(sender),
|
||||
workers: handles,
|
||||
}
|
||||
}
|
||||
|
||||
/// Submit a blocking, runtime-free write closure to the pool and await its
|
||||
/// result.
|
||||
///
|
||||
/// The closure runs on a pool worker thread with no ambient tokio runtime,
|
||||
/// so it may ship WAL segments over gRPC. Its `Result` is bridged back
|
||||
/// through a oneshot the caller awaits, so the reactor is never blocked.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// * [`TidalError::Backpressure`] (→ 429) when the queue is full — the work
|
||||
/// was never enqueued, so the caller may safely retry after backing off.
|
||||
/// * [`ServerError::Cluster`] (→ 500) if the pool has been shut down, or a
|
||||
/// worker dropped the job without responding (e.g. mid-shutdown).
|
||||
pub async fn submit<F, T>(&self, f: F) -> Result<T>
|
||||
where
|
||||
F: FnOnce() -> Result<T> + Send + 'static,
|
||||
T: Send + 'static,
|
||||
{
|
||||
let sender = self
|
||||
.sender
|
||||
.as_ref()
|
||||
.ok_or_else(|| ServerError::Cluster("cluster write pool is shut down".into()))?;
|
||||
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
let job: WriteJob = Box::new(move || {
|
||||
// Ignore send errors: the receiver is only gone if the request
|
||||
// future was cancelled, in which case nobody is waiting.
|
||||
let _ = tx.send(f());
|
||||
});
|
||||
|
||||
match sender.try_send(job) {
|
||||
Ok(()) => {}
|
||||
Err(TrySendError::Full(_)) => {
|
||||
// Bounded queue saturated: degrade to engine-native backpressure
|
||||
// (429) instead of growing threads/queue without limit. The work
|
||||
// was never enqueued, so it is safe to retry.
|
||||
return Err(ServerError::Tidal(TidalError::Backpressure {
|
||||
retry_after_ms: 50,
|
||||
}));
|
||||
}
|
||||
Err(TrySendError::Disconnected(_)) => {
|
||||
return Err(ServerError::Cluster(
|
||||
"cluster write pool has no live workers".into(),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
rx.await.map_err(|_| {
|
||||
ServerError::Cluster("cluster write worker dropped the job without responding".into())
|
||||
})?
|
||||
}
|
||||
}
|
||||
|
||||
/// Drain jobs until the [`Sender`] is dropped, then exit so the thread can be
|
||||
/// joined.
|
||||
fn worker_loop(rx: &Receiver<WriteJob>) {
|
||||
// `recv` blocks until a job arrives, and returns `Err` only once every
|
||||
// `Sender` has been dropped — that is the shutdown signal.
|
||||
while let Ok(job) = rx.recv() {
|
||||
job();
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for ClusterWritePool {
|
||||
/// Drop the [`Sender`] so workers see a disconnected queue and exit, then
|
||||
/// join every worker so no in-flight job is abandoned mid-ship.
|
||||
fn drop(&mut self) {
|
||||
// Dropping the sender disconnects the channel; each worker's `recv`
|
||||
// returns `Err` after draining what it already holds, and the loop exits.
|
||||
self.sender = None;
|
||||
for handle in self.workers.drain(..) {
|
||||
// A worker only panics if a submitted closure panics; that is the
|
||||
// closure author's bug, not the pool's. Log and continue joining the
|
||||
// rest so shutdown still completes.
|
||||
if handle.join().is_err() {
|
||||
tracing::error!("cluster write worker thread panicked during shutdown");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[allow(clippy::unwrap_used)]
|
||||
mod tests {
|
||||
use std::sync::{
|
||||
Arc,
|
||||
atomic::{AtomicUsize, Ordering},
|
||||
};
|
||||
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn pool_processes_submitted_work() {
|
||||
let pool = ClusterWritePool::new(ClusterWritePoolConfig {
|
||||
workers: 2,
|
||||
queue_depth: 8,
|
||||
});
|
||||
let counter = Arc::new(AtomicUsize::new(0));
|
||||
|
||||
let mut results = Vec::new();
|
||||
for i in 0..16u64 {
|
||||
let c = Arc::clone(&counter);
|
||||
results.push(
|
||||
pool.submit(move || {
|
||||
c.fetch_add(1, Ordering::SeqCst);
|
||||
Ok::<u64, ServerError>(i * 2)
|
||||
})
|
||||
.await,
|
||||
);
|
||||
}
|
||||
|
||||
for (i, r) in results.into_iter().enumerate() {
|
||||
assert_eq!(r.unwrap(), (i as u64) * 2);
|
||||
}
|
||||
assert_eq!(counter.load(Ordering::SeqCst), 16);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pool_runs_off_any_tokio_runtime() {
|
||||
// The whole point of the pool: closures execute with NO ambient runtime
|
||||
// handle, which is what `GrpcTransport::send_segment` asserts.
|
||||
let pool = ClusterWritePool::new(ClusterWritePoolConfig {
|
||||
workers: 1,
|
||||
queue_depth: 4,
|
||||
});
|
||||
let outside = pool
|
||||
.submit(|| Ok::<bool, ServerError>(tokio::runtime::Handle::try_current().is_err()))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
outside,
|
||||
"pool worker must not carry an ambient tokio runtime"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn full_queue_yields_backpressure() {
|
||||
// One worker, queue depth of 1: occupy the worker with a job that parks
|
||||
// until released, fill the single queue slot directly, then prove the
|
||||
// next submit is rejected with Backpressure (→ 429) — not a 500 and not
|
||||
// an unbounded thread/queue grow.
|
||||
let pool = ClusterWritePool::new(ClusterWritePoolConfig {
|
||||
workers: 1,
|
||||
queue_depth: 1,
|
||||
});
|
||||
|
||||
// Block the sole worker on a channel until the test releases it. We
|
||||
// submit the parking job directly through the pool's own sender so the
|
||||
// worker picks it up, then wait for an explicit "started" signal — no
|
||||
// sleep-based timing.
|
||||
let sender = pool.sender.as_ref().unwrap().clone();
|
||||
let (started_tx, started_rx) = std::sync::mpsc::channel::<()>();
|
||||
let (release_tx, release_rx) = std::sync::mpsc::channel::<()>();
|
||||
let park: WriteJob = Box::new(move || {
|
||||
// Tell the test the worker is now busy, then block until released.
|
||||
started_tx.send(()).unwrap();
|
||||
release_rx.recv().unwrap();
|
||||
});
|
||||
sender
|
||||
.try_send(park)
|
||||
.expect("worker accepts the parking job");
|
||||
// Deterministic handoff: the worker has dequeued and started `park`, so
|
||||
// the single queue slot is empty again and the worker is occupied.
|
||||
started_rx
|
||||
.recv_timeout(std::time::Duration::from_secs(5))
|
||||
.expect("worker started the parking job");
|
||||
|
||||
// Job B fills the one free queue slot (worker is busy on `park`).
|
||||
let job_b: WriteJob = Box::new(|| {});
|
||||
sender.try_send(job_b).expect("queue slot accepts job B");
|
||||
|
||||
// Job C: queue is now full -> Backpressure (429), via the real submit path.
|
||||
let c = pool.submit(|| Ok::<(), ServerError>(())).await;
|
||||
match c {
|
||||
Err(ServerError::Tidal(TidalError::Backpressure { .. })) => {}
|
||||
other => panic!("expected Backpressure, got {other:?}"),
|
||||
}
|
||||
|
||||
// Release the worker so the pool drains and shuts down cleanly.
|
||||
release_tx.send(()).unwrap();
|
||||
}
|
||||
}
|
||||
@ -81,11 +81,13 @@ impl MakeRequestId for SequentialRequestId {
|
||||
/// Keeping `/health` outside the timeout/concurrency layers means health probes
|
||||
/// are never queued or timed out under saturation, preventing false liveness failures.
|
||||
pub fn build_router(state: Arc<ServerState>, api_key: Option<Arc<str>>) -> Router {
|
||||
// Public routes — exempt from auth so health probes always work.
|
||||
// Public routes — exempt from auth so health probes always work. The
|
||||
// startup/live probes are shared with the cluster router via [`crate::health`]
|
||||
// so the two modes can never advertise a different probe contract.
|
||||
let public = Router::new()
|
||||
.route("/health", get(health))
|
||||
.route("/health/startup", get(health_startup))
|
||||
.route("/health/live", get(health_live))
|
||||
.route("/health/startup", get(crate::health::health_startup))
|
||||
.route("/health/live", get(crate::health::health_live))
|
||||
.with_state(Arc::clone(&state));
|
||||
|
||||
// Protected routes — gated by Bearer token when a key is configured.
|
||||
@ -232,9 +234,17 @@ async fn feed(
|
||||
}
|
||||
let retrieve = builder.build().map_err(|e| TidalErrorWrapper(e.into()))?;
|
||||
|
||||
let result = state
|
||||
.retrieve(query.region.as_deref(), &retrieve)
|
||||
.map_err(AppError)?;
|
||||
// `retrieve` is a synchronous, CPU-bound engine call (candidate scan +
|
||||
// scoring + diversity). Running it inline would pin this reactor worker for
|
||||
// the whole query, so a burst of feeds could starve every other request.
|
||||
// Offload to the blocking pool — the same treatment the cluster path gives
|
||||
// the identical work (see [`crate::offload`]).
|
||||
let offload_state = Arc::clone(&state);
|
||||
let region = query.region.clone();
|
||||
let result =
|
||||
crate::offload::offload_read(move || offload_state.retrieve(region.as_deref(), &retrieve))
|
||||
.await
|
||||
.map_err(AppError)?;
|
||||
|
||||
Ok(Json(FeedResponse {
|
||||
items: feed_items(&result.items),
|
||||
@ -252,12 +262,19 @@ async fn search(
|
||||
builder = builder.for_user(user_id);
|
||||
}
|
||||
let search = builder.build().map_err(|e| TidalErrorWrapper(e.into()))?;
|
||||
state
|
||||
.reload_text_index(query.region.as_deref())
|
||||
.map_err(AppError)?;
|
||||
let result = state
|
||||
.search(query.region.as_deref(), &search)
|
||||
.map_err(AppError)?;
|
||||
|
||||
// The text-index reload and the search are both synchronous, CPU/IO-bound
|
||||
// engine calls (tantivy reload + BM25/ANN/RRF pipeline + scoring +
|
||||
// diversity). Offload the pair together so the reactor stays free — see
|
||||
// [`crate::offload`].
|
||||
let offload_state = Arc::clone(&state);
|
||||
let region = query.region.clone();
|
||||
let result = crate::offload::offload_read(move || {
|
||||
offload_state.reload_text_index(region.as_deref())?;
|
||||
offload_state.search(region.as_deref(), &search)
|
||||
})
|
||||
.await
|
||||
.map_err(AppError)?;
|
||||
|
||||
Ok(Json(SearchResponse {
|
||||
items: search_items(&result.items),
|
||||
@ -266,16 +283,6 @@ async fn search(
|
||||
}))
|
||||
}
|
||||
|
||||
/// Startup probe: always 200 (no migrations to check).
|
||||
async fn health_startup() -> Json<serde_json::Value> {
|
||||
Json(serde_json::json!({ "ok": true, "service": "tidaldb" }))
|
||||
}
|
||||
|
||||
/// Liveness probe: always 200 (process is alive).
|
||||
async fn health_live() -> Json<serde_json::Value> {
|
||||
Json(serde_json::json!({ "ok": true, "service": "tidaldb" }))
|
||||
}
|
||||
|
||||
/// Readiness probe: 200 when ready, 503 when shutting down.
|
||||
async fn health(
|
||||
State(state): State<Arc<ServerState>>,
|
||||
@ -330,6 +337,7 @@ impl IntoResponse for AppError {
|
||||
pub const fn status_from_error(err: &ServerError) -> StatusCode {
|
||||
match err {
|
||||
ServerError::BadRequest(_) | ServerError::SchemaConfig(_) => StatusCode::BAD_REQUEST,
|
||||
ServerError::Unavailable(_) => StatusCode::SERVICE_UNAVAILABLE,
|
||||
ServerError::Tidal(tidal_err) => match tidal_err {
|
||||
tidaldb::TidalError::NotFound { .. } => StatusCode::NOT_FOUND,
|
||||
tidaldb::TidalError::Schema(_) | tidaldb::TidalError::InvalidInput(_) => {
|
||||
@ -345,3 +353,29 @@ pub const fn status_from_error(err: &ServerError) -> StatusCode {
|
||||
_ => StatusCode::INTERNAL_SERVER_ERROR,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// A post-shutdown cluster access surfaces as [`ServerError::Unavailable`],
|
||||
/// which must map to 503 (clean degradation) rather than a 500 or a panic —
|
||||
/// the load-bearing behavior of the `cluster_ref`/`cluster_arc` Result change.
|
||||
#[test]
|
||||
fn unavailable_maps_to_503() {
|
||||
let err = ServerError::Unavailable("server shutting down".into());
|
||||
assert_eq!(status_from_error(&err), StatusCode::SERVICE_UNAVAILABLE);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bad_request_maps_to_400() {
|
||||
let err = ServerError::BadRequest("bad".into());
|
||||
assert_eq!(status_from_error(&err), StatusCode::BAD_REQUEST);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cluster_error_maps_to_500() {
|
||||
let err = ServerError::Cluster("boom".into());
|
||||
assert_eq!(status_from_error(&err), StatusCode::INTERNAL_SERVER_ERROR);
|
||||
}
|
||||
}
|
||||
|
||||
@ -9,6 +9,16 @@
|
||||
//! - **Scatter-gather SEARCH**: fan out to all shards, merge by score
|
||||
//! - **Deadline propagation**: configurable budget with per-hop overhead subtracted
|
||||
//! - **Partial failure**: degraded results when some shards are unreachable
|
||||
//!
|
||||
//! # Size / cohesion note (M0–M10 review Maintainability-S)
|
||||
//!
|
||||
//! This file is large (~1.5k lines), but roughly half is its in-module test
|
||||
//! suite and the production code is a single tightly-coupled concern: the
|
||||
//! fan-out → gather → dedup → reconcile → diversity → truncate merge pipeline,
|
||||
//! whose helpers share the [`Sourced`]/[`GatherState`]/[`MergeItem`] types and
|
||||
//! must stay consistent. Splitting it would fragment that pipeline without
|
||||
//! reducing real complexity, so the split is deferred (not forced mid-campaign).
|
||||
//! Tracked for a dedicated follow-up.
|
||||
|
||||
use std::{
|
||||
collections::{HashMap, HashSet},
|
||||
@ -34,6 +44,38 @@ const DEFAULT_DEADLINE_MS: u64 = 50;
|
||||
/// Estimated network overhead per shard hop (subtracted from deadline).
|
||||
const NETWORK_OVERHEAD_MS: u64 = 5;
|
||||
|
||||
/// Server-side ceiling on a client-supplied `deadline_ms`.
|
||||
///
|
||||
/// The scatter-gather budget is fully client-controlled (the `?deadline_ms=`
|
||||
/// query param flows straight into [`scatter_gather_retrieve`] /
|
||||
/// [`scatter_gather_search`]). An unbounded value lets a single request pin a
|
||||
/// blocking-pool worker for an arbitrarily long time — a trivial resource-
|
||||
/// exhaustion / slow-loris vector against the whole node. We clamp every
|
||||
/// request to this ceiling (10s), which is already two orders of magnitude
|
||||
/// above the 50ms spec target, so it never constrains a legitimate query while
|
||||
/// capping the worst case. See [`clamp_deadline_ms`].
|
||||
const MAX_DEADLINE_MS: u64 = 10_000;
|
||||
|
||||
/// Clamp a client-supplied scatter-gather deadline to [`MAX_DEADLINE_MS`].
|
||||
///
|
||||
/// `None` keeps the [`DEFAULT_DEADLINE_MS`] default. Any explicit value above
|
||||
/// the ceiling is logged once and reduced, so a malicious or buggy client
|
||||
/// cannot hold a worker indefinitely.
|
||||
fn clamp_deadline_ms(requested: Option<u64>) -> u64 {
|
||||
match requested {
|
||||
None => DEFAULT_DEADLINE_MS,
|
||||
Some(ms) if ms > MAX_DEADLINE_MS => {
|
||||
tracing::warn!(
|
||||
requested_ms = ms,
|
||||
cap_ms = MAX_DEADLINE_MS,
|
||||
"scatter-gather deadline_ms exceeds server cap; clamping"
|
||||
);
|
||||
MAX_DEADLINE_MS
|
||||
}
|
||||
Some(ms) => ms,
|
||||
}
|
||||
}
|
||||
|
||||
/// Metadata about scatter-gather query execution.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ScatterGatherMeta {
|
||||
@ -162,9 +204,29 @@ struct ShardOutcome<T> {
|
||||
total_candidates: usize,
|
||||
}
|
||||
|
||||
/// A merged item tagged with the shard ([`RegionId`]) that actually returned
|
||||
/// it.
|
||||
///
|
||||
/// Coordinator-level diversity resolves each item's creator from `region` —
|
||||
/// the node that owns the item in an entity-sharded topology — rather than from
|
||||
/// a single fixed replica. In a replicated topology every shard holds every
|
||||
/// entity, so `region` is simply one healthy replica that has the item; in an
|
||||
/// entity-sharded topology it is the *only* node that has it. Either way the
|
||||
/// creator lookup hits a node that actually stores the item. See
|
||||
/// [`enforce_max_per_creator`].
|
||||
struct Sourced<T> {
|
||||
/// The shard that returned this item (and therefore can resolve its
|
||||
/// metadata locally).
|
||||
region: RegionId,
|
||||
item: T,
|
||||
}
|
||||
|
||||
/// Per-shard query state accumulated by [`dispatch_shards`].
|
||||
struct GatherState<T> {
|
||||
items: Vec<T>,
|
||||
/// Each merged item paired with the shard that returned it, so the merge
|
||||
/// stage can resolve creators from the owning node (entity-sharded) or a
|
||||
/// holding replica (replicated). See [`Sourced`].
|
||||
items: Vec<Sourced<T>>,
|
||||
/// Each contributing shard's `total_candidates`, kept separately so the
|
||||
/// merge stage can reconcile the count instead of blindly summing — a sum
|
||||
/// double-counts the candidate universe across REPLICATED shards (every
|
||||
@ -339,7 +401,15 @@ fn fold_outcome<T>(
|
||||
match outcome {
|
||||
Ok(outcome) => {
|
||||
state.per_shard_totals.push(outcome.total_candidates);
|
||||
state.items.extend(outcome.items);
|
||||
// Tag every item with the shard that returned it so coordinator-
|
||||
// level diversity can resolve its creator from a node that actually
|
||||
// stores the item (the owning shard in entity-sharded mode).
|
||||
state
|
||||
.items
|
||||
.extend(outcome.items.into_iter().map(|item| Sourced {
|
||||
region: shard,
|
||||
item,
|
||||
}));
|
||||
state.shards_queried += 1;
|
||||
}
|
||||
Err(e) => {
|
||||
@ -385,24 +455,27 @@ impl MergeItem for tidaldb::query::search::SearchResultItem {
|
||||
///
|
||||
/// In a replicated topology every shard holds every entity, so the same entity
|
||||
/// can appear once per shard. Without this the merged result would list the
|
||||
/// same item up to `num_shards` times. Returns the deduped item vector and
|
||||
/// `true` if any duplicate was collapsed (i.e. shards overlapped) — the caller
|
||||
/// uses that signal to reconcile `total_candidates`.
|
||||
fn dedup_by_entity<T: MergeItem>(items: Vec<T>) -> (Vec<T>, bool) {
|
||||
/// same item up to `num_shards` times. The surviving copy carries its own
|
||||
/// source [`RegionId`] (the shard that returned the best-scoring copy), which a
|
||||
/// later creator lookup uses — that node demonstrably holds the entity, so the
|
||||
/// metadata read can never miss. Returns the deduped item vector and `true` if
|
||||
/// any duplicate was collapsed (i.e. shards overlapped) — the caller uses that
|
||||
/// signal to reconcile `total_candidates`.
|
||||
fn dedup_by_entity<T: MergeItem>(items: Vec<Sourced<T>>) -> (Vec<Sourced<T>>, bool) {
|
||||
// entity_id → index of the best-scoring copy seen so far in `out`.
|
||||
let mut best: HashMap<u64, usize> = HashMap::with_capacity(items.len());
|
||||
let mut out: Vec<T> = Vec::with_capacity(items.len());
|
||||
let mut out: Vec<Sourced<T>> = Vec::with_capacity(items.len());
|
||||
let mut overlap = false;
|
||||
for item in items {
|
||||
let key = item.entity_id().as_u64();
|
||||
for sourced in items {
|
||||
let key = sourced.item.entity_id().as_u64();
|
||||
if let Some(idx) = best.get(&key).copied() {
|
||||
overlap = true;
|
||||
if item.score() > out[idx].score() {
|
||||
out[idx] = item;
|
||||
if sourced.item.score() > out[idx].item.score() {
|
||||
out[idx] = sourced;
|
||||
}
|
||||
} else {
|
||||
best.insert(key, out.len());
|
||||
out.push(item);
|
||||
out.push(sourced);
|
||||
}
|
||||
}
|
||||
(out, overlap)
|
||||
@ -445,20 +518,23 @@ fn reconcile_total_candidates(
|
||||
/// hit the cap, mirroring the engine's per-shard `DiversitySelector` but at the
|
||||
/// coordinator level.
|
||||
///
|
||||
/// Creators are resolved from the `creator_id` metadata of the read region
|
||||
/// (any healthy replica has all entities; an entity-sharded coordinator reads
|
||||
/// each item from whichever live shard returned it via the cluster's item
|
||||
/// store). Items whose creator cannot be resolved (no `creator_id`, or a
|
||||
/// metadata read error) are treated as having a unique, uncapped creator —
|
||||
/// never dropped — so a metadata gap can only ever UNDER-enforce, never hide a
|
||||
/// result.
|
||||
/// Creators are resolved from the `creator_id` metadata of **the shard that
|
||||
/// actually returned each item** ([`Sourced::region`]), not from one fixed
|
||||
/// replica. This is the fix for the entity-sharded under-enforcement bug: when
|
||||
/// items live on disjoint shards, the leader does not hold the ones owned by
|
||||
/// other shards, so a leader-only lookup resolved their creator to `None` and
|
||||
/// silently treated a prolific creator's items as uncapped. Reading from the
|
||||
/// returning shard guarantees the metadata read hits a node that stores the
|
||||
/// item, in both replicated and entity-sharded topologies. Items whose creator
|
||||
/// still cannot be resolved (no `creator_id`, or a metadata read error) are
|
||||
/// treated as having a unique, uncapped creator — never dropped — so a genuine
|
||||
/// metadata gap can only ever UNDER-enforce, never hide a result.
|
||||
///
|
||||
/// Returns the number of items dropped, so the caller can flag the result as
|
||||
/// not fully constraint-satisfied.
|
||||
fn enforce_max_per_creator<T: MergeItem>(
|
||||
cluster: &SimulatedCluster,
|
||||
read_region: RegionId,
|
||||
items: &mut Vec<T>,
|
||||
items: &mut Vec<Sourced<T>>,
|
||||
max_per_creator: usize,
|
||||
) -> usize {
|
||||
if max_per_creator == 0 {
|
||||
@ -467,13 +543,16 @@ fn enforce_max_per_creator<T: MergeItem>(
|
||||
return 0;
|
||||
}
|
||||
|
||||
let db = &cluster.node(read_region).db;
|
||||
let mut per_creator: HashMap<u64, usize> = HashMap::new();
|
||||
let mut dropped = 0usize;
|
||||
|
||||
items.retain(|item| {
|
||||
let creator = db
|
||||
.get_item_metadata(item.entity_id())
|
||||
items.retain(|sourced| {
|
||||
// Resolve from the node that returned this item, so an entity-sharded
|
||||
// owner's metadata is always reachable.
|
||||
let creator = cluster
|
||||
.node(sourced.region)
|
||||
.db
|
||||
.get_item_metadata(sourced.item.entity_id())
|
||||
.ok()
|
||||
.flatten()
|
||||
.and_then(|meta| meta.get("creator_id").and_then(|c| c.parse::<u64>().ok()));
|
||||
@ -507,7 +586,8 @@ fn enforce_max_per_creator<T: MergeItem>(
|
||||
/// over-counted the candidate universe by up to `num_shards`x.
|
||||
/// 3. **Re-enforces `max_per_creator`** across the merged set when the query
|
||||
/// declares it (see [`enforce_max_per_creator`]) — each shard only enforces
|
||||
/// diversity over its own slice.
|
||||
/// diversity over its own slice, and each item's creator is resolved from
|
||||
/// the shard that returned it so entity-sharded owners are reachable.
|
||||
/// 4. Sorts by score descending and takes the top-K.
|
||||
///
|
||||
/// Shards that error or miss the deadline are reported as degraded — never
|
||||
@ -529,10 +609,9 @@ pub fn scatter_gather_retrieve(
|
||||
deadline_ms: Option<u64>,
|
||||
) -> Result<(RetrieveResults, ScatterGatherMeta)> {
|
||||
let start = Instant::now();
|
||||
let total_deadline = Duration::from_millis(deadline_ms.unwrap_or(DEFAULT_DEADLINE_MS));
|
||||
let shard_deadline_ms = deadline_ms
|
||||
.unwrap_or(DEFAULT_DEADLINE_MS)
|
||||
.saturating_sub(NETWORK_OVERHEAD_MS);
|
||||
let budget_ms = clamp_deadline_ms(deadline_ms);
|
||||
let total_deadline = Duration::from_millis(budget_ms);
|
||||
let shard_deadline_ms = budget_ms.saturating_sub(NETWORK_OVERHEAD_MS);
|
||||
|
||||
// Each detached worker holds an owned clone of the query, so it can outlive
|
||||
// this stack frame if the shard's blocking read exceeds the deadline.
|
||||
@ -563,18 +642,18 @@ pub fn scatter_gather_retrieve(
|
||||
|
||||
// Sort the deduped set by score descending.
|
||||
all_items.sort_by(|a, b| {
|
||||
b.score
|
||||
.partial_cmp(&a.score)
|
||||
b.item
|
||||
.score
|
||||
.partial_cmp(&a.item.score)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
});
|
||||
|
||||
// Re-enforce max-per-creator across the merged set (each shard only
|
||||
// enforced over its own slice). Resolve creators from the leader replica.
|
||||
// enforced over its own slice). Each item's creator is resolved from the
|
||||
// shard that returned it, so an entity-sharded owner's metadata is reachable.
|
||||
let mut constraints_satisfied = true;
|
||||
if let Some(max_per_creator) = query.diversity.as_ref().and_then(|d| d.max_per_creator) {
|
||||
let read_region = cluster.leader_region();
|
||||
let dropped =
|
||||
enforce_max_per_creator(cluster, read_region, &mut all_items, max_per_creator);
|
||||
let dropped = enforce_max_per_creator(cluster, &mut all_items, max_per_creator);
|
||||
if dropped > 0 {
|
||||
constraints_satisfied = false;
|
||||
}
|
||||
@ -585,9 +664,11 @@ pub fn scatter_gather_retrieve(
|
||||
let total_candidates =
|
||||
reconcile_total_candidates(&per_shard_totals, all_items.len(), overlap_detected);
|
||||
|
||||
// Take top limit and re-rank (assign 1-based ranks after merge).
|
||||
// Take top limit, drop the source-region tags, and re-rank (assign 1-based
|
||||
// ranks after merge).
|
||||
let limit = query.limit.min(all_items.len());
|
||||
all_items.truncate(limit);
|
||||
let mut all_items: Vec<RetrieveResult> = all_items.into_iter().map(|s| s.item).collect();
|
||||
for (i, item) in all_items.iter_mut().enumerate() {
|
||||
item.rank = i + 1;
|
||||
}
|
||||
@ -652,10 +733,9 @@ pub fn scatter_gather_search(
|
||||
deadline_ms: Option<u64>,
|
||||
) -> Result<(SearchResults, ScatterGatherMeta)> {
|
||||
let start = Instant::now();
|
||||
let total_deadline = Duration::from_millis(deadline_ms.unwrap_or(DEFAULT_DEADLINE_MS));
|
||||
let shard_deadline_ms = deadline_ms
|
||||
.unwrap_or(DEFAULT_DEADLINE_MS)
|
||||
.saturating_sub(NETWORK_OVERHEAD_MS);
|
||||
let budget_ms = clamp_deadline_ms(deadline_ms);
|
||||
let total_deadline = Duration::from_millis(budget_ms);
|
||||
let shard_deadline_ms = budget_ms.saturating_sub(NETWORK_OVERHEAD_MS);
|
||||
|
||||
// Each detached worker holds an owned clone of the query, so it can outlive
|
||||
// this stack frame if the shard's blocking search exceeds the deadline.
|
||||
@ -690,18 +770,18 @@ pub fn scatter_gather_search(
|
||||
|
||||
// Sort the deduped set by score descending.
|
||||
all_items.sort_by(|a, b| {
|
||||
b.score
|
||||
.partial_cmp(&a.score)
|
||||
b.item
|
||||
.score
|
||||
.partial_cmp(&a.item.score)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
});
|
||||
|
||||
// Re-enforce max-per-creator across the merged set (each shard only
|
||||
// enforced over its own slice). Resolve creators from the leader replica.
|
||||
// enforced over its own slice). Each item's creator is resolved from the
|
||||
// shard that returned it, so an entity-sharded owner's metadata is reachable.
|
||||
let mut constraints_satisfied = true;
|
||||
if let Some(max_per_creator) = query.diversity.as_ref().and_then(|d| d.max_per_creator) {
|
||||
let read_region = cluster.leader_region();
|
||||
let dropped =
|
||||
enforce_max_per_creator(cluster, read_region, &mut all_items, max_per_creator);
|
||||
let dropped = enforce_max_per_creator(cluster, &mut all_items, max_per_creator);
|
||||
if dropped > 0 {
|
||||
constraints_satisfied = false;
|
||||
}
|
||||
@ -714,6 +794,8 @@ pub fn scatter_gather_search(
|
||||
|
||||
let limit = query.limit as usize;
|
||||
all_items.truncate(limit);
|
||||
let mut all_items: Vec<tidaldb::query::search::SearchResultItem> =
|
||||
all_items.into_iter().map(|s| s.item).collect();
|
||||
for (i, item) in all_items.iter_mut().enumerate() {
|
||||
item.rank = i + 1;
|
||||
}
|
||||
@ -754,8 +836,12 @@ pub fn scatter_gather_search(
|
||||
|
||||
#[cfg(test)]
|
||||
// Test exemptions: unwrap on known-good fixtures + loop-counter casts in
|
||||
// distribution math are idiomatic here.
|
||||
#[allow(clippy::unwrap_used, clippy::cast_precision_loss)]
|
||||
// distribution / sharding math are idiomatic here.
|
||||
#[allow(
|
||||
clippy::unwrap_used,
|
||||
clippy::cast_precision_loss,
|
||||
clippy::cast_possible_truncation
|
||||
)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
@ -943,6 +1029,44 @@ mod tests {
|
||||
assert_eq!(meta.shard_deadline_ms, 0, "should saturate at 0");
|
||||
}
|
||||
|
||||
/// A client-supplied `deadline_ms` is clamped to the server-side ceiling so
|
||||
/// a single request cannot pin a worker indefinitely.
|
||||
#[test]
|
||||
fn deadline_ms_is_clamped_to_server_cap() {
|
||||
// None → default budget.
|
||||
assert_eq!(clamp_deadline_ms(None), DEFAULT_DEADLINE_MS);
|
||||
// Below the cap → passed through unchanged.
|
||||
assert_eq!(clamp_deadline_ms(Some(250)), 250);
|
||||
// Exactly at the cap → unchanged.
|
||||
assert_eq!(clamp_deadline_ms(Some(MAX_DEADLINE_MS)), MAX_DEADLINE_MS);
|
||||
// Above the cap → clamped down.
|
||||
assert_eq!(
|
||||
clamp_deadline_ms(Some(MAX_DEADLINE_MS + 1)),
|
||||
MAX_DEADLINE_MS
|
||||
);
|
||||
assert_eq!(clamp_deadline_ms(Some(u64::MAX)), MAX_DEADLINE_MS);
|
||||
}
|
||||
|
||||
/// An over-budget `deadline_ms` flowing through the full retrieve path is
|
||||
/// clamped: the reported `shard_deadline_ms` reflects the capped budget, not
|
||||
/// the (absurd) requested one.
|
||||
#[test]
|
||||
fn scatter_gather_retrieve_clamps_oversized_deadline() {
|
||||
let (cluster, shards, names) = four_region_cluster();
|
||||
let retrieve = tidaldb::query::retrieve::Retrieve::builder()
|
||||
.profile("trending")
|
||||
.limit(5)
|
||||
.build()
|
||||
.unwrap();
|
||||
let (_result, meta) =
|
||||
scatter_gather_retrieve(&cluster, &retrieve, &shards, &names, Some(u64::MAX)).unwrap();
|
||||
assert_eq!(
|
||||
meta.shard_deadline_ms,
|
||||
MAX_DEADLINE_MS - NETWORK_OVERHEAD_MS,
|
||||
"oversized deadline must be clamped before overhead subtraction"
|
||||
);
|
||||
}
|
||||
|
||||
/// Scatter-gather search works across shards.
|
||||
#[test]
|
||||
fn scatter_gather_search_merges_across_shards() {
|
||||
@ -1098,6 +1222,15 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
/// Wrap a [`RetrieveResult`] as if it were returned by `region`, for the
|
||||
/// diversity/merge unit tests that exercise [`Sourced`]-keyed paths.
|
||||
fn sourced(region: RegionId, entity_id: u64, score: f64) -> Sourced<RetrieveResult> {
|
||||
Sourced {
|
||||
region,
|
||||
item: retrieve_result(entity_id, score),
|
||||
}
|
||||
}
|
||||
|
||||
/// `entity_shard` must agree with the engine `ShardRouter::hash` (FNV-1a)
|
||||
/// for EVERY entity, so server-side write/read routing can never disagree
|
||||
/// with the engine's own mapping. A divergent hash silently sends the same
|
||||
@ -1122,26 +1255,29 @@ mod tests {
|
||||
#[test]
|
||||
fn dedup_collapses_replicated_copies_keeping_best_score() {
|
||||
// Entity 1 returned by three replicas with different scores; entity 2
|
||||
// by two replicas; entity 3 once.
|
||||
// by two replicas; entity 3 once. The best-scoring copy's source region
|
||||
// must survive (it is the one a later creator lookup reads from).
|
||||
let items = vec![
|
||||
retrieve_result(1, 0.5),
|
||||
retrieve_result(1, 0.9), // best for entity 1
|
||||
retrieve_result(1, 0.7),
|
||||
retrieve_result(2, 0.3),
|
||||
retrieve_result(2, 0.4), // best for entity 2
|
||||
retrieve_result(3, 0.8),
|
||||
sourced(RegionId(0), 1, 0.5),
|
||||
sourced(RegionId(1), 1, 0.9), // best for entity 1, from region 1
|
||||
sourced(RegionId(2), 1, 0.7),
|
||||
sourced(RegionId(0), 2, 0.3),
|
||||
sourced(RegionId(3), 2, 0.4), // best for entity 2, from region 3
|
||||
sourced(RegionId(2), 3, 0.8),
|
||||
];
|
||||
let (deduped, overlap) = dedup_by_entity(items);
|
||||
assert!(overlap, "duplicates across replicas must set overlap=true");
|
||||
assert_eq!(deduped.len(), 3, "one row per distinct entity");
|
||||
|
||||
let mut by_id: HashMap<u64, f64> = HashMap::new();
|
||||
for item in &deduped {
|
||||
by_id.insert(item.entity_id.as_u64(), item.score);
|
||||
let mut by_id: HashMap<u64, (f64, RegionId)> = HashMap::new();
|
||||
for s in &deduped {
|
||||
by_id.insert(s.item.entity_id.as_u64(), (s.item.score, s.region));
|
||||
}
|
||||
assert!((by_id[&1] - 0.9).abs() < f64::EPSILON, "kept best for e1");
|
||||
assert!((by_id[&2] - 0.4).abs() < f64::EPSILON, "kept best for e2");
|
||||
assert!((by_id[&3] - 0.8).abs() < f64::EPSILON);
|
||||
assert!((by_id[&1].0 - 0.9).abs() < f64::EPSILON, "kept best for e1");
|
||||
assert_eq!(by_id[&1].1, RegionId(1), "kept the winning copy's region");
|
||||
assert!((by_id[&2].0 - 0.4).abs() < f64::EPSILON, "kept best for e2");
|
||||
assert_eq!(by_id[&2].1, RegionId(3), "kept the winning copy's region");
|
||||
assert!((by_id[&3].0 - 0.8).abs() < f64::EPSILON);
|
||||
}
|
||||
|
||||
/// Disjoint (entity-sharded) shards never return the same entity, so dedup
|
||||
@ -1149,9 +1285,9 @@ mod tests {
|
||||
#[test]
|
||||
fn dedup_no_overlap_for_disjoint_shards() {
|
||||
let items = vec![
|
||||
retrieve_result(10, 0.5),
|
||||
retrieve_result(20, 0.6),
|
||||
retrieve_result(30, 0.7),
|
||||
sourced(RegionId(0), 10, 0.5),
|
||||
sourced(RegionId(1), 20, 0.6),
|
||||
sourced(RegionId(2), 30, 0.7),
|
||||
];
|
||||
let (deduped, overlap) = dedup_by_entity(items);
|
||||
assert!(!overlap, "disjoint shards must report overlap=false");
|
||||
@ -1267,6 +1403,7 @@ mod tests {
|
||||
#[test]
|
||||
fn enforce_max_per_creator_zero_is_no_cap() {
|
||||
let (cluster, _shards, _names) = four_region_cluster();
|
||||
let leader = cluster.leader_region();
|
||||
let mut meta = HashMap::new();
|
||||
meta.insert("creator_id".to_string(), "7".to_string());
|
||||
for i in 1..=3u64 {
|
||||
@ -1275,11 +1412,11 @@ mod tests {
|
||||
.unwrap();
|
||||
}
|
||||
let mut items = vec![
|
||||
retrieve_result(1, 0.9),
|
||||
retrieve_result(2, 0.8),
|
||||
retrieve_result(3, 0.7),
|
||||
sourced(leader, 1, 0.9),
|
||||
sourced(leader, 2, 0.8),
|
||||
sourced(leader, 3, 0.7),
|
||||
];
|
||||
let dropped = enforce_max_per_creator(&cluster, cluster.leader_region(), &mut items, 0);
|
||||
let dropped = enforce_max_per_creator(&cluster, &mut items, 0);
|
||||
assert_eq!(dropped, 0, "zero cap must not drop anything");
|
||||
assert_eq!(items.len(), 3);
|
||||
}
|
||||
@ -1289,17 +1426,106 @@ mod tests {
|
||||
#[test]
|
||||
fn enforce_max_per_creator_keeps_unattributed_items() {
|
||||
let (cluster, _shards, _names) = four_region_cluster();
|
||||
let leader = cluster.leader_region();
|
||||
// Write items WITHOUT creator_id.
|
||||
for i in 1..=5u64 {
|
||||
cluster
|
||||
.write_item_with_metadata(EntityId::new(i), &HashMap::new())
|
||||
.unwrap();
|
||||
}
|
||||
let mut items: Vec<RetrieveResult> = (1..=5u64)
|
||||
.map(|i| retrieve_result(i, 1.0 / i as f64))
|
||||
let mut items: Vec<Sourced<RetrieveResult>> = (1..=5u64)
|
||||
.map(|i| sourced(leader, i, 1.0 / i as f64))
|
||||
.collect();
|
||||
let dropped = enforce_max_per_creator(&cluster, cluster.leader_region(), &mut items, 1);
|
||||
let dropped = enforce_max_per_creator(&cluster, &mut items, 1);
|
||||
assert_eq!(dropped, 0, "unattributed items must never be capped");
|
||||
assert_eq!(items.len(), 5);
|
||||
}
|
||||
|
||||
/// Entity-sharded under-enforcement regression: items owned by different
|
||||
/// shards must be capped by resolving each item's creator from the shard
|
||||
/// that returned it. Writing the same prolific creator's items to disjoint
|
||||
/// shards (so the leader does NOT hold the non-leader ones) and merging them
|
||||
/// must still respect the cap — the old leader-only lookup resolved the
|
||||
/// non-leader items' creator to `None` and let them through uncapped.
|
||||
#[test]
|
||||
fn enforce_max_per_creator_resolves_from_owning_shard() {
|
||||
let (cluster, _shards, _names) = four_region_cluster();
|
||||
let regions = [RegionId(0), RegionId(1), RegionId(2), RegionId(3)];
|
||||
|
||||
// Write 8 items from creator 7, each to a DIFFERENT shard's local store
|
||||
// only (no replication) — exactly the entity-sharded layout where the
|
||||
// leader holds only its own slice.
|
||||
let mut meta = HashMap::new();
|
||||
meta.insert("creator_id".to_string(), "7".to_string());
|
||||
let mut items: Vec<Sourced<RetrieveResult>> = Vec::new();
|
||||
for i in 1..=8u64 {
|
||||
let region = regions[(i as usize - 1) % regions.len()];
|
||||
cluster
|
||||
.node(region)
|
||||
.db
|
||||
.write_item_with_metadata(EntityId::new(i), &meta)
|
||||
.unwrap();
|
||||
// Source-tag each item with the shard that "returned" it.
|
||||
items.push(sourced(region, i, 1.0 / i as f64));
|
||||
}
|
||||
|
||||
let dropped = enforce_max_per_creator(&cluster, &mut items, 3);
|
||||
assert_eq!(
|
||||
items.len(),
|
||||
3,
|
||||
"creator cap must hold across shards; got {} items",
|
||||
items.len()
|
||||
);
|
||||
assert_eq!(dropped, 5, "8 items from one creator, cap 3 → drop 5");
|
||||
}
|
||||
|
||||
/// Counter-test proving the bug the fix removes: resolving every item's
|
||||
/// creator from the LEADER only (the old behavior) under-enforces, because
|
||||
/// the leader does not hold the non-leader shards' items.
|
||||
#[test]
|
||||
fn leader_only_creator_lookup_under_enforces_when_sharded() {
|
||||
let (cluster, _shards, _names) = four_region_cluster();
|
||||
let regions = [RegionId(0), RegionId(1), RegionId(2), RegionId(3)];
|
||||
let leader = cluster.leader_region();
|
||||
|
||||
let mut meta = HashMap::new();
|
||||
meta.insert("creator_id".to_string(), "7".to_string());
|
||||
let mut leader_only: HashMap<u64, usize> = HashMap::new();
|
||||
let mut kept = 0usize;
|
||||
for i in 1..=8u64 {
|
||||
let region = regions[(i as usize - 1) % regions.len()];
|
||||
cluster
|
||||
.node(region)
|
||||
.db
|
||||
.write_item_with_metadata(EntityId::new(i), &meta)
|
||||
.unwrap();
|
||||
// Old behavior: always read metadata from the leader.
|
||||
let creator = cluster
|
||||
.node(leader)
|
||||
.db
|
||||
.get_item_metadata(EntityId::new(i))
|
||||
.ok()
|
||||
.flatten()
|
||||
.and_then(|m| m.get("creator_id").and_then(|c| c.parse::<u64>().ok()));
|
||||
let pass = creator.is_none_or(|cid| {
|
||||
let c = leader_only.entry(cid).or_insert(0);
|
||||
if *c >= 3 {
|
||||
false
|
||||
} else {
|
||||
*c += 1;
|
||||
true
|
||||
}
|
||||
});
|
||||
if pass {
|
||||
kept += 1;
|
||||
}
|
||||
}
|
||||
// The leader holds only its own ~2 items, so the other ~6 resolve to
|
||||
// None and slip through uncapped — proving the silent under-enforcement
|
||||
// the per-shard lookup fixes.
|
||||
assert!(
|
||||
kept > 3,
|
||||
"leader-only lookup should under-enforce (kept {kept} > cap 3)"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@ -34,6 +34,7 @@ fn three_region_topology() -> TopologySpec {
|
||||
},
|
||||
],
|
||||
leader: "us-east".into(),
|
||||
write_workers: None,
|
||||
}
|
||||
}
|
||||
|
||||
@ -61,7 +62,7 @@ fn grpc_replication_converges_across_followers() {
|
||||
// servers via GrpcTransport::new, which blocks on its own runtime.
|
||||
let state = ClusterState::new(&three_region_topology(), view_schema(), Vec::new())
|
||||
.expect("cluster builds with gRPC transports");
|
||||
let cluster = state.cluster();
|
||||
let cluster = state.cluster().expect("cluster is live before shutdown");
|
||||
|
||||
let leader = cluster.leader_region();
|
||||
let followers: Vec<_> = cluster
|
||||
|
||||
291
tidal-server/tests/standalone.rs
Normal file
291
tidal-server/tests/standalone.rs
Normal file
@ -0,0 +1,291 @@
|
||||
// Integration-test exemption (same posture as the other tidal-server tests):
|
||||
// unwrap on known-good fixtures + rank index casts are idiomatic here.
|
||||
#![allow(clippy::unwrap_used, clippy::cast_possible_truncation)]
|
||||
//! End-to-end integration coverage for the STANDALONE data surface.
|
||||
//!
|
||||
//! `middleware.rs` covers auth / body-limit / request-ID and
|
||||
//! `standalone_offload.rs` covers the `/feed` offload round trip, but neither
|
||||
//! drives the full data path through the HTTP surface nor pins the
|
||||
//! engine-result → DTO mapping. This test:
|
||||
//!
|
||||
//! 1. Drives the real handlers in-process via `tower::ServiceExt::oneshot`
|
||||
//! (no TCP bind): POST /items → POST /embeddings → POST /signals → GET /feed
|
||||
//! → GET /search.
|
||||
//! 2. Asserts the [`crate::dto`] shape: `signals` is OMITTED when empty and
|
||||
//! PRESENT when a profile populates the snapshot; search items carry
|
||||
//! `bm25_score`.
|
||||
//! 3. Asserts the standalone region-rejection path: any `?region=` param is a
|
||||
//! 400 in both `/feed` and `/search` (region routing is a cluster concept).
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{
|
||||
body::Body,
|
||||
http::{Method, Request, StatusCode},
|
||||
};
|
||||
use tidal_server::{router::build_router, state::ServerState};
|
||||
use tidaldb::TidalDb;
|
||||
use tower::ServiceExt;
|
||||
|
||||
fn make_app() -> axum::Router {
|
||||
// Default schema + built-in profiles (incl. `for_you`/`trending`), same
|
||||
// setup as `run_standalone` in main.rs. No API key so requests need no auth.
|
||||
let (schema, profiles) = tidal_server::config::load_schema(None).unwrap();
|
||||
let db = TidalDb::builder()
|
||||
.ephemeral()
|
||||
.with_schema(schema)
|
||||
.with_profiles(profiles)
|
||||
.open()
|
||||
.unwrap();
|
||||
let state = Arc::new(ServerState::new(db));
|
||||
build_router(state, None)
|
||||
}
|
||||
|
||||
async fn post_json(app: &axum::Router, uri: &str, body: serde_json::Value) -> StatusCode {
|
||||
app.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method(Method::POST)
|
||||
.uri(uri)
|
||||
.header("Content-Type", "application/json")
|
||||
.body(Body::from(serde_json::to_vec(&body).unwrap()))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.status()
|
||||
}
|
||||
|
||||
async fn get_json(app: &axum::Router, uri: &str) -> (StatusCode, serde_json::Value) {
|
||||
let response = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method(Method::GET)
|
||||
.uri(uri)
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let status = response.status();
|
||||
let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.unwrap();
|
||||
let json: serde_json::Value = serde_json::from_slice(&bytes).unwrap_or(serde_json::Value::Null);
|
||||
(status, json)
|
||||
}
|
||||
|
||||
/// Full standalone round trip: items + embeddings + signals, then a ranked
|
||||
/// /feed and a /search over the indexed titles. Pins the DTO shape in both
|
||||
/// directions.
|
||||
#[tokio::test]
|
||||
async fn items_signals_feed_search_round_trip() {
|
||||
let app = make_app();
|
||||
|
||||
// Two items with searchable titles + a created_at for the Hot sort age.
|
||||
for (id, title) in [(1u64, "jazz piano nocturne"), (2u64, "ambient jazz drift")] {
|
||||
let status = post_json(
|
||||
&app,
|
||||
"/items",
|
||||
serde_json::json!({
|
||||
"entity_id": id,
|
||||
"metadata": { "created_at": "1700000000", "title": title, "category": "music" }
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::CREATED, "item {id} create");
|
||||
|
||||
let status = post_json(
|
||||
&app,
|
||||
"/embeddings",
|
||||
serde_json::json!({ "entity_id": id, "values": vec![0.1f32; 128] }),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::NO_CONTENT, "item {id} embedding");
|
||||
}
|
||||
|
||||
// Drive view signals so the ranking profile has a signal snapshot to expose.
|
||||
// Item 2 gets more weight so it should out-rank item 1.
|
||||
for (id, weight) in [(1u64, 1.0), (2u64, 5.0)] {
|
||||
let status = post_json(
|
||||
&app,
|
||||
"/signals",
|
||||
serde_json::json!({ "entity_id": id, "signal": "view", "weight": weight }),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::NO_CONTENT, "signal for item {id}");
|
||||
}
|
||||
|
||||
// ── GET /feed ────────────────────────────────────────────────────────────
|
||||
// `for_you` boosts the `view` signal, so the feed items must carry a
|
||||
// populated `signals` array (the DTO present-when-populated branch).
|
||||
let (status, body) = get_json(&app, "/feed?profile=for_you&limit=10").await;
|
||||
assert_eq!(status, StatusCode::OK, "feed body: {body}");
|
||||
let items = body
|
||||
.get("items")
|
||||
.and_then(serde_json::Value::as_array)
|
||||
.expect("feed has items array");
|
||||
assert!(
|
||||
!items.is_empty(),
|
||||
"feed should rank the written items: {body}"
|
||||
);
|
||||
|
||||
// Ranks contiguous from 1; the heavier-weighted item 2 ranks first.
|
||||
for (idx, item) in items.iter().enumerate() {
|
||||
let rank = item
|
||||
.get("rank")
|
||||
.and_then(serde_json::Value::as_u64)
|
||||
.unwrap();
|
||||
assert_eq!(rank as usize, idx + 1, "ranks contiguous from 1");
|
||||
assert!(
|
||||
item.get("score")
|
||||
.and_then(serde_json::Value::as_f64)
|
||||
.is_some(),
|
||||
"each item carries a score"
|
||||
);
|
||||
}
|
||||
assert_eq!(
|
||||
items[0]
|
||||
.get("entity_id")
|
||||
.and_then(serde_json::Value::as_u64),
|
||||
Some(2),
|
||||
"heavier-weighted item ranks first: {body}"
|
||||
);
|
||||
|
||||
// DTO: at least one ranked item must expose a populated `signals` array
|
||||
// (the `for_you` view boost), proving the present-when-populated branch.
|
||||
let any_with_signals = items.iter().any(|item| {
|
||||
item.get("signals")
|
||||
.and_then(serde_json::Value::as_array)
|
||||
.is_some_and(|s| !s.is_empty())
|
||||
});
|
||||
assert!(
|
||||
any_with_signals,
|
||||
"a boosted profile must populate the signals snapshot: {body}"
|
||||
);
|
||||
|
||||
// ── GET /search ──────────────────────────────────────────────────────────
|
||||
// The background text syncer commits on a ~2s cadence (ephemeral index,
|
||||
// manual reload), and the `/search` handler reloads the reader on every
|
||||
// call. Poll the real endpoint until the committed docs are visible — a
|
||||
// deterministic end-to-end wait on the real syncer, not a fixed sleep.
|
||||
let body = poll_search_until_nonempty(&app, "/search?query=jazz&limit=10").await;
|
||||
let results = body
|
||||
.get("items")
|
||||
.and_then(serde_json::Value::as_array)
|
||||
.expect("search has items array");
|
||||
assert!(
|
||||
!results.is_empty(),
|
||||
"search should match the jazz titles within the syncer window: {body}"
|
||||
);
|
||||
for (idx, item) in results.iter().enumerate() {
|
||||
let rank = item
|
||||
.get("rank")
|
||||
.and_then(serde_json::Value::as_u64)
|
||||
.unwrap();
|
||||
assert_eq!(rank as usize, idx + 1, "search ranks contiguous from 1");
|
||||
}
|
||||
// A pure-text query must surface a BM25 component in the DTO for the
|
||||
// matched items (the `skip_serializing_if = Option::is_none` field is
|
||||
// present when scored).
|
||||
let any_bm25 = results.iter().any(|item| {
|
||||
item.get("bm25_score")
|
||||
.and_then(serde_json::Value::as_f64)
|
||||
.is_some()
|
||||
});
|
||||
assert!(any_bm25, "text search must expose bm25_score: {body}");
|
||||
}
|
||||
|
||||
/// Re-issue a `/search` against the real handler until it returns a non-empty
|
||||
/// `items` array or a bounded deadline elapses. Each call asserts 200 and
|
||||
/// reloads the text reader (the handler does), so this waits on the real
|
||||
/// background syncer commit (≈2s cadence) without a blind fixed sleep.
|
||||
async fn poll_search_until_nonempty(app: &axum::Router, uri: &str) -> serde_json::Value {
|
||||
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(8);
|
||||
loop {
|
||||
let (status, body) = get_json(app, uri).await;
|
||||
assert_eq!(status, StatusCode::OK, "search body: {body}");
|
||||
let nonempty = body
|
||||
.get("items")
|
||||
.and_then(serde_json::Value::as_array)
|
||||
.is_some_and(|a| !a.is_empty());
|
||||
if nonempty || std::time::Instant::now() >= deadline {
|
||||
return body;
|
||||
}
|
||||
// The text syncer runs on its own OS thread, so a brief blocking sleep
|
||||
// here (between awaits, not during one) does not stall the commit — it
|
||||
// just paces the retry. `tokio::time` is not enabled for this crate, so
|
||||
// a std sleep is the right tool.
|
||||
std::thread::sleep(std::time::Duration::from_millis(100));
|
||||
}
|
||||
}
|
||||
|
||||
/// DTO: an item with NO signals must serialize with `signals` ABSENT (the
|
||||
/// skip-when-empty branch). Uses the `new` profile (sort-by-recency, no boosts),
|
||||
/// so no signal snapshot is attached even though the item exists.
|
||||
#[tokio::test]
|
||||
async fn feed_omits_signals_when_empty() {
|
||||
let app = make_app();
|
||||
|
||||
let status = post_json(
|
||||
&app,
|
||||
"/items",
|
||||
serde_json::json!({
|
||||
"entity_id": 42,
|
||||
"metadata": { "created_at": "1700000000", "title": "no signals here" }
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::CREATED);
|
||||
|
||||
// `new` is a built-in sort-by-recency profile with no signal boosts, so the
|
||||
// ranked item carries an empty snapshot → `signals` omitted from the JSON.
|
||||
let (status, body) = get_json(&app, "/feed?profile=new&limit=10").await;
|
||||
assert_eq!(status, StatusCode::OK, "feed body: {body}");
|
||||
let items = body
|
||||
.get("items")
|
||||
.and_then(serde_json::Value::as_array)
|
||||
.expect("feed has items array");
|
||||
assert!(!items.is_empty(), "the item should still rank: {body}");
|
||||
for item in items {
|
||||
assert!(
|
||||
item.get("signals").is_none(),
|
||||
"empty signal snapshot must be OMITTED, not [] : {item}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Standalone region rejection: any `?region=` param is a 400 on `/feed`.
|
||||
/// Region routing is a cluster-only concept; a standalone server must reject it
|
||||
/// loudly rather than silently ignore it.
|
||||
#[tokio::test]
|
||||
async fn feed_rejects_region_param() {
|
||||
let app = make_app();
|
||||
let (status, body) = get_json(&app, "/feed?profile=for_you®ion=us-east").await;
|
||||
assert_eq!(status, StatusCode::BAD_REQUEST, "body: {body}");
|
||||
let err = body
|
||||
.get("error")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.unwrap_or("");
|
||||
assert!(
|
||||
err.contains("region routing requires cluster mode"),
|
||||
"error must explain region routing is cluster-only: {body}"
|
||||
);
|
||||
}
|
||||
|
||||
/// Standalone region rejection also applies to `/search`.
|
||||
#[tokio::test]
|
||||
async fn search_rejects_region_param() {
|
||||
let app = make_app();
|
||||
let (status, body) = get_json(&app, "/search?query=jazz®ion=us-east").await;
|
||||
assert_eq!(status, StatusCode::BAD_REQUEST, "body: {body}");
|
||||
let err = body
|
||||
.get("error")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.unwrap_or("");
|
||||
assert!(
|
||||
err.contains("region routing requires cluster mode"),
|
||||
"error must explain region routing is cluster-only: {body}"
|
||||
);
|
||||
}
|
||||
159
tidal-server/tests/standalone_offload.rs
Normal file
159
tidal-server/tests/standalone_offload.rs
Normal file
@ -0,0 +1,159 @@
|
||||
// Integration-test exemption (same posture as the other tidal-server tests):
|
||||
// unwrap on known-good fixtures is idiomatic here.
|
||||
#![allow(clippy::unwrap_used)]
|
||||
//! End-to-end integration test for the standalone data path through the
|
||||
//! blocking-work offload (`crate::offload::offload_read`).
|
||||
//!
|
||||
//! CRITICAL 8 moved the standalone `feed`/`search` engine calls onto the tokio
|
||||
//! blocking pool. This test drives the real handlers in-process via
|
||||
//! `tower::ServiceExt::oneshot` (no TCP bind) — POST /items + /signals, then GET
|
||||
//! /feed — and asserts a ranked result comes back, proving the offload path is
|
||||
//! wired correctly and returns the same answers as an inline call would.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{
|
||||
body::Body,
|
||||
http::{Method, Request, StatusCode},
|
||||
};
|
||||
use tidal_server::{router::build_router, state::ServerState};
|
||||
use tidaldb::TidalDb;
|
||||
use tower::ServiceExt;
|
||||
|
||||
fn make_app() -> axum::Router {
|
||||
// Default schema + built-in profiles (incl. `for_you`), same setup as
|
||||
// `run_standalone` in main.rs. No API key so requests need no auth.
|
||||
let (schema, profiles) = tidal_server::config::load_schema(None).unwrap();
|
||||
let db = TidalDb::builder()
|
||||
.ephemeral()
|
||||
.with_schema(schema)
|
||||
.with_profiles(profiles)
|
||||
.open()
|
||||
.unwrap();
|
||||
let state = Arc::new(ServerState::new(db));
|
||||
build_router(state, None)
|
||||
}
|
||||
|
||||
async fn post_json(app: &axum::Router, uri: &str, body: serde_json::Value) -> StatusCode {
|
||||
let response = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method(Method::POST)
|
||||
.uri(uri)
|
||||
.header("Content-Type", "application/json")
|
||||
.body(Body::from(serde_json::to_vec(&body).unwrap()))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
response.status()
|
||||
}
|
||||
|
||||
async fn get_json(app: &axum::Router, uri: &str) -> (StatusCode, serde_json::Value) {
|
||||
let response = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method(Method::GET)
|
||||
.uri(uri)
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let status = response.status();
|
||||
let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.unwrap();
|
||||
let json: serde_json::Value = serde_json::from_slice(&bytes).unwrap_or(serde_json::Value::Null);
|
||||
(status, json)
|
||||
}
|
||||
|
||||
/// POST /items + /signals then GET /feed returns ranked items via the offload
|
||||
/// path. Proves the blocking engine call runs on the blocking pool and produces
|
||||
/// a correct ranked result end to end.
|
||||
#[tokio::test]
|
||||
async fn feed_returns_ranked_items_through_offload() {
|
||||
let app = make_app();
|
||||
|
||||
// Write three items with a created_at so the `for_you` Hot sort has an age.
|
||||
for id in 1u64..=3 {
|
||||
let status = post_json(
|
||||
&app,
|
||||
"/items",
|
||||
serde_json::json!({
|
||||
"entity_id": id,
|
||||
"metadata": { "created_at": "1700000000", "category": "tech" }
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::CREATED, "item {id} create");
|
||||
}
|
||||
|
||||
// Drive view signals so the offloaded scoring path has something to rank.
|
||||
// Item 3 gets the most weight, so it should out-rank the rest.
|
||||
for (id, weight) in [(1u64, 1.0), (2, 2.0), (3, 5.0)] {
|
||||
let status = post_json(
|
||||
&app,
|
||||
"/signals",
|
||||
serde_json::json!({
|
||||
"entity_id": id,
|
||||
"signal": "view",
|
||||
"weight": weight
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::NO_CONTENT, "signal for item {id}");
|
||||
}
|
||||
|
||||
// GET /feed — this is the handler that now offloads `retrieve` to the
|
||||
// blocking pool. Assert it returns a ranked, non-empty result set.
|
||||
let (status, body) = get_json(&app, "/feed?profile=for_you&limit=10").await;
|
||||
assert_eq!(status, StatusCode::OK, "feed body: {body}");
|
||||
|
||||
let items = body
|
||||
.get("items")
|
||||
.and_then(serde_json::Value::as_array)
|
||||
.expect("feed response has an items array");
|
||||
assert!(
|
||||
!items.is_empty(),
|
||||
"feed should rank the written items: {body}"
|
||||
);
|
||||
|
||||
// Ranks must be contiguous starting at 1 (the engine's rank is 1-based) — a
|
||||
// real ordered result, not a degenerate empty/echo response.
|
||||
for (idx, item) in items.iter().enumerate() {
|
||||
let rank = item
|
||||
.get("rank")
|
||||
.and_then(serde_json::Value::as_u64)
|
||||
.unwrap();
|
||||
assert_eq!(rank as usize, idx + 1, "ranks contiguous from 1");
|
||||
assert!(
|
||||
item.get("score")
|
||||
.and_then(serde_json::Value::as_f64)
|
||||
.is_some(),
|
||||
"each item carries a score"
|
||||
);
|
||||
}
|
||||
|
||||
let total = body
|
||||
.get("total_candidates")
|
||||
.and_then(serde_json::Value::as_u64)
|
||||
.expect("total_candidates present");
|
||||
assert!(total >= 3, "all three items were candidates: {body}");
|
||||
}
|
||||
|
||||
/// GET /feed succeeds even with no data — the offload path must not hang or 500
|
||||
/// on an empty engine; it returns an empty ranked result.
|
||||
#[tokio::test]
|
||||
async fn empty_feed_succeeds_through_offload() {
|
||||
let app = make_app();
|
||||
let (status, body) = get_json(&app, "/feed?profile=for_you&limit=10").await;
|
||||
assert_eq!(status, StatusCode::OK, "feed body: {body}");
|
||||
let items = body
|
||||
.get("items")
|
||||
.and_then(serde_json::Value::as_array)
|
||||
.expect("items array present");
|
||||
assert!(items.is_empty(), "no items written, feed is empty");
|
||||
}
|
||||
@ -138,9 +138,17 @@ name = "m7p3_social_scale"
|
||||
[[test]]
|
||||
name = "m8p2_replication"
|
||||
|
||||
[[test]]
|
||||
name = "m8p2_replication_durability"
|
||||
required-features = ["test-utils"]
|
||||
|
||||
[[test]]
|
||||
name = "m8p3_crdt"
|
||||
|
||||
[[test]]
|
||||
name = "m8p3_reconcile_production"
|
||||
required-features = ["test-utils"]
|
||||
|
||||
[[test]]
|
||||
name = "m8p4_session"
|
||||
|
||||
@ -158,6 +166,10 @@ name = "vector_usearch"
|
||||
name = "m0m10_recovery"
|
||||
required-features = ["test-utils"]
|
||||
|
||||
[[test]]
|
||||
name = "m0m10_durability"
|
||||
required-features = ["test-utils"]
|
||||
|
||||
[[bench]]
|
||||
name = "signals"
|
||||
harness = false
|
||||
|
||||
@ -48,7 +48,7 @@ fn make_index(n: u64) -> TextIndex {
|
||||
w.index_item(EntityId::new(i), &meta).unwrap();
|
||||
}
|
||||
|
||||
w.commit(n).unwrap();
|
||||
w.commit().unwrap();
|
||||
drop(w);
|
||||
|
||||
idx.reload_reader().unwrap();
|
||||
|
||||
@ -140,10 +140,19 @@ Never hardcode a single filtering strategy. Estimate selectivity, then branch:
|
||||
The entity store is the source of truth. Tantivy is a materialized view. If the Tantivy index is corrupted or lost, it can be rebuilt from the entity store.
|
||||
|
||||
Consistency pattern:
|
||||
1. Write to entity store (within transaction / WAL)
|
||||
2. Background indexer reads outbox and feeds Tantivy
|
||||
3. On each Tantivy commit, store last-processed sequence number in commit payload
|
||||
4. On crash recovery, replay from that sequence number
|
||||
1. Write to entity store (within transaction / WAL) — the durable source of truth.
|
||||
2. Background indexer reads the outbox and feeds Tantivy (best-effort, async).
|
||||
3. On crash recovery, unconditionally rebuild the text index from the entity
|
||||
store at open — the same proven pattern as the vector index, which is rebuilt
|
||||
from the durable embeddings on every reopen.
|
||||
|
||||
> Crash recovery is NOT driven by a Tantivy commit-payload sequence number.
|
||||
> Because the indexer is best-effort and async, an item can be durably persisted
|
||||
> to the entity store and then lost in a crash before the syncer commits it to
|
||||
> Tantivy. A stored sequence number does not heal that gap by itself and is
|
||||
> easy to leave un-wired (it once silently was, hiding the lost write). The
|
||||
> unconditional rebuild-from-store at open is deterministic and cannot silently
|
||||
> regress: any entity in the store but missing from Tantivy is re-indexed.
|
||||
|
||||
### Hybrid fusion starts with RRF
|
||||
|
||||
|
||||
1270
tidal/docs/reviews/CODE_REVIEW_m0-m10.md
Normal file
1270
tidal/docs/reviews/CODE_REVIEW_m0-m10.md
Normal file
File diff suppressed because it is too large
Load Diff
@ -1090,21 +1090,34 @@ struct TextIndexer {
|
||||
|
||||
### 12.3 Crash Recovery
|
||||
|
||||
On startup, the text indexer:
|
||||
> **Implementation note (recovery is rebuild-at-open, not seqno replay).**
|
||||
> The shipped engine does NOT recover the text index by reading a
|
||||
> `last_committed_seqno` from a Tantivy commit payload. The async syncer is
|
||||
> best-effort, so an item can be durably persisted to the entity store and then
|
||||
> lost in a crash before the syncer commits it to Tantivy; a stored sequence
|
||||
> number does not heal that gap on its own and proved easy to leave un-wired,
|
||||
> hiding the lost write. Instead, on every open the engine **unconditionally
|
||||
> rebuilds the item and creator text indexes from the durable entity stores**
|
||||
> (`db::state_rebuild::rebuild_text_indexes_at_open`), exactly as the vector
|
||||
> index is rebuilt from the durable embeddings. Any entity present in the store
|
||||
> but missing from Tantivy is re-indexed. This is deterministic and cannot
|
||||
> silently regress. The item-side rebuild shares a single scan with the
|
||||
> suggestion-index rebuild.
|
||||
|
||||
1. Opens the Tantivy index.
|
||||
2. Reads the last commit's payload to recover `last_committed_seqno`.
|
||||
3. Replays all outbox entries with `seqno > last_committed_seqno`.
|
||||
4. Resumes normal polling.
|
||||
On startup, the engine:
|
||||
|
||||
1. Opens the Tantivy index (item and creator).
|
||||
2. Unconditionally rebuilds each text index from its durable entity store.
|
||||
3. Resumes normal async syncing of subsequent writes.
|
||||
|
||||
**Failure modes and recovery:**
|
||||
|
||||
| Failure | State After Crash | Recovery |
|
||||
|---------|-------------------|----------|
|
||||
| Crash before Tantivy commit | Entity store ahead of text index. Outbox entries exist for uncommitted docs. | Replay from `last_committed_seqno`. Documents appear in search after recovery. |
|
||||
| Crash during Tantivy commit | Tantivy rolls back to last successful commit. | Same as above -- replay from last committed seqno. |
|
||||
| Crash after Tantivy commit but before outbox cleanup | Outbox may re-deliver entries. | Tantivy silently handles duplicate deletes. Duplicate adds create duplicate documents briefly until the next merge consolidates them. The `_entity_id` field provides deduplication at query time. |
|
||||
| Tantivy index corruption | Text index is unusable. | Full rebuild from entity store (Section 12.5). |
|
||||
| Crash before Tantivy commit | Entity store ahead of text index; recently-written items not yet in Tantivy. | Unconditional rebuild-from-store at open re-indexes them; they appear in search after recovery. |
|
||||
| Crash during Tantivy commit | Tantivy rolls back to last successful commit. | Same as above -- rebuild-from-store at open. |
|
||||
| Crash after Tantivy commit | Text index may briefly hold duplicate documents from a re-applied write. | The delete-then-add update path and the `entity_id` term make re-indexing idempotent; rebuild-from-store re-establishes the exact set. |
|
||||
| Tantivy index corruption | Text index is unusable. | Full rebuild from entity store (Section 12.5) — the same primitive run at open. |
|
||||
|
||||
### 12.4 Outbox Key Encoding
|
||||
|
||||
|
||||
@ -36,6 +36,7 @@ fn random_unit_vector(dim: usize, rng: &mut impl Rng) -> Vec<f32> {
|
||||
v.iter().map(|x| x / norm).collect()
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_lines)] // Linear walkthrough script; splitting would hurt readability.
|
||||
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
// Initialize tracing so spans emitted by tidalDB are visible.
|
||||
tracing_subscriber::fmt()
|
||||
|
||||
@ -4,9 +4,10 @@
|
||||
//!
|
||||
//! `CohortSignalLedger::checkpoint()` serializes all in-memory cohort signal
|
||||
//! state to the `StorageEngine` as a single atomic `WriteBatch`. Each entry
|
||||
//! is serialized using the same 983-byte fixed-format record as the global
|
||||
//! signal ledger (via `signals::checkpoint::format`), but keyed under
|
||||
//! `Tag::Cohort` instead of `Tag::Sig`.
|
||||
//! is serialized using the same fixed-format record as the global signal
|
||||
//! ledger (via `signals::checkpoint::format`, currently the V2 1116-byte
|
||||
//! layout — `format::ENTRY_SIZE`), but keyed under `Tag::Cohort` instead of
|
||||
//! `Tag::Sig`.
|
||||
//!
|
||||
//! # Restore
|
||||
//!
|
||||
@ -26,7 +27,7 @@
|
||||
|
||||
use super::CohortSignalLedger;
|
||||
use crate::{
|
||||
cohort::types::MAX_COHORT_NAME_LEN,
|
||||
cohort::types::{CohortName, MAX_COHORT_NAME_LEN},
|
||||
schema::{EntityId, TidalError},
|
||||
signals::{
|
||||
SignalTypeId,
|
||||
@ -42,9 +43,10 @@ use crate::{
|
||||
impl CohortSignalLedger {
|
||||
/// Write all in-memory cohort signal state to the storage engine atomically.
|
||||
///
|
||||
/// Each entry is serialized as a 983-byte record (same format as the global
|
||||
/// signal ledger) and keyed under `Tag::Cohort` with a suffix encoding the
|
||||
/// signal type ID and cohort name.
|
||||
/// Each entry is serialized with the global signal ledger's fixed-format
|
||||
/// record (currently the V2 1116-byte layout, `format::ENTRY_SIZE`) and
|
||||
/// keyed under `Tag::Cohort` with a suffix encoding the signal type ID and
|
||||
/// cohort name.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
@ -68,6 +70,7 @@ impl CohortSignalLedger {
|
||||
batch.put(meta_key, serialize_meta(&meta));
|
||||
|
||||
// Write all cohort-entity-signal entries.
|
||||
let mut entry_count = 0usize;
|
||||
for entry_ref in self.iter_entries() {
|
||||
let (cohort_name, entity_id, signal_type_id) = entry_ref.key();
|
||||
let entry = entry_ref.value();
|
||||
@ -76,10 +79,17 @@ impl CohortSignalLedger {
|
||||
let key = encode_key(*entity_id, Tag::Cohort, &suffix);
|
||||
let value = serialize_entry(*entity_id, *signal_type_id, entry);
|
||||
batch.put(key, value);
|
||||
entry_count += 1;
|
||||
}
|
||||
|
||||
storage.write_batch(batch)?;
|
||||
storage.flush()?;
|
||||
|
||||
tracing::debug!(
|
||||
entry_count,
|
||||
wal_sequence = meta.wal_sequence,
|
||||
"cohort signal checkpoint written"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@ -93,35 +103,53 @@ impl CohortSignalLedger {
|
||||
///
|
||||
/// # Corruption policy
|
||||
///
|
||||
/// The cohort aggregate cannot be rebuilt from the WAL — this checkpoint is
|
||||
/// its only durability mechanism — so a partial restore would silently
|
||||
/// resurrect a wrong aggregate. We therefore treat *any* structurally
|
||||
/// invalid `Tag::Cohort` entry as fatal corruption and abort the whole
|
||||
/// restore with `TidalError::Internal`, whether the damage is in the key
|
||||
/// suffix (malformed `signal_type_id`/name length) or in the value
|
||||
/// (un-deserializable entry). The two used to be handled inconsistently
|
||||
/// (a corrupt value aborted but a malformed suffix only warned-and-skipped);
|
||||
/// they are now uniform. Failing loudly on corruption beats serving a
|
||||
/// silently-truncated aggregate — "don't ship what you wouldn't trust at 3am."
|
||||
/// The cohort aggregate is *derived* state (named user segments' decaying
|
||||
/// signal counts). It cannot be rebuilt from the WAL — this checkpoint is its
|
||||
/// only durability mechanism — so a *partial* restore would silently resurrect
|
||||
/// a wrong aggregate. We therefore treat the checkpoint as ALL-OR-NOTHING: if
|
||||
/// any structurally invalid `Tag::Cohort` record is encountered (corrupt meta,
|
||||
/// malformed key suffix, or un-deserializable value), we log a
|
||||
/// `tracing::error!` naming the corrupt record and start from an EMPTY cohort
|
||||
/// ledger (`Ok(None)`), committing none of the buffered entries.
|
||||
///
|
||||
/// We deliberately do NOT abort the whole database open on cohort-checkpoint
|
||||
/// corruption: cohort attribution is a non-critical ranking input that
|
||||
/// re-accumulates from live signal traffic, so refusing to open the entire DB
|
||||
/// for one corrupt derived aggregate is a worse 3am outcome than serving fresh
|
||||
/// cohort state with a loud error in the log. The escalation is the `error!`
|
||||
/// record (an operator alarm), not a process abort — and the caller in
|
||||
/// `TidalDb::from_parts` no longer silently downgrades an abort to a `warn!`:
|
||||
/// code and this contract now agree. (The sibling `CommunityLedger::restore`
|
||||
/// shares the "not WAL-rebuildable" property and the same start-fresh-on-
|
||||
/// corruption posture.)
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// - `TidalError::Storage` on I/O failure
|
||||
/// - `TidalError::Internal` on any corrupt checkpoint record (meta, key
|
||||
/// suffix, or entry value)
|
||||
/// - `TidalError::Storage` on a genuine I/O failure reading the checkpoint.
|
||||
/// Checkpoint *corruption* is NOT returned as an error: it is logged at
|
||||
/// `error!` and yields `Ok(None)` (start empty).
|
||||
pub fn restore(&self, storage: &dyn StorageEngine) -> crate::Result<Option<CheckpointMeta>> {
|
||||
// Read checkpoint metadata first.
|
||||
// Read checkpoint metadata first. Corrupt meta -> start empty (loud).
|
||||
let meta_key = encode_key(EntityId::new(0), Tag::Cohort, META_SUFFIX);
|
||||
let meta = match storage.get(&meta_key)? {
|
||||
None => None,
|
||||
Some(meta_bytes) => Some(deserialize_meta(&meta_bytes).map_err(|e| {
|
||||
TidalError::internal(
|
||||
"cohort_checkpoint_restore",
|
||||
format!("corrupt cohort checkpoint meta: {e}"),
|
||||
)
|
||||
})?),
|
||||
Some(meta_bytes) => match deserialize_meta(&meta_bytes) {
|
||||
Ok(m) => Some(m),
|
||||
Err(e) => {
|
||||
tracing::error!(
|
||||
error = %e,
|
||||
"corrupt cohort checkpoint meta; starting from EMPTY cohort \
|
||||
state (derived aggregate re-accumulates from live signals)"
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
},
|
||||
};
|
||||
|
||||
// Buffer every entry and commit only if the ENTIRE checkpoint is intact:
|
||||
// a partial cohort aggregate is a wrong aggregate, so any corruption
|
||||
// discards everything and starts empty (see the Corruption policy above).
|
||||
//
|
||||
// TECH DEBT: scan_prefix(&[]) iterates the entire keyspace. This mirrors
|
||||
// the global signal-ledger restore scan (signals::checkpoint::restore) and
|
||||
// is correct but not selective: it touches every key, not just Tag::Cohort
|
||||
@ -129,6 +157,7 @@ impl CohortSignalLedger {
|
||||
// `scan_tag(Tag::Cohort)`) once the storage engine exposes a tag-bounded
|
||||
// iterator, so reopen cost scales with cohort state rather than total
|
||||
// keyspace size. Documented here for parity with the global-ledger note.
|
||||
let mut buffered = Vec::new();
|
||||
for item in storage.scan_prefix(&[]) {
|
||||
let (key, value) = item?;
|
||||
if let Some((entity_id, Tag::Cohort, suffix)) = parse_key(&key) {
|
||||
@ -137,33 +166,47 @@ impl CohortSignalLedger {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Parse the suffix to extract signal_type_id and cohort_name.
|
||||
// A malformed suffix on a Tag::Cohort key is checkpoint corruption,
|
||||
// not a foreign key — abort, consistent with the value-corruption
|
||||
// path below (see the Corruption policy section above).
|
||||
// not a foreign key. Start empty (loud) rather than commit a wrong
|
||||
// partial aggregate (see the Corruption policy section above).
|
||||
let Some((signal_type_id, cohort_name)) = parse_cohort_entry_suffix(suffix) else {
|
||||
return Err(TidalError::internal(
|
||||
"cohort_checkpoint_restore",
|
||||
format!(
|
||||
"corrupt cohort checkpoint entry: malformed suffix \
|
||||
({} bytes) for entity_id={}",
|
||||
suffix.len(),
|
||||
entity_id.as_u64(),
|
||||
),
|
||||
));
|
||||
tracing::error!(
|
||||
entity_id = entity_id.as_u64(),
|
||||
suffix_len = suffix.len(),
|
||||
"corrupt cohort checkpoint entry (malformed suffix); \
|
||||
starting from EMPTY cohort state"
|
||||
);
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let (_eid, _stid, entry) = deserialize_entry(&value).map_err(|e| {
|
||||
TidalError::internal(
|
||||
"cohort_checkpoint_restore",
|
||||
format!("corrupt cohort checkpoint entry: {e}"),
|
||||
)
|
||||
})?;
|
||||
let entry = match deserialize_entry(&value) {
|
||||
Ok((_eid, _stid, entry)) => entry,
|
||||
Err(e) => {
|
||||
tracing::error!(
|
||||
error = %e,
|
||||
entity_id = entity_id.as_u64(),
|
||||
"corrupt cohort checkpoint entry (bad value); \
|
||||
starting from EMPTY cohort state"
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
|
||||
self.insert_entry(cohort_name, entity_id, signal_type_id, entry);
|
||||
buffered.push((cohort_name, entity_id, signal_type_id, entry));
|
||||
}
|
||||
}
|
||||
|
||||
// Whole checkpoint validated: commit every buffered entry.
|
||||
let restored_count = buffered.len();
|
||||
for (cohort_name, entity_id, signal_type_id, entry) in buffered {
|
||||
self.insert_entry(cohort_name, entity_id, signal_type_id, entry);
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
restored_count,
|
||||
wal_sequence = meta.as_ref().map_or(0, |m| m.wal_sequence),
|
||||
"cohort signal ledger restored from checkpoint"
|
||||
);
|
||||
Ok(meta)
|
||||
}
|
||||
}
|
||||
@ -203,7 +246,10 @@ fn cohort_entry_suffix(
|
||||
}
|
||||
|
||||
/// Parse a cohort entry suffix back into `(SignalTypeId, cohort_name)`.
|
||||
fn parse_cohort_entry_suffix(suffix: &[u8]) -> Option<(SignalTypeId, String)> {
|
||||
///
|
||||
/// The cohort name is returned as a [`CohortName`] (`Arc<str>`) so it can be
|
||||
/// inserted into the ledger's `Arc<str>`-keyed map without a second copy.
|
||||
fn parse_cohort_entry_suffix(suffix: &[u8]) -> Option<(SignalTypeId, CohortName)> {
|
||||
// Minimum: 2 (signal_type_id) + 2 (name_len) = 4 bytes
|
||||
if suffix.len() < 4 {
|
||||
return None;
|
||||
@ -214,7 +260,7 @@ fn parse_cohort_entry_suffix(suffix: &[u8]) -> Option<(SignalTypeId, String)> {
|
||||
return None;
|
||||
}
|
||||
let name = std::str::from_utf8(&suffix[4..4 + name_len]).ok()?;
|
||||
Some((signal_type_id, name.to_owned()))
|
||||
Some((signal_type_id, CohortName::from(name)))
|
||||
}
|
||||
|
||||
// ── Tests ───────────────────────────────────────────────────────────────────
|
||||
@ -230,6 +276,13 @@ mod tests {
|
||||
storage::InMemoryBackend,
|
||||
};
|
||||
|
||||
/// Build a `CohortName` (`Arc<str>`) for `CohortSignalLedger::record`, which
|
||||
/// keys by `&CohortName` (the production fan-out clones an `Arc` rather than
|
||||
/// allocating a `String` per matched cohort).
|
||||
fn cn(name: &str) -> CohortName {
|
||||
CohortName::from(name)
|
||||
}
|
||||
|
||||
fn test_schema() -> crate::schema::Schema {
|
||||
let mut builder = SchemaBuilder::new();
|
||||
let _ = builder
|
||||
@ -265,9 +318,15 @@ mod tests {
|
||||
let now = Timestamp::now();
|
||||
let type_id = ledger.resolve_signal_type("view").unwrap();
|
||||
for i in 1..=5u64 {
|
||||
ledger.record("tech", EntityId::new(i), type_id, 1.0, now.as_nanos());
|
||||
ledger.record(&cn("tech"), EntityId::new(i), type_id, 1.0, now.as_nanos());
|
||||
}
|
||||
ledger.record("sports", EntityId::new(1), type_id, 1.0, now.as_nanos());
|
||||
ledger.record(
|
||||
&cn("sports"),
|
||||
EntityId::new(1),
|
||||
type_id,
|
||||
1.0,
|
||||
now.as_nanos(),
|
||||
);
|
||||
|
||||
let storage = InMemoryBackend::new();
|
||||
let meta = CheckpointMeta {
|
||||
@ -308,12 +367,18 @@ mod tests {
|
||||
|
||||
// Record signals across cohorts.
|
||||
for i in 1..=10u64 {
|
||||
ledger.record("tech", EntityId::new(i), view_id, 1.0, ts_ns);
|
||||
ledger.record(&cn("tech"), EntityId::new(i), view_id, 1.0, ts_ns);
|
||||
}
|
||||
for i in 1..=5u64 {
|
||||
ledger.record("sports", EntityId::new(i), like_id, 2.0, ts_ns);
|
||||
ledger.record(&cn("sports"), EntityId::new(i), like_id, 2.0, ts_ns);
|
||||
}
|
||||
ledger.record("tech", EntityId::new(1), view_id, 1.0, ts_ns + 1_000_000);
|
||||
ledger.record(
|
||||
&cn("tech"),
|
||||
EntityId::new(1),
|
||||
view_id,
|
||||
1.0,
|
||||
ts_ns + 1_000_000,
|
||||
);
|
||||
|
||||
let storage = InMemoryBackend::new();
|
||||
let meta = CheckpointMeta {
|
||||
@ -360,7 +425,7 @@ mod tests {
|
||||
|
||||
// First checkpoint: 3 entries.
|
||||
for i in 1..=3u64 {
|
||||
ledger.record("cohort_a", EntityId::new(i), view_id, 1.0, ts_ns);
|
||||
ledger.record(&cn("cohort_a"), EntityId::new(i), view_id, 1.0, ts_ns);
|
||||
}
|
||||
ledger
|
||||
.checkpoint(
|
||||
@ -375,7 +440,7 @@ mod tests {
|
||||
|
||||
// Add 2 more entries, then second checkpoint: 5 entries total.
|
||||
for i in 4..=5u64 {
|
||||
ledger.record("cohort_a", EntityId::new(i), view_id, 1.0, ts_ns);
|
||||
ledger.record(&cn("cohort_a"), EntityId::new(i), view_id, 1.0, ts_ns);
|
||||
}
|
||||
ledger
|
||||
.checkpoint(
|
||||
@ -405,7 +470,7 @@ mod tests {
|
||||
let (parsed_type_id, parsed_cohort) =
|
||||
parse_cohort_entry_suffix(&suffix).expect("valid suffix");
|
||||
assert_eq!(parsed_type_id, type_id);
|
||||
assert_eq!(parsed_cohort, cohort);
|
||||
assert_eq!(&*parsed_cohort, cohort);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@ -421,7 +486,7 @@ mod tests {
|
||||
let at_limit = "y".repeat(MAX_COHORT_NAME_LEN);
|
||||
let suffix = cohort_entry_suffix(type_id, &at_limit).expect("at-limit name is valid");
|
||||
let (_id, parsed) = parse_cohort_entry_suffix(&suffix).expect("valid suffix");
|
||||
assert_eq!(parsed, at_limit);
|
||||
assert_eq!(&*parsed, at_limit.as_str());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@ -432,7 +497,7 @@ mod tests {
|
||||
let ledger = CohortSignalLedger::new(&schema);
|
||||
let type_id = ledger.resolve_signal_type("view").unwrap();
|
||||
let now = Timestamp::now();
|
||||
let over_long = "z".repeat(MAX_COHORT_NAME_LEN + 1);
|
||||
let over_long = cn(&"z".repeat(MAX_COHORT_NAME_LEN + 1));
|
||||
ledger.record(&over_long, EntityId::new(1), type_id, 1.0, now.as_nanos());
|
||||
|
||||
let storage = InMemoryBackend::new();
|
||||
@ -448,10 +513,11 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn restore_aborts_on_malformed_suffix() {
|
||||
fn restore_starts_empty_on_malformed_suffix() {
|
||||
// A Tag::Cohort key with a too-short suffix is checkpoint corruption.
|
||||
// The restore must abort (not warn-and-skip), consistent with the
|
||||
// value-corruption path.
|
||||
// Restore must start from an EMPTY ledger (committing nothing), not commit
|
||||
// a wrong partial aggregate and not abort the whole DB open. The escalation
|
||||
// is the error! log; here we assert the all-or-nothing outcome.
|
||||
let storage = InMemoryBackend::new();
|
||||
|
||||
// Write a valid meta key so restore proceeds past the meta read.
|
||||
@ -474,16 +540,24 @@ mod tests {
|
||||
|
||||
let schema = test_schema();
|
||||
let ledger = CohortSignalLedger::new(&schema);
|
||||
let err = ledger
|
||||
let meta = ledger
|
||||
.restore(&storage)
|
||||
.expect_err("malformed suffix must abort restore");
|
||||
assert!(matches!(err, TidalError::Internal(_)));
|
||||
.expect("corruption is not an error: restore must return Ok and start empty");
|
||||
assert!(
|
||||
meta.is_none(),
|
||||
"malformed suffix must yield Ok(None) (start empty)"
|
||||
);
|
||||
assert_eq!(
|
||||
ledger.entry_count(),
|
||||
0,
|
||||
"no entries may be committed from a corrupt checkpoint"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn restore_aborts_on_corrupt_value() {
|
||||
fn restore_starts_empty_on_corrupt_value() {
|
||||
// A Tag::Cohort key with a valid suffix but a truncated/garbage value
|
||||
// is checkpoint corruption and must abort the restore.
|
||||
// is checkpoint corruption: restore starts empty (all-or-nothing).
|
||||
let storage = InMemoryBackend::new();
|
||||
|
||||
let meta_key = encode_key(EntityId::new(0), Tag::Cohort, META_SUFFIX);
|
||||
@ -506,10 +580,18 @@ mod tests {
|
||||
|
||||
let schema = test_schema();
|
||||
let ledger = CohortSignalLedger::new(&schema);
|
||||
let err = ledger
|
||||
let meta = ledger
|
||||
.restore(&storage)
|
||||
.expect_err("corrupt value must abort restore");
|
||||
assert!(matches!(err, TidalError::Internal(_)));
|
||||
.expect("corruption is not an error: restore must return Ok and start empty");
|
||||
assert!(
|
||||
meta.is_none(),
|
||||
"corrupt value must yield Ok(None) (start empty)"
|
||||
);
|
||||
assert_eq!(
|
||||
ledger.entry_count(),
|
||||
0,
|
||||
"no entries may be committed from a corrupt checkpoint"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@ -7,11 +7,12 @@
|
||||
//! Cohort-level signal writes are best-effort: they do not go through the WAL
|
||||
//! and do not fail the primary signal write path.
|
||||
|
||||
use std::{collections::HashMap, fmt};
|
||||
use std::{collections::HashMap, fmt, sync::Arc};
|
||||
|
||||
use dashmap::DashMap;
|
||||
|
||||
use crate::{
|
||||
cohort::types::CohortName,
|
||||
schema::{EntityId, Schema, SchemaError, TidalError, Window},
|
||||
signals::{
|
||||
SignalTypeId, hot::HotSignalState, ledger::types::EntitySignalEntry, warm::BucketedCounter,
|
||||
@ -20,12 +21,18 @@ use crate::{
|
||||
|
||||
/// Per-cohort signal state.
|
||||
///
|
||||
/// The key type is `(String, EntityId, SignalTypeId)` where `String` is the
|
||||
/// cohort name. Using `String` avoids lifetime issues with `Arc<str>` in
|
||||
/// `DashMap` while keeping the performance impact negligible (cohort names
|
||||
/// are short strings, and signal writes are already I/O-bound).
|
||||
/// The key type is `(CohortName, EntityId, SignalTypeId)` where
|
||||
/// `CohortName = Arc<str>`. The attribution fan-out (`db::signals` resolves a
|
||||
/// `Vec<CohortName>` and records the signal into each one) holds the cohort
|
||||
/// name as an `Arc<str>` already, so `record` keys the map with an `Arc::clone`
|
||||
/// — a refcount bump, not a fresh allocation. The previous `String` key forced
|
||||
/// a `cohort.to_owned()` heap allocation per matched cohort on every signal
|
||||
/// write, which violates `CODING_GUIDELINES §1` (no `String` in the hot-path).
|
||||
/// Read paths (which take a borrowed `&str` from the query executor) build a
|
||||
/// single lookup `Arc<str>`; that is the same single allocation the old `String`
|
||||
/// key incurred, and reads are off the signal-write fan-out path.
|
||||
pub struct CohortSignalLedger {
|
||||
entries: DashMap<(String, EntityId, SignalTypeId), EntitySignalEntry>,
|
||||
entries: DashMap<(CohortName, EntityId, SignalTypeId), EntitySignalEntry>,
|
||||
signal_lambdas: HashMap<SignalTypeId, Vec<f64>>,
|
||||
signal_name_to_id: HashMap<String, SignalTypeId>,
|
||||
}
|
||||
@ -65,9 +72,14 @@ impl CohortSignalLedger {
|
||||
///
|
||||
/// This is best-effort: it does not write to WAL and never returns an error.
|
||||
/// If recording fails (e.g., unknown signal type), it is silently dropped.
|
||||
///
|
||||
/// Takes the cohort name as `&CohortName` (`&Arc<str>`) so the per-write key
|
||||
/// is built with `Arc::clone` (a refcount bump) rather than a fresh `String`
|
||||
/// allocation. A single signal can fan out to every matching cohort, so this
|
||||
/// is the hot path the `String` key previously taxed once per matched cohort.
|
||||
pub fn record(
|
||||
&self,
|
||||
cohort: &str,
|
||||
cohort: &CohortName,
|
||||
entity_id: EntityId,
|
||||
type_id: SignalTypeId,
|
||||
weight: f64,
|
||||
@ -78,7 +90,7 @@ impl CohortSignalLedger {
|
||||
.get(&type_id)
|
||||
.map_or_else(<&[f64]>::default, Vec::as_slice);
|
||||
|
||||
let key = (cohort.to_owned(), entity_id, type_id);
|
||||
let key = (Arc::clone(cohort), entity_id, type_id);
|
||||
let entry = self
|
||||
.entries
|
||||
.entry(key)
|
||||
@ -120,7 +132,12 @@ impl CohortSignalLedger {
|
||||
.copied()
|
||||
.unwrap_or(0.0);
|
||||
|
||||
let key = (cohort.to_owned(), entity_id, type_id);
|
||||
// Reads take a borrowed name from the query executor, so we build a
|
||||
// single lookup `Arc<str>` (the tuple key element is `Arc<str>`, which is
|
||||
// not `Copy`, so the key must own it). This is the same single allocation
|
||||
// the former `String` key incurred — reads are off the signal-write
|
||||
// fan-out hot path that the `Arc::clone` change in `record` optimizes.
|
||||
let key: (CohortName, EntityId, SignalTypeId) = (Arc::from(cohort), entity_id, type_id);
|
||||
match self.entries.get(&key) {
|
||||
None => Ok(None),
|
||||
Some(entry) => Ok(Some(entry.hot.current_score(
|
||||
@ -144,10 +161,16 @@ impl CohortSignalLedger {
|
||||
window: Window,
|
||||
) -> crate::Result<u64> {
|
||||
let type_id = self.resolve_signal_type(signal_type_name)?;
|
||||
let key = (cohort.to_owned(), entity_id, type_id);
|
||||
// Read the clock here (matching `read_decay_score`) and thread it into the
|
||||
// warm tier so a quiet cohort-entity's windowed/velocity count ages
|
||||
// correctly at read time — without it CohortTrending stays frozen at its
|
||||
// peak forever (CRITICAL #6, CODING_GUIDELINES §8).
|
||||
let now_ns = crate::schema::Timestamp::now().as_nanos();
|
||||
// Build the lookup key with a single `Arc<str>` (see `read_decay_score`).
|
||||
let key: (CohortName, EntityId, SignalTypeId) = (Arc::from(cohort), entity_id, type_id);
|
||||
match self.entries.get(&key) {
|
||||
None => Ok(0),
|
||||
Some(entry) => Ok(entry.warm.windowed_count(window)),
|
||||
Some(entry) => Ok(entry.warm.windowed_count(window, now_ns)),
|
||||
}
|
||||
}
|
||||
|
||||
@ -193,7 +216,7 @@ impl CohortSignalLedger {
|
||||
) -> impl Iterator<
|
||||
Item = dashmap::mapref::multiple::RefMulti<
|
||||
'_,
|
||||
(String, EntityId, SignalTypeId),
|
||||
(CohortName, EntityId, SignalTypeId),
|
||||
EntitySignalEntry,
|
||||
>,
|
||||
> {
|
||||
@ -203,7 +226,7 @@ impl CohortSignalLedger {
|
||||
/// Insert a restored entry (for checkpoint restore).
|
||||
pub fn insert_entry(
|
||||
&self,
|
||||
cohort: String,
|
||||
cohort: CohortName,
|
||||
entity_id: EntityId,
|
||||
type_id: SignalTypeId,
|
||||
entry: EntitySignalEntry,
|
||||
@ -237,6 +260,13 @@ mod tests {
|
||||
use super::*;
|
||||
use crate::schema::{DecaySpec, EntityKind, SchemaBuilder, Timestamp, Window};
|
||||
|
||||
/// Build a `CohortName` (`Arc<str>`) from a string literal for `record`,
|
||||
/// which keys by `&CohortName` so the production fan-out clones an `Arc`
|
||||
/// (refcount bump) instead of allocating a `String` per matched cohort.
|
||||
fn cn(name: &str) -> CohortName {
|
||||
Arc::from(name)
|
||||
}
|
||||
|
||||
fn test_schema() -> Schema {
|
||||
let mut builder = SchemaBuilder::new();
|
||||
let _ = builder
|
||||
@ -277,7 +307,7 @@ mod tests {
|
||||
let type_id = ledger.resolve_signal_type("view").unwrap();
|
||||
let ts_ns = Timestamp::now().as_nanos();
|
||||
|
||||
ledger.record("tech", entity_id, type_id, 1.0, ts_ns);
|
||||
ledger.record(&cn("tech"), entity_id, type_id, 1.0, ts_ns);
|
||||
|
||||
let score = ledger
|
||||
.read_decay_score("tech", entity_id, "view", 0)
|
||||
@ -294,8 +324,8 @@ mod tests {
|
||||
let type_id = ledger.resolve_signal_type("view").unwrap();
|
||||
let ts_ns = Timestamp::now().as_nanos();
|
||||
|
||||
ledger.record("sports", entity_id, type_id, 1.0, ts_ns);
|
||||
ledger.record("sports", entity_id, type_id, 1.0, ts_ns + 1_000_000);
|
||||
ledger.record(&cn("sports"), entity_id, type_id, 1.0, ts_ns);
|
||||
ledger.record(&cn("sports"), entity_id, type_id, 1.0, ts_ns + 1_000_000);
|
||||
|
||||
let count = ledger
|
||||
.read_windowed_count("sports", entity_id, "view", Window::AllTime)
|
||||
@ -303,6 +333,39 @@ mod tests {
|
||||
assert_eq!(count, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quiet_cohort_entity_windowed_count_decays_to_zero_on_read() {
|
||||
// CRITICAL #6 regression for the cohort tier: a cohort-entity that went
|
||||
// quiet must report a decaying 1h windowed count on a pure read so
|
||||
// CohortTrending ages correctly. Record events stamped ~2h in the past,
|
||||
// then read at the present clock — the read-time rotation must zero the
|
||||
// stale minute buckets.
|
||||
let schema = test_schema();
|
||||
let ledger = CohortSignalLedger::new(&schema);
|
||||
let entity_id = EntityId::new(3);
|
||||
let type_id = ledger.resolve_signal_type("view").unwrap();
|
||||
|
||||
let two_hours_ns = 2 * 3_600_000_000_000_u64;
|
||||
let past = Timestamp::now().as_nanos().saturating_sub(two_hours_ns);
|
||||
for _ in 0..20 {
|
||||
ledger.record(&cn("tech"), entity_id, type_id, 1.0, past);
|
||||
}
|
||||
|
||||
let one_hour = ledger
|
||||
.read_windowed_count("tech", entity_id, "view", Window::OneHour)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
one_hour, 0,
|
||||
"a quiet cohort-entity's 1h window must decay to 0 once its events fall outside it"
|
||||
);
|
||||
|
||||
// AllTime is unbounded and must still reflect every recorded event.
|
||||
let all_time = ledger
|
||||
.read_windowed_count("tech", entity_id, "view", Window::AllTime)
|
||||
.unwrap();
|
||||
assert_eq!(all_time, 20);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn different_cohorts_independent() {
|
||||
let schema = test_schema();
|
||||
@ -311,9 +374,9 @@ mod tests {
|
||||
let type_id = ledger.resolve_signal_type("view").unwrap();
|
||||
let ts_ns = Timestamp::now().as_nanos();
|
||||
|
||||
ledger.record("cohort_a", entity_id, type_id, 1.0, ts_ns);
|
||||
ledger.record("cohort_a", entity_id, type_id, 1.0, ts_ns + 1_000_000);
|
||||
ledger.record("cohort_b", entity_id, type_id, 1.0, ts_ns);
|
||||
ledger.record(&cn("cohort_a"), entity_id, type_id, 1.0, ts_ns);
|
||||
ledger.record(&cn("cohort_a"), entity_id, type_id, 1.0, ts_ns + 1_000_000);
|
||||
ledger.record(&cn("cohort_b"), entity_id, type_id, 1.0, ts_ns);
|
||||
|
||||
let count_a = ledger
|
||||
.read_windowed_count("cohort_a", entity_id, "view", Window::AllTime)
|
||||
@ -361,7 +424,7 @@ mod tests {
|
||||
let type_id = ledger.resolve_signal_type("view").unwrap();
|
||||
let ts_ns = Timestamp::now().as_nanos();
|
||||
|
||||
ledger.record("c", entity_id, type_id, 1.0, ts_ns);
|
||||
ledger.record(&cn("c"), entity_id, type_id, 1.0, ts_ns);
|
||||
let vel = ledger
|
||||
.read_velocity("c", entity_id, "view", Window::AllTime)
|
||||
.unwrap();
|
||||
@ -376,15 +439,15 @@ mod tests {
|
||||
|
||||
let type_id = ledger.resolve_signal_type("view").unwrap();
|
||||
let ts_ns = Timestamp::now().as_nanos();
|
||||
ledger.record("c", EntityId::new(1), type_id, 1.0, ts_ns);
|
||||
ledger.record(&cn("c"), EntityId::new(1), type_id, 1.0, ts_ns);
|
||||
assert_eq!(ledger.entry_count(), 1);
|
||||
|
||||
// Same key increments count, does not add new entry.
|
||||
ledger.record("c", EntityId::new(1), type_id, 1.0, ts_ns + 1);
|
||||
ledger.record(&cn("c"), EntityId::new(1), type_id, 1.0, ts_ns + 1);
|
||||
assert_eq!(ledger.entry_count(), 1);
|
||||
|
||||
// Different cohort = different key.
|
||||
ledger.record("d", EntityId::new(1), type_id, 1.0, ts_ns);
|
||||
ledger.record(&cn("d"), EntityId::new(1), type_id, 1.0, ts_ns);
|
||||
assert_eq!(ledger.entry_count(), 2);
|
||||
}
|
||||
|
||||
@ -397,7 +460,7 @@ mod tests {
|
||||
hot: HotSignalState::new(42, type_id.as_u16()),
|
||||
warm: BucketedCounter::new(),
|
||||
};
|
||||
ledger.insert_entry("test".to_owned(), EntityId::new(42), type_id, entry);
|
||||
ledger.insert_entry(cn("test"), EntityId::new(42), type_id, entry);
|
||||
assert_eq!(ledger.entry_count(), 1);
|
||||
}
|
||||
|
||||
|
||||
@ -110,6 +110,29 @@ impl CohortResolver {
|
||||
cache.pop(&user_id);
|
||||
}
|
||||
|
||||
/// Invalidate every cached membership at once.
|
||||
///
|
||||
/// Call this whenever the cohort *registry* changes (a cohort is defined or
|
||||
/// removed at runtime): every cached entry was computed against the previous
|
||||
/// cohort set and is now stale. Per-user `invalidate` is insufficient — a
|
||||
/// newly-defined cohort affects all users whose membership was already
|
||||
/// resolved, and there is no cheap way to know which of them now match the
|
||||
/// new predicate without re-resolving them.
|
||||
///
|
||||
/// A full clear is correct and cheap precisely because the cache is a pure
|
||||
/// accelerator: every dropped entry is recomputed from the registry on the
|
||||
/// next `resolve()`. The alternative (leaving the cache intact) silently
|
||||
/// drops attribution for the new cohort until each user is LRU-evicted or
|
||||
/// their metadata is rewritten — a no-error, no-log wrong answer.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Panics only if the internal cache mutex is poisoned.
|
||||
pub fn invalidate_all(&self) {
|
||||
let mut cache = self.cache.lock().expect("cohort resolver cache poisoned");
|
||||
cache.clear();
|
||||
}
|
||||
|
||||
/// Number of cached user resolutions (for diagnostics).
|
||||
///
|
||||
/// # Panics
|
||||
@ -266,6 +289,28 @@ mod tests {
|
||||
assert_eq!(resolver.cache_len(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalidate_all_clears_every_entry() {
|
||||
let reg = make_registry();
|
||||
let resolver = CohortResolver::new(reg);
|
||||
|
||||
// Populate several distinct users.
|
||||
let _ = resolver.resolve(1, &tech_user_metadata());
|
||||
let _ = resolver.resolve(2, &sports_user_metadata());
|
||||
let _ = resolver.resolve(3, &tech_user_metadata());
|
||||
assert_eq!(resolver.cache_len(), 3);
|
||||
|
||||
// A registry change invalidates every cached membership at once.
|
||||
resolver.invalidate_all();
|
||||
assert_eq!(resolver.cache_len(), 0);
|
||||
|
||||
// The cache still works after a full clear: the next resolve recomputes
|
||||
// and re-caches, since the cache is a pure accelerator.
|
||||
let memberships = resolver.resolve(1, &tech_user_metadata());
|
||||
assert!(memberships.iter().any(|n| n.to_string() == "tech_en"));
|
||||
assert_eq!(resolver.cache_len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_is_bounded_and_evicts_lru() {
|
||||
let reg = make_registry();
|
||||
@ -296,8 +341,7 @@ mod tests {
|
||||
// A capacity-1 cache still functions: resolves succeed and the result
|
||||
// is correct even though only one entry is retained at a time.
|
||||
let memberships = resolver.resolve(1, &tech_user_metadata());
|
||||
let names: Vec<String> = memberships.iter().map(|n| n.to_string()).collect();
|
||||
assert!(names.contains(&"tech_en".to_string()));
|
||||
assert!(memberships.iter().any(|n| n.to_string() == "tech_en"));
|
||||
assert_eq!(resolver.cache_len(), 1);
|
||||
}
|
||||
|
||||
|
||||
@ -110,10 +110,15 @@ impl CohortRegistry {
|
||||
/// Look up a cohort by name.
|
||||
///
|
||||
/// Returns a `DashMap` reference guard or `None` if not found.
|
||||
///
|
||||
/// The lookup probes by borrowed `&str`: `Arc<str>: Borrow<str>`, so
|
||||
/// `DashMap::get` hashes and compares the key in place with no per-probe
|
||||
/// allocation. (The previous `Arc::from(name)` allocated a throwaway
|
||||
/// `Arc<str>` on every call — including the resolver's per-cohort probes in
|
||||
/// its membership-resolution loop — purely to satisfy the key type.)
|
||||
#[must_use]
|
||||
pub fn get(&self, name: &str) -> Option<dashmap::mapref::one::Ref<'_, CohortName, CohortDef>> {
|
||||
let key: CohortName = Arc::from(name);
|
||||
self.cohorts.get(&key)
|
||||
self.cohorts.get(name)
|
||||
}
|
||||
|
||||
/// List all registered cohort names.
|
||||
@ -152,12 +157,25 @@ impl std::fmt::Debug for CohortRegistry {
|
||||
// ── Serialization ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// Serialize a cohort definition as JSON bytes for durable storage.
|
||||
#[must_use]
|
||||
pub fn serialize_cohort_def(def: &CohortDef) -> Vec<u8> {
|
||||
// serde_json::to_vec will not fail for these types (no maps with
|
||||
// non-string keys, no recursion limits exceeded). Falling back to
|
||||
// an empty vec is defensive but should never happen.
|
||||
serde_json::to_vec(def).unwrap_or_default()
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TidalError::Internal` if serialization fails. `serde_json::to_vec`
|
||||
/// cannot fail for the current `CohortDef`/`Predicate` shape (no maps with
|
||||
/// non-string keys, no recursion-limit overflow), so this is effectively
|
||||
/// infallible today. We propagate the error rather than swallowing it because
|
||||
/// the previous `unwrap_or_default()` turned the impossible failure into an
|
||||
/// empty `Vec` — and an empty `Vec` persisted to `Tag::CohortDef` deserializes
|
||||
/// back to `None` on the next open, silently dropping the cohort definition
|
||||
/// after a restart. Surfacing the error keeps that footgun out of the durable
|
||||
/// write path if a future `Predicate` variant ever makes serialization fallible.
|
||||
pub fn serialize_cohort_def(def: &CohortDef) -> Result<Vec<u8>, TidalError> {
|
||||
serde_json::to_vec(def).map_err(|e| {
|
||||
TidalError::internal(
|
||||
"serialize_cohort_def",
|
||||
format!("failed to serialize cohort '{}': {e}", def.name),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
/// Deserialize a cohort definition from JSON bytes.
|
||||
@ -323,7 +341,7 @@ mod tests {
|
||||
value: "en".into(),
|
||||
},
|
||||
};
|
||||
let bytes = serialize_cohort_def(&def);
|
||||
let bytes = serialize_cohort_def(&def).unwrap();
|
||||
let restored = deserialize_cohort_def(&bytes).unwrap();
|
||||
assert_eq!(restored.name, "tech_en");
|
||||
match &restored.predicate {
|
||||
@ -357,7 +375,7 @@ mod tests {
|
||||
]),
|
||||
]),
|
||||
};
|
||||
let bytes = serialize_cohort_def(&def);
|
||||
let bytes = serialize_cohort_def(&def).unwrap();
|
||||
let restored = deserialize_cohort_def(&bytes).unwrap();
|
||||
assert_eq!(restored.name, "compound");
|
||||
// Verify the structure survived.
|
||||
|
||||
@ -210,19 +210,10 @@ impl super::TidalDb {
|
||||
})
|
||||
})?;
|
||||
|
||||
// Cohort and co-engagement checkpoints are secondary state; log
|
||||
// and continue rather than aborting the backup if they fail.
|
||||
if self.cohort_ledger.entry_count() > 0
|
||||
&& let Err(e) = self.cohort_ledger.checkpoint(storage.items_engine(), meta)
|
||||
{
|
||||
tracing::warn!(error = %e, "backup: cohort checkpoint failed (non-fatal)");
|
||||
}
|
||||
|
||||
if self.co_engagement.edge_count() > 0
|
||||
&& let Err(e) = self.co_engagement.checkpoint(storage.items_engine())
|
||||
{
|
||||
tracing::warn!(error = %e, "backup: co-engagement checkpoint failed (non-fatal)");
|
||||
}
|
||||
// Cohort / co-engagement / community / preference checkpoints are
|
||||
// secondary derived state with no WAL backstop; log and continue
|
||||
// rather than aborting the backup if any fails.
|
||||
self.checkpoint_secondary_state(storage, meta);
|
||||
}
|
||||
|
||||
// 5. Flush text indexes (non-fatal: search indexes are rebuilt on
|
||||
@ -234,6 +225,37 @@ impl super::TidalDb {
|
||||
tracing::warn!(error = %e, "backup: creator text index flush failed (non-fatal)");
|
||||
}
|
||||
|
||||
// 5b. WAL writer in-flight events: why the copy below is still consistent.
|
||||
//
|
||||
// `backup_in_progress` (set by the CAS above) makes every NEW `signal()`
|
||||
// return Backpressure before it enqueues. The only writes it does NOT
|
||||
// cover are events a concurrent `signal()` already enqueued before it
|
||||
// observed the flag — those sit in the writer's crossbeam channel,
|
||||
// possibly mid-fsync, while we copy. We deliberately do NOT block on them,
|
||||
// and the resulting backup is still crash-consistent, for two reasons:
|
||||
//
|
||||
// 1. Torn tail batch in a copied segment — `WalHandle::open` runs
|
||||
// `reader::recover`, which two-phase-validates every batch (magic +
|
||||
// bounds, then BLAKE3) and truncates any corrupted/half-written tail.
|
||||
// A copy that captured a partially-fsynced tail batch is repaired on
|
||||
// the next open of the backup, exactly as a crash mid-write would be.
|
||||
//
|
||||
// 2. An enqueued-but-unflushed event missing from the copy — that event
|
||||
// belongs to a `signal()` call still blocked in `WalHandle::append`
|
||||
// (which only returns AFTER its batch is fsynced). The caller has not
|
||||
// yet been told the write succeeded, so its absence from the backup is
|
||||
// identical to a write that lost a race with a crash — an acceptable,
|
||||
// un-acked loss, not an inconsistency. Any in-flight event that DID
|
||||
// reach a copied segment carries seq > the WAL checkpoint marker
|
||||
// written in step 8, so backup replay re-applies it without
|
||||
// double-counting against the ledger checkpoint from step 4.
|
||||
//
|
||||
// A *guaranteed* drain (capturing acked-but-just-enqueued writes too)
|
||||
// would need a reply-after-fsync barrier the WAL does not yet expose;
|
||||
// `truncate_before` cannot serve as one because its reply is sent before
|
||||
// the in-progress batch is flushed. Tracked for a future
|
||||
// `WalCommand::Flush`; the current behavior is already crash-consistent.
|
||||
|
||||
// 6. Flush fjall storage (all keyspaces).
|
||||
//
|
||||
// This is fatal: if the memtable cannot be flushed, the backup
|
||||
@ -291,6 +313,39 @@ impl super::TidalDb {
|
||||
size_bytes,
|
||||
})
|
||||
}
|
||||
|
||||
/// Checkpoint the secondary derived ledgers (cohort, co-engagement,
|
||||
/// community, preference vectors) into `storage`.
|
||||
///
|
||||
/// Each is non-fatal on its own: these are derived aggregates whose absence
|
||||
/// degrades but never corrupts ranking, so a failure is logged and the others
|
||||
/// still run. Shared by the backup path; the shutdown path has the same set
|
||||
/// inline because it must additionally fold any failure into its first-error
|
||||
/// return.
|
||||
fn checkpoint_secondary_state(
|
||||
&self,
|
||||
storage: &super::storage_box::StorageBox,
|
||||
meta: crate::signals::checkpoint::CheckpointMeta,
|
||||
) {
|
||||
if self.cohort_ledger.entry_count() > 0
|
||||
&& let Err(e) = self.cohort_ledger.checkpoint(storage.items_engine(), meta)
|
||||
{
|
||||
tracing::warn!(error = %e, "backup: cohort checkpoint failed (non-fatal)");
|
||||
}
|
||||
if self.co_engagement.edge_count() > 0
|
||||
&& let Err(e) = self.co_engagement.checkpoint(storage.items_engine())
|
||||
{
|
||||
tracing::warn!(error = %e, "backup: co-engagement checkpoint failed (non-fatal)");
|
||||
}
|
||||
if let Err(e) = self.community_ledger.checkpoint(storage.items_engine()) {
|
||||
tracing::warn!(error = %e, "backup: community checkpoint failed (non-fatal)");
|
||||
}
|
||||
if !self.preference_vectors.is_empty()
|
||||
&& let Err(e) = self.preference_vectors.checkpoint(storage.items_engine())
|
||||
{
|
||||
tracing::warn!(error = %e, "backup: preference-vector checkpoint failed (non-fatal)");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@ -342,4 +397,76 @@ mod tests {
|
||||
"expected persistent error, got: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
/// Minimal schema: one Item "view" signal over the `AllTime` window
|
||||
/// (ephemeral WITHOUT a schema has no storage/ledger).
|
||||
fn minimal_schema() -> crate::schema::Schema {
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::schema::{DecaySpec, EntityKind, SchemaBuilder, Window};
|
||||
|
||||
let mut builder = SchemaBuilder::new();
|
||||
let _ = builder
|
||||
.signal(
|
||||
"view",
|
||||
EntityKind::Item,
|
||||
DecaySpec::Exponential {
|
||||
half_life: Duration::from_secs(3600),
|
||||
},
|
||||
)
|
||||
.windows(&[Window::AllTime])
|
||||
.velocity(false)
|
||||
.add();
|
||||
builder.build().expect("schema must be valid")
|
||||
}
|
||||
|
||||
/// Every signal whose `signal()` call has RETURNED before the backup must be
|
||||
/// present in the restored backup — this is the consistency guarantee that
|
||||
/// makes the no-explicit-WAL-flush copy safe (`WalHandle::append` only returns
|
||||
/// after its batch is fsynced, so an acked write is already on disk when the
|
||||
/// copy runs). The restored DB is opened fresh, exercising the full
|
||||
/// WAL-replay + torn-tail-recovery path on the copied data dir.
|
||||
#[test]
|
||||
fn backup_captures_all_acked_signals() {
|
||||
use crate::schema::{EntityId, Timestamp};
|
||||
|
||||
let src = tempfile::tempdir().unwrap();
|
||||
let bak = tempfile::tempdir().unwrap();
|
||||
let backup_path = bak.path().join("backup");
|
||||
|
||||
let db = crate::TidalDb::builder()
|
||||
.with_data_dir(src.path())
|
||||
.with_schema(minimal_schema())
|
||||
.open()
|
||||
.unwrap();
|
||||
|
||||
// Each of these returns only after its WAL batch is durably fsynced, so by
|
||||
// the time create_backup runs they are all on disk.
|
||||
for i in 1u64..=5 {
|
||||
db.signal("view", EntityId::new(i), 1.0, Timestamp::now())
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
let info = db.create_backup(&backup_path).unwrap();
|
||||
assert!(info.size_bytes > 0, "backup must copy non-empty data dir");
|
||||
db.close().unwrap();
|
||||
|
||||
// Reopen the backup from scratch: WAL replay + torn-tail recovery rebuild
|
||||
// the ledger. Every acked signal's decay score must be present.
|
||||
let restored = crate::TidalDb::builder()
|
||||
.with_data_dir(&backup_path)
|
||||
.with_schema(minimal_schema())
|
||||
.open()
|
||||
.unwrap();
|
||||
for i in 1u64..=5 {
|
||||
let score = restored
|
||||
.read_decay_score(EntityId::new(i), "view", 0)
|
||||
.unwrap();
|
||||
assert!(
|
||||
score.is_some_and(|s| s > 0.0),
|
||||
"acked signal for entity {i} must survive the backup/restore round-trip"
|
||||
);
|
||||
}
|
||||
restored.close().unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
@ -19,7 +19,12 @@ impl TidalDb {
|
||||
/// Grant an agent a capability token with the given scope permissions and a
|
||||
/// TTL. Returns the new token id. Persisted under `Tag::Capability`.
|
||||
///
|
||||
/// A `ttl` of `None` grants a non-expiring token.
|
||||
/// A `ttl` of `None` grants a non-expiring token. The expiry is stored as a
|
||||
/// nanosecond `u64` wall-clock timestamp, so the maximum representable TTL is
|
||||
/// `u64::MAX` ns (~584.9 years from now). A `ttl` whose nanoseconds exceed
|
||||
/// `u64::MAX`, or whose addition would overflow, saturates to `u64::MAX`
|
||||
/// (effectively non-expiring) rather than wrapping to a near-future instant
|
||||
/// that would expire the token immediately.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
@ -33,7 +38,13 @@ impl TidalDb {
|
||||
) -> crate::Result<String> {
|
||||
self.require_writeable("grant_capability")?;
|
||||
let now_ns = Timestamp::now().as_nanos();
|
||||
let expires_at_ns = ttl.map_or(u64::MAX, |d| now_ns.saturating_add(d.as_nanos() as u64));
|
||||
// `Duration::as_nanos` is u128; clamp to u64 BEFORE adding so a TTL above
|
||||
// `u64::MAX` ns (~584.9 years) does not truncate-then-wrap into a
|
||||
// near-future expiry. `saturating_add` then guards the addition itself.
|
||||
let expires_at_ns = ttl.map_or(u64::MAX, |d| {
|
||||
let ttl_ns = u64::try_from(d.as_nanos()).unwrap_or(u64::MAX);
|
||||
now_ns.saturating_add(ttl_ns)
|
||||
});
|
||||
let token_id = self.capability_registry.next_token_id(agent_id);
|
||||
let token = CapabilityToken {
|
||||
token_id: token_id.clone(),
|
||||
@ -230,3 +241,81 @@ fn parse_revocation_record(bytes: &[u8]) -> Option<(String, String, u64)> {
|
||||
let revoked_at_ns = v.get("revoked_at_ns")?.as_u64()?;
|
||||
Some((token_id, agent_id, revoked_at_ns))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::{
|
||||
db::TidalDb,
|
||||
governance::{ScopeClass, ScopePermission},
|
||||
};
|
||||
|
||||
/// A TTL whose nanoseconds exceed `u64::MAX` must saturate the stored expiry
|
||||
/// to `u64::MAX` (effectively non-expiring) instead of truncating the u128
|
||||
/// nanos to u64 and wrapping into a near-future instant that would expire the
|
||||
/// token almost immediately.
|
||||
#[test]
|
||||
fn very_long_ttl_saturates_instead_of_expiring_early() {
|
||||
let db = TidalDb::builder()
|
||||
.ephemeral()
|
||||
.open()
|
||||
.expect("ephemeral open should succeed");
|
||||
|
||||
// `Duration::MAX` is ~5.85e11 years; its nanos vastly exceed `u64::MAX`.
|
||||
// Before the fix `d.as_nanos() as u64` truncated this to a tiny value,
|
||||
// producing an expiry only a few seconds out.
|
||||
let ttl = std::time::Duration::MAX;
|
||||
let token_id = db
|
||||
.grant_capability(
|
||||
"agent-long-lived",
|
||||
vec![ScopePermission::read_write(ScopeClass::Community)],
|
||||
Some(ttl),
|
||||
)
|
||||
.expect("grant should succeed");
|
||||
|
||||
let tokens = db.agent_capabilities("agent-long-lived");
|
||||
let token = tokens
|
||||
.iter()
|
||||
.find(|t| t.token_id == token_id)
|
||||
.expect("granted token must be retrievable");
|
||||
|
||||
assert_eq!(
|
||||
token.expires_at_ns,
|
||||
u64::MAX,
|
||||
"an over-u64 TTL must saturate the expiry to u64::MAX, not wrap"
|
||||
);
|
||||
// The capability must still be live right now (and far into the future),
|
||||
// proving it did not expire early due to a truncated/wrapped expiry.
|
||||
assert!(
|
||||
db.check_capability("agent-long-lived", ScopeClass::Community, true),
|
||||
"saturated-expiry token must remain a live community-write capability"
|
||||
);
|
||||
assert!(
|
||||
!token.is_expired(u64::MAX - 1),
|
||||
"token must not be expired even near the far end of the u64 ns range"
|
||||
);
|
||||
}
|
||||
|
||||
/// A `None` TTL grants a non-expiring token (`expires_at_ns == u64::MAX`).
|
||||
#[test]
|
||||
fn none_ttl_is_non_expiring() {
|
||||
let db = TidalDb::builder()
|
||||
.ephemeral()
|
||||
.open()
|
||||
.expect("ephemeral open should succeed");
|
||||
|
||||
let token_id = db
|
||||
.grant_capability(
|
||||
"agent-forever",
|
||||
vec![ScopePermission::read_write(ScopeClass::Community)],
|
||||
None,
|
||||
)
|
||||
.expect("grant should succeed");
|
||||
|
||||
let tokens = db.agent_capabilities("agent-forever");
|
||||
let token = tokens
|
||||
.iter()
|
||||
.find(|t| t.token_id == token_id)
|
||||
.expect("granted token must be retrievable");
|
||||
assert_eq!(token.expires_at_ns, u64::MAX, "None TTL must never expire");
|
||||
}
|
||||
}
|
||||
|
||||
@ -27,7 +27,10 @@ impl TidalDb {
|
||||
// definition. (This is the one persist caller whose contract is
|
||||
// warn-and-continue rather than propagate.)
|
||||
let key = encode_key(EntityId::new(0), Tag::CohortDef, def.name.as_bytes());
|
||||
let value = serialize_cohort_def(&def);
|
||||
// A serialization failure is (today) impossible for these types, but if a
|
||||
// future Predicate variant makes it possible we fail loudly rather than
|
||||
// persist an empty Vec that silently drops the definition on restart.
|
||||
let value = serialize_cohort_def(&def)?;
|
||||
if let Err(e) = self.persist_to_items(&key, &value, "define_cohort") {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
@ -36,6 +39,20 @@ impl TidalDb {
|
||||
);
|
||||
}
|
||||
|
||||
self.cohort_registry.define(def)
|
||||
self.cohort_registry.define(def)?;
|
||||
|
||||
// The resolver caches per-user memberships computed against the cohort
|
||||
// set as it was when each user was last resolved. A newly-defined cohort
|
||||
// changes that set for everyone, so every cached membership is now stale
|
||||
// and would silently omit this cohort — its ledger would accumulate zero
|
||||
// signal for any already-resolved user until they were LRU-evicted or
|
||||
// their metadata was rewritten. Clear the whole cache so the next signal
|
||||
// for each user re-resolves against the new registry. We do this only
|
||||
// after `define` succeeds, so a rejected definition (e.g. duplicate
|
||||
// name) does not needlessly flush the accelerator. The clear is cheap:
|
||||
// the cache is a pure accelerator and recomputes every entry on miss.
|
||||
self.cohort_resolver.invalidate_all();
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@ -296,6 +296,26 @@ impl TidalDb {
|
||||
effective_weight,
|
||||
timestamp.as_nanos(),
|
||||
);
|
||||
|
||||
// ZERO loss window for an accepted contribution: the community
|
||||
// aggregate is the SOLE durable copy of per-contributor identity (the
|
||||
// WAL community event carries none), so a periodic checkpoint alone
|
||||
// would still lose this contribution on a crash before the next tick.
|
||||
// The `signal_for_community` rustdoc promises `Ok(true)` means
|
||||
// "durably logged" — to honor that for the per-contributor aggregate we
|
||||
// synchronously checkpoint it here BEFORE returning Ok(true). The
|
||||
// checkpoint stages the full snapshot into one atomic WriteBatch, so
|
||||
// the durable Tag::Community state is never behind an accepted write.
|
||||
if let Some(ref storage) = self.storage
|
||||
&& let Err(e) = self.community_ledger.checkpoint(storage.items_engine())
|
||||
{
|
||||
// A failed durability flush must not be reported as a durable
|
||||
// accept: surface it so the caller can retry rather than believing
|
||||
// the contribution survived a crash.
|
||||
return Err(TidalError::Durability(crate::schema::DurabilityError {
|
||||
message: format!("community contribution checkpoint failed: {e}"),
|
||||
}));
|
||||
}
|
||||
}
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
@ -5,11 +5,7 @@ use std::collections::HashMap;
|
||||
use super::TidalDb;
|
||||
use crate::{
|
||||
schema::{EntityId, EntityKind, TidalError},
|
||||
storage::vector::{
|
||||
BruteForceIndex, QuantizationLevel, VectorIndexConfig, deserialize_embedding,
|
||||
insert_embedding,
|
||||
registry::{EmbeddingSlotState, EmbeddingSource, HnswParams},
|
||||
},
|
||||
storage::vector::deserialize_embedding,
|
||||
};
|
||||
|
||||
impl TidalDb {
|
||||
@ -28,14 +24,17 @@ impl TidalDb {
|
||||
let storage = self.storage()?;
|
||||
super::users::write_entity_meta(storage.creators_engine(), id, metadata)?;
|
||||
|
||||
// Enqueue for creator text index (best-effort).
|
||||
// Enqueue for the async creator text syncer (best-effort). The creator
|
||||
// entity store is the source of truth; the text index is derived state,
|
||||
// and any write the syncer has not yet committed before a crash is
|
||||
// recovered by the unconditional rebuild-from-store at open (see
|
||||
// `db::mod::rebuild_text_indexes_at_open`).
|
||||
if let Ok(guard) = self.creator_text_tx.lock()
|
||||
&& let Some(tx) = guard.as_ref()
|
||||
{
|
||||
let _ = tx.send(crate::text::PendingWrite {
|
||||
entity_id: id,
|
||||
metadata: metadata.clone(),
|
||||
seq: 0,
|
||||
deleted: false,
|
||||
});
|
||||
}
|
||||
@ -71,84 +70,17 @@ impl TidalDb {
|
||||
/// - `TidalError::Internal` if no storage backend is wired or lock is poisoned.
|
||||
/// - `TidalError::Storage` on storage engine failure.
|
||||
/// - `TidalError::Internal` if the embedding has zero norm.
|
||||
#[allow(clippy::significant_drop_tightening)] // lock must be held across insert_embedding call
|
||||
pub fn write_creator_embedding(&self, id: EntityId, embedding: &[f32]) -> crate::Result<()> {
|
||||
self.require_writeable("write_creator_embedding")?;
|
||||
let storage = self.storage()?;
|
||||
|
||||
// Auto-register the creator "content" slot if absent.
|
||||
{
|
||||
let mut registry = self.embedding_registry.write().map_err(|_| {
|
||||
TidalError::internal(
|
||||
"write_creator_embedding",
|
||||
"embedding_registry write lock poisoned",
|
||||
)
|
||||
})?;
|
||||
if registry.get(EntityKind::Creator, "content").is_none() {
|
||||
let config = VectorIndexConfig {
|
||||
dimensions: embedding.len(),
|
||||
..VectorIndexConfig::default()
|
||||
};
|
||||
let state = EmbeddingSlotState {
|
||||
index: Box::new(BruteForceIndex::new(config)),
|
||||
dimensions: embedding.len(),
|
||||
quantization: QuantizationLevel::F32,
|
||||
source: EmbeddingSource::External,
|
||||
params: HnswParams::default(),
|
||||
};
|
||||
registry
|
||||
.register(EntityKind::Creator, "content".to_string(), state)
|
||||
.map_err(|e| {
|
||||
TidalError::internal(
|
||||
"write_creator_embedding",
|
||||
format!("slot registration failed: {e}"),
|
||||
)
|
||||
})?;
|
||||
}
|
||||
}
|
||||
|
||||
// Insert into the HNSW index and store in entity engine.
|
||||
// Read the slot dimensions first, dropping the lock before the storage write.
|
||||
let dimensions = {
|
||||
let registry = self.embedding_registry.read().map_err(|_| {
|
||||
TidalError::internal(
|
||||
"write_creator_embedding",
|
||||
"embedding_registry read lock poisoned",
|
||||
)
|
||||
})?;
|
||||
registry
|
||||
.get(EntityKind::Creator, "content")
|
||||
.ok_or_else(|| {
|
||||
TidalError::internal("write_creator_embedding", "creator content slot missing")
|
||||
})?
|
||||
.dimensions
|
||||
};
|
||||
// Re-acquire read lock to get the index reference, perform the insert,
|
||||
// then drop the lock.
|
||||
{
|
||||
let registry = self.embedding_registry.read().map_err(|_| {
|
||||
TidalError::internal(
|
||||
"write_creator_embedding",
|
||||
"embedding_registry read lock poisoned",
|
||||
)
|
||||
})?;
|
||||
let slot = registry
|
||||
.get(EntityKind::Creator, "content")
|
||||
.ok_or_else(|| {
|
||||
TidalError::internal("write_creator_embedding", "creator content slot missing")
|
||||
})?;
|
||||
insert_embedding(
|
||||
id,
|
||||
"content",
|
||||
embedding,
|
||||
dimensions,
|
||||
slot.index.as_ref(),
|
||||
storage.creators_engine(),
|
||||
)
|
||||
.map_err(|e| TidalError::internal("write_creator_embedding", format!("{e}")))?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
self.write_entity_embedding(
|
||||
id,
|
||||
EntityKind::Creator,
|
||||
"content",
|
||||
embedding,
|
||||
storage.creators_engine(),
|
||||
"write_creator_embedding",
|
||||
)
|
||||
}
|
||||
|
||||
/// Read a stored creator embedding for a given entity ID.
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
//! RLHF training data export and cross-session aggregation types.
|
||||
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::collections::{BinaryHeap, HashMap, HashSet};
|
||||
|
||||
use crate::schema::{SchemaError, SignalTypeDef, TidalError};
|
||||
|
||||
@ -256,12 +256,120 @@ impl super::TidalDb {
|
||||
// Batch WAL events: skip entirely when user_id filter is set because
|
||||
// batch WAL records have no user context — including them would pollute
|
||||
// user-filtered results with anonymous events.
|
||||
//
|
||||
// Bounded scan (was: load every WAL event into one unbounded Vec, then
|
||||
// truncate). `read_all_events` materialized the *entire* pre-compaction
|
||||
// WAL — tens of millions of events under write load — regardless of the
|
||||
// requested limit, so a `limit: Some(10)` export could OOM. We instead
|
||||
// stream segments and retain at most `effective_limit` events in memory
|
||||
// via a bounded heap keyed by timestamp (see `scan_batch_wal_bounded`).
|
||||
let mut results = Vec::new();
|
||||
if request.user_id.is_none() {
|
||||
let all_events = crate::wal::reader::read_all_events(&wal_dir)
|
||||
.map_err(|e| TidalError::internal("export_signals", e.to_string()))?;
|
||||
let batch_signals = scan_batch_wal_bounded(
|
||||
&wal_dir,
|
||||
request,
|
||||
&type_filter,
|
||||
&type_names,
|
||||
effective_limit,
|
||||
)
|
||||
.map_err(|e| TidalError::internal("export_signals", e.to_string()))?;
|
||||
results.extend(batch_signals);
|
||||
}
|
||||
|
||||
for event in all_events {
|
||||
// Session journal events: always read (they carry user context).
|
||||
let session_signals = read_session_journal_signals(&wal_dir, request);
|
||||
results.extend(session_signals);
|
||||
|
||||
// Sort merged results by timestamp ascending for deterministic output.
|
||||
results.sort_by_key(|s| s.timestamp_ns);
|
||||
|
||||
// Apply limit across the merged, sorted result.
|
||||
results.truncate(effective_limit);
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
}
|
||||
|
||||
// ── scan_batch_wal_bounded ────────────────────────────────────────────────────
|
||||
|
||||
/// Stream the batch WAL segments and return at most `limit` matching events,
|
||||
/// holding no more than `limit` events in memory at any point.
|
||||
///
|
||||
/// # Why this is not `read_all_events`
|
||||
///
|
||||
/// [`read_all_events`](crate::wal::reader::read_all_events) decodes EVERY event
|
||||
/// in EVERY segment into one `Vec` before any filter or limit applies. Between
|
||||
/// checkpoints the WAL can hold tens of millions of events (online compaction
|
||||
/// only deletes segments fully covered by a checkpoint), so a small-limit export
|
||||
/// on a large pre-compaction WAL allocated the whole tail and could OOM. This
|
||||
/// scanner pushes the time-range + signal-type filters and the `limit` *down*
|
||||
/// into the segment-by-segment scan and retains only the surviving window.
|
||||
///
|
||||
/// # Which events survive
|
||||
///
|
||||
/// To preserve the prior output (the earliest-N events in the range — the merge
|
||||
/// step then re-sorts ascending by timestamp and truncates), this keeps the
|
||||
/// `limit` events with the SMALLEST `timestamp_ns` seen across all segments. A
|
||||
/// max-heap keyed by timestamp evicts the current largest once the heap exceeds
|
||||
/// `limit`, so the working set is capped at `limit` regardless of WAL size. With
|
||||
/// `limit == 0` no events are retained.
|
||||
///
|
||||
/// Segments are scanned read-only ([`scan_segment_readonly`] never truncates a
|
||||
/// torn tail), so this is safe to run concurrently with the live writer thread.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns [`WalError`](crate::wal::error::WalError) if a segment file cannot be
|
||||
/// listed or opened. A torn tail in the active segment is ignored, not an error.
|
||||
fn scan_batch_wal_bounded(
|
||||
wal_dir: &std::path::Path,
|
||||
request: &ExportRequest,
|
||||
type_filter: &HashSet<u8>,
|
||||
type_names: &HashMap<u8, String>,
|
||||
limit: usize,
|
||||
) -> Result<Vec<ExportedSignal>, crate::wal::error::WalError> {
|
||||
use crate::wal::{reader::scan_segment_readonly, segment::list_segments};
|
||||
|
||||
/// Heap element ordered ONLY by its `(timestamp, ordinal)` key.
|
||||
///
|
||||
/// `ExportedSignal` is not `Ord` (it carries floats), so we cannot put it in
|
||||
/// a `BinaryHeap` directly. This wrapper derives ordering from the key alone:
|
||||
/// a plain max-heap of `HeapEntry` therefore yields the LARGEST timestamp at
|
||||
/// `peek`/`pop` — exactly the element to evict when the heap exceeds `limit`,
|
||||
/// leaving the `limit` smallest-timestamp (earliest) survivors. The `ordinal`
|
||||
/// breaks timestamp ties deterministically so eviction is stable across runs.
|
||||
struct HeapEntry {
|
||||
key: (u64, u64),
|
||||
signal: ExportedSignal,
|
||||
}
|
||||
impl PartialEq for HeapEntry {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
self.key == other.key
|
||||
}
|
||||
}
|
||||
impl Eq for HeapEntry {}
|
||||
impl PartialOrd for HeapEntry {
|
||||
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
|
||||
Some(self.cmp(other))
|
||||
}
|
||||
}
|
||||
impl Ord for HeapEntry {
|
||||
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
|
||||
self.key.cmp(&other.key)
|
||||
}
|
||||
}
|
||||
|
||||
let mut heap: BinaryHeap<HeapEntry> = BinaryHeap::new();
|
||||
let mut ordinal: u64 = 0;
|
||||
|
||||
// A zero limit cannot retain anything; skip the scan entirely.
|
||||
if limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
for (_seq, path) in &list_segments(wal_dir)? {
|
||||
for (_header, events) in scan_segment_readonly(path)? {
|
||||
for event in events {
|
||||
// Time range filter: since is inclusive, until is exclusive.
|
||||
if let Some(since) = request.since
|
||||
&& event.timestamp_nanos < since
|
||||
@ -279,13 +387,12 @@ impl super::TidalDb {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Resolve signal type name from the u8 ID.
|
||||
let signal_type_name = type_names
|
||||
.get(&event.signal_type)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| format!("unknown:{}", event.signal_type));
|
||||
|
||||
results.push(ExportedSignal {
|
||||
let exported = ExportedSignal {
|
||||
entity_id: event.entity_id,
|
||||
signal_type: signal_type_name,
|
||||
weight: event.weight,
|
||||
@ -293,22 +400,24 @@ impl super::TidalDb {
|
||||
user_id: None,
|
||||
session_id: None,
|
||||
annotation: None,
|
||||
};
|
||||
|
||||
heap.push(HeapEntry {
|
||||
key: (event.timestamp_nanos, ordinal),
|
||||
signal: exported,
|
||||
});
|
||||
ordinal += 1;
|
||||
|
||||
// Cap the working set: evict the current largest timestamp so the
|
||||
// heap never exceeds `limit` events.
|
||||
if heap.len() > limit {
|
||||
heap.pop();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Session journal events: always read (they carry user context).
|
||||
let session_signals = read_session_journal_signals(&wal_dir, request);
|
||||
results.extend(session_signals);
|
||||
|
||||
// Sort merged results by timestamp ascending for deterministic output.
|
||||
results.sort_by_key(|s| s.timestamp_ns);
|
||||
|
||||
// Apply limit across the merged, sorted result.
|
||||
results.truncate(effective_limit);
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
Ok(heap.into_iter().map(|entry| entry.signal).collect())
|
||||
}
|
||||
|
||||
// ── read_session_journal_signals ──────────────────────────────────────────────
|
||||
@ -318,6 +427,9 @@ impl super::TidalDb {
|
||||
///
|
||||
/// Returns an empty vec if the session journal does not exist or cannot be read
|
||||
/// (non-fatal — the batch WAL results are still valid).
|
||||
// The export wire intentionally narrows the f64 session weight to f32 for the
|
||||
// RLHF/training output format; the durable journal keeps full precision.
|
||||
#[allow(clippy::cast_possible_truncation)]
|
||||
fn read_session_journal_signals(
|
||||
wal_dir: &std::path::Path,
|
||||
request: &ExportRequest,
|
||||
@ -394,8 +506,11 @@ fn read_session_journal_signals(
|
||||
|
||||
results.push(ExportedSignal {
|
||||
entity_id: *entity_id,
|
||||
// The RLHF export wire keeps weight as f32 (a display/training
|
||||
// format, not the durable score); the session journal stores it
|
||||
// as f64, so narrow explicitly here at the export boundary.
|
||||
weight: *weight as f32,
|
||||
signal_type: signal_name.clone(),
|
||||
weight: *weight,
|
||||
timestamp_ns: *ts_ns,
|
||||
user_id: mapped_user_id,
|
||||
session_id: Some(*session_id),
|
||||
|
||||
@ -170,6 +170,99 @@ fn exported_signal_json_lines_stream_stays_one_record_per_line() {
|
||||
}
|
||||
}
|
||||
|
||||
/// The bounded WAL scanner holds at most `limit` events in memory even when the
|
||||
/// WAL is far larger than the limit, and returns the EARLIEST-`limit` events in
|
||||
/// the range (the merge step in `export_signals` re-sorts ascending and
|
||||
/// truncates, so earliest-N preserves the prior output shape).
|
||||
///
|
||||
/// Regression: `export_signals` used to call `read_all_events`, which decoded
|
||||
/// EVERY event in EVERY segment into one `Vec` before the limit applied — a
|
||||
/// `limit: Some(10)` export over a multi-million-event pre-compaction WAL could
|
||||
/// OOM. This proves the working set is now bounded by the limit.
|
||||
#[test]
|
||||
fn scan_batch_wal_bounded_caps_working_set_and_keeps_earliest() {
|
||||
use std::collections::{HashMap, HashSet};
|
||||
|
||||
use crate::replication::ShardId;
|
||||
use crate::wal::{format::EventRecord, segment::segment_filename};
|
||||
|
||||
// Write a WAL FAR larger than the export limit: 50 segments * 100 events =
|
||||
// 5_000 events, each timestamped by its global ordinal so "earliest" is the
|
||||
// lowest ordinals. The limit is 10, so 4_990 events must never be retained.
|
||||
const SEGMENTS: u64 = 50;
|
||||
const EVENTS_PER_SEG: u64 = 100;
|
||||
const TOTAL: u64 = SEGMENTS * EVENTS_PER_SEG;
|
||||
const LIMIT: usize = 10;
|
||||
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
||||
let mut ordinal: u64 = 0;
|
||||
for seg in 0..SEGMENTS {
|
||||
let first_seq = seg * EVENTS_PER_SEG + 1;
|
||||
let events: Vec<EventRecord> = (0..EVENTS_PER_SEG)
|
||||
.map(|_| {
|
||||
// signal_type 0, weight 1.0, timestamp = global ordinal.
|
||||
let ev = EventRecord::signal(ordinal, 0, 1.0, ordinal);
|
||||
ordinal += 1;
|
||||
ev
|
||||
})
|
||||
.collect();
|
||||
let bytes = crate::wal::format::encode_batch(&events, first_seq, first_seq).unwrap();
|
||||
let name = segment_filename(ShardId::SINGLE, first_seq);
|
||||
std::fs::write(dir.path().join(name), &bytes).unwrap();
|
||||
}
|
||||
assert_eq!(ordinal, TOTAL);
|
||||
|
||||
// No type filter; map id 0 -> "view".
|
||||
let type_filter: HashSet<u8> = HashSet::new();
|
||||
let mut type_names: HashMap<u8, String> = HashMap::new();
|
||||
type_names.insert(0, "view".to_string());
|
||||
|
||||
let req = ExportRequest::time_range(0, u64::MAX);
|
||||
let got = super::scan_batch_wal_bounded(dir.path(), &req, &type_filter, &type_names, LIMIT)
|
||||
.expect("bounded scan should succeed");
|
||||
|
||||
// Exactly `limit` events survive — never the whole 5_000-event WAL.
|
||||
assert_eq!(
|
||||
got.len(),
|
||||
LIMIT,
|
||||
"bounded scan must cap the result at the limit"
|
||||
);
|
||||
|
||||
// The survivors are the EARLIEST `limit` events (ordinals/timestamps 0..LIMIT).
|
||||
let mut timestamps: Vec<u64> = got.iter().map(|s| s.timestamp_ns).collect();
|
||||
timestamps.sort_unstable();
|
||||
let expected: Vec<u64> = (0..LIMIT as u64).collect();
|
||||
assert_eq!(
|
||||
timestamps, expected,
|
||||
"bounded scan must retain the earliest-N events, not an arbitrary window"
|
||||
);
|
||||
}
|
||||
|
||||
/// A zero limit retains nothing and never touches the heap, even with a
|
||||
/// populated WAL — the early-return guard short-circuits the scan.
|
||||
#[test]
|
||||
fn scan_batch_wal_bounded_zero_limit_returns_empty() {
|
||||
use std::collections::{HashMap, HashSet};
|
||||
|
||||
use crate::replication::ShardId;
|
||||
use crate::wal::{format::EventRecord, segment::segment_filename};
|
||||
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let events = vec![EventRecord::signal(1, 0, 1.0, 100)];
|
||||
let bytes = crate::wal::format::encode_batch(&events, 1, 100).unwrap();
|
||||
std::fs::write(
|
||||
dir.path().join(segment_filename(ShardId::SINGLE, 1)),
|
||||
&bytes,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let req = ExportRequest::time_range(0, u64::MAX);
|
||||
let got = super::scan_batch_wal_bounded(dir.path(), &req, &HashSet::new(), &HashMap::new(), 0)
|
||||
.expect("bounded scan should succeed");
|
||||
assert!(got.is_empty(), "zero limit must retain no events");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn export_limit_exceeds_max_returns_error() {
|
||||
use crate::TidalDb;
|
||||
|
||||
@ -5,14 +5,7 @@ use std::collections::HashMap;
|
||||
use super::{TidalDb, metadata::deserialize_metadata};
|
||||
use crate::{
|
||||
schema::{EntityId, EntityKind, TidalError, Timestamp},
|
||||
storage::{
|
||||
Tag, encode_key,
|
||||
vector::{
|
||||
BruteForceIndex, QuantizationLevel, VectorIndexConfig, deserialize_embedding,
|
||||
insert_embedding,
|
||||
registry::{EmbeddingSlotState, EmbeddingSource, HnswParams},
|
||||
},
|
||||
},
|
||||
storage::{Tag, encode_key, vector::deserialize_embedding},
|
||||
};
|
||||
|
||||
impl TidalDb {
|
||||
@ -22,6 +15,20 @@ impl TidalDb {
|
||||
/// to storage AND inserts the entity into the bitmap, range, and universe
|
||||
/// indexes so it is discoverable by RETRIEVE queries.
|
||||
///
|
||||
/// # Item-universe limit
|
||||
///
|
||||
/// The in-memory candidate-generation indexes (universe bitmap, category /
|
||||
/// format / creator / tag bitmaps, range indexes) are keyed on a `u32`
|
||||
/// entity ID, which bounds the item universe to `u32::MAX` (~4 billion)
|
||||
/// items per instance. This is well beyond the single-node target of 10M
|
||||
/// items, but it is a hard contract: an `EntityId` whose `as_u64()` exceeds
|
||||
/// `u32::MAX` would alias a lower ID in every bitmap, silently colliding
|
||||
/// query results across distinct items. Such an ID is therefore rejected at
|
||||
/// the write boundary with `TidalError::invalid_input` rather than truncated
|
||||
/// — embedders must allocate dense item IDs within the `u32` range. If a
|
||||
/// larger universe is ever required, widen the bitmap key type behind the
|
||||
/// index abstraction instead of relaxing this guard.
|
||||
///
|
||||
/// Recognized metadata keys for indexing:
|
||||
/// - `"category"` -> category bitmap index
|
||||
/// - `"format"` -> format bitmap index
|
||||
@ -32,6 +39,8 @@ impl TidalDb {
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// - `TidalError::InvalidInput` if `id.as_u64()` exceeds the `u32`
|
||||
/// item-universe limit.
|
||||
/// - `TidalError::Internal` if no storage backend is wired.
|
||||
/// - `TidalError::Storage` on storage engine failure.
|
||||
#[allow(clippy::too_many_lines)]
|
||||
@ -42,6 +51,25 @@ impl TidalDb {
|
||||
metadata: &HashMap<String, String>,
|
||||
) -> crate::Result<()> {
|
||||
self.require_writeable("write_item_with_metadata")?;
|
||||
|
||||
// Enforce the u32 item-universe limit BEFORE any durable write. The
|
||||
// in-memory candidate-generation indexes key on a u32 id, so an id past
|
||||
// u32::MAX would alias a lower id in every bitmap (silent cross-id
|
||||
// collision in query results). Reject it here rather than truncate so a
|
||||
// rejected item never lands in storage in an unindexable state. NOTE:
|
||||
// the sibling write paths in signals.rs / relationships.rs /
|
||||
// collections.rs / user_state.rs still perform a bare `as u32` narrowing
|
||||
// on item-side ids; they should be routed through a shared guard so the
|
||||
// limit is enforced uniformly (out of scope for this change).
|
||||
if id.as_u64() > u64::from(u32::MAX) {
|
||||
let raw = id.as_u64();
|
||||
return Err(TidalError::invalid_input(format!(
|
||||
"item entity id {raw} exceeds the u32 item-universe limit ({}); \
|
||||
allocate dense item IDs within the u32 range",
|
||||
u32::MAX
|
||||
)));
|
||||
}
|
||||
|
||||
// Validate metadata size limits.
|
||||
const MAX_METADATA_KEYS: usize = 64;
|
||||
const MAX_KEY_BYTES: usize = 512;
|
||||
@ -108,16 +136,10 @@ impl TidalDb {
|
||||
// durable-write step; all indexing happens below via the shared path.
|
||||
self.persist_item_metadata(id, &stored_metadata)?;
|
||||
|
||||
// Truncate entity ID to u32 for roaring bitmap. This limits the universe
|
||||
// to ~4 billion entities per instance, which is well beyond the single-node
|
||||
// target of 10M items.
|
||||
// Narrow the entity ID to u32 for the roaring bitmap / range indexes.
|
||||
// The boundary guard above already rejected any id > u32::MAX, so this
|
||||
// cast is lossless: no truncation, no cross-id collision is possible.
|
||||
let id_u32 = id.as_u64() as u32;
|
||||
if id.as_u64() > u64::from(u32::MAX) {
|
||||
tracing::warn!(
|
||||
entity_id = id.as_u64(),
|
||||
"entity ID exceeds u32::MAX; universe bitmap entry will collide with a lower ID"
|
||||
);
|
||||
}
|
||||
|
||||
// Index via the SAME shared path the restart rebuild uses, so the two
|
||||
// can never diverge (created_at default, creator_items population, etc.).
|
||||
@ -150,14 +172,17 @@ impl TidalDb {
|
||||
}
|
||||
}
|
||||
|
||||
// Enqueue for text index (best-effort -- channel missing or dropped is non-fatal).
|
||||
// Enqueue for the async text syncer (best-effort -- channel missing or
|
||||
// dropped is non-fatal). The durable entity store is the source of
|
||||
// truth; the text index is derived state, and any write the syncer has
|
||||
// not yet committed before a crash is recovered by the unconditional
|
||||
// rebuild-from-store at open (see `db::mod::rebuild_text_indexes_at_open`).
|
||||
if let Ok(guard) = self.text_tx.lock()
|
||||
&& let Some(tx) = guard.as_ref()
|
||||
{
|
||||
let _ = tx.send(crate::text::PendingWrite {
|
||||
entity_id: id,
|
||||
metadata: (*stored_metadata).clone(),
|
||||
seq: 0, // seq tracking is best-effort for now
|
||||
deleted: false,
|
||||
});
|
||||
}
|
||||
@ -247,80 +272,18 @@ impl TidalDb {
|
||||
/// - `TidalError::Internal` if no storage backend is wired or lock is poisoned.
|
||||
/// - `TidalError::Storage` on storage engine failure.
|
||||
/// - `TidalError::Internal` if the embedding has zero norm.
|
||||
#[allow(clippy::significant_drop_tightening)] // lock must be held across insert_embedding call
|
||||
pub fn write_item_embedding(&self, id: EntityId, embedding: &[f32]) -> crate::Result<()> {
|
||||
self.require_writeable("write_item_embedding")?;
|
||||
let storage = self.storage()?;
|
||||
let slot_name = self.item_embedding_slot();
|
||||
|
||||
// Auto-register the item embedding slot if absent.
|
||||
{
|
||||
let mut registry = self.embedding_registry.write().map_err(|_| {
|
||||
TidalError::internal(
|
||||
"write_item_embedding",
|
||||
"embedding_registry write lock poisoned",
|
||||
)
|
||||
})?;
|
||||
if registry.get(EntityKind::Item, slot_name).is_none() {
|
||||
let config = VectorIndexConfig {
|
||||
dimensions: embedding.len(),
|
||||
..VectorIndexConfig::default()
|
||||
};
|
||||
let state = EmbeddingSlotState {
|
||||
index: Box::new(BruteForceIndex::new(config)),
|
||||
dimensions: embedding.len(),
|
||||
quantization: QuantizationLevel::F32,
|
||||
source: EmbeddingSource::External,
|
||||
params: HnswParams::default(),
|
||||
};
|
||||
registry
|
||||
.register(EntityKind::Item, slot_name.to_string(), state)
|
||||
.map_err(|e| {
|
||||
TidalError::internal(
|
||||
"write_item_embedding",
|
||||
format!("slot registration failed: {e}"),
|
||||
)
|
||||
})?;
|
||||
}
|
||||
}
|
||||
|
||||
// Insert into the HNSW index and store in entity engine.
|
||||
let dimensions = {
|
||||
let registry = self.embedding_registry.read().map_err(|_| {
|
||||
TidalError::internal(
|
||||
"write_item_embedding",
|
||||
"embedding_registry read lock poisoned",
|
||||
)
|
||||
})?;
|
||||
registry
|
||||
.get(EntityKind::Item, slot_name)
|
||||
.ok_or_else(|| {
|
||||
TidalError::internal("write_item_embedding", "item embedding slot missing")
|
||||
})?
|
||||
.dimensions
|
||||
};
|
||||
{
|
||||
let registry = self.embedding_registry.read().map_err(|_| {
|
||||
TidalError::internal(
|
||||
"write_item_embedding",
|
||||
"embedding_registry read lock poisoned",
|
||||
)
|
||||
})?;
|
||||
let slot = registry.get(EntityKind::Item, slot_name).ok_or_else(|| {
|
||||
TidalError::internal("write_item_embedding", "item embedding slot missing")
|
||||
})?;
|
||||
insert_embedding(
|
||||
id,
|
||||
slot_name,
|
||||
embedding,
|
||||
dimensions,
|
||||
slot.index.as_ref(),
|
||||
storage.items_engine(),
|
||||
)
|
||||
.map_err(|e| TidalError::internal("write_item_embedding", format!("{e}")))?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
self.write_entity_embedding(
|
||||
id,
|
||||
EntityKind::Item,
|
||||
slot_name,
|
||||
embedding,
|
||||
storage.items_engine(),
|
||||
"write_item_embedding",
|
||||
)
|
||||
}
|
||||
|
||||
/// Read the stored content embedding for an item.
|
||||
@ -472,12 +435,11 @@ impl TidalDb {
|
||||
let Some(idx) = self.text_index.as_ref() else {
|
||||
return Ok(()); // no text fields declared -> nothing to rebuild
|
||||
};
|
||||
let last_seq = self.last_wal_seq.load(std::sync::atomic::Ordering::Acquire);
|
||||
// Materialize first so the durable scan completes before touching the
|
||||
// Tantivy writer (rebuild_from deletes all docs up front).
|
||||
// `all_item_metadata` already excludes `EMB:` embedding rows.
|
||||
let items = self.all_item_metadata()?;
|
||||
idx.rebuild_from(items.into_iter(), last_seq)?;
|
||||
idx.rebuild_from(items.into_iter())?;
|
||||
idx.reload_reader()?;
|
||||
// Resynced: clear the latch (the live is_healthy() flag is owned and
|
||||
// reset by the syncer itself once it recovers).
|
||||
@ -512,6 +474,25 @@ mod tests {
|
||||
|
||||
use crate::{TidalDb, schema::EntityId};
|
||||
|
||||
/// Minimal valid schema so ephemeral mode wires in-memory storage (needed by
|
||||
/// any test that actually persists an item).
|
||||
fn minimal_schema() -> crate::schema::Schema {
|
||||
use crate::schema::{DecaySpec, EntityKind, SchemaBuilder, Window};
|
||||
let mut b = SchemaBuilder::new();
|
||||
let _ = b
|
||||
.signal(
|
||||
"view",
|
||||
EntityKind::Item,
|
||||
DecaySpec::Exponential {
|
||||
half_life: std::time::Duration::from_secs(3600),
|
||||
},
|
||||
)
|
||||
.windows(&[Window::AllTime])
|
||||
.velocity(false)
|
||||
.add();
|
||||
b.build().expect("schema must be valid")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_item_rejects_oversized_metadata_value() {
|
||||
let db = TidalDb::builder().ephemeral().open().unwrap();
|
||||
@ -562,4 +543,45 @@ mod tests {
|
||||
);
|
||||
db.close().unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_item_rejects_id_above_u32_universe_limit() {
|
||||
let db = TidalDb::builder()
|
||||
.ephemeral()
|
||||
.with_schema(minimal_schema())
|
||||
.open()
|
||||
.unwrap();
|
||||
let meta = HashMap::new();
|
||||
// u32::MAX + 1: the smallest id that would alias a lower id once
|
||||
// narrowed to u32 for the candidate-generation bitmaps.
|
||||
let oversized = EntityId::new(u64::from(u32::MAX) + 1);
|
||||
let err = db.write_item_with_metadata(oversized, &meta).unwrap_err();
|
||||
assert!(
|
||||
err.to_string().contains("u32 item-universe limit"),
|
||||
"expected u32 item-universe rejection, got: {err}"
|
||||
);
|
||||
// The rejected item must not have leaked into storage or the universe
|
||||
// bitmap — a rejected write leaves no trace.
|
||||
assert_eq!(db.item_count(), 0, "rejected item must not enter universe");
|
||||
assert!(
|
||||
db.get_item_metadata(oversized).unwrap().is_none(),
|
||||
"rejected item must not be persisted"
|
||||
);
|
||||
db.close().unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_item_accepts_id_at_u32_universe_boundary() {
|
||||
let db = TidalDb::builder()
|
||||
.ephemeral()
|
||||
.with_schema(minimal_schema())
|
||||
.open()
|
||||
.unwrap();
|
||||
let meta = HashMap::new();
|
||||
// u32::MAX is the largest in-bounds id: narrowing is lossless.
|
||||
let boundary = EntityId::new(u64::from(u32::MAX));
|
||||
db.write_item_with_metadata(boundary, &meta).unwrap();
|
||||
assert_eq!(db.item_count(), 1, "boundary id must be accepted");
|
||||
db.close().unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
@ -189,6 +189,26 @@ impl TidalDb {
|
||||
tracing::error!(error = %e, "community ledger checkpoint failed during shutdown");
|
||||
first_err.get_or_insert(e);
|
||||
}
|
||||
// Checkpoint preference vectors (derived state, no WAL backstop).
|
||||
if !self.preference_vectors.is_empty()
|
||||
&& let Err(e) = self.preference_vectors.checkpoint(storage.items_engine())
|
||||
{
|
||||
tracing::error!(error = %e, "preference-vector checkpoint failed during shutdown");
|
||||
first_err.get_or_insert(e);
|
||||
}
|
||||
// BLOCKER 6: persist the follower replication high-water-mark so a
|
||||
// clean restart restores its idempotency boundary instead of
|
||||
// resetting to 0. The receiver thread was already joined at the top
|
||||
// of this method, so no apply can advance the HWM past what we write
|
||||
// here. Non-fatal: a failure is captured (surfaced from close()) but
|
||||
// the next open simply restores an older, still-correct boundary.
|
||||
if let Err(e) = crate::db::state_rebuild::checkpoint_replication_state(
|
||||
storage.items_engine(),
|
||||
self.replication_state.as_ref(),
|
||||
) {
|
||||
tracing::error!(error = %e, "replication-state checkpoint failed during shutdown");
|
||||
first_err.get_or_insert(e);
|
||||
}
|
||||
if let Err(e) = storage.flush() {
|
||||
tracing::error!(error = %e, "storage flush failed during shutdown");
|
||||
first_err.get_or_insert(e);
|
||||
|
||||
@ -2,6 +2,46 @@
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Every length in the metadata format is a little-endian `u32`, so the entry
|
||||
/// count and each key/value byte length must fit in `u32::MAX` (~4.29e9). Item
|
||||
/// metadata never legitimately approaches this; the bound exists only to keep a
|
||||
/// pathological/oversized input from silently truncating its length prefix.
|
||||
const MAX_METADATA_LEN: usize = u32::MAX as usize;
|
||||
|
||||
/// Saturate a `usize` length to `u32`, clamping to `u32::MAX` on overflow.
|
||||
///
|
||||
/// Pure, assertion-free, total: the saturation half of [`len_as_u32`], split
|
||||
/// out so the over-limit branch can be unit-tested directly without tripping the
|
||||
/// debug assertion that fires on the (production-unexpected) oversized input.
|
||||
const fn saturate_len(len: usize) -> u32 {
|
||||
if len > MAX_METADATA_LEN {
|
||||
u32::MAX
|
||||
} else {
|
||||
// In-range by the guard above, so this cast cannot truncate.
|
||||
len as u32
|
||||
}
|
||||
}
|
||||
|
||||
/// Cast a `usize` length to the `u32` the wire format uses, saturating to
|
||||
/// `u32::MAX` on overflow instead of truncating.
|
||||
///
|
||||
/// On a 64-bit target `usize` can exceed `u32::MAX`, and a bare `as u32` would
|
||||
/// WRAP (e.g. `u32::MAX as usize + 1` -> `0`), writing a length prefix that
|
||||
/// mis-frames the record and silently corrupts every following field on read.
|
||||
/// Saturating instead writes `u32::MAX`, which the bounds-checked
|
||||
/// [`deserialize_metadata`] treats as a truncated/malformed record and stops on
|
||||
/// cleanly — a lossy-but-safe degradation rather than silent cross-field
|
||||
/// corruption. The `debug_assert!` surfaces any such input in test/debug builds
|
||||
/// so an oversized write is caught at its source rather than only on read-back.
|
||||
fn len_as_u32(len: usize) -> u32 {
|
||||
debug_assert!(
|
||||
len <= MAX_METADATA_LEN,
|
||||
"metadata length {len} exceeds the u32 wire limit ({MAX_METADATA_LEN}); \
|
||||
saturating to u32::MAX (read-back will treat the record as truncated)"
|
||||
);
|
||||
saturate_len(len)
|
||||
}
|
||||
|
||||
/// Serialize `HashMap<String, String>` as length-prefixed binary pairs.
|
||||
///
|
||||
/// This is the **items-only** metadata format (no embedding header).
|
||||
@ -19,13 +59,19 @@ use std::collections::HashMap;
|
||||
/// [key_len: 4 bytes][key bytes]
|
||||
/// [val_len: 4 bytes][value bytes]
|
||||
/// ```
|
||||
///
|
||||
/// Each length is bounded by [`MAX_METADATA_LEN`] (`u32::MAX`). An entry count
|
||||
/// or key/value length that exceeds it saturates to `u32::MAX` (see
|
||||
/// [`len_as_u32`]) rather than truncating, so an oversized input degrades to a
|
||||
/// record [`deserialize_metadata`] reads as truncated instead of one that
|
||||
/// silently mis-frames the following fields.
|
||||
pub(super) fn serialize_metadata(map: &HashMap<String, String>) -> Vec<u8> {
|
||||
let mut buf = Vec::new();
|
||||
buf.extend_from_slice(&(map.len() as u32).to_le_bytes());
|
||||
buf.extend_from_slice(&len_as_u32(map.len()).to_le_bytes());
|
||||
for (k, v) in map {
|
||||
buf.extend_from_slice(&(k.len() as u32).to_le_bytes());
|
||||
buf.extend_from_slice(&len_as_u32(k.len()).to_le_bytes());
|
||||
buf.extend_from_slice(k.as_bytes());
|
||||
buf.extend_from_slice(&(v.len() as u32).to_le_bytes());
|
||||
buf.extend_from_slice(&len_as_u32(v.len()).to_le_bytes());
|
||||
buf.extend_from_slice(v.as_bytes());
|
||||
}
|
||||
buf
|
||||
@ -102,6 +148,32 @@ pub(super) fn is_positive_engagement_signal(signal_type: &str) -> bool {
|
||||
)
|
||||
}
|
||||
|
||||
/// The `UserStateIndex` map a signal type drives, if any.
|
||||
///
|
||||
/// This is the single classifier that wires `signal_with_context` to the
|
||||
/// completion / saved / liked maps the `InProgress` filter and `DateSaved` sort
|
||||
/// read. Keeping it here (next to the other signal classifiers) means the write
|
||||
/// path and any rebuild path classify identically.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(super) enum UserStateSignal {
|
||||
/// `completion`: record a completion ratio (`record_completion`).
|
||||
Completion,
|
||||
/// `save` / `bookmark`: record a timestamped save (`add_save_timestamped`).
|
||||
Save,
|
||||
/// `like`: record a like (`add_like`).
|
||||
Like,
|
||||
}
|
||||
|
||||
/// Classify a signal type into the `UserStateIndex` map it populates, or `None`.
|
||||
pub(super) fn classify_user_state_signal(signal_type: &str) -> Option<UserStateSignal> {
|
||||
match signal_type {
|
||||
"completion" => Some(UserStateSignal::Completion),
|
||||
"save" | "bookmark" => Some(UserStateSignal::Save),
|
||||
"like" => Some(UserStateSignal::Like),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@ -145,4 +217,26 @@ mod tests {
|
||||
let out = deserialize_metadata(&bytes);
|
||||
assert!(out.is_empty(), "oversize value length must not slice OOB");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn saturate_len_clamps_over_u32_max() {
|
||||
// In-range lengths pass through unchanged (also the path len_as_u32 takes
|
||||
// for every legitimate write).
|
||||
assert_eq!(saturate_len(0), 0);
|
||||
assert_eq!(saturate_len(42), 42);
|
||||
assert_eq!(saturate_len(MAX_METADATA_LEN), u32::MAX);
|
||||
assert_eq!(len_as_u32(MAX_METADATA_LEN), u32::MAX);
|
||||
|
||||
// Over the u32 limit, the conversion must SATURATE to u32::MAX, never
|
||||
// wrap. (Only reachable where usize is wider than u32, i.e. 64-bit
|
||||
// targets; on a 32-bit target usize cannot exceed u32::MAX.) We exercise
|
||||
// `saturate_len` directly rather than `len_as_u32` because the latter's
|
||||
// debug_assert deliberately fires on this production-unexpected input.
|
||||
#[cfg(target_pointer_width = "64")]
|
||||
{
|
||||
// A bare `as u32` would wrap this to 0; saturation keeps it u32::MAX.
|
||||
assert_eq!(saturate_len(MAX_METADATA_LEN + 1), u32::MAX);
|
||||
assert_eq!(saturate_len(usize::MAX), u32::MAX);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -71,13 +71,33 @@ mod replication_ops;
|
||||
mod sweeper;
|
||||
mod text_syncer;
|
||||
|
||||
use self::text_syncer::spawn_text_syncer;
|
||||
use self::text_syncer::open_text_syncer;
|
||||
|
||||
/// A tidalDB database instance.
|
||||
///
|
||||
/// Created via [`TidalDb::builder()`]. Call [`close`](Self::close) for
|
||||
/// explicit shutdown; [`Drop`] runs best-effort cleanup if not called.
|
||||
/// `Send + Sync` -- wrap in `Arc` for multi-threaded access.
|
||||
///
|
||||
/// # Pending refactor: field-cluster extraction (`CODING_GUIDELINES` §9)
|
||||
///
|
||||
/// This struct carries ~60 fields spanning M0–M10 concerns and `mod.rs` exceeds
|
||||
/// the §9 file-size guideline. The shared-default initializers were extracted
|
||||
/// into [`SharedDefaults`] to de-duplicate the two constructors, but the fields
|
||||
/// themselves still live flat on `TidalDb`. The next structural step is to group
|
||||
/// cohesive clusters into sub-structs — candidates: `replication` (state +
|
||||
/// receiver/shipper handles + bridge + tenancy + control plane), `governance`
|
||||
/// (membership/community/governance/capability registries), `text` (both text
|
||||
/// indexes, channels, syncer threads, flush channels, health latch), and
|
||||
/// `lifecycle` (checkpoint/sweeper handles + their shutdown flags).
|
||||
///
|
||||
/// Each cluster's blast radius reaches OUTSIDE `mod.rs`, so the extraction must
|
||||
/// be done as one whole-crate change per cluster (it rewrites every
|
||||
/// `self.<field>` to `self.<cluster>.<field>` across `db/lifecycle.rs`,
|
||||
/// `db/sweeper.rs`, `db/diagnostics.rs`, `db/communities.rs`, `db/backup.rs`,
|
||||
/// `db/remove_scope.rs`, `db/capabilities.rs`, `db/replication_ops.rs`, ...).
|
||||
/// It cannot be applied safely from a `mod.rs`-only edit. Do it incrementally,
|
||||
/// one cluster per change, with the whole crate compiling between steps.
|
||||
pub struct TidalDb {
|
||||
config: Config,
|
||||
closed: AtomicBool,
|
||||
@ -188,6 +208,29 @@ pub struct TidalDb {
|
||||
lock_file: Option<std::fs::File>,
|
||||
}
|
||||
|
||||
/// The subset of [`TidalDb`] fields that both constructors initialize
|
||||
/// identically. Produced once by [`TidalDb::shared_defaults`] and destructured
|
||||
/// into each constructor's struct literal, so the verbatim default initializers
|
||||
/// live in exactly one place. Field names mirror [`TidalDb`] exactly so the
|
||||
/// destructure at each call site is a one-to-one move with no renaming.
|
||||
struct SharedDefaults {
|
||||
sessions: dashmap::DashMap<SessionId, Arc<SessionState>>,
|
||||
next_session_id: AtomicU64,
|
||||
closed_sessions: dashmap::DashMap<SessionId, SessionSnapshot>,
|
||||
collection_index: Arc<CollectionIndex>,
|
||||
notification_tracker: Arc<notification_tracker::NotificationTracker>,
|
||||
load_detector: Arc<crate::load::LoadDetector>,
|
||||
backpressure_config: crate::load::BackpressureConfig,
|
||||
rate_limiter: Arc<crate::load::RateLimiter>,
|
||||
shutdown_sweeper: Arc<AtomicBool>,
|
||||
sweeper_thread: std::sync::Mutex<Option<std::thread::JoinHandle<()>>>,
|
||||
backup_in_progress: Arc<AtomicBool>,
|
||||
receiver_handle: std::sync::Mutex<Option<crate::replication::receiver::SegmentReceiverHandle>>,
|
||||
shipper_handle: Option<crate::replication::WalShipperHandle>,
|
||||
session_bridge: Arc<crate::replication::SessionReplicationBridge>,
|
||||
lock_file: Option<std::fs::File>,
|
||||
}
|
||||
|
||||
impl TidalDb {
|
||||
/// Returns a new [`TidalDbBuilder`] with default (ephemeral) configuration.
|
||||
#[must_use]
|
||||
@ -255,6 +298,44 @@ impl TidalDb {
|
||||
(control_plane, router)
|
||||
}
|
||||
|
||||
/// The fields whose initializers are byte-for-byte identical between
|
||||
/// `from_config` and `from_parts`.
|
||||
///
|
||||
/// Both constructors used to inline these ~18 default initializers verbatim,
|
||||
/// so the two struct literals had to be kept in lock-step by hand (DRY tax,
|
||||
/// `CODING_GUIDELINES` §9). Centralizing them here means each constructor only
|
||||
/// spells out the fields that genuinely differ (opened storage/ledger/WAL,
|
||||
/// the in-memory indexes, and the control-plane/replication wiring); the
|
||||
/// shared remainder is destructured from one `SharedDefaults::new()` call.
|
||||
///
|
||||
/// This deliberately does NOT carry session/text channel state: `from_parts`
|
||||
/// installs real text-syncer bundles where `from_config` installs `None`, so
|
||||
/// those fields stay at the call site. Everything here is a pure default that
|
||||
/// is constructed identically regardless of how the DB was opened.
|
||||
fn shared_defaults() -> SharedDefaults {
|
||||
SharedDefaults {
|
||||
sessions: dashmap::DashMap::new(),
|
||||
next_session_id: AtomicU64::new(1),
|
||||
closed_sessions: dashmap::DashMap::new(),
|
||||
collection_index: Arc::new(CollectionIndex::new()),
|
||||
notification_tracker: Arc::new(notification_tracker::NotificationTracker::new()),
|
||||
load_detector: Arc::new(crate::load::LoadDetector::new(
|
||||
crate::load::DegradationThresholds::default(),
|
||||
)),
|
||||
backpressure_config: crate::load::BackpressureConfig::default(),
|
||||
rate_limiter: Arc::new(crate::load::RateLimiter::new(
|
||||
crate::load::RateLimiterConfig::default(),
|
||||
)),
|
||||
shutdown_sweeper: Arc::new(AtomicBool::new(false)),
|
||||
sweeper_thread: std::sync::Mutex::new(None),
|
||||
backup_in_progress: Arc::new(AtomicBool::new(false)),
|
||||
receiver_handle: std::sync::Mutex::new(None),
|
||||
shipper_handle: None,
|
||||
session_bridge: Self::make_single_node_session_bridge(),
|
||||
lock_file: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Construct a `TidalDb` without a schema (M0 compatibility mode).
|
||||
///
|
||||
/// `metrics` must be uniquely owned (no other `Arc` clones yet): this
|
||||
@ -294,6 +375,24 @@ impl TidalDb {
|
||||
m.set_control_plane(Arc::clone(&control_plane));
|
||||
}
|
||||
|
||||
let SharedDefaults {
|
||||
sessions,
|
||||
next_session_id,
|
||||
closed_sessions,
|
||||
collection_index,
|
||||
notification_tracker,
|
||||
load_detector,
|
||||
backpressure_config,
|
||||
rate_limiter,
|
||||
shutdown_sweeper,
|
||||
sweeper_thread,
|
||||
backup_in_progress,
|
||||
receiver_handle,
|
||||
shipper_handle,
|
||||
session_bridge,
|
||||
lock_file,
|
||||
} = Self::shared_defaults();
|
||||
|
||||
Self {
|
||||
config,
|
||||
closed: AtomicBool::new(false),
|
||||
@ -330,32 +429,35 @@ impl TidalDb {
|
||||
creator_text_syncer_thread: std::sync::Mutex::new(None),
|
||||
text_flush_tx: None,
|
||||
creator_text_flush_tx: None,
|
||||
sessions: dashmap::DashMap::new(),
|
||||
next_session_id: AtomicU64::new(1),
|
||||
closed_sessions: dashmap::DashMap::new(),
|
||||
sessions,
|
||||
next_session_id,
|
||||
closed_sessions,
|
||||
co_engagement: Arc::new(CoEngagementIndex::new()),
|
||||
user_signal_index: Arc::new(UserSignalIndex::new()),
|
||||
cohort_registry,
|
||||
cohort_resolver,
|
||||
cohort_ledger: Arc::new(crate::cohort::CohortSignalLedger::empty()),
|
||||
collection_index: Arc::new(CollectionIndex::new()),
|
||||
collection_index,
|
||||
suggestion_index: Arc::new(crate::query::suggest::SuggestionIndex::new()),
|
||||
notification_tracker: Arc::new(notification_tracker::NotificationTracker::new()),
|
||||
load_detector: Arc::new(crate::load::LoadDetector::new(
|
||||
crate::load::DegradationThresholds::default(),
|
||||
)),
|
||||
backpressure_config: crate::load::BackpressureConfig::default(),
|
||||
rate_limiter: Arc::new(crate::load::RateLimiter::new(
|
||||
crate::load::RateLimiterConfig::default(),
|
||||
)),
|
||||
shutdown_sweeper: Arc::new(AtomicBool::new(false)),
|
||||
sweeper_thread: std::sync::Mutex::new(None),
|
||||
backup_in_progress: Arc::new(AtomicBool::new(false)),
|
||||
replication_state: Arc::new(crate::replication::state::ReplicationState::single()),
|
||||
receiver_handle: std::sync::Mutex::new(None),
|
||||
shipper_handle: None,
|
||||
session_bridge: Self::make_single_node_session_bridge(),
|
||||
lock_file: None,
|
||||
notification_tracker,
|
||||
load_detector,
|
||||
backpressure_config,
|
||||
rate_limiter,
|
||||
shutdown_sweeper,
|
||||
sweeper_thread,
|
||||
backup_in_progress,
|
||||
// obs-REPL-1: share the SAME Arc the control-plane lag gauge was
|
||||
// built from (above) instead of constructing a second, divergent
|
||||
// ReplicationState. No-schema mode never starts replication today, but
|
||||
// a divergent Arc here is a latent failover/lag-reporting landmine.
|
||||
// The `from_config_shares_one_replication_state_arc` test pins that
|
||||
// this field and the control-plane lag gauge observe one
|
||||
// `ReplicationState`.
|
||||
replication_state: rep_state,
|
||||
receiver_handle,
|
||||
shipper_handle,
|
||||
session_bridge,
|
||||
lock_file,
|
||||
tenant_router,
|
||||
control_plane,
|
||||
membership_registry: Arc::new(crate::governance::MembershipRegistry::new()),
|
||||
@ -408,6 +510,11 @@ impl TidalDb {
|
||||
let ledger = Some(ledger);
|
||||
let storage = Some(storage);
|
||||
|
||||
// Wrap the (already-restored) preference vectors in an Arc up-front so the
|
||||
// periodic checkpoint thread can share the SAME instance the read path uses
|
||||
// — restored in open.rs, then checkpointed here so the loss window is bounded.
|
||||
let preference_vectors = Arc::new(preference_vectors);
|
||||
|
||||
// M6 co-engagement index: restore from durable storage if available.
|
||||
let co_engagement = Arc::new(CoEngagementIndex::new());
|
||||
if let Some(ref sb) = storage
|
||||
@ -485,9 +592,28 @@ impl TidalDb {
|
||||
replayed += 1;
|
||||
}
|
||||
}
|
||||
if replayed > 0 {
|
||||
|
||||
// Replay agent-scope purge tombstones (M10p3). Same crash window as
|
||||
// the user tombstones above: a remove-by-agent that crashed before the
|
||||
// next community checkpoint must not resurrect the revoked agent's
|
||||
// contributions. `apply_agent_purge_tombstone` re-purges by writer and
|
||||
// is idempotent, so a purge already reflected in the checkpoint is a
|
||||
// no-op.
|
||||
let agent_prefix = crate::storage::entity_tag_prefix(
|
||||
crate::schema::EntityId::new(0),
|
||||
crate::storage::Tag::AgentPurgeTombstone,
|
||||
);
|
||||
let mut agent_replayed = 0u64;
|
||||
for (_, value) in sb.items_engine().scan_prefix(&agent_prefix).flatten() {
|
||||
if let Some(t) = crate::governance::AgentPurgeTombstone::from_bytes(&value) {
|
||||
community_ledger.apply_agent_purge_tombstone(&t);
|
||||
agent_replayed += 1;
|
||||
}
|
||||
}
|
||||
if replayed > 0 || agent_replayed > 0 {
|
||||
tracing::info!(
|
||||
tombstones = replayed,
|
||||
user_tombstones = replayed,
|
||||
agent_tombstones = agent_replayed,
|
||||
"replayed purge tombstones over restored community aggregate"
|
||||
);
|
||||
}
|
||||
@ -538,30 +664,79 @@ impl TidalDb {
|
||||
);
|
||||
}
|
||||
Ok(None) => {
|
||||
// First boot or no cohort checkpoint yet -- nothing to restore.
|
||||
// First boot, no cohort checkpoint yet, OR a corrupt checkpoint
|
||||
// that restore() already escalated at error! and chose to start
|
||||
// empty for (cohort state is a re-accumulating derived aggregate;
|
||||
// see CohortSignalLedger::restore's corruption policy). Nothing
|
||||
// more to do here.
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
// Genuine I/O failure reading the checkpoint (NOT corruption,
|
||||
// which restore() handles internally). Start empty and surface
|
||||
// it loudly so the lost cohort state is observable.
|
||||
tracing::error!(
|
||||
error = %e,
|
||||
"cohort signal ledger restore failed; starting from empty state"
|
||||
"cohort signal ledger restore hit a storage I/O error; \
|
||||
starting from empty cohort state"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let text_bundle = spawn_text_syncer(
|
||||
// ── Text-index open + crash-recovery rebuild + syncer spawn ─────────
|
||||
//
|
||||
// BLOCKER 7 deadlock fix: the syncer thread holds the Tantivy writer lock
|
||||
// for its ENTIRE lifetime (single-writer design). The rebuild-from-store
|
||||
// ALSO needs that lock (`rebuild_from` -> `writer_guard()`). If we spawned
|
||||
// the syncer first, the rebuild would block forever and open would hang.
|
||||
//
|
||||
// So open is split into three steps with the writer lock free for the
|
||||
// middle one:
|
||||
// 1. open the item + creator indexes and create their channels
|
||||
// (`open_text_syncer`), NO thread spawned yet;
|
||||
// 2. run `rebuild_text_indexes_at_open` on the freshly-opened indexes
|
||||
// while no thread holds the writer lock — this also folds the
|
||||
// suggestion-index rebuild into the single item scan, and runs the
|
||||
// suggestion rebuild even when there is no item text index;
|
||||
// 3. spawn both syncer threads (`.spawn()`); each now grabs the writer
|
||||
// lock for life, which is fine because the rebuild already finished.
|
||||
let text_pending = open_text_syncer(
|
||||
schema_def.text_fields(),
|
||||
&config,
|
||||
"text_index",
|
||||
"tidaldb-text-syncer",
|
||||
);
|
||||
let creator_text_bundle = spawn_text_syncer(
|
||||
let creator_text_pending = open_text_syncer(
|
||||
schema_def.creator_text_fields(),
|
||||
&config,
|
||||
"creator_text_index",
|
||||
"tidaldb-creator-text-syncer",
|
||||
);
|
||||
|
||||
// The suggestion index is rebuilt in the SAME item scan as the item text
|
||||
// index, so it must exist before the rebuild and be moved (not recreated)
|
||||
// into the struct below.
|
||||
let suggestion_index = Arc::new(crate::query::suggest::SuggestionIndex::new());
|
||||
|
||||
// Step 2: rebuild from the durable store BEFORE any syncer thread takes
|
||||
// the writer lock (persistent mode only; ephemeral storage starts empty,
|
||||
// so the existing guard inside the rebuild makes this a no-op there).
|
||||
// The suggestion-index rebuild runs regardless of whether an item text
|
||||
// index exists.
|
||||
if let Some(ref sb) = storage {
|
||||
crate::db::state_rebuild::rebuild_text_indexes_at_open(
|
||||
sb,
|
||||
text_pending.index().map(|a| a.as_ref()),
|
||||
creator_text_pending.index().map(|a| a.as_ref()),
|
||||
&suggestion_index,
|
||||
);
|
||||
}
|
||||
|
||||
// Step 3: now that the rebuild is done, spawn both syncer threads. The
|
||||
// syncer grabbing the writer lock after the rebuild is correct.
|
||||
let text_bundle = text_pending.spawn();
|
||||
let creator_text_bundle = creator_text_pending.spawn();
|
||||
|
||||
// Wrap the embedding registry early so the checkpoint thread can share it.
|
||||
let embedding_registry = Arc::new(RwLock::new(embedding_registry));
|
||||
|
||||
@ -619,10 +794,21 @@ impl TidalDb {
|
||||
// receiver advances it as it applies segments, so
|
||||
// `ControlPlane::lag_for(SINGLE)` reflects true follower lag rather than
|
||||
// staying 0).
|
||||
// BLOCKER 6: restore the per-shard applied-seqno high-water-mark from its
|
||||
// durable checkpoint so a restarted follower keeps its idempotency
|
||||
// boundary instead of resetting to 0 (which would re-apply and
|
||||
// double-count every re-shipped segment). On first boot / no checkpoint /
|
||||
// ephemeral mode this restores all shards at 0, identical to before.
|
||||
// This is the SINGLE Arc the receiver advances, the checkpoint thread
|
||||
// persists, and the lag gauge reads (obs-REPL-1: one Arc, never two).
|
||||
let replication_state = {
|
||||
let mut shards: Vec<crate::replication::ShardId> = vec![config.cluster.shard_id];
|
||||
shards.extend_from_slice(&config.cluster.peer_shards);
|
||||
Arc::new(crate::replication::state::ReplicationState::new(&shards))
|
||||
let restored = storage.as_ref().map_or_else(
|
||||
|| crate::replication::state::ReplicationState::new(&shards),
|
||||
|sb| state_rebuild::restore_replication_state(sb.items_engine(), &shards),
|
||||
);
|
||||
Arc::new(restored)
|
||||
};
|
||||
// The tenant router shares the control plane's topology Arc (see
|
||||
// `build_control_plane`) so a topology update is never invisible to the
|
||||
@ -656,6 +842,13 @@ impl TidalDb {
|
||||
let shutdown = Arc::clone(&shutdown_checkpoint);
|
||||
let ledger_clone = Arc::clone(ledger_arc);
|
||||
let cohort_clone = Arc::clone(&cohort_ledger);
|
||||
let community_clone = Arc::clone(&community_ledger);
|
||||
let co_engagement_clone = Arc::clone(&co_engagement);
|
||||
let preference_clone = Arc::clone(&preference_vectors);
|
||||
// Share the SAME replication_state Arc the receiver advances
|
||||
// and the lag gauge reads, so the persisted high-water-mark is
|
||||
// always this node's real applied seqno (BLOCKER 6 / obs-REPL-1).
|
||||
let replication_state_clone = Arc::clone(&replication_state);
|
||||
let seq_clone = Arc::clone(&last_seq);
|
||||
let wal_dir = config.data_dir.as_ref().map(|d| d.join("wal"));
|
||||
let metrics_clone = Arc::clone(&metrics);
|
||||
@ -694,6 +887,10 @@ impl TidalDb {
|
||||
shutdown,
|
||||
ledger_clone,
|
||||
cohort_clone,
|
||||
community_clone,
|
||||
co_engagement_clone,
|
||||
preference_clone,
|
||||
replication_state_clone,
|
||||
items,
|
||||
seq_clone,
|
||||
wal_dir,
|
||||
@ -718,6 +915,24 @@ impl TidalDb {
|
||||
std::sync::Mutex::new(handle)
|
||||
};
|
||||
|
||||
let SharedDefaults {
|
||||
sessions,
|
||||
next_session_id,
|
||||
closed_sessions,
|
||||
collection_index,
|
||||
notification_tracker,
|
||||
load_detector,
|
||||
backpressure_config,
|
||||
rate_limiter,
|
||||
shutdown_sweeper,
|
||||
sweeper_thread,
|
||||
backup_in_progress,
|
||||
receiver_handle,
|
||||
shipper_handle,
|
||||
session_bridge,
|
||||
lock_file,
|
||||
} = Self::shared_defaults();
|
||||
|
||||
let db = Self {
|
||||
config,
|
||||
closed: AtomicBool::new(false),
|
||||
@ -745,7 +960,7 @@ impl TidalDb {
|
||||
user_state: Arc::new(user_state),
|
||||
hard_negatives: Arc::new(hard_negatives),
|
||||
interaction_ledger: Arc::new(interaction_ledger),
|
||||
preference_vectors: Arc::new(preference_vectors),
|
||||
preference_vectors,
|
||||
text_index: text_bundle.index,
|
||||
text_tx: std::sync::Mutex::new(text_bundle.write_tx),
|
||||
text_syncer_thread: text_bundle.thread,
|
||||
@ -754,32 +969,31 @@ impl TidalDb {
|
||||
creator_text_syncer_thread: creator_text_bundle.thread,
|
||||
text_flush_tx: text_bundle.flush_tx,
|
||||
creator_text_flush_tx: creator_text_bundle.flush_tx,
|
||||
sessions: dashmap::DashMap::new(),
|
||||
next_session_id: AtomicU64::new(1),
|
||||
closed_sessions: dashmap::DashMap::new(),
|
||||
sessions,
|
||||
next_session_id,
|
||||
closed_sessions,
|
||||
co_engagement,
|
||||
user_signal_index,
|
||||
cohort_registry,
|
||||
cohort_resolver,
|
||||
cohort_ledger,
|
||||
collection_index: Arc::new(CollectionIndex::new()),
|
||||
suggestion_index: Arc::new(crate::query::suggest::SuggestionIndex::new()),
|
||||
notification_tracker: Arc::new(notification_tracker::NotificationTracker::new()),
|
||||
load_detector: Arc::new(crate::load::LoadDetector::new(
|
||||
crate::load::DegradationThresholds::default(),
|
||||
)),
|
||||
backpressure_config: crate::load::BackpressureConfig::default(),
|
||||
rate_limiter: Arc::new(crate::load::RateLimiter::new(
|
||||
crate::load::RateLimiterConfig::default(),
|
||||
)),
|
||||
shutdown_sweeper: Arc::new(AtomicBool::new(false)),
|
||||
sweeper_thread: std::sync::Mutex::new(None),
|
||||
backup_in_progress: Arc::new(AtomicBool::new(false)),
|
||||
// Empty here; rebuilt from durable storage just below via
|
||||
// `rebuild_collections` (the shared default is the same empty index
|
||||
// both constructors start from).
|
||||
collection_index,
|
||||
suggestion_index,
|
||||
notification_tracker,
|
||||
load_detector,
|
||||
backpressure_config,
|
||||
rate_limiter,
|
||||
shutdown_sweeper,
|
||||
sweeper_thread,
|
||||
backup_in_progress,
|
||||
replication_state,
|
||||
receiver_handle: std::sync::Mutex::new(None),
|
||||
shipper_handle: None,
|
||||
session_bridge: Self::make_single_node_session_bridge(),
|
||||
lock_file: None,
|
||||
receiver_handle,
|
||||
shipper_handle,
|
||||
session_bridge,
|
||||
lock_file,
|
||||
tenant_router,
|
||||
control_plane,
|
||||
membership_registry,
|
||||
@ -789,12 +1003,22 @@ impl TidalDb {
|
||||
};
|
||||
|
||||
// M6p4: rebuild collection index from durable storage.
|
||||
// M6p5: rebuild suggestion index title terms from durable item metadata.
|
||||
//
|
||||
// BLOCKER 7: the item AND creator Tantivy text indexes (and the
|
||||
// suggestion index, folded into the item scan) are rebuilt from the
|
||||
// durable entity stores ABOVE — BEFORE the syncer threads are spawned —
|
||||
// so the rebuild's writer-lock acquisition does not deadlock against the
|
||||
// syncer, which holds that lock for its whole lifetime. Only the
|
||||
// collection index remains to rebuild here; it needs no Tantivy writer
|
||||
// lock and is independent of the text syncers.
|
||||
if let Some(ref sb) = db.storage {
|
||||
crate::db::collections::rebuild_collections(&db.collection_index, sb.items_engine());
|
||||
crate::db::state_rebuild::rebuild_suggestion_index(sb, &db.suggestion_index);
|
||||
}
|
||||
|
||||
// Restore durable tenant-migration routing so a crash mid-migration does
|
||||
// not silently revert a finalized tenant to jump-hash routing.
|
||||
db.restore_migration_routing();
|
||||
|
||||
if !session_events.is_empty() {
|
||||
db.restore_session_wal_events(&session_events);
|
||||
}
|
||||
@ -819,3 +1043,90 @@ impl TidalDb {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::replication::ShardId;
|
||||
|
||||
/// Pins the obs-REPL-1 invariant: in BOTH constructors the
|
||||
/// `replication_state` field and the control-plane lag gauge must observe ONE
|
||||
/// `ReplicationState`, never two divergent Arcs.
|
||||
///
|
||||
/// Earlier `from_config` built the lag gauge from one `ReplicationState::single()`
|
||||
/// and the field from a SECOND one, so a receiver advancing the field would
|
||||
/// leave the lag gauge reading an orphaned Arc — replication lag stuck forever.
|
||||
/// We prove the Arcs are shared *behaviorally* rather than poking private
|
||||
/// internals: advancing the field's high-water-mark must be observable through
|
||||
/// the gauge. `ReplicationLagGauge::applied_seqno()` reads its own `state` Arc,
|
||||
/// so the assertion can only pass when both point at the same allocation
|
||||
/// (equivalent to `Arc::ptr_eq`, but without widening the gauge's interface).
|
||||
fn assert_shared_replication_state(db: &TidalDb) {
|
||||
let gauge = db.control_plane().lag_gauge();
|
||||
// Fresh DB: both start at 0.
|
||||
assert_eq!(gauge.applied_seqno(), 0, "gauge must start at applied 0");
|
||||
assert_eq!(
|
||||
db.replication_state().applied_seqno(ShardId::SINGLE),
|
||||
Some(0),
|
||||
"field must start at applied 0"
|
||||
);
|
||||
|
||||
// Advance ONLY the field's high-water-mark (what the receiver does).
|
||||
db.replication_state().advance(ShardId::SINGLE, 7);
|
||||
|
||||
// The gauge must observe it. If the field and gauge held divergent Arcs
|
||||
// this would still read 0 — the exact stuck-lag bug obs-REPL-1 prevents.
|
||||
assert_eq!(
|
||||
gauge.applied_seqno(),
|
||||
7,
|
||||
"lag gauge and replication_state field must share ONE Arc; gauge \
|
||||
did not observe the field's advance (divergent-Arc regression)"
|
||||
);
|
||||
assert_eq!(
|
||||
db.replication_state().applied_seqno(ShardId::SINGLE),
|
||||
Some(7),
|
||||
"field must reflect its own advance"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_config_shares_one_replication_state_arc() {
|
||||
// No-schema (M0 compatibility) open routes through `from_config`.
|
||||
let db = TidalDb::builder()
|
||||
.ephemeral()
|
||||
.open()
|
||||
.expect("no-schema ephemeral open should succeed");
|
||||
assert_shared_replication_state(&db);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_parts_shares_one_replication_state_arc() {
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::schema::{DecaySpec, EntityKind, SchemaBuilder, Window};
|
||||
|
||||
// Persistent-with-schema open routes through `from_parts`.
|
||||
let mut builder = SchemaBuilder::new();
|
||||
let _ = builder
|
||||
.signal(
|
||||
"view",
|
||||
EntityKind::Item,
|
||||
DecaySpec::Exponential {
|
||||
half_life: Duration::from_secs(3600),
|
||||
},
|
||||
)
|
||||
.windows(&[Window::AllTime])
|
||||
.velocity(false)
|
||||
.add();
|
||||
let schema = builder.build().expect("schema must be valid");
|
||||
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let db = TidalDb::builder()
|
||||
.with_data_dir(dir.path())
|
||||
.with_schema(schema)
|
||||
.open()
|
||||
.expect("persistent open with schema should succeed");
|
||||
assert_shared_replication_state(&db);
|
||||
db.close().expect("close should succeed");
|
||||
}
|
||||
}
|
||||
|
||||
@ -196,6 +196,26 @@ impl super::TidalDb {
|
||||
let interaction_ledger = InteractionLedger::new();
|
||||
rebuild_entity_state(&storage, &user_state, &interaction_ledger)?;
|
||||
|
||||
// Restore the signal-driven UserStateIndex maps (completion /
|
||||
// saved-timestamp / liked) from their durable Tag::UserState rows
|
||||
// so the InProgress filter and DateSaved sort survive a restart.
|
||||
if let Err(e) = user_state.restore(storage.items_engine()) {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
"user-state restore failed; InProgress/DateSaved start empty"
|
||||
);
|
||||
}
|
||||
|
||||
// Restore hard negatives from durable Tag::HardNeg rows so a
|
||||
// skip/dislike-rejected item stays excluded after a crash.
|
||||
let hard_negatives = HardNegIndex::new();
|
||||
if let Err(e) = hard_negatives.restore(storage.items_engine()) {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
"hard-negative restore failed; rejected items may resurface"
|
||||
);
|
||||
}
|
||||
|
||||
// Rebuild item indexes (universe bitmap + bitmap/range indexes +
|
||||
// creator-items) from persisted metadata so queries and /health
|
||||
// work immediately. The shared `ItemIndexes::index` is the SAME
|
||||
@ -221,6 +241,16 @@ impl super::TidalDb {
|
||||
.first()
|
||||
.map_or(128, |s| s.dimensions);
|
||||
|
||||
// Restore preference vectors from their Tag::Preference checkpoint
|
||||
// so personalized ranking does not snap to cold-start after a crash.
|
||||
let preference_vectors = PreferenceVectors::new(pref_dim);
|
||||
if let Err(e) = preference_vectors.restore(storage.items_engine()) {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
"preference vector restore failed; starting from cold-start"
|
||||
);
|
||||
}
|
||||
|
||||
Ok(OpenResult {
|
||||
storage,
|
||||
ledger,
|
||||
@ -238,9 +268,9 @@ impl super::TidalDb {
|
||||
schema_def,
|
||||
creator_items,
|
||||
user_state,
|
||||
hard_negatives: HardNegIndex::new(),
|
||||
hard_negatives,
|
||||
interaction_ledger,
|
||||
preference_vectors: PreferenceVectors::new(pref_dim),
|
||||
preference_vectors,
|
||||
session_events,
|
||||
})
|
||||
}
|
||||
|
||||
@ -48,7 +48,9 @@ impl TidalDb {
|
||||
.add_block_creator(from.as_u64(), to.as_u64());
|
||||
}
|
||||
RelationshipType::Hide => {
|
||||
self.user_state.add_hide(from.as_u64(), to.as_u64() as u32);
|
||||
// Route through the observable narrowing helper so an item id
|
||||
// above the u32 universe limit is warned, not silently truncated.
|
||||
self.user_state.add_hide_item(from.as_u64(), to.as_u64());
|
||||
}
|
||||
RelationshipType::Follows => {
|
||||
// Forward: user -> followed creator
|
||||
@ -57,7 +59,12 @@ impl TidalDb {
|
||||
self.user_state
|
||||
.add_creator_follower(to.as_u64(), from.as_u64());
|
||||
}
|
||||
_ => {}
|
||||
// InteractionWeight edges carry the (user, creator) running score
|
||||
// (maintained on the signal path via `persist_interaction_edge`, not
|
||||
// here), and Mute is durably persisted but currently inert (no
|
||||
// muted-set is applied to retrieval yet — tracked). Explicit arm so
|
||||
// the no-op is deliberate and greppable, not absorbed by a catch-all.
|
||||
RelationshipType::InteractionWeight | RelationshipType::Mute => {}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@ -95,15 +102,15 @@ impl TidalDb {
|
||||
.remove_block_creator(from.as_u64(), to.as_u64());
|
||||
}
|
||||
RelationshipType::Hide => {
|
||||
self.user_state
|
||||
.remove_hide(from.as_u64(), to.as_u64() as u32);
|
||||
self.user_state.remove_hide_item(from.as_u64(), to.as_u64());
|
||||
}
|
||||
RelationshipType::Follows => {
|
||||
self.user_state.remove_follow(from.as_u64(), to.as_u64());
|
||||
self.user_state
|
||||
.remove_creator_follower(to.as_u64(), from.as_u64());
|
||||
}
|
||||
_ => {}
|
||||
// See `write_relationship` for why these are explicit no-ops.
|
||||
RelationshipType::InteractionWeight | RelationshipType::Mute => {}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
|
||||
@ -10,9 +10,11 @@
|
||||
use super::TidalDb;
|
||||
use crate::{
|
||||
governance::{
|
||||
PurgeReceipt, RemoveScope, community_ledger::ContributionProvenance, scope::CommunityId,
|
||||
AgentPurgeTombstone, PurgeReceipt, RemoveScope, community_ledger::ContributionProvenance,
|
||||
scope::CommunityId,
|
||||
},
|
||||
schema::{EntityId, Timestamp},
|
||||
schema::{EntityId, TidalError, Timestamp},
|
||||
storage::{Tag, encode_key},
|
||||
};
|
||||
|
||||
impl TidalDb {
|
||||
@ -45,10 +47,33 @@ impl TidalDb {
|
||||
} => {
|
||||
self.require_writeable("remove_from_personalization")?;
|
||||
let watermark_ns = Timestamp::now().as_nanos();
|
||||
|
||||
// Durability-before-mutation, mirroring the user-purge path:
|
||||
// persist an agent-scoped tombstone BEFORE mutating the in-memory
|
||||
// aggregate so a crash before the next community checkpoint replays
|
||||
// to the same outcome (the revoked agent's contributions stay gone)
|
||||
// rather than restoring the pre-purge snapshot.
|
||||
let tombstone =
|
||||
AgentPurgeTombstone::new(community, writer_agent.clone(), watermark_ns);
|
||||
self.persist_agent_tombstone(&tombstone)?;
|
||||
|
||||
let contributions_removed =
|
||||
self.community_ledger.purge_writer(community, &writer_agent);
|
||||
self.community_ledger
|
||||
.set_purge_watermark(community, watermark_ns);
|
||||
|
||||
// Re-checkpoint so the durable Tag::Community snapshot already
|
||||
// reflects the removal (belt-and-suspenders with the tombstone
|
||||
// replay: either alone keeps the agent gone after a crash).
|
||||
if let Some(ref storage) = self.storage
|
||||
&& let Err(e) = self.community_ledger.checkpoint(storage.items_engine())
|
||||
{
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
"agent-purge re-checkpoint failed; tombstone replay still covers the crash window"
|
||||
);
|
||||
}
|
||||
|
||||
Ok(PurgeReceipt {
|
||||
contributions_removed,
|
||||
watermark_ns,
|
||||
@ -57,6 +82,23 @@ impl TidalDb {
|
||||
}
|
||||
}
|
||||
|
||||
/// Persist an agent-scope purge tombstone under `Tag::AgentPurgeTombstone`.
|
||||
/// No-op in ephemeral mode.
|
||||
fn persist_agent_tombstone(&self, t: &AgentPurgeTombstone) -> crate::Result<()> {
|
||||
let Some(ref storage) = self.storage else {
|
||||
return Ok(());
|
||||
};
|
||||
let key = encode_key(
|
||||
EntityId::new(0),
|
||||
Tag::AgentPurgeTombstone,
|
||||
&t.storage_suffix(),
|
||||
);
|
||||
storage
|
||||
.items_engine()
|
||||
.put(&key, &t.to_bytes()?)
|
||||
.map_err(TidalError::from)
|
||||
}
|
||||
|
||||
/// Attribute a community aggregate to its contributions (M10p3
|
||||
/// explainability): each surviving `(user, writer_agent, epoch, score, count)`
|
||||
/// that composes the `(community, entity, signal_type)` aggregate, sorted
|
||||
|
||||
@ -39,6 +39,222 @@ impl TidalDb {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Test-only: durably checkpoint the follower's replicated state into a
|
||||
/// crash-consistent on-disk snapshot.
|
||||
///
|
||||
/// Writes, atomically-consistent with each other:
|
||||
/// 1. the signal-ledger aggregate (fjall checkpoint at the current WAL seq),
|
||||
/// 2. the WAL checkpoint marker at the SAME seq (so a subsequent crash replays
|
||||
/// only post-checkpoint follower-WAL events — no double-counting of events
|
||||
/// already folded into the aggregate), and
|
||||
/// 3. the replication high-water-mark ([`ReplicationState`] checkpoint).
|
||||
///
|
||||
/// This lands the exact on-disk state a clean checkpoint cycle would, without
|
||||
/// the 30s wait, so a crash-recovery test can then simulate a hard crash (no
|
||||
/// graceful shutdown) and prove the follower restores its replicated state and
|
||||
/// applied-seqno boundary instead of resetting to a fresh node (BLOCKER 6).
|
||||
///
|
||||
/// No-op in ephemeral mode (no durable storage / no WAL). Returns the same
|
||||
/// errors the periodic and shutdown checkpoints would.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TidalError::Storage`/`TidalError::Durability` if any of the
|
||||
/// signal-ledger, WAL-marker, or replication-state checkpoints fail.
|
||||
///
|
||||
/// [`ReplicationState`]: crate::replication::state::ReplicationState
|
||||
#[cfg(any(test, feature = "test-utils"))]
|
||||
pub fn force_replication_checkpoint(&self) -> crate::Result<()> {
|
||||
let (Some(ledger), Some(storage)) = (&self.ledger, &self.storage) else {
|
||||
return Ok(());
|
||||
};
|
||||
let seq = self.last_wal_seq.load(std::sync::atomic::Ordering::Acquire);
|
||||
let meta = crate::signals::checkpoint::CheckpointMeta {
|
||||
checkpoint_time_ns: crate::schema::Timestamp::now().as_nanos(),
|
||||
wal_sequence: seq,
|
||||
payload_hash: [0u8; 32],
|
||||
};
|
||||
ledger.checkpoint(storage.items_engine(), meta)?;
|
||||
crate::db::state_rebuild::checkpoint_replication_state(
|
||||
storage.items_engine(),
|
||||
self.replication_state.as_ref(),
|
||||
)?;
|
||||
// Advance the WAL checkpoint marker so reopen replays only follower-WAL
|
||||
// events strictly after `seq` — keeping the WAL marker consistent with the
|
||||
// fjall ledger checkpoint at the same seq (the same invariant shutdown
|
||||
// maintains), so the recovered aggregate is not double-counted. The
|
||||
// marker write is delegated out of the lock scope below to keep the WAL
|
||||
// mutex held for only the borrow.
|
||||
if seq > 0 {
|
||||
let marker_result = {
|
||||
let guard = self
|
||||
.wal
|
||||
.lock()
|
||||
.map_err(|_| crate::TidalError::internal("wal_access", "WAL mutex poisoned"))?;
|
||||
guard.as_ref().map(|wal| wal.checkpoint(seq))
|
||||
};
|
||||
if let Some(Err(e)) = marker_result {
|
||||
return Err(crate::TidalError::Durability(
|
||||
crate::schema::DurabilityError {
|
||||
message: format!("force_replication_checkpoint WAL marker failed: {e}"),
|
||||
},
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Snapshot this node's live CRDT state for reconciliation after a partition.
|
||||
///
|
||||
/// Walks the signal ledger and the hard-negative index into a
|
||||
/// [`StateSnapshot`](crate::replication::reconcile::StateSnapshot):
|
||||
///
|
||||
/// - **Signals:** every `(entity, signal_type)` entry becomes a per-node
|
||||
/// [`CrdtSignalState`](crate::replication::crdt::CrdtSignalState) attributed
|
||||
/// to this node's shard, carrying its decay score, `last_update_ns`, and —
|
||||
/// crucially — its live `AllTime` windowed count as the per-node bucket
|
||||
/// count, so the windowed total round-trips through merge instead of
|
||||
/// collapsing to 0.
|
||||
/// - **Hard negatives:** every in-memory `(user, item)` pair becomes an
|
||||
/// `LWWRegister<HardNegAction::Hide>` stamped at the current HLC. The index
|
||||
/// stores only resolved membership, so "currently hidden as of now" is the
|
||||
/// honest local assertion; the reconciliation engine's HLC LWW then resolves
|
||||
/// it against any remote hide/unhide.
|
||||
///
|
||||
/// Pair this with [`reconcile_with`](Self::reconcile_with) on the *other*
|
||||
/// node's snapshot (exchanged over the control plane after a partition heals)
|
||||
/// to deterministically merge diverged state. Producing a snapshot does not
|
||||
/// mutate any state.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TidalError::Internal` if this node was opened without a schema
|
||||
/// (no signal ledger to snapshot).
|
||||
pub fn take_crdt_snapshot(
|
||||
&self,
|
||||
) -> crate::Result<crate::replication::reconcile::StateSnapshot> {
|
||||
use crate::replication::{
|
||||
crdt::{Hlc, LWWRegister},
|
||||
reconcile::{HardNegAction, StateSnapshot},
|
||||
};
|
||||
|
||||
let ledger = self.ledger()?;
|
||||
let local_shard = self.config.cluster.shard_id;
|
||||
let now_ns = crate::schema::Timestamp::now().as_nanos();
|
||||
|
||||
let mut snapshot = StateSnapshot::new();
|
||||
|
||||
// Signal contributions (decay score + windowed bucket count per entry).
|
||||
for (entity_id, type_id, state) in ledger.crdt_contributions(local_shard, now_ns) {
|
||||
snapshot.add_signal_state(entity_id, type_id, state);
|
||||
}
|
||||
|
||||
// Hard-negative membership → per-pair Hide registers stamped at now().
|
||||
// One HLC for this node attributes every register to this shard with a
|
||||
// strictly-monotonic timestamp, so the reconcile engine can LWW-order
|
||||
// them against remote hide/unhide decisions.
|
||||
let hlc = Hlc::for_shard(local_shard);
|
||||
for (user_id, item_id) in self.hard_negatives.all_pairs() {
|
||||
let mut reg = LWWRegister::empty();
|
||||
reg.write(HardNegAction::Hide, hlc.now());
|
||||
snapshot.add_hardneg_register(
|
||||
crate::schema::EntityId::new(user_id),
|
||||
crate::schema::EntityId::new(u64::from(item_id)),
|
||||
reg,
|
||||
);
|
||||
}
|
||||
|
||||
Ok(snapshot)
|
||||
}
|
||||
|
||||
/// Reconcile this node's live state against a remote node's CRDT snapshot
|
||||
/// after a partition heals.
|
||||
///
|
||||
/// Builds a [`ReconciliationEngine`](crate::replication::reconcile::ReconciliationEngine)
|
||||
/// bound to this node's signal ledger and hard-negative index, takes this
|
||||
/// node's current [`take_crdt_snapshot`](Self::take_crdt_snapshot) as the
|
||||
/// local side, computes the deterministic
|
||||
/// [`MergePlan`](crate::replication::reconcile::MergePlan) against
|
||||
/// `remote_snapshot`, and applies it:
|
||||
///
|
||||
/// - signal states are CRDT-merged per `(entity, signal_type)` (per-node
|
||||
/// contributions preserved, windowed counts summed) and folded into both
|
||||
/// the hot and warm tiers via `apply_crdt_state`;
|
||||
/// - hard negatives are LWW-resolved per `(user, item)` by HLC timestamp and
|
||||
/// committed (Hide → add, Unhide → remove).
|
||||
///
|
||||
/// Applying is idempotent: re-running with the same `remote_snapshot`
|
||||
/// produces identical state, so this is safe under at-least-once delivery of
|
||||
/// the snapshot exchange. Returns the number of operations applied (signal
|
||||
/// merges + hard-negative resolutions) for observability.
|
||||
///
|
||||
/// This is the production entry point that heals a partition — without it the
|
||||
/// CRDT subsystem never runs outside tests. Wire it into the anti-entropy /
|
||||
/// reconnection flow by exchanging snapshots between the two sides and calling
|
||||
/// this on each. (No standing reconnection hook exists yet; the control plane
|
||||
/// tracks health only.)
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TidalError::Internal` if this node has no schema/ledger, or if a
|
||||
/// signal type in the plan is unknown to this ledger, or a hard-negative item
|
||||
/// ID exceeds the `u32` bitmap range.
|
||||
pub fn reconcile_with(
|
||||
&self,
|
||||
remote_snapshot: &crate::replication::reconcile::StateSnapshot,
|
||||
) -> crate::Result<usize> {
|
||||
use crate::replication::reconcile::ReconciliationEngine;
|
||||
|
||||
self.require_writeable("reconcile_with")?;
|
||||
|
||||
let ledger = Arc::clone(self.ledger()?);
|
||||
let hard_neg = Arc::clone(&self.hard_negatives);
|
||||
let local_snapshot = self.take_crdt_snapshot()?;
|
||||
|
||||
let engine = ReconciliationEngine::new(ledger, hard_neg);
|
||||
let plan = engine.plan(&local_snapshot, remote_snapshot);
|
||||
let op_count = plan.operation_count();
|
||||
engine.apply(&plan)?;
|
||||
|
||||
tracing::info!(
|
||||
ops = op_count,
|
||||
signal_merges = plan.signal_merges.len(),
|
||||
hardneg_resolutions = plan.hardneg_resolutions.len(),
|
||||
"reconcile_with: applied CRDT merge plan after partition heal"
|
||||
);
|
||||
|
||||
Ok(op_count)
|
||||
}
|
||||
|
||||
/// Restore the tenant router's migration routing state (dual-write pairs +
|
||||
/// shard pins) from its durable checkpoint.
|
||||
///
|
||||
/// Idempotent and safe to call once at open: on first boot / no checkpoint /
|
||||
/// ephemeral mode it restores nothing (every tenant uses jump-hash routing).
|
||||
/// Without it, a crash mid-migration silently reverts a finalized tenant to
|
||||
/// jump-hash routing (Accuracy-W). A read error is non-fatal: the node starts
|
||||
/// with empty routing and logs the failure.
|
||||
///
|
||||
/// **Open-time wiring (db/mod.rs, out of scope):** call this once in
|
||||
/// `from_parts` immediately after `build_control_plane`, e.g.
|
||||
/// `db.restore_migration_routing();` after the `TidalDb` value is built — or,
|
||||
/// before the value exists, `crate::replication::tenant::restore_migration_routing(sb.items_engine(), &tenant_router)`.
|
||||
pub(crate) fn restore_migration_routing(&self) {
|
||||
let Ok(storage) = self.storage() else {
|
||||
return; // ephemeral / no schema: nothing durable to restore.
|
||||
};
|
||||
if let Err(e) = crate::replication::tenant::restore_migration_routing(
|
||||
storage.items_engine(),
|
||||
&self.tenant_router,
|
||||
) {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
"migration routing checkpoint read failed; tenant routing starts empty \
|
||||
(tenants fall back to jump-hash routing)"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the shared `TenantRouter` for multi-tenant routing.
|
||||
#[must_use]
|
||||
#[allow(clippy::missing_const_for_fn)]
|
||||
@ -63,14 +279,28 @@ impl TidalDb {
|
||||
///
|
||||
/// Dispatches to each assigned shard. On single-node deployments all
|
||||
/// assignments resolve to the local shard, so behavior is identical to
|
||||
/// calling `signal()` directly. During a dual-write migration, the signal
|
||||
/// is written to both the source and target shard; writes to remote shards
|
||||
/// are logged for visibility until a transport is wired in by the caller.
|
||||
/// calling `signal()` directly.
|
||||
///
|
||||
/// # Remote-shard dispatch is NOT yet wired (fails closed)
|
||||
///
|
||||
/// During a dual-write migration the tenant router resolves two assignments
|
||||
/// (source and target shard). Cross-node dispatch of the remote half over
|
||||
/// the replication transport is **not yet implemented** (tracked as the
|
||||
/// "Task 5" shipper-integrated remote dispatch in `docs/specs/14`, per
|
||||
/// `CODING_GUIDELINES.md` §12 "no TODO without a tracking note"). Rather than
|
||||
/// silently dropping the remote write — which would lose signals during a
|
||||
/// migration with no error visible to the caller — this method **fails
|
||||
/// closed**: if any assignment targets a shard that is not local to this node
|
||||
/// it returns `TidalError::Internal`. A misconfigured dual-write is therefore
|
||||
/// loud, not silently lossy. The local half (if any) is still written before
|
||||
/// the error is returned.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// - `TidalError::QuotaExceeded` — tenant rate limit exceeded.
|
||||
/// - `TidalError::Internal` — no eligible shards for the tenant.
|
||||
/// - `TidalError::Internal` — no eligible shards for the tenant, OR an
|
||||
/// assignment targets a non-local shard while cross-node dispatch is
|
||||
/// unwired (see above).
|
||||
/// - Any error from the underlying `signal()` call.
|
||||
pub fn signal_for_tenant(
|
||||
&self,
|
||||
@ -92,16 +322,30 @@ impl TidalDb {
|
||||
// Local shard: write directly to the signal ledger.
|
||||
self.signal(signal_type, entity_id, weight, timestamp)?;
|
||||
} else {
|
||||
// Remote shard (dual-write migration path): a transport-based dispatch
|
||||
// will be wired here once the shipper is integrated (Task 5).
|
||||
// For now we record the pending dispatch so it is visible to operators.
|
||||
tracing::debug!(
|
||||
// Remote shard (dual-write migration path). Cross-node dispatch
|
||||
// over the replication transport is not yet wired (Task 5 in
|
||||
// docs/specs/14 per CODING_GUIDELINES §12). Fail CLOSED instead of
|
||||
// silently dropping the write: a dual-write that cannot reach its
|
||||
// remote half would otherwise lose signals during migration with
|
||||
// no error surfaced. Log at warn so the misconfiguration is
|
||||
// observable, then return a typed error.
|
||||
tracing::warn!(
|
||||
tenant_id = tenant_id.0,
|
||||
shard_id = assignment.shard_id.0,
|
||||
entity_id = entity_id.as_u64(),
|
||||
signal_type,
|
||||
"pending remote-shard dispatch (transport not yet wired)"
|
||||
"signal_for_tenant: remote-shard dispatch is not wired (Task 5); \
|
||||
failing closed to avoid silently dropping the dual-write"
|
||||
);
|
||||
return Err(crate::TidalError::internal(
|
||||
"signal_for_tenant",
|
||||
format!(
|
||||
"remote-shard dispatch not yet wired: assignment targets shard {} \
|
||||
but this node is shard {}; cross-node dual-write requires the \
|
||||
transport-integrated shipper (Task 5, docs/specs/14)",
|
||||
assignment.shard_id.0, local_shard.0
|
||||
),
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
|
||||
@ -24,6 +24,33 @@
|
||||
///
|
||||
/// Hashing `Duration::as_nanos()` (u128) instead of the derived `f64` lambda
|
||||
/// avoids IEEE 754 representation differences across platforms.
|
||||
///
|
||||
/// # Deliberately excluded fields (and the drift risk they carry)
|
||||
///
|
||||
/// A signal's **windows**, **velocity-enabled flag**, and **positive-engagement
|
||||
/// flag** are intentionally NOT hashed. The fingerprint exists to protect
|
||||
/// *on-disk integrity* — to reject a reopen whose schema change would make the
|
||||
/// persisted bytes uninterpretable (a different signal set, decay model, decay
|
||||
/// parameter, or embedding dimensionality). Windows/velocity/positive-engagement
|
||||
/// do not change how already-written data is laid out or decoded; they only
|
||||
/// change how the engine *aggregates and folds future signals*. Including them
|
||||
/// is therefore not required for byte-level correctness.
|
||||
///
|
||||
/// The accepted trade-off is **silent semantic drift on reopen**: an operator
|
||||
/// who edits only a window set, the velocity flag, or the positive-engagement
|
||||
/// flag will reopen successfully, and the engine will quietly start aggregating
|
||||
/// under the new declaration while pre-change durable aggregates were built
|
||||
/// under the old one — no `SchemaMismatch` is raised. This is a deliberate
|
||||
/// choice, not an oversight.
|
||||
///
|
||||
/// These fields are *excluded rather than included* because adding them now
|
||||
/// would change the fingerprint of every signal and would brick reopen of every
|
||||
/// existing data directory with a spurious `SchemaMismatch` (the stored
|
||||
/// fingerprint was computed without them). If a future migration is willing to
|
||||
/// pay that back-compat cost (or versions the fingerprint), folding these in —
|
||||
/// using a stable encoding for each `Window`, a single flag byte for velocity,
|
||||
/// and a single flag byte for positive-engagement — is the way to turn the
|
||||
/// silent semantic drift above into a loud, caught mismatch.
|
||||
pub fn compute_schema_fingerprint(schema: &crate::schema::Schema) -> [u8; 32] {
|
||||
use blake3::Hasher;
|
||||
|
||||
@ -39,13 +66,8 @@ pub fn compute_schema_fingerprint(schema: &crate::schema::Schema) -> [u8; 32] {
|
||||
// Signal name.
|
||||
hasher.update(sig.name().as_bytes());
|
||||
|
||||
// Entity kind.
|
||||
let kind_byte = match sig.target() {
|
||||
crate::schema::EntityKind::Item => 0u8,
|
||||
crate::schema::EntityKind::User => 1u8,
|
||||
crate::schema::EntityKind::Creator => 2u8,
|
||||
};
|
||||
hasher.update(&[kind_byte]);
|
||||
// Entity kind — single canonical encoding (Item=0/User=1/Creator=2).
|
||||
hasher.update(&[sig.target().fingerprint_byte()]);
|
||||
|
||||
// Decay model.
|
||||
match sig.decay() {
|
||||
@ -65,16 +87,11 @@ pub fn compute_schema_fingerprint(schema: &crate::schema::Schema) -> [u8; 32] {
|
||||
|
||||
// Embedding slots. The accessor returns `(name, kind, dims)` already sorted
|
||||
// deterministically by `(name, kind)`, so no extra sort is needed here. The
|
||||
// kind byte uses the same Item=0/User=1/Creator=2 mapping as the signal
|
||||
// section above, keeping one canonical kind encoding across the fingerprint.
|
||||
// kind byte uses the same canonical `EntityKind::fingerprint_byte` mapping as
|
||||
// the signal section above, keeping one source of truth for the encoding.
|
||||
for (name, kind, dims) in schema.embedding_slot_fingerprint() {
|
||||
hasher.update(name.as_bytes());
|
||||
let kind_byte = match kind {
|
||||
crate::schema::EntityKind::Item => 0u8,
|
||||
crate::schema::EntityKind::User => 1u8,
|
||||
crate::schema::EntityKind::Creator => 2u8,
|
||||
};
|
||||
hasher.update(&[kind_byte]);
|
||||
hasher.update(&[kind.fingerprint_byte()]);
|
||||
hasher.update(&(dims as u64).to_le_bytes());
|
||||
}
|
||||
|
||||
|
||||
@ -10,7 +10,7 @@ use std::{
|
||||
|
||||
use super::TidalDb;
|
||||
use crate::{
|
||||
schema::{EntityId, Timestamp},
|
||||
schema::EntityId,
|
||||
session::{self as session_mod, AgentId, SessionId, SessionState},
|
||||
storage::{Tag, encode_key, parse_key},
|
||||
wal::format::session::SessionSeqNo,
|
||||
@ -86,7 +86,7 @@ impl TidalDb {
|
||||
/// Called from `from_parts()` during startup. The storage-backed
|
||||
/// `restore_sessions()` is additive and handles archived snapshots; this
|
||||
/// method handles in-flight sessions that were active at crash time.
|
||||
#[allow(clippy::too_many_lines, clippy::cast_possible_truncation)]
|
||||
#[allow(clippy::too_many_lines)]
|
||||
pub(super) fn restore_session_wal_events(
|
||||
&self,
|
||||
events: &[crate::wal::format::SessionWalEvent],
|
||||
@ -102,8 +102,6 @@ impl TidalDb {
|
||||
return;
|
||||
};
|
||||
|
||||
let now_ns = Timestamp::now().as_nanos();
|
||||
|
||||
for (session_id, (user_id, started_at_ns, agent_id_str, policy_name)) in open_sessions {
|
||||
if schema.session_policy(&policy_name).is_none() {
|
||||
tracing::warn!(
|
||||
@ -159,19 +157,37 @@ impl TidalDb {
|
||||
|
||||
// Replay signals and annotations into the restored session state.
|
||||
// Advance the seqno HWM so incoming replication events are correctly deduped.
|
||||
//
|
||||
// CRITICAL (§8 "WAL replay produces identical state"): replay events
|
||||
// in ascending ts_ns order and seed each per-signal-type state with
|
||||
// that type's *earliest* event timestamp, not the restore-time wall
|
||||
// clock. The live path seeds `SessionSignalState::new(ts_ns, ..)`
|
||||
// with the first signal's own timestamp (db/sessions.rs), so its
|
||||
// BucketedCounter rotations fire as real time advances. If we seeded
|
||||
// with `now_ns` and fed past timestamps, `maybe_rotate` would early-
|
||||
// return on every event and collapse all replayed signals into one
|
||||
// current-minute bucket, diverging the restored window counts. By
|
||||
// sorting ascending, the first occurrence of each signal_name in the
|
||||
// iteration carries that name's earliest ts, so the `or_insert_with`
|
||||
// seed is exactly the live seed and rotations advance identically.
|
||||
let mut max_seqno: Option<u64> = None;
|
||||
if let Some(signals) = session_signals.get(&session_id) {
|
||||
for (entity_id, weight, ts_ns, signal_name, annotation, seqno_raw) in signals {
|
||||
let mut ordered: Vec<_> = signals.iter().collect();
|
||||
ordered.sort_by_key(|(_, _, ts_ns, _, _, _)| *ts_ns);
|
||||
for (entity_id, weight, ts_ns, signal_name, annotation, seqno_raw) in ordered {
|
||||
let lambda = schema
|
||||
.signal(signal_name)
|
||||
.and_then(|def| def.decay().lambda())
|
||||
.unwrap_or(session_mod::DEFAULT_SESSION_LAMBDA);
|
||||
|
||||
// Seed with this event's own ts_ns: because events are sorted
|
||||
// ascending, the first signal of each name carries its
|
||||
// earliest timestamp, matching the live `new(ts_ns, lambda)`.
|
||||
let ss = state
|
||||
.signals
|
||||
.entry(signal_name.clone())
|
||||
.or_insert_with(|| session_mod::SessionSignalState::new(now_ns, lambda));
|
||||
ss.on_signal(f64::from(*weight), *ts_ns);
|
||||
.or_insert_with(|| session_mod::SessionSignalState::new(*ts_ns, lambda));
|
||||
ss.on_signal(*weight, *ts_ns);
|
||||
drop(ss);
|
||||
|
||||
// Route through the capped insert so WAL replay honors the
|
||||
@ -245,14 +261,14 @@ impl TidalDb {
|
||||
events: &[crate::wal::format::SessionWalEvent],
|
||||
) -> (
|
||||
HashMap<u64, (u64, u64, String, String)>,
|
||||
HashMap<u64, Vec<(u64, f32, u64, String, Option<String>, Option<u64>)>>,
|
||||
HashMap<u64, Vec<(u64, f64, u64, String, Option<String>, Option<u64>)>>,
|
||||
) {
|
||||
use crate::wal::format::SessionWalEvent;
|
||||
|
||||
let mut open_sessions: HashMap<u64, (u64, u64, String, String)> = HashMap::new();
|
||||
let mut session_signals: HashMap<
|
||||
u64,
|
||||
Vec<(u64, f32, u64, String, Option<String>, Option<u64>)>,
|
||||
Vec<(u64, f64, u64, String, Option<String>, Option<u64>)>,
|
||||
> = HashMap::new();
|
||||
|
||||
for event in events {
|
||||
|
||||
@ -193,19 +193,6 @@ impl TidalDb {
|
||||
tracing::warn!(error = %e, session_id = %session_id, "session journal append failed");
|
||||
}
|
||||
|
||||
// Evict oldest closed sessions if the cap is exceeded.
|
||||
if self.closed_sessions.len() >= session_mod::MAX_CLOSED_SESSIONS {
|
||||
let mut entries: Vec<_> = self
|
||||
.closed_sessions
|
||||
.iter()
|
||||
.map(|r| (*r.key(), r.value().closed_at_ns))
|
||||
.collect();
|
||||
entries.sort_unstable_by_key(|(_, ts)| *ts);
|
||||
for (id, _) in entries.iter().take(session_mod::EVICT_BATCH_SIZE) {
|
||||
self.closed_sessions.remove(id);
|
||||
}
|
||||
}
|
||||
|
||||
// Resolve user_id before state is dropped.
|
||||
let user_id = state.user_id;
|
||||
|
||||
@ -218,7 +205,15 @@ impl TidalDb {
|
||||
// embedding into the user's global preference vector with adaptive learning rate.
|
||||
self.apply_session_preference_update(user_id, &snapshot);
|
||||
|
||||
// Insert first, then enforce the cap, so the post-condition is always
|
||||
// `len() <= MAX_CLOSED_SESSIONS` once this returns. The previous
|
||||
// check-then-insert ordering tested the cap before the insert, so N
|
||||
// concurrent closes that each saw `len() < MAX` would all skip eviction
|
||||
// and then all insert, leaving the map at `MAX - 1 + N` — a persistent
|
||||
// overshoot. `enforce_closed_session_cap` trims down to the cap after
|
||||
// the insert instead, so any overshoot is transient and self-healing.
|
||||
self.closed_sessions.insert(session_id, snapshot);
|
||||
self.enforce_closed_session_cap();
|
||||
|
||||
#[cfg(feature = "metrics")]
|
||||
{
|
||||
@ -251,6 +246,38 @@ impl TidalDb {
|
||||
})
|
||||
}
|
||||
|
||||
/// Trim `closed_sessions` down to `MAX_CLOSED_SESSIONS` by evicting the
|
||||
/// oldest snapshots (smallest `closed_at_ns`).
|
||||
///
|
||||
/// Called after an insert so the cap is enforced as a post-condition rather
|
||||
/// than a pre-condition: when this returns, `closed_sessions.len()` is at
|
||||
/// most `MAX_CLOSED_SESSIONS` from this thread's view. The trim removes a
|
||||
/// bounded `EVICT_BATCH_SIZE` extra entries below the cap so steady-state
|
||||
/// closes do not re-sort the whole map on every single close.
|
||||
///
|
||||
/// Concurrency: `DashMap::remove` is atomic, so two threads trimming at once
|
||||
/// cooperate — a double-remove of the same oldest id simply returns `None`.
|
||||
/// Any overshoot from concurrent inserts is therefore transient and
|
||||
/// self-healing, never a persistent cap violation.
|
||||
fn enforce_closed_session_cap(&self) {
|
||||
if self.closed_sessions.len() <= session_mod::MAX_CLOSED_SESSIONS {
|
||||
return;
|
||||
}
|
||||
// Snapshot (id, closed_at_ns) and evict the oldest down to a low-water
|
||||
// mark (cap - EVICT_BATCH_SIZE) so the cap holds with headroom.
|
||||
let mut entries: Vec<_> = self
|
||||
.closed_sessions
|
||||
.iter()
|
||||
.map(|r| (*r.key(), r.value().closed_at_ns))
|
||||
.collect();
|
||||
entries.sort_unstable_by_key(|(_, ts)| *ts);
|
||||
let target = session_mod::MAX_CLOSED_SESSIONS.saturating_sub(session_mod::EVICT_BATCH_SIZE);
|
||||
let evict = entries.len().saturating_sub(target);
|
||||
for (id, _) in entries.iter().take(evict) {
|
||||
self.closed_sessions.remove(id);
|
||||
}
|
||||
}
|
||||
|
||||
/// List all currently active sessions.
|
||||
#[must_use]
|
||||
pub fn active_sessions(&self) -> Vec<SessionInfo> {
|
||||
@ -436,7 +463,7 @@ impl TidalDb {
|
||||
&& let Err(e) = wal.session_signal(
|
||||
handle.id.as_u64(),
|
||||
entity_id.as_u64(),
|
||||
weight as f32,
|
||||
weight,
|
||||
ts_ns,
|
||||
signal_type,
|
||||
ann_for_wal.as_deref(),
|
||||
|
||||
@ -4,10 +4,13 @@ use std::collections::HashMap;
|
||||
|
||||
use super::{
|
||||
TidalDb,
|
||||
metadata::{is_positive_engagement_signal, serialize_metadata},
|
||||
metadata::{
|
||||
UserStateSignal, classify_user_state_signal, is_positive_engagement_signal,
|
||||
serialize_metadata,
|
||||
},
|
||||
};
|
||||
use crate::{
|
||||
entities::HardNegIndex,
|
||||
entities::{HardNegIndex, RelationshipType},
|
||||
governance::SignalScope,
|
||||
schema::{EntityId, EntityKind, TidalError, Timestamp, Window},
|
||||
};
|
||||
@ -294,9 +297,20 @@ impl TidalDb {
|
||||
/// # Durability
|
||||
///
|
||||
/// The base signal (entity, type, weight, timestamp) is WAL-backed and survives crashes.
|
||||
/// User-context side effects (seen state, interaction weights, preference vector updates)
|
||||
/// are reconstructed from durable storage on restart via `rebuild_entity_state`.
|
||||
/// Hard negatives (hide/block) are durably written via `write_relationship()`.
|
||||
/// The per-user side effects do NOT live in the WAL (the signal record carries no
|
||||
/// `for_user` identity), so each one is made crash-recoverable on its own path:
|
||||
///
|
||||
/// - **Hard negatives** (skip/hide/dislike/block) are written write-through as durable
|
||||
/// `Tag::HardNeg` rows (zero loss window) and restored on open via
|
||||
/// [`HardNegIndex::restore`](crate::entities::HardNegIndex::restore).
|
||||
/// - **Interaction weights** are persisted write-through as a merged
|
||||
/// `RelationshipType::InteractionWeight` edge carrying the decayed running score and
|
||||
/// last-update timestamp, so `rebuild_entity_state` reproduces the ledger on open.
|
||||
/// - **Completion / save / like** state is written write-through as durable
|
||||
/// `Tag::UserState` rows and restored on open, so the `InProgress` filter and
|
||||
/// `DateSaved` sort survive a crash.
|
||||
/// - **Preference vectors** are periodically checkpointed (bounded post-crash loss
|
||||
/// window) and restored on open.
|
||||
///
|
||||
/// Seen state from regular views/likes is intentionally ephemeral -- users should see
|
||||
/// content again after a restart. Only explicit hides (via `write_relationship` with
|
||||
@ -321,20 +335,53 @@ impl TidalDb {
|
||||
// Signal dispatch: side effects based on signal type and context.
|
||||
if let Some(user_id) = for_user {
|
||||
let item_u32 = entity_id.as_u64() as u32;
|
||||
// Durable-write target for the write-through side effects below.
|
||||
// `None` in ephemeral mode (no persistence; in-memory state is the
|
||||
// whole story and there is nothing to crash-recover).
|
||||
let storage = self.storage.as_ref();
|
||||
|
||||
// 1. Hard negatives.
|
||||
// 1. Hard negatives. Write-through: a durable Tag::HardNeg row is
|
||||
// persisted with the SAME durability class as the base signal so a
|
||||
// skip/dislike survives a hard crash with zero loss window, then
|
||||
// restored into the in-memory bitmap on open.
|
||||
if HardNegIndex::is_hard_neg_signal(signal_type) {
|
||||
self.hard_negatives.add(user_id, item_u32);
|
||||
if let Some(sb) = storage
|
||||
&& let Err(e) =
|
||||
self.hard_negatives
|
||||
.persist(sb.items_engine(), user_id, item_u32)
|
||||
{
|
||||
tracing::error!(
|
||||
user_id,
|
||||
item = item_u32,
|
||||
error = %e,
|
||||
"failed to persist hard-negative durable row; \
|
||||
it would not survive a crash"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// 2. Seen tracking.
|
||||
self.user_state.mark_seen(user_id, item_u32);
|
||||
|
||||
// 3. Interaction weight: if creator is known, update the
|
||||
// (user, creator) interaction strength.
|
||||
// (user, creator) interaction strength, then persist the decayed
|
||||
// running score + last-update as a merged InteractionWeight edge so
|
||||
// `rebuild_entity_state` reproduces the ledger on open.
|
||||
if let Some(cid) = creator_id {
|
||||
self.interaction_ledger
|
||||
.record(user_id, cid, weight, timestamp.as_nanos());
|
||||
self.persist_interaction_edge(user_id, cid);
|
||||
}
|
||||
|
||||
// 3b. Completion / save / like: drive the UserStateIndex maps the
|
||||
// InProgress filter and DateSaved sort read. Write-through so they
|
||||
// survive a crash. Signals do not carry a completion ratio, so the
|
||||
// signal weight IS the completion ratio (documented choice).
|
||||
if let Some(kind) = classify_user_state_signal(signal_type) {
|
||||
self.persist_user_state_signal(
|
||||
kind, user_id, entity_id, item_u32, weight, timestamp,
|
||||
);
|
||||
}
|
||||
|
||||
// 4. Preference vector + co-engagement: for positive engagement signals,
|
||||
@ -361,6 +408,100 @@ impl TidalDb {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Persist the current `(user, creator)` interaction weight as a durable
|
||||
/// `RelationshipType::InteractionWeight` edge so `rebuild_entity_state`
|
||||
/// reproduces the ledger on open.
|
||||
///
|
||||
/// Reads back the raw running score + last-update timestamp (post-`record`)
|
||||
/// and overwrites the edge with those exact values. Because the rebuild path
|
||||
/// replays a stored edge through `InteractionLedger::record` into a fresh
|
||||
/// ledger — folding the persisted score as one in-order event at
|
||||
/// `last_update_ns` — the rebuilt entry equals this one within decay
|
||||
/// tolerance. Best-effort: a storage error is logged, not propagated, so the
|
||||
/// base signal still succeeds.
|
||||
fn persist_interaction_edge(&self, user_id: u64, creator_id: u64) {
|
||||
use crate::entities::relationship::{
|
||||
encode_relationship_key, serialize_relationship_value,
|
||||
};
|
||||
let Some((score, last_update_ns)) = self.interaction_ledger.raw_entry(user_id, creator_id)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let Some(storage) = self.storage.as_ref() else {
|
||||
return;
|
||||
};
|
||||
let key = encode_relationship_key(
|
||||
EntityId::new(user_id),
|
||||
RelationshipType::InteractionWeight,
|
||||
EntityId::new(creator_id),
|
||||
);
|
||||
let value = serialize_relationship_value(score, Timestamp::from_nanos(last_update_ns));
|
||||
if let Err(e) = storage.users_engine().put(&key, &value) {
|
||||
tracing::error!(
|
||||
user_id,
|
||||
creator_id,
|
||||
error = %e,
|
||||
"failed to persist interaction-weight edge; \
|
||||
it would reset to zero on restart"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Drive (and durably persist) the `UserStateIndex` map for a
|
||||
/// completion / save / like signal.
|
||||
///
|
||||
/// The signal weight is used as the completion ratio (signals carry no
|
||||
/// explicit ratio). Write-through so the `InProgress` filter / `DateSaved`
|
||||
/// sort survive a crash; best-effort persistence (logged, not propagated).
|
||||
fn persist_user_state_signal(
|
||||
&self,
|
||||
kind: UserStateSignal,
|
||||
user_id: u64,
|
||||
entity_id: EntityId,
|
||||
item_u32: u32,
|
||||
weight: f64,
|
||||
timestamp: Timestamp,
|
||||
) {
|
||||
let Some(storage) = self.storage.as_ref() else {
|
||||
// Ephemeral mode: update in-memory state only (nothing to recover).
|
||||
match kind {
|
||||
UserStateSignal::Completion => {
|
||||
self.user_state
|
||||
.record_completion(user_id, entity_id.as_u64(), weight);
|
||||
}
|
||||
UserStateSignal::Save => {
|
||||
self.user_state
|
||||
.add_save_timestamped(user_id, item_u32, timestamp.as_nanos());
|
||||
}
|
||||
UserStateSignal::Like => self.user_state.add_like(user_id, item_u32),
|
||||
}
|
||||
return;
|
||||
};
|
||||
let engine = storage.items_engine();
|
||||
let result = match kind {
|
||||
UserStateSignal::Completion => self.user_state.record_completion_durable(
|
||||
engine,
|
||||
user_id,
|
||||
entity_id.as_u64(),
|
||||
weight,
|
||||
),
|
||||
UserStateSignal::Save => {
|
||||
self.user_state
|
||||
.add_save_durable(engine, user_id, item_u32, timestamp.as_nanos())
|
||||
}
|
||||
UserStateSignal::Like => self.user_state.add_like_durable(engine, user_id, item_u32),
|
||||
};
|
||||
if let Err(e) = result {
|
||||
tracing::error!(
|
||||
user_id,
|
||||
item = item_u32,
|
||||
error = %e,
|
||||
"failed to persist user-state durable row; \
|
||||
InProgress/DateSaved would not survive a crash"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Best-effort cohort attribution: resolves user cohorts and records the signal
|
||||
/// into each matching cohort's signal ledger.
|
||||
///
|
||||
|
||||
@ -16,14 +16,89 @@ use super::{metadata::deserialize_metadata, metrics::MetricsState, storage_box::
|
||||
use crate::{
|
||||
cohort::CohortSignalLedger,
|
||||
query::suggest::SuggestionIndex,
|
||||
schema::{TidalError, Timestamp},
|
||||
replication::{ShardId, state::ReplicationState},
|
||||
schema::{EntityId, TidalError, Timestamp},
|
||||
signals::{DEFAULT_MAX_SIGNAL_ENTRIES, SignalLedger, trim_cold_entries},
|
||||
storage::{
|
||||
StorageEngine, Tag,
|
||||
StorageEngine, Tag, encode_key,
|
||||
indexes::{bitmap::BitmapIndex, range::RangeIndex},
|
||||
},
|
||||
};
|
||||
|
||||
/// Suffix for the single follower-replication-state checkpoint row.
|
||||
const REPLICATION_STATE_SUFFIX: &[u8] = b"hwm";
|
||||
|
||||
/// Storage key for the follower replication high-water-mark checkpoint.
|
||||
///
|
||||
/// One row at entity ID 0 under [`Tag::ReplicationState`], holding the
|
||||
/// JSON-serialized per-shard `applied_seqno`.
|
||||
#[must_use]
|
||||
fn replication_state_key() -> Vec<u8> {
|
||||
encode_key(
|
||||
EntityId::new(0),
|
||||
Tag::ReplicationState,
|
||||
REPLICATION_STATE_SUFFIX,
|
||||
)
|
||||
}
|
||||
|
||||
/// Durably persist the follower's per-shard replication high-water-mark.
|
||||
///
|
||||
/// Serializes [`ReplicationState::to_checkpoint_bytes`] and commits it (with an
|
||||
/// fsync via `flush`) so a follower crash cannot reset its idempotency boundary
|
||||
/// to 0. Called from the periodic checkpoint thread and from shutdown. Failure
|
||||
/// is non-fatal at the call site: a missed checkpoint means the next open
|
||||
/// restores an older high-water-mark (correct, just re-applies a few re-shipped
|
||||
/// segments idempotently), but it still signals a failing disk.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TidalError::Storage` if the write or flush fails.
|
||||
pub(super) fn checkpoint_replication_state(
|
||||
storage: &dyn StorageEngine,
|
||||
replication_state: &ReplicationState,
|
||||
) -> crate::Result<()> {
|
||||
let bytes = replication_state.to_checkpoint_bytes();
|
||||
let mut batch = crate::storage::WriteBatch::new();
|
||||
batch.put(replication_state_key(), bytes);
|
||||
storage.write_batch(batch)?;
|
||||
storage.flush()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Restore a follower [`ReplicationState`] from its durable checkpoint.
|
||||
///
|
||||
/// Reads the persisted per-shard high-water-mark and reconstructs a
|
||||
/// `ReplicationState` tracking `shards`. Shards absent from the checkpoint (or
|
||||
/// the first-boot / corrupt-bytes cases) start at 0 — replaying from 0 is
|
||||
/// always correct, just slower. The returned state is the SINGLE Arc the
|
||||
/// receiver advances and the lag gauge reads (obs-REPL-1): the caller must not
|
||||
/// build a second `ReplicationState` for the same node.
|
||||
#[must_use]
|
||||
pub(super) fn restore_replication_state(
|
||||
storage: &dyn StorageEngine,
|
||||
shards: &[ShardId],
|
||||
) -> ReplicationState {
|
||||
match storage.get(&replication_state_key()) {
|
||||
Ok(Some(bytes)) => {
|
||||
let restored = ReplicationState::from_checkpoint_bytes(&bytes, shards);
|
||||
tracing::info!(
|
||||
shards = shards.len(),
|
||||
"follower replication high-water-mark restored from checkpoint"
|
||||
);
|
||||
restored
|
||||
}
|
||||
Ok(None) => ReplicationState::new(shards),
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
"replication-state checkpoint read failed; \
|
||||
follower starts at applied seqno 0 (re-applies re-shipped segments)"
|
||||
);
|
||||
ReplicationState::new(shards)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Index health metrics handles ────────────────────────────────────────────
|
||||
|
||||
/// Handles to live index instances for the periodic checkpoint-thread refresh.
|
||||
@ -111,7 +186,7 @@ pub(super) fn rebuild_entity_state(
|
||||
rel_count += 1;
|
||||
}
|
||||
RelationshipType::Hide => {
|
||||
user_state.add_hide(from_id_u64, to_id.as_u64() as u32);
|
||||
user_state.add_hide_item(from_id_u64, to_id.as_u64());
|
||||
rel_count += 1;
|
||||
}
|
||||
RelationshipType::Follows => {
|
||||
@ -146,34 +221,6 @@ pub(super) fn rebuild_entity_state(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Rebuild `SuggestionIndex` title terms from durable item metadata on restart.
|
||||
///
|
||||
/// Scans the items keyspace for `Tag::Meta` keys, deserializes metadata, and
|
||||
/// calls `suggestion_index.index_title(title)` for each item that has a `"title"`
|
||||
/// field. This ensures autocomplete works correctly after a restart without
|
||||
/// requiring all items to be re-written.
|
||||
///
|
||||
/// For ephemeral mode the engine is empty, so this is a no-op.
|
||||
pub(super) fn rebuild_suggestion_index(storage: &StorageBox, suggestion_index: &SuggestionIndex) {
|
||||
let mut indexed = 0u64;
|
||||
for entry in storage.items_engine().scan_prefix(&[]) {
|
||||
let Ok((key, value)) = entry else { continue };
|
||||
if let Some((_entity_id, Tag::Meta, _suffix)) = crate::storage::keys::parse_key(&key) {
|
||||
let meta = deserialize_metadata(&value);
|
||||
if let Some(title) = meta.get("title") {
|
||||
suggestion_index.index_title(title);
|
||||
indexed += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
if indexed > 0 {
|
||||
tracing::info!(
|
||||
items = indexed,
|
||||
"suggestion index rebuilt from durable storage"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Borrowed references to every in-memory item index, bundled so the write path
|
||||
/// and the restart-rebuild path share one indexing implementation.
|
||||
///
|
||||
@ -329,9 +376,12 @@ pub(super) fn rebuild_item_indexes(
|
||||
|
||||
/// Background thread body: checkpoint signal state to storage every 30 seconds.
|
||||
///
|
||||
/// Checkpoints both the global signal ledger and the cohort signal ledger
|
||||
/// atomically (each writes its own `WriteBatch`). The cohort checkpoint uses
|
||||
/// the same storage engine and the same `CheckpointMeta` as the global ledger.
|
||||
/// Checkpoints the global signal ledger, the cohort signal ledger, the
|
||||
/// per-contributor community ledger, the co-engagement index, and the per-user
|
||||
/// preference vectors (each writes its own `WriteBatch`). Every one of these is
|
||||
/// derived state with no WAL backing for its per-user/per-contributor identity,
|
||||
/// so a periodic checkpoint is what bounds their post-crash loss window to the
|
||||
/// checkpoint interval instead of "everything since the last clean shutdown."
|
||||
///
|
||||
/// After each successful checkpoint, compacts WAL segments that are fully
|
||||
/// covered by the checkpoint. Compaction failure is non-fatal: a warning is
|
||||
@ -344,11 +394,19 @@ pub(super) fn rebuild_item_indexes(
|
||||
/// The `Arc` arguments are intentionally passed by value: the thread must own
|
||||
/// them for its entire lifetime (references cannot satisfy the `'static` bound
|
||||
/// required by `std::thread::spawn`).
|
||||
#[allow(clippy::needless_pass_by_value, clippy::too_many_arguments)]
|
||||
#[allow(
|
||||
clippy::needless_pass_by_value,
|
||||
clippy::too_many_arguments,
|
||||
clippy::too_many_lines
|
||||
)]
|
||||
pub(super) fn run_checkpoint_thread(
|
||||
shutdown: Arc<AtomicBool>,
|
||||
ledger: Arc<SignalLedger>,
|
||||
cohort_ledger: Arc<CohortSignalLedger>,
|
||||
community_ledger: Arc<crate::governance::CommunityLedger>,
|
||||
co_engagement: Arc<crate::entities::CoEngagementIndex>,
|
||||
preference_vectors: Arc<crate::entities::PreferenceVectors>,
|
||||
replication_state: Arc<ReplicationState>,
|
||||
storage: Box<dyn StorageEngine + Send + Sync>,
|
||||
last_wal_seq: Arc<AtomicU64>,
|
||||
wal_dir: Option<PathBuf>,
|
||||
@ -377,7 +435,7 @@ pub(super) fn run_checkpoint_thread(
|
||||
// feature so search liveness is always observed.
|
||||
if index_metrics_elapsed >= INDEX_METRICS_INTERVAL {
|
||||
index_metrics_elapsed = Duration::ZERO;
|
||||
poll_text_index_health(&index_handles, storage.as_ref(), last_wal_seq.as_ref());
|
||||
poll_text_index_health(&index_handles, storage.as_ref());
|
||||
#[cfg(feature = "metrics")]
|
||||
refresh_index_metrics(&index_handles, &metrics);
|
||||
}
|
||||
@ -447,6 +505,19 @@ pub(super) fn run_checkpoint_thread(
|
||||
// never deletes the segment the writer is appending to (deleting
|
||||
// it would unlink the inode out from under the open FD and lose
|
||||
// post-checkpoint writes on Linux).
|
||||
//
|
||||
// NOTE: the WAL checkpoint *marker* is deliberately NOT advanced
|
||||
// here — only on clean shutdown (lifecycle.rs) and via the
|
||||
// test-only force_replication_checkpoint. Advancing it on the racy
|
||||
// periodic snapshot would trade the existing safe over-count
|
||||
// behavior (a crash replays a few already-materialized events) for
|
||||
// a potential lost-write (the marker would suppress replay of an
|
||||
// event the concurrent snapshot missed). For eventually-consistent
|
||||
// signal state, over-count self-heals via decay; an acked lost
|
||||
// write does not. The follower's replicated aggregate is therefore
|
||||
// fully recoverable from checkpoint + WAL replay, and its persisted
|
||||
// high-water-mark (checkpointed below) is the leader-seqno
|
||||
// idempotency boundary that re-shipped segments are gated against.
|
||||
if let Some(ref dir) = wal_dir
|
||||
&& seq > 0
|
||||
{
|
||||
@ -465,10 +536,13 @@ pub(super) fn run_checkpoint_thread(
|
||||
}
|
||||
}
|
||||
|
||||
// Update WAL lag bytes: sum remaining segment file sizes.
|
||||
// Update WAL lag bytes: sum only segments NOT yet compacted
|
||||
// by the checkpoint, which is exactly what the gauge's doc
|
||||
// ("Bytes of WAL segments not yet compacted") promises. `seq`
|
||||
// is the checkpoint sequence we just compacted at.
|
||||
#[cfg(feature = "metrics")]
|
||||
{
|
||||
let lag = compute_wal_lag_bytes(dir);
|
||||
let lag = compute_wal_lag_bytes(dir, seq);
|
||||
metrics.wal_lag_bytes.store(lag, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
@ -479,6 +553,28 @@ pub(super) fn run_checkpoint_thread(
|
||||
{
|
||||
tracing::error!(error = %e, "periodic cohort checkpoint failed");
|
||||
}
|
||||
|
||||
// Checkpoint the derived per-user / per-contributor ledgers that have
|
||||
// no WAL backstop. A periodic checkpoint bounds their post-crash loss
|
||||
// window to the interval instead of "everything since the last clean
|
||||
// shutdown" (the community ledger is the SOLE durable copy of
|
||||
// per-contributor identity, so this is its only periodic durability).
|
||||
checkpoint_derived_ledgers(
|
||||
storage.as_ref(),
|
||||
community_ledger.as_ref(),
|
||||
co_engagement.as_ref(),
|
||||
preference_vectors.as_ref(),
|
||||
);
|
||||
|
||||
// Persist the follower replication high-water-mark (BLOCKER 6): a
|
||||
// restarted follower must restore its idempotency boundary, not
|
||||
// reset to 0 and double-count every re-shipped segment. Non-fatal —
|
||||
// a missed checkpoint restores an older (still-correct) boundary.
|
||||
if let Err(e) =
|
||||
checkpoint_replication_state(storage.as_ref(), replication_state.as_ref())
|
||||
{
|
||||
tracing::error!(error = %e, "periodic replication-state checkpoint failed");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -487,6 +583,33 @@ pub(super) fn run_checkpoint_thread(
|
||||
let _ = &metrics;
|
||||
}
|
||||
|
||||
/// Checkpoint the derived per-user / per-contributor ledgers (community,
|
||||
/// co-engagement, preference vectors) into `storage`.
|
||||
///
|
||||
/// None of these has a WAL backstop for its per-user/per-contributor identity, so
|
||||
/// this periodic checkpoint is what bounds their post-crash loss window. Each is
|
||||
/// independent and non-fatal: a failure is logged and the others still run.
|
||||
fn checkpoint_derived_ledgers(
|
||||
storage: &dyn StorageEngine,
|
||||
community_ledger: &crate::governance::CommunityLedger,
|
||||
co_engagement: &crate::entities::CoEngagementIndex,
|
||||
preference_vectors: &crate::entities::PreferenceVectors,
|
||||
) {
|
||||
if let Err(e) = community_ledger.checkpoint(storage) {
|
||||
tracing::error!(error = %e, "periodic community checkpoint failed");
|
||||
}
|
||||
if co_engagement.edge_count() > 0
|
||||
&& let Err(e) = co_engagement.checkpoint(storage)
|
||||
{
|
||||
tracing::error!(error = %e, "periodic co-engagement checkpoint failed");
|
||||
}
|
||||
if !preference_vectors.is_empty()
|
||||
&& let Err(e) = preference_vectors.checkpoint(storage)
|
||||
{
|
||||
tracing::error!(error = %e, "periodic preference-vector checkpoint failed");
|
||||
}
|
||||
}
|
||||
|
||||
/// Poll text-index liveness and schedule a resync rebuild when unhealthy.
|
||||
///
|
||||
/// Runs every 10s on the checkpoint thread regardless of the `metrics` feature
|
||||
@ -502,11 +625,7 @@ pub(super) fn run_checkpoint_thread(
|
||||
/// A rebuild that itself fails leaves the latch set (the next reopen or a manual
|
||||
/// resync recovers); a rebuild that succeeds clears the latch for the item side,
|
||||
/// and the syncer's own `is_healthy()` flag governs from then on.
|
||||
fn poll_text_index_health(
|
||||
handles: &IndexMetricsHandles,
|
||||
items_engine: &dyn StorageEngine,
|
||||
last_wal_seq: &AtomicU64,
|
||||
) {
|
||||
fn poll_text_index_health(handles: &IndexMetricsHandles, items_engine: &dyn StorageEngine) {
|
||||
let item_ok = handles
|
||||
.text_index
|
||||
.as_ref()
|
||||
@ -530,11 +649,7 @@ fn poll_text_index_health(
|
||||
|
||||
// Resync the item index from the durable store this thread owns.
|
||||
if !item_ok && let Some(idx) = handles.text_index.as_ref() {
|
||||
match rebuild_text_index_from_engine(
|
||||
idx,
|
||||
items_engine,
|
||||
last_wal_seq.load(Ordering::Acquire),
|
||||
) {
|
||||
match rebuild_text_index_from_engine(idx, items_engine) {
|
||||
Ok(indexed) => {
|
||||
tracing::info!(indexed, "item text index rebuilt; clearing unhealthy latch");
|
||||
// Clear only if the creator side is also healthy; otherwise leave
|
||||
@ -559,7 +674,6 @@ fn poll_text_index_health(
|
||||
fn rebuild_text_index_from_engine(
|
||||
idx: &crate::text::TextIndex,
|
||||
items_engine: &dyn StorageEngine,
|
||||
last_seq: u64,
|
||||
) -> crate::Result<usize> {
|
||||
use crate::storage::keys::parse_key;
|
||||
|
||||
@ -574,11 +688,160 @@ fn rebuild_text_index_from_engine(
|
||||
}
|
||||
}
|
||||
let count = items.len();
|
||||
idx.rebuild_from(items.into_iter(), last_seq)?;
|
||||
idx.rebuild_from(items.into_iter())?;
|
||||
idx.reload_reader()?;
|
||||
Ok(count)
|
||||
}
|
||||
|
||||
/// Unconditionally rebuild the item and creator Tantivy text indexes from the
|
||||
/// durable entity stores at open, and fold the suggestion-index rebuild into the
|
||||
/// same item scan.
|
||||
///
|
||||
/// # Why this exists (BLOCKER 7)
|
||||
///
|
||||
/// The text index is DERIVED state fed by a best-effort async syncer. Items are
|
||||
/// persisted durably to the entity store and *then* enqueued to the syncer; a
|
||||
/// hard crash within the syncer's batch window (default 2s) leaves those items
|
||||
/// in the durable store but never committed to Tantivy. The old design relied on
|
||||
/// a Tantivy commit-payload sequence number for recovery, but that machinery was
|
||||
/// never wired (every write enqueued `seq: 0`) and never read at open, so those
|
||||
/// items vanished from BM25 search permanently — silently, with health GREEN.
|
||||
///
|
||||
/// The fix mirrors the vector index, which is unconditionally rebuilt from the
|
||||
/// durable embeddings at open: here we unconditionally rebuild the text indexes
|
||||
/// from the durable entity stores (the source of truth) on every reopen. Any
|
||||
/// item/creator present in the store but missing from Tantivy is re-indexed.
|
||||
///
|
||||
/// # Efficiency
|
||||
///
|
||||
/// The item side performs a SINGLE scan of the items engine and drives BOTH the
|
||||
/// suggestion index and the item text index from it (the suggestion rebuild was
|
||||
/// previously its own separate scan). The creator side is a separate scan of the
|
||||
/// creators engine because creators live in their own store and are absent from
|
||||
/// every item-keyed pass. `EMB:` embedding rows (items) are skipped, matching
|
||||
/// the live write path's text exclusions.
|
||||
///
|
||||
/// Failures are non-fatal and logged: a failed text rebuild leaves search
|
||||
/// degraded for that index, but the durable store is intact and the
|
||||
/// health-driven resync (`poll_text_index_health`) remains a backstop.
|
||||
pub(super) fn rebuild_text_indexes_at_open(
|
||||
storage: &StorageBox,
|
||||
item_text_index: Option<&crate::text::TextIndex>,
|
||||
creator_text_index: Option<&crate::text::TextIndex>,
|
||||
suggestion_index: &SuggestionIndex,
|
||||
) {
|
||||
rebuild_item_text_and_suggestions_at_open(storage, item_text_index, suggestion_index);
|
||||
rebuild_creator_text_at_open(storage, creator_text_index);
|
||||
}
|
||||
|
||||
/// Single item scan that rebuilds the suggestion index and (if present) the item
|
||||
/// text index. See [`rebuild_text_indexes_at_open`] for rationale.
|
||||
fn rebuild_item_text_and_suggestions_at_open(
|
||||
storage: &StorageBox,
|
||||
item_text_index: Option<&crate::text::TextIndex>,
|
||||
suggestion_index: &SuggestionIndex,
|
||||
) {
|
||||
use crate::storage::keys::parse_key;
|
||||
|
||||
// One pass over every persisted item, feeding both derived indexes.
|
||||
let mut text_items: Vec<(EntityId, HashMap<String, String>)> = Vec::new();
|
||||
let mut suggestions = 0u64;
|
||||
for entry in storage.items_engine().scan_prefix(&[]) {
|
||||
let Ok((key, value)) = entry else { continue };
|
||||
let Some((entity_id, Tag::Meta, suffix)) = parse_key(&key) else {
|
||||
continue;
|
||||
};
|
||||
// Skip embedding rows: they share Tag::Meta but carry a serialized
|
||||
// vector, not searchable metadata.
|
||||
if suffix.starts_with(b"EMB:") {
|
||||
continue;
|
||||
}
|
||||
let meta = deserialize_metadata(&value);
|
||||
if let Some(title) = meta.get("title") {
|
||||
suggestion_index.index_title(title);
|
||||
suggestions += 1;
|
||||
}
|
||||
if item_text_index.is_some() {
|
||||
text_items.push((entity_id, meta));
|
||||
}
|
||||
}
|
||||
|
||||
if suggestions > 0 {
|
||||
tracing::info!(
|
||||
items = suggestions,
|
||||
"suggestion index rebuilt from durable storage"
|
||||
);
|
||||
}
|
||||
|
||||
if let Some(idx) = item_text_index {
|
||||
let count = text_items.len();
|
||||
// The health flag is freshly `true` on a just-opened index, so a
|
||||
// successful rebuild needs no explicit clear — search is in sync.
|
||||
match idx
|
||||
.rebuild_from(text_items.into_iter())
|
||||
.and_then(|()| idx.reload_reader())
|
||||
{
|
||||
Ok(()) => {
|
||||
if count > 0 {
|
||||
tracing::info!(
|
||||
items = count,
|
||||
"item text index rebuilt from durable storage at open"
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(e) => tracing::error!(
|
||||
error = %e,
|
||||
"item text index rebuild at open failed; BM25 search degraded \
|
||||
until the next health-driven resync"
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Scan the creators engine and rebuild the creator text index from it.
|
||||
///
|
||||
/// Creators are stored via `serialize_entity(None, metadata)` (like users), so
|
||||
/// each value is deserialized with [`deserialize_entity`](crate::entities::deserialize_entity)
|
||||
/// rather than the raw-metadata codec used for items.
|
||||
fn rebuild_creator_text_at_open(
|
||||
storage: &StorageBox,
|
||||
creator_text_index: Option<&crate::text::TextIndex>,
|
||||
) {
|
||||
use crate::storage::keys::parse_key;
|
||||
|
||||
let Some(idx) = creator_text_index else {
|
||||
return; // no creator text fields declared
|
||||
};
|
||||
|
||||
let mut creators: Vec<(EntityId, HashMap<String, String>)> = Vec::new();
|
||||
for entry in storage.creators_engine().scan_prefix(&[]) {
|
||||
let Ok((key, value)) = entry else { continue };
|
||||
if let Some((entity_id, Tag::Meta, _suffix)) = parse_key(&key) {
|
||||
let (_emb, meta) = crate::entities::deserialize_entity(&value);
|
||||
creators.push((entity_id, meta));
|
||||
}
|
||||
}
|
||||
|
||||
let count = creators.len();
|
||||
match idx
|
||||
.rebuild_from(creators.into_iter())
|
||||
.and_then(|()| idx.reload_reader())
|
||||
{
|
||||
Ok(()) => {
|
||||
if count > 0 {
|
||||
tracing::info!(
|
||||
creators = count,
|
||||
"creator text index rebuilt from durable storage at open"
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(e) => tracing::error!(
|
||||
error = %e,
|
||||
"creator text index rebuild at open failed; creator BM25 search degraded"
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
/// Refresh index health metrics from the live index handles.
|
||||
///
|
||||
/// Called once per checkpoint cycle (~30s). Reads current stats from the
|
||||
@ -620,18 +883,114 @@ fn refresh_index_metrics(handles: &IndexMetricsHandles, metrics: &MetricsState)
|
||||
.store(cardinality, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Sum the file sizes of all remaining WAL segment files in the directory.
|
||||
/// Sum the file sizes of WAL segments NOT yet compacted by the checkpoint at
|
||||
/// `checkpoint_seq` — i.e. the true "lag behind the checkpoint."
|
||||
///
|
||||
/// # Why not "all segments"
|
||||
///
|
||||
/// This previously summed EVERY segment, including ones already covered by the
|
||||
/// checkpoint (deletable on the next compaction) and the active segment, so the
|
||||
/// number it stored into `metrics.wal_lag_bytes` did not match that gauge's
|
||||
/// documentation ("Bytes of WAL segments not yet compacted"). A segment is
|
||||
/// "not yet compacted" exactly when [`compact_wal_online`](crate::wal::compaction::compact_wal_online)
|
||||
/// would retain it: that compactor deletes segments whose `first_seq < floor`,
|
||||
/// where `floor = min(checkpoint_seq, active_first_seq)` (the active segment is
|
||||
/// always protected because it still holds uncheckpointed tail events). We
|
||||
/// mirror that floor here and sum only the segments with `first_seq >= floor`,
|
||||
/// which makes the gauge faithful to its own HELP string.
|
||||
///
|
||||
/// Run immediately AFTER `compact_wal_online`, the covered segments are already
|
||||
/// gone, so in steady state this equals the on-disk retained bytes — but the
|
||||
/// floor computation keeps the gauge honest even if compaction was skipped (e.g.
|
||||
/// `seq == 0`) or fell behind.
|
||||
///
|
||||
/// Returns 0 if the directory cannot be read or contains no segments.
|
||||
/// Errors on individual file metadata reads are treated as 0 bytes
|
||||
/// (non-fatal: this is a best-effort monitoring metric).
|
||||
#[cfg(feature = "metrics")]
|
||||
fn compute_wal_lag_bytes(wal_dir: &std::path::Path) -> u64 {
|
||||
fn compute_wal_lag_bytes(wal_dir: &std::path::Path, checkpoint_seq: u64) -> u64 {
|
||||
let Ok(segments) = crate::wal::segment::list_segments(wal_dir) else {
|
||||
return 0;
|
||||
};
|
||||
|
||||
// The active segment (max `first_seq`) is never compacted away while the
|
||||
// writer is live; clamp the floor so it is always counted as lag. This is
|
||||
// the identical floor `compact_wal_online` uses, so "lag" == "what survived
|
||||
// compaction" by construction.
|
||||
let floor = match segments.iter().map(|(first_seq, _)| *first_seq).max() {
|
||||
Some(active_first_seq) => checkpoint_seq.min(active_first_seq),
|
||||
None => checkpoint_seq,
|
||||
};
|
||||
|
||||
segments
|
||||
.iter()
|
||||
.filter(|(first_seq, _)| *first_seq >= floor)
|
||||
.map(|(_, path)| std::fs::metadata(path).map(|m| m.len()).unwrap_or(0))
|
||||
.sum()
|
||||
}
|
||||
|
||||
#[cfg(all(test, feature = "metrics"))]
|
||||
#[allow(clippy::unwrap_used)]
|
||||
mod wal_lag_tests {
|
||||
use super::compute_wal_lag_bytes;
|
||||
use crate::{replication::ShardId, wal::segment::segment_filename};
|
||||
|
||||
/// Write a segment file of `len` bytes whose name encodes `first_seq`, and
|
||||
/// return its byte length so callers can build expected sums.
|
||||
fn write_segment(dir: &std::path::Path, first_seq: u64, len: usize) -> u64 {
|
||||
let path = dir.join(segment_filename(ShardId::SINGLE, first_seq));
|
||||
std::fs::write(&path, vec![0u8; len]).unwrap();
|
||||
len as u64
|
||||
}
|
||||
|
||||
/// The lag gauge must EXCLUDE segments fully covered by the checkpoint and
|
||||
/// INCLUDE both the straddling/active segment and any post-checkpoint
|
||||
/// segment — matching the gauge's "not yet compacted" documentation. The old
|
||||
/// sum-everything implementation over-counted the covered segments.
|
||||
#[test]
|
||||
fn lag_excludes_covered_segments_and_counts_the_active_tail() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
||||
// Three segments. Checkpoint at seq=200 means:
|
||||
// seg first_seq=1 -> covered (1 < floor) -> EXCLUDED
|
||||
// seg first_seq=200 -> active/straddling tail -> INCLUDED
|
||||
// (200 is the max first_seq, so floor = min(200,200) = 200)
|
||||
let _covered = write_segment(dir.path(), 1, 4096);
|
||||
let active = write_segment(dir.path(), 200, 1024);
|
||||
|
||||
let lag = compute_wal_lag_bytes(dir.path(), 200);
|
||||
assert_eq!(
|
||||
lag, active,
|
||||
"lag must count only the uncompacted active segment, not the covered one"
|
||||
);
|
||||
}
|
||||
|
||||
/// A second post-checkpoint segment is also lag; only the truly-covered
|
||||
/// segment is dropped from the sum.
|
||||
#[test]
|
||||
fn lag_counts_all_uncompacted_segments() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
||||
// floor = min(checkpoint_seq=50, active_first_seq=300) = 50.
|
||||
// first_seq=1 -> 1 < 50 -> EXCLUDED (covered)
|
||||
// first_seq=100 -> >= 50 -> INCLUDED
|
||||
// first_seq=300 -> active -> INCLUDED
|
||||
let _covered = write_segment(dir.path(), 1, 8192);
|
||||
let mid = write_segment(dir.path(), 100, 2048);
|
||||
let active = write_segment(dir.path(), 300, 512);
|
||||
|
||||
let lag = compute_wal_lag_bytes(dir.path(), 50);
|
||||
assert_eq!(lag, mid + active, "every uncompacted segment counts as lag");
|
||||
}
|
||||
|
||||
/// With no checkpoint yet (`checkpoint_seq == 0`) nothing is covered, so the
|
||||
/// whole on-disk WAL is lag — there is nothing the checkpoint has caught up to.
|
||||
#[test]
|
||||
fn lag_with_zero_checkpoint_counts_everything() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let a = write_segment(dir.path(), 1, 1000);
|
||||
let b = write_segment(dir.path(), 50, 2000);
|
||||
let lag = compute_wal_lag_bytes(dir.path(), 0);
|
||||
assert_eq!(lag, a + b, "no checkpoint => all WAL bytes are lag");
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,7 +1,20 @@
|
||||
//! Background text-index syncer infrastructure.
|
||||
//!
|
||||
//! [`TextSyncerBundle`] holds the handles for a running `TextIndexSyncer`
|
||||
//! thread; [`spawn_text_syncer`] creates one; [`shutdown_text_syncer`] joins it.
|
||||
//! Opening a text index and running its syncer is split into two phases so the
|
||||
//! crash-recovery rebuild-from-store can run while NO thread holds the Tantivy
|
||||
//! writer lock:
|
||||
//!
|
||||
//! 1. [`open_text_syncer`] opens the `TextIndex` and creates the channels, but
|
||||
//! does NOT spawn the syncer thread. The caller now holds the only handle to
|
||||
//! the index and the writer lock is free.
|
||||
//! 2. The caller runs `rebuild_from_store` (see `state_rebuild`) on the open
|
||||
//! index — it acquires and releases the writer lock for the rebuild.
|
||||
//! 3. [`TextSyncerPending::spawn`] starts the syncer thread, which acquires the
|
||||
//! writer lock for its entire lifetime (Tantivy single-writer design).
|
||||
//!
|
||||
//! Spawning the thread *before* the rebuild deadlocks: the syncer holds the
|
||||
//! writer lock forever, so `rebuild_from`'s `writer_guard()` blocks forever and
|
||||
//! the open hangs. [`shutdown_text_syncer`] joins the thread on shutdown.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
@ -16,24 +29,128 @@ pub struct TextSyncerBundle {
|
||||
pub thread: std::sync::Mutex<Option<std::thread::JoinHandle<crate::Result<()>>>>,
|
||||
}
|
||||
|
||||
/// Spawn a background text syncer thread for a text index.
|
||||
/// Phase 1 of text-syncer setup: an opened [`TextIndex`] plus its channels, with
|
||||
/// NO syncer thread yet running.
|
||||
///
|
||||
/// Creates a `TextIndex` (on-disk if `config.data_dir` is set, ephemeral otherwise),
|
||||
/// a crossbeam channel for pending writes, an optional flush channel for synchronous
|
||||
/// commit requests, and a syncer thread that commits batches.
|
||||
pub fn spawn_text_syncer(
|
||||
/// While this value is held, the Tantivy writer lock is free, so the caller can
|
||||
/// safely run the crash-recovery rebuild-from-store on [`index`](Self::index)
|
||||
/// before [`spawn`](Self::spawn)ing the thread that takes the lock for life.
|
||||
///
|
||||
/// An empty `text_fields` set (no text index) yields a [`TextSyncerPending`] with
|
||||
/// `index: None`; calling [`spawn`](Self::spawn) on it produces an empty
|
||||
/// [`TextSyncerBundle`] (no thread, no channels).
|
||||
pub struct TextSyncerPending {
|
||||
index: Option<Arc<TextIndex>>,
|
||||
/// `(write_rx, flush_tx, flush_rx, commit_n, commit_secs, thread_name)`, all
|
||||
/// captured at open time and consumed by [`spawn`](Self::spawn). `None` when
|
||||
/// there is no index.
|
||||
spawn_inputs: Option<SpawnInputs>,
|
||||
write_tx: Option<crossbeam::channel::Sender<PendingWrite>>,
|
||||
flush_tx: Option<crossbeam::channel::Sender<crossbeam::channel::Sender<()>>>,
|
||||
}
|
||||
|
||||
/// Inputs captured at open time and consumed when the syncer thread is spawned.
|
||||
struct SpawnInputs {
|
||||
write_rx: crossbeam::channel::Receiver<PendingWrite>,
|
||||
flush_rx: crossbeam::channel::Receiver<crossbeam::channel::Sender<()>>,
|
||||
commit_n: usize,
|
||||
commit_secs: u64,
|
||||
thread_name: String,
|
||||
}
|
||||
|
||||
impl TextSyncerPending {
|
||||
/// Borrow the opened index so the caller can run the rebuild-from-store while
|
||||
/// no syncer thread holds the writer lock. `None` when no text fields were
|
||||
/// declared (no index exists).
|
||||
#[must_use]
|
||||
pub const fn index(&self) -> Option<&Arc<TextIndex>> {
|
||||
self.index.as_ref()
|
||||
}
|
||||
|
||||
/// Phase 2: spawn the syncer thread for the opened index and return the
|
||||
/// running [`TextSyncerBundle`].
|
||||
///
|
||||
/// MUST be called only after any rebuild-from-store on [`index`](Self::index)
|
||||
/// has completed: the spawned syncer acquires the Tantivy writer lock for its
|
||||
/// entire lifetime, so a rebuild attempted afterward would deadlock.
|
||||
#[must_use]
|
||||
pub fn spawn(self) -> TextSyncerBundle {
|
||||
let (Some(index), Some(inputs)) = (self.index, self.spawn_inputs) else {
|
||||
// No text index — nothing to spawn.
|
||||
return TextSyncerBundle {
|
||||
index: None,
|
||||
write_tx: None,
|
||||
flush_tx: None,
|
||||
thread: std::sync::Mutex::new(None),
|
||||
};
|
||||
};
|
||||
|
||||
let idx_clone = Arc::clone(&index);
|
||||
let SpawnInputs {
|
||||
write_rx,
|
||||
flush_rx,
|
||||
commit_n,
|
||||
commit_secs,
|
||||
thread_name,
|
||||
} = inputs;
|
||||
|
||||
let handle = match std::thread::Builder::new()
|
||||
.name(thread_name)
|
||||
.spawn(move || {
|
||||
TextIndexSyncer::new(idx_clone, write_rx, commit_n, commit_secs)
|
||||
.with_flush_rx(flush_rx)
|
||||
.run()
|
||||
}) {
|
||||
Ok(h) => Some(h),
|
||||
Err(e) => {
|
||||
// Thread spawn failed (e.g. EAGAIN / resource limits). The bundle
|
||||
// still carries the open index and live channels, but with no
|
||||
// syncer thread draining them, queued writes never reach Tantivy.
|
||||
// `.ok()` discarded this silently; log it so the disabled syncer
|
||||
// is observable rather than a mystery stale index.
|
||||
tracing::error!(
|
||||
error = %e,
|
||||
"text syncer thread spawn failed; index writes will not be committed"
|
||||
);
|
||||
None
|
||||
}
|
||||
};
|
||||
|
||||
TextSyncerBundle {
|
||||
index: Some(index),
|
||||
write_tx: self.write_tx,
|
||||
flush_tx: self.flush_tx,
|
||||
thread: std::sync::Mutex::new(handle),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Phase 1 of background text-syncer setup: open the index and create channels,
|
||||
/// WITHOUT spawning the syncer thread.
|
||||
///
|
||||
/// Creates a `TextIndex` (on-disk if `config.data_dir` is set, ephemeral
|
||||
/// otherwise), a crossbeam channel for pending writes, and an optional flush
|
||||
/// channel for synchronous commit requests. The syncer thread is NOT started —
|
||||
/// call [`TextSyncerPending::spawn`] after running the crash-recovery rebuild on
|
||||
/// [`TextSyncerPending::index`].
|
||||
///
|
||||
/// On empty `text_fields` (no text index) or an index-open failure, returns a
|
||||
/// pending handle with `index: None` whose `spawn()` is an empty no-op bundle.
|
||||
pub fn open_text_syncer(
|
||||
text_fields: &[crate::schema::TextFieldDef],
|
||||
config: &Config,
|
||||
index_name: &str,
|
||||
thread_name: &str,
|
||||
) -> TextSyncerBundle {
|
||||
) -> TextSyncerPending {
|
||||
let empty = || TextSyncerPending {
|
||||
index: None,
|
||||
spawn_inputs: None,
|
||||
write_tx: None,
|
||||
flush_tx: None,
|
||||
};
|
||||
|
||||
if text_fields.is_empty() {
|
||||
return TextSyncerBundle {
|
||||
index: None,
|
||||
write_tx: None,
|
||||
flush_tx: None,
|
||||
thread: std::sync::Mutex::new(None),
|
||||
};
|
||||
return empty();
|
||||
}
|
||||
|
||||
let text_config = config
|
||||
@ -51,37 +168,35 @@ pub fn spawn_text_syncer(
|
||||
|
||||
let index = match index {
|
||||
Ok(idx) => Arc::new(idx),
|
||||
Err(_) => {
|
||||
return TextSyncerBundle {
|
||||
index: None,
|
||||
write_tx: None,
|
||||
flush_tx: None,
|
||||
thread: std::sync::Mutex::new(None),
|
||||
};
|
||||
Err(e) => {
|
||||
// A corrupt Tantivy directory, a permission error, or a full disk all
|
||||
// surface here. Swallowing the error returned an empty bundle
|
||||
// (`index: None`) that is INDISTINGUISHABLE from "no text fields
|
||||
// declared", so a hard text-index failure silently disabled BM25
|
||||
// search with no operator signal. Log it loudly before degrading.
|
||||
tracing::error!(
|
||||
index = index_name,
|
||||
error = %e,
|
||||
"text index open failed; text search disabled for this index"
|
||||
);
|
||||
return empty();
|
||||
}
|
||||
};
|
||||
|
||||
let (tx, rx) = crossbeam::channel::unbounded();
|
||||
let (flush_tx, flush_rx) = crossbeam::channel::bounded::<crossbeam::channel::Sender<()>>(4);
|
||||
let idx_clone = Arc::clone(&index);
|
||||
let commit_n = text_config.commit_every_n_docs;
|
||||
let commit_secs = text_config.commit_every_secs;
|
||||
let tname = thread_name.to_owned();
|
||||
|
||||
let handle = std::thread::Builder::new()
|
||||
.name(tname)
|
||||
.spawn(move || {
|
||||
TextIndexSyncer::new(idx_clone, rx, commit_n, commit_secs)
|
||||
.with_flush_rx(flush_rx)
|
||||
.run()
|
||||
})
|
||||
.ok();
|
||||
|
||||
TextSyncerBundle {
|
||||
TextSyncerPending {
|
||||
index: Some(index),
|
||||
spawn_inputs: Some(SpawnInputs {
|
||||
write_rx: rx,
|
||||
flush_rx,
|
||||
commit_n: text_config.commit_every_n_docs,
|
||||
commit_secs: text_config.commit_every_secs,
|
||||
thread_name: thread_name.to_owned(),
|
||||
}),
|
||||
write_tx: Some(tx),
|
||||
flush_tx: Some(flush_tx),
|
||||
thread: std::sync::Mutex::new(handle),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -10,8 +10,14 @@ use std::collections::HashMap;
|
||||
|
||||
use super::TidalDb;
|
||||
use crate::{
|
||||
schema::{EntityId, TidalError},
|
||||
storage::{StorageEngine, Tag, encode_key},
|
||||
schema::{EntityId, EntityKind, TidalError},
|
||||
storage::{
|
||||
StorageEngine, Tag, encode_key,
|
||||
vector::{
|
||||
BruteForceIndex, QuantizationLevel, VectorIndexConfig, insert_embedding,
|
||||
registry::{EmbeddingSlotState, EmbeddingSource, HnswParams},
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
// ── Shared entity-metadata persistence helpers ───────────────────────────────
|
||||
@ -91,6 +97,99 @@ impl TidalDb {
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Auto-register an embedding slot (if absent), then L2-normalize the
|
||||
/// embedding, store it in `engine` (the entity-kind's source of truth), and
|
||||
/// insert it into that slot's vector index.
|
||||
///
|
||||
/// This is the shared body of [`TidalDb::write_item_embedding`] and
|
||||
/// [`TidalDb::write_creator_embedding`], which differ only in their target
|
||||
/// `kind`, `slot_name`, storage `engine`, and the `op` label used for error
|
||||
/// diagnostics. Centralizing it keeps the registry lock discipline defined
|
||||
/// in exactly one place: the write lock is taken only for the one-shot slot
|
||||
/// registration and dropped before the read phase; a read lock is then taken
|
||||
/// to resolve the slot dimensions, dropped, and re-taken across the
|
||||
/// `insert_embedding` call (which must hold it for the duration of the index
|
||||
/// mutation). Behavior is identical to the prior per-method bodies.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// - `TidalError::Internal` if the embedding registry lock is poisoned, the
|
||||
/// slot cannot be registered, the slot is missing after registration, or
|
||||
/// the embedding has zero norm.
|
||||
/// - `TidalError::Storage` on storage engine failure.
|
||||
#[allow(clippy::significant_drop_tightening)] // lock must be held across insert_embedding call
|
||||
pub(super) fn write_entity_embedding(
|
||||
&self,
|
||||
id: EntityId,
|
||||
kind: EntityKind,
|
||||
slot_name: &str,
|
||||
embedding: &[f32],
|
||||
engine: &dyn StorageEngine,
|
||||
op: &'static str,
|
||||
) -> crate::Result<()> {
|
||||
// Auto-register the slot if absent. The write lock is scoped to this
|
||||
// block and dropped before the read phase below.
|
||||
{
|
||||
let mut registry = self
|
||||
.embedding_registry
|
||||
.write()
|
||||
.map_err(|_| TidalError::internal(op, "embedding_registry write lock poisoned"))?;
|
||||
if registry.get(kind, slot_name).is_none() {
|
||||
let config = VectorIndexConfig {
|
||||
dimensions: embedding.len(),
|
||||
..VectorIndexConfig::default()
|
||||
};
|
||||
let state = EmbeddingSlotState {
|
||||
index: Box::new(BruteForceIndex::new(config)),
|
||||
dimensions: embedding.len(),
|
||||
quantization: QuantizationLevel::F32,
|
||||
source: EmbeddingSource::External,
|
||||
params: HnswParams::default(),
|
||||
};
|
||||
registry
|
||||
.register(kind, slot_name.to_string(), state)
|
||||
.map_err(|e| {
|
||||
TidalError::internal(op, format!("slot registration failed: {e}"))
|
||||
})?;
|
||||
}
|
||||
}
|
||||
|
||||
// Read the slot dimensions, dropping the lock before the storage write.
|
||||
let dimensions = {
|
||||
let registry = self
|
||||
.embedding_registry
|
||||
.read()
|
||||
.map_err(|_| TidalError::internal(op, "embedding_registry read lock poisoned"))?;
|
||||
registry
|
||||
.get(kind, slot_name)
|
||||
.ok_or_else(|| TidalError::internal(op, "embedding slot missing"))?
|
||||
.dimensions
|
||||
};
|
||||
|
||||
// Re-acquire the read lock to get the index reference, perform the
|
||||
// insert, then drop the lock.
|
||||
{
|
||||
let registry = self
|
||||
.embedding_registry
|
||||
.read()
|
||||
.map_err(|_| TidalError::internal(op, "embedding_registry read lock poisoned"))?;
|
||||
let slot = registry
|
||||
.get(kind, slot_name)
|
||||
.ok_or_else(|| TidalError::internal(op, "embedding slot missing"))?;
|
||||
insert_embedding(
|
||||
id,
|
||||
slot_name,
|
||||
embedding,
|
||||
dimensions,
|
||||
slot.index.as_ref(),
|
||||
engine,
|
||||
)
|
||||
.map_err(|e| TidalError::internal(op, format!("{e}")))?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl TidalDb {
|
||||
@ -149,7 +248,16 @@ impl TidalDb {
|
||||
let mut out = Vec::new();
|
||||
for entry in storage.users_engine().scan_prefix(&[]) {
|
||||
let (key, value) = entry.map_err(TidalError::from)?;
|
||||
if let Some((entity_id, Tag::Meta, _suffix)) = parse_key(&key) {
|
||||
if let Some((entity_id, Tag::Meta, suffix)) = parse_key(&key) {
|
||||
// Skip embedding rows: they share Tag::Meta but their value is a
|
||||
// serialized vector (persisted under the `EMB:{slot}` suffix by
|
||||
// `embedding_store_key`), not a serialized entity. Deserializing
|
||||
// those as an entity yields garbage that must never reach callers
|
||||
// (embedder key-mapping reconstruction). Mirrors the same skip in
|
||||
// `all_item_metadata`.
|
||||
if suffix.starts_with(b"EMB:") {
|
||||
continue;
|
||||
}
|
||||
let (_emb, meta) = crate::entities::deserialize_entity(&value);
|
||||
out.push((entity_id, meta));
|
||||
}
|
||||
@ -157,3 +265,70 @@ impl TidalDb {
|
||||
Ok(out)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[allow(clippy::unwrap_used)]
|
||||
mod tests {
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::{
|
||||
TidalDb,
|
||||
schema::EntityId,
|
||||
storage::vector::{embedding_store_key, serialize_embedding},
|
||||
};
|
||||
|
||||
/// Minimal valid schema so ephemeral mode wires in-memory storage.
|
||||
fn minimal_schema() -> crate::schema::Schema {
|
||||
use crate::schema::{DecaySpec, EntityKind, SchemaBuilder, Window};
|
||||
let mut b = SchemaBuilder::new();
|
||||
let _ = b
|
||||
.signal(
|
||||
"view",
|
||||
EntityKind::Item,
|
||||
DecaySpec::Exponential {
|
||||
half_life: std::time::Duration::from_secs(3600),
|
||||
},
|
||||
)
|
||||
.windows(&[Window::AllTime])
|
||||
.velocity(false)
|
||||
.add();
|
||||
b.build().expect("schema must be valid")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn all_user_metadata_skips_embedding_rows() {
|
||||
let db = TidalDb::builder()
|
||||
.ephemeral()
|
||||
.with_schema(minimal_schema())
|
||||
.open()
|
||||
.unwrap();
|
||||
|
||||
// A real user with plain metadata.
|
||||
let mut meta = HashMap::new();
|
||||
meta.insert("name".to_string(), "alice".to_string());
|
||||
db.write_user(EntityId::new(1), &meta).unwrap();
|
||||
|
||||
// An embedding row written under the SAME Tag::Meta key space the user
|
||||
// embedding durability path uses (`EMB:{slot}` suffix). Its value is a
|
||||
// serialized vector, not a serialized entity — `all_user_metadata` must
|
||||
// not surface it as user metadata.
|
||||
let emb_key = embedding_store_key(EntityId::new(1), "content");
|
||||
let emb_value = serialize_embedding(&[0.1, 0.2, 0.3]);
|
||||
db.storage()
|
||||
.unwrap()
|
||||
.users_engine()
|
||||
.put(&emb_key, &emb_value)
|
||||
.unwrap();
|
||||
|
||||
let all = db.all_user_metadata().unwrap();
|
||||
assert_eq!(
|
||||
all.len(),
|
||||
1,
|
||||
"embedding row must be skipped; expected only the plain user row, got {all:?}"
|
||||
);
|
||||
assert_eq!(all[0].0, EntityId::new(1));
|
||||
assert_eq!(all[0].1.get("name").map(String::as_str), Some("alice"));
|
||||
|
||||
db.close().unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,23 +1,38 @@
|
||||
//! Hard-negative tracking: per-user explicit rejections.
|
||||
//!
|
||||
//! Tracks items a user has explicitly rejected (skip, hide, dislike).
|
||||
//! Tracks items a user has explicitly rejected (skip, hide, dislike, block).
|
||||
//! Hard negatives are used to exclude items from personalized ranking
|
||||
//! and to downweight similar items in the preference vector.
|
||||
//!
|
||||
//! The index is in-memory (rebuilt from signals on startup) and provides
|
||||
//! O(1) lookup via `RoaringBitmap`.
|
||||
//! # Durability
|
||||
//!
|
||||
//! The in-memory index is a `DashMap<u64, RoaringBitmap>` for O(1) lookup, but
|
||||
//! it is NOT the source of truth: each `(user, item)` rejection is also written
|
||||
//! write-through as a durable `Tag::HardNeg` row by the signal path
|
||||
//! ([`TidalDb::signal_with_context`](crate::TidalDb::signal_with_context)), with
|
||||
//! the same durability class as the base signal (zero loss window). On open,
|
||||
//! [`HardNegIndex::restore`] scans those rows and repopulates the bitmaps. The
|
||||
//! WAL signal record carries no per-user identity, so these durable rows — not
|
||||
//! WAL replay — are what makes a `skip`/`dislike` survive a hard crash.
|
||||
|
||||
use dashmap::DashMap;
|
||||
use roaring::RoaringBitmap;
|
||||
|
||||
use crate::{
|
||||
schema::EntityId,
|
||||
storage::{StorageEngine, Tag, encode_key, entity_tag_prefix, parse_key},
|
||||
};
|
||||
|
||||
/// Signal types that constitute a hard negative.
|
||||
pub const HARD_NEG_SIGNALS: &[&str] = &["skip", "hide", "dislike", "block"];
|
||||
|
||||
/// Per-user hard-negative bitmap index.
|
||||
///
|
||||
/// Each user has a `RoaringBitmap` of item IDs they have explicitly
|
||||
/// rejected. The bitmap is populated from signal events during startup
|
||||
/// and updated in real-time as new signals arrive.
|
||||
/// rejected. The bitmap is repopulated on open by [`HardNegIndex::restore`]
|
||||
/// from the durable `Tag::HardNeg` rows the signal path writes write-through,
|
||||
/// and updated in real-time as new signals arrive (the in-memory `add` is paired
|
||||
/// with a synchronous durable write at the call site).
|
||||
///
|
||||
/// Thread-safe via `DashMap` -- concurrent reads and writes to different
|
||||
/// users never contend.
|
||||
@ -94,6 +109,103 @@ impl HardNegIndex {
|
||||
pub fn is_hard_neg_signal(signal_type: &str) -> bool {
|
||||
HARD_NEG_SIGNALS.contains(&signal_type)
|
||||
}
|
||||
|
||||
/// Collect every `(user_id, item_id)` hard-negative pair currently held in
|
||||
/// memory.
|
||||
///
|
||||
/// Used by `TidalDb::take_crdt_snapshot` to materialize the local
|
||||
/// hard-negative membership into per-pair LWW registers for reconciliation.
|
||||
/// This index stores only resolved membership (no per-item HLC); the caller
|
||||
/// stamps each pair with a timestamp when building the snapshot. The returned
|
||||
/// vector is a point-in-time copy, so it does not hold any `DashMap` shard
|
||||
/// lock.
|
||||
#[must_use]
|
||||
pub(crate) fn all_pairs(&self) -> Vec<(u64, u32)> {
|
||||
let mut pairs = Vec::new();
|
||||
for entry in &self.inner {
|
||||
let user_id = *entry.key();
|
||||
for item_id in entry.value() {
|
||||
pairs.push((user_id, item_id));
|
||||
}
|
||||
}
|
||||
pairs
|
||||
}
|
||||
}
|
||||
|
||||
// ── Durable rows ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// The durable storage key for one `(user, item)` hard-negative row.
|
||||
///
|
||||
/// All hard-neg rows share the sentinel entity id (`0`) so [`HardNegIndex::restore`]
|
||||
/// can target them with a single prefix scan rather than a full table scan.
|
||||
///
|
||||
/// Key format: `encode_key(EntityId(0), Tag::HardNeg, [user: 8B BE, item: 4B BE])`.
|
||||
/// The value is empty — presence of the key is the fact. The item is stored as
|
||||
/// `u32` because the in-memory `RoaringBitmap` is u32-keyed; this matches the
|
||||
/// truncation the signal path already applies at `entity_id.as_u64() as u32`.
|
||||
#[must_use]
|
||||
pub fn hard_neg_key(user_id: u64, item_id: u32) -> Vec<u8> {
|
||||
let mut suffix = [0u8; 12];
|
||||
suffix[..8].copy_from_slice(&user_id.to_be_bytes());
|
||||
suffix[8..].copy_from_slice(&item_id.to_be_bytes());
|
||||
encode_key(EntityId::new(0), Tag::HardNeg, &suffix)
|
||||
}
|
||||
|
||||
impl HardNegIndex {
|
||||
/// Persist a `(user, item)` hard-negative as a durable `Tag::HardNeg` row.
|
||||
///
|
||||
/// Called write-through by the signal path so a `skip`/`dislike` survives a
|
||||
/// hard crash with zero loss window. The in-memory [`add`](Self::add) is the
|
||||
/// paired hot-path mutation; this is its durable backing.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns storage errors from the underlying engine.
|
||||
pub fn persist(
|
||||
&self,
|
||||
storage: &dyn StorageEngine,
|
||||
user_id: u64,
|
||||
item_id: u32,
|
||||
) -> crate::Result<()> {
|
||||
storage
|
||||
.put(&hard_neg_key(user_id, item_id), &[])
|
||||
.map_err(crate::schema::TidalError::from)
|
||||
}
|
||||
|
||||
/// Restore the in-memory bitmaps from durable `Tag::HardNeg` rows on open.
|
||||
///
|
||||
/// Uses a targeted prefix scan over the sentinel entity (`0`), so restore is
|
||||
/// `O(hard_neg_rows)` rather than `O(total_keys)`.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns storage errors from the underlying engine.
|
||||
pub fn restore(&self, storage: &dyn StorageEngine) -> crate::Result<()> {
|
||||
let prefix = entity_tag_prefix(EntityId::new(0), Tag::HardNeg);
|
||||
let mut count = 0u64;
|
||||
for entry in storage.scan_prefix(&prefix) {
|
||||
let (key, _) = entry.map_err(crate::schema::TidalError::from)?;
|
||||
if let Some((_, Tag::HardNeg, suffix)) = parse_key(&key) {
|
||||
if suffix.len() < 12 {
|
||||
continue;
|
||||
}
|
||||
let user_id = u64::from_be_bytes([
|
||||
suffix[0], suffix[1], suffix[2], suffix[3], suffix[4], suffix[5], suffix[6],
|
||||
suffix[7],
|
||||
]);
|
||||
let item_id = u32::from_be_bytes([suffix[8], suffix[9], suffix[10], suffix[11]]);
|
||||
self.add(user_id, item_id);
|
||||
count += 1;
|
||||
}
|
||||
}
|
||||
if count > 0 {
|
||||
tracing::info!(
|
||||
rows = count,
|
||||
"hard-negative index restored from durable rows"
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for HardNegIndex {
|
||||
|
||||
@ -134,6 +134,22 @@ impl InteractionLedger {
|
||||
}
|
||||
}
|
||||
|
||||
/// The raw stored running score and last-update timestamp for a
|
||||
/// `(user, creator)` pair, with **no** decay applied. Returns `None` if no
|
||||
/// interaction has been recorded.
|
||||
///
|
||||
/// This is the durability projection: persisting `(score, last_update_ns)`
|
||||
/// into a `RelationshipType::InteractionWeight` edge and replaying it through
|
||||
/// [`record`](Self::record) into a *fresh* ledger reproduces this entry
|
||||
/// exactly (the empty ledger folds the stored score as a single in-order
|
||||
/// event at `last_update_ns`, yielding `score` with timestamp `last_update_ns`).
|
||||
#[must_use]
|
||||
pub fn raw_entry(&self, user_id: u64, creator_id: u64) -> Option<(f64, u64)> {
|
||||
let user_map = self.inner.get(&user_id)?;
|
||||
let entry = user_map.get(&creator_id)?;
|
||||
Some((entry.score, entry.last_update_ns))
|
||||
}
|
||||
|
||||
/// Read the current interaction weight for a (user, creator) pair.
|
||||
///
|
||||
/// Applies lazy decay to the current time. Returns 0.0 if no interaction
|
||||
@ -239,6 +255,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
|
||||
fn out_of_order_event_is_pre_decayed_not_full_weight() {
|
||||
// Regression: the old inline form clamped dt to 0 for out-of-order
|
||||
// events, adding a late event's weight at FULL value and silently
|
||||
|
||||
@ -3,6 +3,15 @@
|
||||
//! Tracks user taste by maintaining a preference vector that evolves with
|
||||
//! interactions. The vector is L2-normalized on every update to ensure
|
||||
//! consistent cosine similarity scoring during personalized ranking.
|
||||
//!
|
||||
//! # Durability
|
||||
//!
|
||||
//! Preference vectors are DERIVED state with no WAL backing (the signal record
|
||||
//! carries no `for_user`), so they are persisted by [`PreferenceVectors::checkpoint`]
|
||||
//! (called on the periodic checkpoint thread, on backup, and on shutdown) and
|
||||
//! restored on open by [`PreferenceVectors::restore`]. This bounds the post-crash
|
||||
//! loss window to the checkpoint interval rather than losing every update since
|
||||
//! the last clean shutdown.
|
||||
|
||||
use dashmap::DashMap;
|
||||
|
||||
@ -98,8 +107,6 @@ impl PreferenceVectors {
|
||||
/// Returns `false` if the interaction embedding dimension does not match.
|
||||
#[must_use]
|
||||
pub fn update(&self, user_id: u64, interaction_embedding: &[f32]) -> bool {
|
||||
use dashmap::mapref::entry::Entry;
|
||||
|
||||
if interaction_embedding.len() != self.dim {
|
||||
return false;
|
||||
}
|
||||
@ -114,20 +121,7 @@ impl PreferenceVectors {
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
let lr = (f64::from(self.learning_rate) / (1.0 + (count as f64).ln_1p())) as f32;
|
||||
|
||||
match self.inner.entry(user_id) {
|
||||
Entry::Occupied(mut occ) => {
|
||||
let pref = occ.get_mut();
|
||||
for (p, &i) in pref.iter_mut().zip(interaction_embedding.iter()) {
|
||||
*p = (1.0 - lr).mul_add(*p, lr * i);
|
||||
}
|
||||
l2_normalize(pref);
|
||||
}
|
||||
Entry::Vacant(vac) => {
|
||||
let mut v = interaction_embedding.to_vec();
|
||||
l2_normalize(&mut v);
|
||||
vac.insert(v);
|
||||
}
|
||||
}
|
||||
self.blend_and_normalize(user_id, interaction_embedding, lr);
|
||||
true
|
||||
}
|
||||
|
||||
@ -145,12 +139,32 @@ impl PreferenceVectors {
|
||||
interaction_embedding: &[f32],
|
||||
lr: f32,
|
||||
) -> bool {
|
||||
use dashmap::mapref::entry::Entry;
|
||||
|
||||
if interaction_embedding.len() != self.dim {
|
||||
return false;
|
||||
}
|
||||
|
||||
self.blend_and_normalize(user_id, interaction_embedding, lr);
|
||||
true
|
||||
}
|
||||
|
||||
/// Blend `interaction_embedding` into the user's preference at learning
|
||||
/// rate `lr` and re-normalize, the shared EMA body for both [`update`] and
|
||||
/// [`update_with_custom_rate`].
|
||||
///
|
||||
/// An existing preference is blended in place via
|
||||
/// `pref = (1 - lr) * pref + lr * interaction` and re-normalized; if no
|
||||
/// preference exists yet, the (normalized) interaction embedding becomes the
|
||||
/// initial preference, avoiding a double-applied blend on first insertion.
|
||||
///
|
||||
/// The caller is responsible for the dimension check — `dim` mismatch is
|
||||
/// validated before this is reached, and the per-element `zip` would
|
||||
/// otherwise silently blend only the overlapping prefix.
|
||||
///
|
||||
/// [`update`]: Self::update
|
||||
/// [`update_with_custom_rate`]: Self::update_with_custom_rate
|
||||
fn blend_and_normalize(&self, user_id: u64, interaction_embedding: &[f32], lr: f32) {
|
||||
use dashmap::mapref::entry::Entry;
|
||||
|
||||
match self.inner.entry(user_id) {
|
||||
Entry::Occupied(mut occ) => {
|
||||
let pref = occ.get_mut();
|
||||
@ -165,7 +179,6 @@ impl PreferenceVectors {
|
||||
vac.insert(v);
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
/// Compute cosine similarity between a user's preference and a candidate embedding.
|
||||
@ -230,6 +243,111 @@ impl PreferenceVectors {
|
||||
}
|
||||
}
|
||||
|
||||
// ── Checkpoint / Restore ─────────────────────────────────────────────────────
|
||||
|
||||
impl PreferenceVectors {
|
||||
/// Checkpoint every per-user preference vector + update count to durable
|
||||
/// storage under `Tag::Preference`, one row per user.
|
||||
///
|
||||
/// All rows share the sentinel entity id (`0`) so [`restore`](Self::restore)
|
||||
/// can target them with a single prefix scan. The swap is staged into one
|
||||
/// `WriteBatch` and committed atomically (matching the co-engagement /
|
||||
/// community checkpoints): stale rows are deleted and current rows written in
|
||||
/// the same batch, so a crash mid-checkpoint never leaves a half-written set.
|
||||
///
|
||||
/// Key suffix: `[user: 8B BE]`. Value: `[update_count: 8B LE][dim: 4B LE][f32 * dim LE]`.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns storage errors from the underlying engine.
|
||||
pub fn checkpoint(&self, storage: &dyn crate::storage::StorageEngine) -> crate::Result<()> {
|
||||
use crate::{
|
||||
schema::EntityId,
|
||||
storage::{Tag, WriteBatch, encode_key, entity_tag_prefix},
|
||||
};
|
||||
|
||||
let prefix = entity_tag_prefix(EntityId::new(0), Tag::Preference);
|
||||
let mut batch = WriteBatch::with_capacity(self.inner.len() + 1);
|
||||
|
||||
// Stage deletion of any stale rows (e.g. a user whose vector was cleared)
|
||||
// so the post-swap snapshot holds exactly the current set.
|
||||
for item in storage.scan_prefix(&prefix) {
|
||||
let (key, _) = item.map_err(crate::schema::TidalError::from)?;
|
||||
batch.delete(key);
|
||||
}
|
||||
|
||||
for entry in &self.inner {
|
||||
let user_id = *entry.key();
|
||||
let vec = entry.value();
|
||||
let count = self.update_count(user_id);
|
||||
let mut value = Vec::with_capacity(8 + 4 + vec.len() * 4);
|
||||
value.extend_from_slice(&count.to_le_bytes());
|
||||
#[allow(clippy::cast_possible_truncation)]
|
||||
value.extend_from_slice(&(vec.len() as u32).to_le_bytes());
|
||||
for v in &**vec {
|
||||
value.extend_from_slice(&v.to_le_bytes());
|
||||
}
|
||||
let key = encode_key(EntityId::new(0), Tag::Preference, &user_id.to_be_bytes());
|
||||
batch.put(key, value);
|
||||
}
|
||||
storage
|
||||
.write_batch(batch)
|
||||
.map_err(crate::schema::TidalError::from)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Restore per-user preference vectors + update counts from a
|
||||
/// `Tag::Preference` checkpoint. Skips rows whose stored dimension does not
|
||||
/// match this store's `dim` (a schema change) and rows that are torn.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns storage errors from the underlying engine.
|
||||
pub fn restore(&self, storage: &dyn crate::storage::StorageEngine) -> crate::Result<()> {
|
||||
use crate::{
|
||||
schema::EntityId,
|
||||
storage::{Tag, entity_tag_prefix, parse_key},
|
||||
};
|
||||
|
||||
let prefix = entity_tag_prefix(EntityId::new(0), Tag::Preference);
|
||||
let mut count = 0u64;
|
||||
for entry in storage.scan_prefix(&prefix) {
|
||||
let (key, value) = entry.map_err(crate::schema::TidalError::from)?;
|
||||
let Some((_, Tag::Preference, suffix)) = parse_key(&key) else {
|
||||
continue;
|
||||
};
|
||||
if suffix.len() < 8 || value.len() < 12 {
|
||||
continue;
|
||||
}
|
||||
let user_id = u64::from_be_bytes(suffix[0..8].try_into().unwrap_or([0u8; 8]));
|
||||
let update_count = u64::from_le_bytes(value[0..8].try_into().unwrap_or([0u8; 8]));
|
||||
let dim = u32::from_le_bytes(value[8..12].try_into().unwrap_or([0u8; 4])) as usize;
|
||||
if dim != self.dim || value.len() < 12 + dim * 4 {
|
||||
continue;
|
||||
}
|
||||
let mut vec = Vec::with_capacity(dim);
|
||||
for i in 0..dim {
|
||||
let off = 12 + i * 4;
|
||||
let f = f32::from_le_bytes(value[off..off + 4].try_into().unwrap_or([0u8; 4]));
|
||||
// Neutralize NaN at the load boundary so a torn row cannot poison
|
||||
// cosine scoring downstream.
|
||||
vec.push(if f.is_nan() { 0.0 } else { f });
|
||||
}
|
||||
// The stored vector is already L2-normalized; insert directly (set()
|
||||
// re-normalizes, which is a no-op for a unit vector) and seed the
|
||||
// adaptive update count so the learning rate resumes where it was.
|
||||
self.inner.insert(user_id, vec);
|
||||
self.update_counts.insert(user_id, update_count);
|
||||
let _ = EntityId::new(user_id);
|
||||
count += 1;
|
||||
}
|
||||
if count > 0 {
|
||||
tracing::info!(users = count, "preference vectors restored from checkpoint");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// L2-normalize a vector in-place. If the vector has zero magnitude, it remains
|
||||
/// as-is (all zeros).
|
||||
fn l2_normalize(vec: &mut [f32]) {
|
||||
|
||||
@ -12,7 +12,24 @@ pub enum RelationshipType {
|
||||
Blocks = 0x02,
|
||||
InteractionWeight = 0x03,
|
||||
Hide = 0x04,
|
||||
/// PERSISTED-BUT-INERT (tracked): a `Mute` edge is durably stored and
|
||||
/// round-trips through [`as_byte`](Self::as_byte) / [`from_byte`](Self::from_byte),
|
||||
/// and the rebuild path counts it, but it has **no effect** on ranking yet —
|
||||
/// there is no muted-set in `UserStateIndex` and no predicate consults it.
|
||||
///
|
||||
/// Implementing it fully requires a muted-set on `UserStateIndex` plus
|
||||
/// application in the unseen/unblocked predicates, and wiring in
|
||||
/// `db/relationships.rs` (`write_relationship` / `delete_relationship`) and
|
||||
/// `db/state_rebuild.rs` (`rebuild_entity_state`). Until then the no-op is
|
||||
/// **deliberate and greppable**: the catch-all arms in `db/relationships.rs`
|
||||
/// and the counter-only arm in `db/state_rebuild.rs` should be (or are)
|
||||
/// explicit `RelationshipType::Mute => { /* tracked: not yet applied */ }`
|
||||
/// arms rather than absorbed by a bare `_`.
|
||||
///
|
||||
/// Tracking: implement muted-set application (M7 production hardening).
|
||||
Mute = 0x05,
|
||||
// `from_byte` returns `None` for any other byte, so adding a variant here
|
||||
// forces every match in this module to be revisited (no silent default).
|
||||
}
|
||||
|
||||
impl RelationshipType {
|
||||
|
||||
@ -12,7 +12,7 @@
|
||||
use dashmap::DashMap;
|
||||
|
||||
use crate::{
|
||||
schema::{EntityId, Window},
|
||||
schema::{EntityId, Timestamp, Window},
|
||||
signals::{SignalTypeId, warm::BucketedCounter},
|
||||
};
|
||||
|
||||
@ -39,7 +39,13 @@ impl UserSignalIndex {
|
||||
.increment(ts_ns);
|
||||
}
|
||||
|
||||
/// Get the windowed count for a `(user, entity, signal)` triplet.
|
||||
/// Get the windowed count for a `(user, entity, signal)` triplet, as of the
|
||||
/// current wall-clock time.
|
||||
///
|
||||
/// Reads the clock here and threads it into the warm tier so a user who went
|
||||
/// quiet has their windowed/velocity contribution age out at read time
|
||||
/// instead of freezing at its peak — social-graph trending depends on this
|
||||
/// (CRITICAL #6, `CODING_GUIDELINES` §8).
|
||||
#[must_use]
|
||||
pub fn windowed_count(
|
||||
&self,
|
||||
@ -48,9 +54,10 @@ impl UserSignalIndex {
|
||||
type_id: SignalTypeId,
|
||||
window: Window,
|
||||
) -> u64 {
|
||||
let now_ns = Timestamp::now().as_nanos();
|
||||
self.entries
|
||||
.get(&(user_id, entity_id.as_u64(), type_id.as_u16()))
|
||||
.map_or(0, |counter| counter.windowed_count(window))
|
||||
.map_or(0, |counter| counter.windowed_count(window, now_ns))
|
||||
}
|
||||
|
||||
/// Get velocity for a `(user, entity, signal)` triplet.
|
||||
@ -124,7 +131,9 @@ mod tests {
|
||||
let index = UserSignalIndex::new();
|
||||
let type_id = SignalTypeId::new(0);
|
||||
let entity = EntityId::new(42);
|
||||
let ts = 1_000_000_000_000_000_000_u64; // 1e18 ns
|
||||
// Base at the real wall clock: `windowed_count` reads `Timestamp::now()`
|
||||
// at read time, so events stamped near now stay inside the 1h window.
|
||||
let ts = Timestamp::now().as_nanos();
|
||||
|
||||
index.record(1, entity, type_id, ts);
|
||||
index.record(1, entity, type_id, ts + 1_000_000_000); // 1 second later
|
||||
@ -141,9 +150,11 @@ mod tests {
|
||||
let index = UserSignalIndex::new();
|
||||
let type_id = SignalTypeId::new(0);
|
||||
let entity = EntityId::new(42);
|
||||
let ts = 1_000_000_000_000_000_000u64;
|
||||
let ts = Timestamp::now().as_nanos();
|
||||
|
||||
// Record 3600 events in 1 second interval -- they all land in the same minute bucket.
|
||||
// Record 3600 events one second apart, spanning the full hour. The read
|
||||
// clock (~ts) precedes the last write's rotation timestamp, so no buckets
|
||||
// age out and all 3600 events stay in the 1h window.
|
||||
for i in 0..3600 {
|
||||
index.record(1, entity, type_id, ts + i * 1_000_000_000);
|
||||
}
|
||||
@ -158,7 +169,7 @@ mod tests {
|
||||
let index = UserSignalIndex::new();
|
||||
let type_id = SignalTypeId::new(0);
|
||||
let entity = EntityId::new(42);
|
||||
let ts = 1_000_000_000_000_000_000u64;
|
||||
let ts = Timestamp::now().as_nanos();
|
||||
|
||||
// User 1: 100 events.
|
||||
for i in 0..100 {
|
||||
@ -180,7 +191,7 @@ mod tests {
|
||||
let index = UserSignalIndex::new();
|
||||
let type_id = SignalTypeId::new(0);
|
||||
let entity = EntityId::new(42);
|
||||
let ts = 1_000_000_000_000_000_000u64;
|
||||
let ts = Timestamp::now().as_nanos();
|
||||
|
||||
index.record(1, entity, type_id, ts);
|
||||
assert_eq!(
|
||||
|
||||
@ -1,8 +1,18 @@
|
||||
//! Per-user state bitmap indexes.
|
||||
//!
|
||||
//! Tracks seen items, blocked creators/items, saved items, liked items,
|
||||
//! and completion progress. All state is in-memory and rebuilt from
|
||||
//! relationships + signals on demand.
|
||||
//! and completion progress.
|
||||
//!
|
||||
//! # Durability
|
||||
//!
|
||||
//! `seen`, `blocked`, `follows`, and `creator_followers` are rebuilt on open
|
||||
//! from durable relationship edges (`rebuild_entity_state`). The
|
||||
//! `completion` / `saved` / `liked` / `save_timestamps` maps have no
|
||||
//! relationship edge — they are driven by `completion` / `save` / `like` signals
|
||||
//! — so they are persisted write-through as durable `Tag::UserState` rows by the
|
||||
//! signal path and restored on open via [`UserStateIndex::restore`]. Without
|
||||
//! that, the `InProgress` filter and `DateSaved` sort silently reset to empty on
|
||||
//! every restart.
|
||||
|
||||
use std::collections::HashSet;
|
||||
|
||||
@ -10,6 +20,49 @@ use dashmap::DashMap;
|
||||
use roaring::{RoaringBitmap, RoaringTreemap};
|
||||
|
||||
use super::CreatorItemsBitmap;
|
||||
use crate::{
|
||||
schema::EntityId,
|
||||
storage::{StorageEngine, Tag, encode_key, entity_tag_prefix, parse_key},
|
||||
};
|
||||
|
||||
/// Narrow a `u64` item-side entity ID into the `u32` slot used by every
|
||||
/// per-user item bitmap (seen / hidden / saved / liked), emitting a
|
||||
/// `tracing::warn!` if the ID exceeds `u32::MAX`.
|
||||
///
|
||||
/// The item universe is bounded to `u32::MAX` (~4 billion) because the
|
||||
/// in-memory candidate-generation indexes key on a `u32` slot. The primary
|
||||
/// write path (`db/items.rs::write_item_with_metadata`) *rejects* an
|
||||
/// over-range item ID before it ever lands in storage, so in a well-formed
|
||||
/// database this narrowing never overflows. The relationship-driven write
|
||||
/// paths (the `Hide` edge in `db/relationships.rs` and its replay in
|
||||
/// `db/state_rebuild.rs`) reach this index with a `to` id that was *not*
|
||||
/// funneled through that guard, so they must route their `u64 -> u32`
|
||||
/// narrowing through this helper to keep every truncation observable and
|
||||
/// consistent with the warning `db/items.rs` already emits. Truncation is
|
||||
/// preserved (the documented universe limit is unchanged); only the silent
|
||||
/// part is fixed.
|
||||
///
|
||||
/// Out-of-scope note: `db/relationships.rs` (the `Hide` relationship write /
|
||||
/// delete) and `db/state_rebuild.rs` (the `Hide` replay) currently perform a
|
||||
/// bare `to.as_u64() as u32`; they should call this helper instead. Those
|
||||
/// files are not edited here.
|
||||
#[must_use]
|
||||
pub(crate) fn narrow_item_slot(item_id: u64) -> u32 {
|
||||
if item_id > u64::from(u32::MAX) {
|
||||
// SLOT-NARROW: mirrors the warn db/items.rs emits when rejecting an
|
||||
// over-range item ID, so a Hide on such an ID is equally observable.
|
||||
tracing::warn!(
|
||||
item_id,
|
||||
limit = u32::MAX,
|
||||
"item id exceeds u32 item-slot universe; narrowing aliases a lower id (silent cross-id collision in per-user bitmaps)"
|
||||
);
|
||||
}
|
||||
// Documented, intentional truncation to the u32 item slot.
|
||||
#[allow(clippy::cast_possible_truncation)]
|
||||
{
|
||||
item_id as u32
|
||||
}
|
||||
}
|
||||
|
||||
/// Blocked state: both individual items and all items from blocked creators.
|
||||
#[derive(Debug, Default)]
|
||||
@ -28,7 +81,15 @@ pub struct UserStateIndex {
|
||||
blocked: DashMap<u64, BlockedState>,
|
||||
saved: DashMap<u64, RoaringBitmap>,
|
||||
liked: DashMap<u64, RoaringBitmap>,
|
||||
completion: DashMap<(u64, u64), f64>,
|
||||
/// `user_id` -> (`item_id` -> completion ratio).
|
||||
///
|
||||
/// Keyed per-user rather than by a flat `(user_id, item_id)` so that
|
||||
/// [`in_progress_items`](Self::in_progress_items) scans only one user's
|
||||
/// completion entries instead of the global map. The threshold is applied
|
||||
/// at query time, so the inner map keeps the raw ratio (it cannot be
|
||||
/// pre-bucketed by an "in-progress" flag that depends on a per-call
|
||||
/// threshold).
|
||||
completion: DashMap<u64, std::collections::HashMap<u64, f64>>,
|
||||
/// `user_id` -> set of followed creator IDs.
|
||||
/// Populated from `RelationshipType::Follows` edges on startup and
|
||||
/// maintained in real-time by `write_relationship` / `delete_relationship`.
|
||||
@ -89,7 +150,7 @@ impl UserStateIndex {
|
||||
|
||||
// ── Blocked ─────────────────────────────────────────────────────────
|
||||
|
||||
/// Hide a specific item for a user.
|
||||
/// Hide a specific item for a user (item already narrowed to a `u32` slot).
|
||||
pub fn add_hide(&self, user_id: u64, item_id: u32) {
|
||||
self.blocked
|
||||
.entry(user_id)
|
||||
@ -98,6 +159,24 @@ impl UserStateIndex {
|
||||
.insert(item_id);
|
||||
}
|
||||
|
||||
/// Hide a specific item for a user, narrowing the `u64` item ID to its
|
||||
/// `u32` slot through [`narrow_item_slot`] so an over-range ID logs a
|
||||
/// warning instead of silently aliasing a lower ID.
|
||||
///
|
||||
/// This is the observable entry point the `Hide` relationship write path
|
||||
/// (`db/relationships.rs`) and its rebuild replay (`db/state_rebuild.rs`)
|
||||
/// should call instead of `add_hide(user_id, to.as_u64() as u32)`.
|
||||
pub fn add_hide_item(&self, user_id: u64, item_id: u64) {
|
||||
self.add_hide(user_id, narrow_item_slot(item_id));
|
||||
}
|
||||
|
||||
/// Un-hide a specific item for a user, narrowing the `u64` item ID to its
|
||||
/// `u32` slot through [`narrow_item_slot`]. Companion to
|
||||
/// [`add_hide_item`](Self::add_hide_item) for the `Hide` delete path.
|
||||
pub fn remove_hide_item(&self, user_id: u64, item_id: u64) {
|
||||
self.remove_hide(user_id, narrow_item_slot(item_id));
|
||||
}
|
||||
|
||||
/// Block an entire creator for a user.
|
||||
pub fn add_block_creator(&self, user_id: u64, creator_id: u64) {
|
||||
self.blocked
|
||||
@ -327,40 +406,198 @@ impl UserStateIndex {
|
||||
// ── Completion ──────────────────────────────────────────────────────
|
||||
|
||||
pub fn record_completion(&self, user_id: u64, item_id: u64, ratio: f64) {
|
||||
self.completion.insert((user_id, item_id), ratio);
|
||||
self.completion
|
||||
.entry(user_id)
|
||||
.or_default()
|
||||
.insert(item_id, ratio);
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn get_completion(&self, user_id: u64, item_id: u64) -> Option<f64> {
|
||||
self.completion.get(&(user_id, item_id)).map(|r| *r)
|
||||
self.completion
|
||||
.get(&user_id)
|
||||
.and_then(|items| items.get(&item_id).copied())
|
||||
}
|
||||
|
||||
/// Check if an item is "in progress" (0.0 < completion < threshold).
|
||||
#[must_use]
|
||||
pub fn is_in_progress(&self, user_id: u64, item_id: u64, threshold: f64) -> bool {
|
||||
self.completion
|
||||
.get(&(user_id, item_id))
|
||||
.is_some_and(|r| *r > 0.0 && *r < threshold)
|
||||
.get(&user_id)
|
||||
.and_then(|items| items.get(&item_id).copied())
|
||||
.is_some_and(|r| r > 0.0 && r < threshold)
|
||||
}
|
||||
|
||||
/// List all item IDs with partial completion for a user (0 < v < threshold).
|
||||
///
|
||||
/// Returns a `RoaringBitmap` of item IDs (truncated to u32) that satisfy
|
||||
/// the in-progress predicate. Used by Stage 2.5 of the executor to build
|
||||
/// the inclusion set for `FilterExpr::InProgress`.
|
||||
/// Returns a `RoaringBitmap` of item IDs (narrowed to a u32 slot via
|
||||
/// [`narrow_item_slot`]) that satisfy the in-progress predicate. Used by
|
||||
/// Stage 2.5 of the executor to build the inclusion set for
|
||||
/// `FilterExpr::InProgress`.
|
||||
///
|
||||
/// Scans only this user's completion entries — the map is keyed per-user,
|
||||
/// so the cost is `O(items completed by this user)`, not `O(global
|
||||
/// completion rows)` as it was when the map was keyed by a flat
|
||||
/// `(user_id, item_id)` tuple.
|
||||
#[must_use]
|
||||
pub fn in_progress_items(&self, user_id: u64, threshold: f64) -> RoaringBitmap {
|
||||
let mut bm = RoaringBitmap::new();
|
||||
for entry in &self.completion {
|
||||
let &(uid, iid) = entry.key();
|
||||
if uid == user_id && *entry.value() > 0.0 && *entry.value() < threshold {
|
||||
bm.insert(iid as u32);
|
||||
if let Some(items) = self.completion.get(&user_id) {
|
||||
for (&iid, &ratio) in items.iter() {
|
||||
if ratio > 0.0 && ratio < threshold {
|
||||
bm.insert(narrow_item_slot(iid));
|
||||
}
|
||||
}
|
||||
}
|
||||
bm
|
||||
}
|
||||
}
|
||||
|
||||
// ── Durable rows (completion / saved-timestamp / liked) ──────────────────────
|
||||
|
||||
/// Discriminator byte for a `Tag::UserState` durable row, distinguishing the
|
||||
/// three signal-driven maps that share the sentinel `(0, Tag::UserState)` prefix.
|
||||
#[repr(u8)]
|
||||
#[derive(Clone, Copy, PartialEq, Eq)]
|
||||
enum UserStateRowKind {
|
||||
/// `(user_be8, item_be8)` -> completion ratio (8B LE f64 value).
|
||||
Completion = 0x01,
|
||||
/// `(user_be8, item_be8)` -> save timestamp ns (8B LE u64 value).
|
||||
Save = 0x02,
|
||||
/// `(user_be8, item_be8)` -> liked (empty value).
|
||||
Like = 0x03,
|
||||
}
|
||||
|
||||
/// Encode a `Tag::UserState` key: `[kind: 1B][user: 8B BE][item: 8B BE]`.
|
||||
fn user_state_key(kind: UserStateRowKind, user_id: u64, item_id: u64) -> Vec<u8> {
|
||||
let mut suffix = [0u8; 17];
|
||||
suffix[0] = kind as u8;
|
||||
suffix[1..9].copy_from_slice(&user_id.to_be_bytes());
|
||||
suffix[9..].copy_from_slice(&item_id.to_be_bytes());
|
||||
encode_key(EntityId::new(0), Tag::UserState, &suffix)
|
||||
}
|
||||
|
||||
/// Decode a `Tag::UserState` suffix into `(kind_byte, user_id, item_id)`.
|
||||
fn decode_user_state_suffix(suffix: &[u8]) -> Option<(u8, u64, u64)> {
|
||||
if suffix.len() < 17 {
|
||||
return None;
|
||||
}
|
||||
let user_id = u64::from_be_bytes(suffix[1..9].try_into().ok()?);
|
||||
let item_id = u64::from_be_bytes(suffix[9..17].try_into().ok()?);
|
||||
Some((suffix[0], user_id, item_id))
|
||||
}
|
||||
|
||||
impl UserStateIndex {
|
||||
/// Persist a completion ratio write-through as a durable `Tag::UserState` row,
|
||||
/// then update the in-memory map. Mirrors [`record_completion`](Self::record_completion).
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns storage errors from the underlying engine.
|
||||
pub fn record_completion_durable(
|
||||
&self,
|
||||
storage: &dyn StorageEngine,
|
||||
user_id: u64,
|
||||
item_id: u64,
|
||||
ratio: f64,
|
||||
) -> crate::Result<()> {
|
||||
let key = user_state_key(UserStateRowKind::Completion, user_id, item_id);
|
||||
storage
|
||||
.put(&key, &ratio.to_le_bytes())
|
||||
.map_err(crate::schema::TidalError::from)?;
|
||||
self.record_completion(user_id, item_id, ratio);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Persist a timestamped save write-through, then update the in-memory map.
|
||||
/// Mirrors [`add_save_timestamped`](Self::add_save_timestamped).
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns storage errors from the underlying engine.
|
||||
pub fn add_save_durable(
|
||||
&self,
|
||||
storage: &dyn StorageEngine,
|
||||
user_id: u64,
|
||||
item_id: u32,
|
||||
timestamp_ns: u64,
|
||||
) -> crate::Result<()> {
|
||||
let key = user_state_key(UserStateRowKind::Save, user_id, u64::from(item_id));
|
||||
storage
|
||||
.put(&key, ×tamp_ns.to_le_bytes())
|
||||
.map_err(crate::schema::TidalError::from)?;
|
||||
self.add_save_timestamped(user_id, item_id, timestamp_ns);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Persist a like write-through, then update the in-memory map.
|
||||
/// Mirrors [`add_like`](Self::add_like).
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns storage errors from the underlying engine.
|
||||
pub fn add_like_durable(
|
||||
&self,
|
||||
storage: &dyn StorageEngine,
|
||||
user_id: u64,
|
||||
item_id: u32,
|
||||
) -> crate::Result<()> {
|
||||
let key = user_state_key(UserStateRowKind::Like, user_id, u64::from(item_id));
|
||||
storage
|
||||
.put(&key, &[])
|
||||
.map_err(crate::schema::TidalError::from)?;
|
||||
self.add_like(user_id, item_id);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Restore the completion / saved-timestamp / liked maps from durable
|
||||
/// `Tag::UserState` rows on open. Targeted prefix scan over the sentinel
|
||||
/// entity (`0`), so restore is `O(user_state_rows)`.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns storage errors from the underlying engine.
|
||||
#[allow(clippy::cast_possible_truncation)]
|
||||
pub fn restore(&self, storage: &dyn StorageEngine) -> crate::Result<()> {
|
||||
let prefix = entity_tag_prefix(EntityId::new(0), Tag::UserState);
|
||||
let mut count = 0u64;
|
||||
for entry in storage.scan_prefix(&prefix) {
|
||||
let (key, value) = entry.map_err(crate::schema::TidalError::from)?;
|
||||
let Some((_, Tag::UserState, suffix)) = parse_key(&key) else {
|
||||
continue;
|
||||
};
|
||||
let Some((kind, user_id, item_id)) = decode_user_state_suffix(suffix) else {
|
||||
continue;
|
||||
};
|
||||
match kind {
|
||||
k if k == UserStateRowKind::Completion as u8 => {
|
||||
if value.len() >= 8 {
|
||||
let ratio = f64::from_le_bytes(value[..8].try_into().unwrap_or([0u8; 8]));
|
||||
self.record_completion(user_id, item_id, ratio);
|
||||
count += 1;
|
||||
}
|
||||
}
|
||||
k if k == UserStateRowKind::Save as u8 => {
|
||||
if value.len() >= 8 {
|
||||
let ts = u64::from_le_bytes(value[..8].try_into().unwrap_or([0u8; 8]));
|
||||
self.add_save_timestamped(user_id, item_id as u32, ts);
|
||||
count += 1;
|
||||
}
|
||||
}
|
||||
k if k == UserStateRowKind::Like as u8 => {
|
||||
self.add_like(user_id, item_id as u32);
|
||||
count += 1;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
if count > 0 {
|
||||
tracing::info!(rows = count, "user-state index restored from durable rows");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for UserStateIndex {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
@ -414,6 +651,42 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn narrow_item_slot_in_range_is_identity() {
|
||||
// Any id within the u32 universe narrows to itself with no warning.
|
||||
assert_eq!(narrow_item_slot(0), 0);
|
||||
assert_eq!(narrow_item_slot(123), 123);
|
||||
assert_eq!(narrow_item_slot(u64::from(u32::MAX)), u32::MAX);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn narrow_item_slot_over_range_truncates() {
|
||||
// Over-range ids are truncated to the low 32 bits (the documented,
|
||||
// intentional universe limit); the warning makes the truncation
|
||||
// observable but the value contract is unchanged.
|
||||
assert_eq!(narrow_item_slot(u64::from(u32::MAX) + 1), 0);
|
||||
assert_eq!(narrow_item_slot((1u64 << 32) | 7), 7);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn add_hide_item_routes_through_narrowing() {
|
||||
let idx = UserStateIndex::new();
|
||||
// In-range id hides exactly that slot.
|
||||
idx.add_hide_item(1, 100);
|
||||
assert!(idx.hidden_items(1).contains(100));
|
||||
|
||||
// Over-range id narrows to its low-32 slot, matching narrow_item_slot.
|
||||
// (Low 32 bits = 42; the high bit is dropped by the narrowing.)
|
||||
idx.add_hide_item(1, (1u64 << 32) + 42);
|
||||
assert!(idx.hidden_items(1).contains(42));
|
||||
|
||||
// remove_hide_item un-hides through the same narrowing.
|
||||
idx.remove_hide_item(1, (1u64 << 32) + 42);
|
||||
assert!(!idx.hidden_items(1).contains(42));
|
||||
// The unrelated in-range hide is untouched.
|
||||
assert!(idx.hidden_items(1).contains(100));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn block_creator() {
|
||||
let idx = UserStateIndex::new();
|
||||
@ -606,4 +879,38 @@ mod tests {
|
||||
assert!(!bm.contains(30)); // == 0.0
|
||||
assert!(!bm.contains(50)); // different user
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn in_progress_items_isolates_users_under_shared_item_ids() {
|
||||
// Regression for the per-user completion keying: many users completing
|
||||
// the *same* item IDs must not bleed into each other's in-progress set.
|
||||
// Previously the global (user_id, item_id) map was scanned in full per
|
||||
// call; the per-user map makes each query touch only its own user.
|
||||
let idx = UserStateIndex::new();
|
||||
for uid in 0..100u64 {
|
||||
// Every user has item 7 in progress and item 8 complete.
|
||||
idx.record_completion(uid, 7, 0.5);
|
||||
idx.record_completion(uid, 8, 0.95);
|
||||
}
|
||||
// User 42 additionally has item 9 in progress.
|
||||
idx.record_completion(42, 9, 0.25);
|
||||
|
||||
let bm = idx.in_progress_items(42, 0.8);
|
||||
assert_eq!(bm.len(), 2, "only user 42's in-progress items");
|
||||
assert!(bm.contains(7));
|
||||
assert!(bm.contains(9));
|
||||
assert!(!bm.contains(8)); // >= threshold
|
||||
|
||||
// A user without the extra item sees only the shared in-progress item.
|
||||
let other = idx.in_progress_items(7, 0.8);
|
||||
assert_eq!(other.len(), 1);
|
||||
assert!(other.contains(7));
|
||||
assert!(!other.contains(9)); // belongs to user 42 only
|
||||
|
||||
// get_completion / is_in_progress agree with the per-user map.
|
||||
assert_eq!(idx.get_completion(42, 9), Some(0.25));
|
||||
assert_eq!(idx.get_completion(7, 9), None);
|
||||
assert!(idx.is_in_progress(42, 9, 0.8));
|
||||
assert!(!idx.is_in_progress(7, 9, 0.8));
|
||||
}
|
||||
}
|
||||
|
||||
@ -15,7 +15,7 @@ use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use dashmap::DashMap;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::scope::SignalScope;
|
||||
use super::scope::{CommunityId, SignalScope};
|
||||
use crate::schema::TidalError;
|
||||
|
||||
/// A scope *class* a capability can grant rights over.
|
||||
@ -49,6 +49,21 @@ impl ScopeClass {
|
||||
}
|
||||
}
|
||||
|
||||
// `ScopeClass` hand-mirrors `SignalScope`'s variant taxonomy and discriminant
|
||||
// ordering (the durable wire byte). These const assertions pin each
|
||||
// `ScopeClass as u8` to the matching [`SignalScope::discriminant`] at *compile
|
||||
// time*, so adding or reordering a scope on either side that breaks the
|
||||
// alignment is a build error — not a silent drift caught (or missed) at runtime.
|
||||
// `Community` uses the `NONE` sentinel because `discriminant` ignores the id.
|
||||
const _: () = {
|
||||
assert!(ScopeClass::Local as u8 == SignalScope::Local.discriminant());
|
||||
assert!(
|
||||
ScopeClass::Community as u8 == SignalScope::Community(CommunityId::NONE).discriminant()
|
||||
);
|
||||
assert!(ScopeClass::Session as u8 == SignalScope::Session.discriminant());
|
||||
assert!(ScopeClass::Agent as u8 == SignalScope::Agent.discriminant());
|
||||
};
|
||||
|
||||
/// A single scope permission within a capability token.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ScopePermission {
|
||||
@ -306,29 +321,67 @@ impl CapabilityRegistry {
|
||||
/// cost is amortized away over the cap. The victim id/agent are cloned at
|
||||
/// most once (only when a new best candidate is found), not per iteration.
|
||||
fn evict_one(&self, now_ns: u64) {
|
||||
// Victim selection state: a `(created_at_ns, token_id, agent_id)` triple
|
||||
// for the oldest live token seen so far. A dead token short-circuits the
|
||||
// scan because dropping it costs no live capability.
|
||||
let mut victim: Option<(u64, String, String)> = None;
|
||||
// Victim selection state: a `(created_at_ns, token_id, agent_id, live)`
|
||||
// tuple for the best candidate seen so far. `live` records whether the
|
||||
// chosen victim is still in force, so the observability warn below needs
|
||||
// no second map lookup. A dead token short-circuits the scan because
|
||||
// dropping it costs no live capability.
|
||||
let mut victim: Option<(u64, String, String, bool)> = None;
|
||||
for entry in &self.tokens {
|
||||
let t = entry.value();
|
||||
if t.is_revoked() || t.is_expired(now_ns) {
|
||||
victim = Some((t.created_at_ns, t.token_id.clone(), t.agent_id.clone()));
|
||||
victim = Some((
|
||||
t.created_at_ns,
|
||||
t.token_id.clone(),
|
||||
t.agent_id.clone(),
|
||||
false,
|
||||
));
|
||||
break;
|
||||
}
|
||||
// Keep the smallest (created_at_ns, token_id) — token_id breaks ties
|
||||
// deterministically. Only clone when this entry becomes the new best.
|
||||
let is_new_best = match &victim {
|
||||
Some((c, id, _)) => (t.created_at_ns, t.token_id.as_str()) < (*c, id.as_str()),
|
||||
Some((c, id, _, _)) => (t.created_at_ns, t.token_id.as_str()) < (*c, id.as_str()),
|
||||
None => true,
|
||||
};
|
||||
if is_new_best {
|
||||
victim = Some((t.created_at_ns, t.token_id.clone(), t.agent_id.clone()));
|
||||
victim = Some((
|
||||
t.created_at_ns,
|
||||
t.token_id.clone(),
|
||||
t.agent_id.clone(),
|
||||
true,
|
||||
));
|
||||
}
|
||||
}
|
||||
let Some((_, victim_id, victim_agent)) = victim else {
|
||||
let Some((victim_created_ns, victim_id, victim_agent, victim_live)) = victim else {
|
||||
return; // registry empty
|
||||
};
|
||||
// Observability: if the victim is still live (non-revoked, non-expired)
|
||||
// we are trimming a *valid* grant out of the in-memory cache. The check
|
||||
// path consults only the in-memory maps, so until the durable
|
||||
// `Tag::Capability` row is rehydrated (today only on the next `open()`)
|
||||
// a `check()` for this agent/scope will return a false denial. Surface
|
||||
// the condition so it is operable rather than silent.
|
||||
//
|
||||
// TODO(rehydrate): make `check()` lazily reload an evicted-but-durable
|
||||
// token from `Tag::Capability` before returning a denial. That requires
|
||||
// a storage-engine handle inside the registry, which neither `check()`
|
||||
// nor `evict_one()` currently hold — both are pure in-memory paths. The
|
||||
// fix belongs where the registry is constructed (`db/mod.rs`), so it is
|
||||
// tracked there rather than papered over here. Until then this warn is
|
||||
// the contract: a live eviction is visible, and the durable log + next
|
||||
// `open()` remain the recovery mechanism.
|
||||
if victim_live {
|
||||
tracing::warn!(
|
||||
token_id = %victim_id,
|
||||
agent_id = %victim_agent,
|
||||
created_at_ns = victim_created_ns,
|
||||
cap = self.cap,
|
||||
"evicting a live capability token under cap pressure; \
|
||||
check() will deny this grant in-memory until it is \
|
||||
rehydrated from durable Tag::Capability on next open()"
|
||||
);
|
||||
}
|
||||
self.tokens.remove(&victim_id);
|
||||
if let Some(mut ids) = self.by_agent.get_mut(&victim_agent) {
|
||||
ids.retain(|id| id != &victim_id);
|
||||
@ -516,7 +569,6 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn scope_class_of_signalscope() {
|
||||
use crate::governance::scope::CommunityId;
|
||||
assert_eq!(ScopeClass::of(SignalScope::Local), ScopeClass::Local);
|
||||
assert_eq!(
|
||||
ScopeClass::of(SignalScope::Community(CommunityId(1))),
|
||||
@ -524,6 +576,27 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scope_class_discriminants_match_signal_scope() {
|
||||
// Runtime twin of the module-level `const _` assertions: pin every
|
||||
// `ScopeClass as u8` to its `SignalScope::discriminant`. The const block
|
||||
// is the load-bearing guard (it fails the build); this test documents
|
||||
// the invariant and keeps it visible in the suite. Iterating every
|
||||
// `SignalScope` variant means a new variant forces an update here too.
|
||||
for scope in [
|
||||
SignalScope::Local,
|
||||
SignalScope::Community(CommunityId::NONE),
|
||||
SignalScope::Session,
|
||||
SignalScope::Agent,
|
||||
] {
|
||||
assert_eq!(
|
||||
ScopeClass::of(scope) as u8,
|
||||
scope.discriminant(),
|
||||
"ScopeClass discriminant drifted from SignalScope for {scope:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a token with an explicit `created_at_ns` for eviction-order tests.
|
||||
fn token_created(agent: &str, id: &str, created_at_ns: u64) -> CapabilityToken {
|
||||
let mut t = token(agent, id, true, u64::MAX);
|
||||
@ -582,6 +655,46 @@ mod tests {
|
||||
assert!(reg.token("a#2").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn live_eviction_keeps_registry_consistent_and_index_clean() {
|
||||
// When every resident token is live and the registry is at cap, the
|
||||
// oldest live token is the victim (the warned-about case). The registry
|
||||
// must stay internally consistent: the evicted id is gone from both the
|
||||
// token map and the agent index, and the survivors are intact.
|
||||
let reg = CapabilityRegistry::with_capacity(2);
|
||||
reg.insert(token_created("a", "a#0", 10)); // oldest live -> victim
|
||||
reg.insert(token_created("a", "a#1", 20));
|
||||
reg.insert(token_created("a", "a#2", 30)); // forces a live eviction
|
||||
|
||||
assert_eq!(reg.len(), 2);
|
||||
assert!(reg.token("a#0").is_none(), "oldest live token is evicted");
|
||||
// The agent index no longer references the evicted id...
|
||||
let ids: Vec<String> = reg
|
||||
.tokens_for("a")
|
||||
.into_iter()
|
||||
.map(|t| t.token_id)
|
||||
.collect();
|
||||
assert!(!ids.contains(&"a#0".to_string()));
|
||||
assert!(ids.contains(&"a#1".to_string()));
|
||||
assert!(ids.contains(&"a#2".to_string()));
|
||||
// ...and check() now denies the evicted-but-still-valid grant purely
|
||||
// from the in-memory maps — the false-denial condition the warn flags.
|
||||
// (a#0 had no successor token granting Community, so this is a clean
|
||||
// InsufficientCapability, not Revoked/Expired.)
|
||||
let reg2 = CapabilityRegistry::with_capacity(1);
|
||||
reg2.insert(token_created("solo", "solo#0", 10));
|
||||
assert!(
|
||||
reg2.check("solo", ScopeClass::Community, true, 100).is_ok(),
|
||||
"live grant is honored while resident"
|
||||
);
|
||||
reg2.insert(token_created("other", "other#0", 20)); // evicts solo#0 (live)
|
||||
assert_eq!(
|
||||
reg2.check("solo", ScopeClass::Community, true, 100),
|
||||
Err(CapabilityDenial::InsufficientCapability),
|
||||
"evicted live grant is denied in-memory until durable rehydrate"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_to_bytes_is_fallible_and_succeeds_for_valid_token() {
|
||||
let t = token("a", "a#0", true, u64::MAX);
|
||||
|
||||
@ -18,7 +18,7 @@ use dashmap::DashMap;
|
||||
|
||||
use super::{membership::UserId, scope::CommunityId};
|
||||
use crate::{
|
||||
schema::{DecayModel, EntityId, Schema, SchemaError, TidalError, Window},
|
||||
schema::{DecayModel, EntityId, Schema, SchemaError, TidalError, Timestamp, Window},
|
||||
signals::{
|
||||
SignalTypeId,
|
||||
checkpoint::format::{deserialize_entry, serialize_entry},
|
||||
@ -243,10 +243,14 @@ impl CommunityLedger {
|
||||
let Some(cell) = self.cells.get(&(community, entity_id, type_id)) else {
|
||||
return Ok(0);
|
||||
};
|
||||
// Read the clock once and age every contributor's warm tier out at read
|
||||
// time, so a quiet community-entity's windowed count decays instead of
|
||||
// freezing at its peak (CRITICAL #6, CODING_GUIDELINES §8).
|
||||
let now_ns = Timestamp::now().as_nanos();
|
||||
Ok(cell
|
||||
.contributors
|
||||
.values()
|
||||
.map(|e| e.warm.windowed_count(window))
|
||||
.map(|e| e.warm.windowed_count(window, now_ns))
|
||||
.sum())
|
||||
}
|
||||
|
||||
@ -374,7 +378,7 @@ impl CommunityLedger {
|
||||
writer_agent: k.1.clone(),
|
||||
epoch: k.2,
|
||||
score: e.hot.current_score(0, now_ns, lambda),
|
||||
count: e.warm.windowed_count(Window::AllTime),
|
||||
count: e.warm.windowed_count(Window::AllTime, now_ns),
|
||||
});
|
||||
}
|
||||
}
|
||||
@ -432,6 +436,20 @@ impl CommunityLedger {
|
||||
}
|
||||
}
|
||||
|
||||
// Stage purge watermarks. These are NOT contributor rows, so they use the
|
||||
// watermark suffix codec (a sentinel `signal_type` discriminant the
|
||||
// contributor decoder rejects) and survive a restart independent of any
|
||||
// surviving contributor cells. Without this a community whose every
|
||||
// contributor was purged would lose its watermark, which queries surface.
|
||||
for w in &self.purge_watermarks {
|
||||
let key = encode_key(
|
||||
EntityId::new(0),
|
||||
Tag::Community,
|
||||
&encode_watermark_suffix(*w.key()),
|
||||
);
|
||||
batch.put(key, w.value().to_le_bytes().to_vec());
|
||||
}
|
||||
|
||||
storage.write_batch(batch).map_err(TidalError::from)?;
|
||||
Ok(())
|
||||
}
|
||||
@ -444,6 +462,15 @@ impl CommunityLedger {
|
||||
let Some(suffix) = key.get(10..) else {
|
||||
continue;
|
||||
};
|
||||
// Watermark rows are shorter than the contributor fixed prefix and are
|
||||
// decoded first so they are not misread as torn contributor rows.
|
||||
if let Some(community) = decode_watermark_suffix(suffix) {
|
||||
if value.len() >= 8 {
|
||||
let w = u64::from_le_bytes(value[..8].try_into().unwrap_or([0u8; 8]));
|
||||
self.set_purge_watermark(community, w);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
let Some(coords) = decode_contributor_suffix(suffix) else {
|
||||
continue;
|
||||
};
|
||||
@ -483,6 +510,22 @@ impl CommunityLedger {
|
||||
self.set_purge_watermark(t.community_id, t.watermark_ns);
|
||||
}
|
||||
|
||||
/// Re-apply a deserialized [`AgentPurgeTombstone`](super::tombstone::AgentPurgeTombstone)
|
||||
/// against the restored ledger, making an agent-scope purge idempotent across
|
||||
/// crash/replay (the agent-path analogue of [`apply_purge_tombstone`](Self::apply_purge_tombstone)).
|
||||
///
|
||||
/// The tombstone is written durably *before* `purge_writer` mutates the
|
||||
/// in-memory aggregate, but the community checkpoint is only flushed on the
|
||||
/// lifecycle/periodic path — so a crash after the purge but before the next
|
||||
/// checkpoint would restore the *pre-purge* snapshot and silently resurrect the
|
||||
/// revoked agent's contributions. On-open replay closes that window: it
|
||||
/// re-purges by writer and re-stamps the watermark. Replaying a purge that
|
||||
/// already survived in the checkpoint is a no-op (purge is idempotent).
|
||||
pub(crate) fn apply_agent_purge_tombstone(&self, t: &super::tombstone::AgentPurgeTombstone) {
|
||||
let _removed = self.purge_writer(t.community_id, &t.writer_agent);
|
||||
self.set_purge_watermark(t.community_id, t.watermark_ns);
|
||||
}
|
||||
|
||||
/// A canonical, order-independent digest of a community's entire aggregate
|
||||
/// state at `now_ns`, for verifying rematerialization determinism (M9p3:
|
||||
/// "rebuilt aggregates identical across repeated replay").
|
||||
@ -509,7 +552,7 @@ impl CommunityLedger {
|
||||
for k in keys {
|
||||
if let Some(e) = cell.value().contributors.get(k) {
|
||||
score += e.hot.current_score(0, now_ns, lambda);
|
||||
count += e.warm.windowed_count(Window::AllTime);
|
||||
count += e.warm.windowed_count(Window::AllTime, now_ns);
|
||||
}
|
||||
}
|
||||
rows.push((entity.as_u64(), type_id.as_u16(), score.to_bits(), count));
|
||||
@ -534,6 +577,28 @@ impl CommunityLedger {
|
||||
/// tag (variable-length UTF-8) follows and runs to the end of the suffix.
|
||||
const CONTRIB_SUFFIX_FIXED_LEN: usize = 8 + 8 + 4 + 8 + 2;
|
||||
|
||||
/// Length of a purge-watermark suffix: a single big-endian community id.
|
||||
///
|
||||
/// Deliberately shorter than [`CONTRIB_SUFFIX_FIXED_LEN`] so a watermark row is
|
||||
/// unambiguously distinguishable from a (possibly torn) contributor row during
|
||||
/// restore — the contributor decoder rejects anything shorter than its fixed
|
||||
/// prefix, so a watermark suffix can never be mis-bucketed as a contribution.
|
||||
const WATERMARK_SUFFIX_LEN: usize = 8;
|
||||
|
||||
/// Encode the durable suffix for a purge-watermark row: `community 8` (BE).
|
||||
fn encode_watermark_suffix(community: CommunityId) -> Vec<u8> {
|
||||
community.as_u64().to_be_bytes().to_vec()
|
||||
}
|
||||
|
||||
/// Decode a purge-watermark suffix, or `None` if it is not exactly the
|
||||
/// watermark length (i.e. it is a contributor row).
|
||||
fn decode_watermark_suffix(suffix: &[u8]) -> Option<CommunityId> {
|
||||
if suffix.len() != WATERMARK_SUFFIX_LEN {
|
||||
return None;
|
||||
}
|
||||
Some(CommunityId(u64::from_be_bytes(suffix.try_into().ok()?)))
|
||||
}
|
||||
|
||||
/// The fully-decoded coordinates of one contributor checkpoint row.
|
||||
struct ContributorCoords {
|
||||
community: CommunityId,
|
||||
@ -689,7 +754,7 @@ mod tests {
|
||||
let digest_before = l.aggregate_digest(c, now);
|
||||
|
||||
// Purge user 3 once, capture digest.
|
||||
l.purge_contributor(c, UserId(3), 0..=0);
|
||||
let _ = l.purge_contributor(c, UserId(3), 0..=0);
|
||||
let digest_after = l.aggregate_digest(c, now);
|
||||
assert_ne!(
|
||||
digest_before, digest_after,
|
||||
@ -698,7 +763,7 @@ mod tests {
|
||||
|
||||
// Re-running the same purge is a no-op: identical digest (idempotent +
|
||||
// deterministic rematerialization).
|
||||
l.purge_contributor(c, UserId(3), 0..=0);
|
||||
let _ = l.purge_contributor(c, UserId(3), 0..=0);
|
||||
assert_eq!(
|
||||
digest_after,
|
||||
l.aggregate_digest(c, now),
|
||||
|
||||
@ -30,7 +30,7 @@ pub use capability::{
|
||||
};
|
||||
pub use community_ledger::{CommunityLedger, ContributionProvenance};
|
||||
pub use policy::{GovernancePolicy, GovernanceRegistry, PolicyRegisterError, WeightingBounds};
|
||||
pub use tombstone::{PurgeReceipt, PurgeTombstone, RemoveScope};
|
||||
pub use tombstone::{AgentPurgeTombstone, PurgeReceipt, PurgeTombstone, RemoveScope};
|
||||
|
||||
pub use membership::{LeaveMode, Membership, MembershipRegistry, MembershipState, UserId};
|
||||
pub use provenance::{SignalProvenance, WRITER_DIRECT};
|
||||
|
||||
@ -16,6 +16,24 @@ use serde::{Deserialize, Serialize};
|
||||
use super::{membership::UserId, scope::CommunityId};
|
||||
use crate::schema::TidalError;
|
||||
|
||||
/// Serialize a durable governance row to bytes, mapping any serde failure to
|
||||
/// [`TidalError::Internal`] tagged with `operation`.
|
||||
///
|
||||
/// Both tombstone types share the same durability invariant: a serialize
|
||||
/// failure must be propagated, never swallowed into empty bytes — an empty
|
||||
/// value deserializes back as a *missing* purge record, which on crash-replay
|
||||
/// silently resurrects purged data (the exact privacy/trust failure these
|
||||
/// durable records exist to prevent). Centralizing the serialize-or-error here
|
||||
/// keeps that single behavior from drifting between the two call sites; each
|
||||
/// caller supplies its own `operation` tag so the error context stays precise.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns [`TidalError::Internal`] if serialization fails.
|
||||
fn to_durable_bytes<T: Serialize>(value: &T, operation: &'static str) -> crate::Result<Vec<u8>> {
|
||||
serde_json::to_vec(value).map_err(|e| TidalError::internal(operation, e.to_string()))
|
||||
}
|
||||
|
||||
/// A durable record of a retroactive purge of one user's community
|
||||
/// contributions over an inclusive epoch range.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
@ -84,8 +102,7 @@ impl PurgeTombstone {
|
||||
///
|
||||
/// Returns [`TidalError::Internal`] if serialization fails.
|
||||
pub fn to_bytes(&self) -> crate::Result<Vec<u8>> {
|
||||
serde_json::to_vec(self)
|
||||
.map_err(|e| TidalError::internal("purge_tombstone_to_bytes", e.to_string()))
|
||||
to_durable_bytes(self, "purge_tombstone_to_bytes")
|
||||
}
|
||||
|
||||
/// Deserialize from durable bytes; malformed bytes return `None`.
|
||||
@ -99,6 +116,71 @@ impl PurgeTombstone {
|
||||
}
|
||||
}
|
||||
|
||||
/// A durable record of a remove-by-agent purge: every contribution written by
|
||||
/// `writer_agent` to `community_id` was removed (M10p3 agent-scope removal).
|
||||
///
|
||||
/// The user-scoped [`PurgeTombstone`] keys on `(community, user, epoch_range)`;
|
||||
/// agent removal cuts across users, so it needs its own tombstone keyed on
|
||||
/// `(community, writer_agent)`. Written under `Tag::AgentPurgeTombstone` *before*
|
||||
/// the in-memory `purge_writer` mutation (durability-before-mutation), and
|
||||
/// replayed on open against the restored community ledger so a crash after the
|
||||
/// purge but before the next community checkpoint does not resurrect the revoked
|
||||
/// agent's contributions.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct AgentPurgeTombstone {
|
||||
/// The community the contributions were removed from.
|
||||
pub community_id: CommunityId,
|
||||
/// The writer agent whose contributions were removed.
|
||||
pub writer_agent: String,
|
||||
/// Wall-clock ns at which the purge was applied; doubles as the purge
|
||||
/// watermark exposed to community queries.
|
||||
pub watermark_ns: u64,
|
||||
}
|
||||
|
||||
impl AgentPurgeTombstone {
|
||||
/// Construct a tombstone for `writer_agent`'s `community` contributions.
|
||||
#[must_use]
|
||||
pub const fn new(community_id: CommunityId, writer_agent: String, watermark_ns: u64) -> Self {
|
||||
Self {
|
||||
community_id,
|
||||
writer_agent,
|
||||
watermark_ns,
|
||||
}
|
||||
}
|
||||
|
||||
/// Durable storage suffix: `community_be8 || writer_agent_utf8`.
|
||||
///
|
||||
/// Keying by `(community, writer_agent)` makes re-purging the same agent
|
||||
/// idempotent (it rewrites the same key) while distinct agents coexist.
|
||||
#[must_use]
|
||||
pub fn storage_suffix(&self) -> Vec<u8> {
|
||||
let mut s = Vec::with_capacity(8 + self.writer_agent.len());
|
||||
s.extend_from_slice(&self.community_id.as_u64().to_be_bytes());
|
||||
s.extend_from_slice(self.writer_agent.as_bytes());
|
||||
s
|
||||
}
|
||||
|
||||
/// Serialize to durable bytes.
|
||||
///
|
||||
/// A serialize failure is propagated rather than swallowed into empty bytes:
|
||||
/// an empty tombstone deserializes back as a *missing* purge, which on
|
||||
/// crash-replay would silently resurrect a revoked agent's data — the exact
|
||||
/// trust failure this durable record exists to prevent.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns [`TidalError::Internal`] if serialization fails.
|
||||
pub fn to_bytes(&self) -> crate::Result<Vec<u8>> {
|
||||
to_durable_bytes(self, "agent_purge_tombstone_to_bytes")
|
||||
}
|
||||
|
||||
/// Deserialize from durable bytes; malformed bytes return `None`.
|
||||
#[must_use]
|
||||
pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
|
||||
serde_json::from_slice(bytes).ok()
|
||||
}
|
||||
}
|
||||
|
||||
/// The scope of a [`remove_from_personalization`] operation (M10p3): which
|
||||
/// contributions to remove, non-globally.
|
||||
///
|
||||
@ -161,4 +243,35 @@ mod tests {
|
||||
let b = PurgeTombstone::new(UserId(1), CommunityId(1), 1, 1, 10);
|
||||
assert_ne!(a.storage_suffix(), b.storage_suffix());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn both_tombstones_serialize_via_shared_helper_and_roundtrip() {
|
||||
// Both `to_bytes` paths now route through `to_durable_bytes`; the shared
|
||||
// helper must still produce non-empty, round-trippable bytes for each
|
||||
// type (the success path the durable log relies on).
|
||||
let user = PurgeTombstone::new(UserId(3), CommunityId(7), 0, 2, 123_456);
|
||||
let user_bytes = user.to_bytes().expect("valid tombstone serializes");
|
||||
assert!(!user_bytes.is_empty());
|
||||
assert_eq!(PurgeTombstone::from_bytes(&user_bytes), Some(user));
|
||||
|
||||
let agent = AgentPurgeTombstone::new(CommunityId(7), "agent-x".to_string(), 999);
|
||||
let agent_bytes = agent.to_bytes().expect("valid tombstone serializes");
|
||||
assert!(!agent_bytes.is_empty());
|
||||
assert_eq!(AgentPurgeTombstone::from_bytes(&agent_bytes), Some(agent));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn durable_helper_tags_error_context_with_operation() {
|
||||
// The helper threads the caller's operation tag into the error context.
|
||||
// A non-serializable value (a map with non-string keys, which JSON
|
||||
// rejects) drives the error path so we can assert the tag is preserved.
|
||||
use std::collections::BTreeMap;
|
||||
let bad: BTreeMap<(u8, u8), u8> = BTreeMap::from([((1, 2), 3)]);
|
||||
let err = to_durable_bytes(&bad, "unit_test_tag")
|
||||
.expect_err("non-string-keyed map is not valid JSON");
|
||||
assert!(
|
||||
err.to_string().contains("unit_test_tag"),
|
||||
"operation tag must survive into the error: {err}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@ -103,6 +103,11 @@ impl RateLimiterConfig {
|
||||
///
|
||||
/// `signals_per_second` is the sustained rate; `burst_capacity` is the
|
||||
/// maximum burst (typically 2x the sustained rate).
|
||||
///
|
||||
/// This constructor is infallible and does **not** validate its inputs;
|
||||
/// validation happens at [`RateLimiter::new`] / [`RateLimiter::try_new`]
|
||||
/// construction time. Prefer [`RateLimiterConfig::try_limited`] to surface
|
||||
/// a degenerate config (zero/NaN/negative rate, sub-unit burst) eagerly.
|
||||
#[must_use]
|
||||
pub const fn limited(signals_per_second: f64, burst_capacity: f64) -> Self {
|
||||
Self {
|
||||
@ -110,6 +115,56 @@ impl RateLimiterConfig {
|
||||
burst_capacity,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a limiting configuration, rejecting degenerate values.
|
||||
///
|
||||
/// Mirrors the validation pattern of
|
||||
/// [`DegradationThresholds::new`](super::DegradationThresholds::new):
|
||||
/// construction fails loudly rather than producing a limiter that can
|
||||
/// never refill (locking an agent out forever and emitting a
|
||||
/// `u64::MAX`-millisecond retry hint).
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `Err` if `signals_per_second` is NaN or `<= 0.0` (a non-positive
|
||||
/// refill rate never adds tokens, so the bucket drains permanently), or if
|
||||
/// `burst_capacity` is NaN or `< 1.0` (a sub-unit bucket can never hold the
|
||||
/// single token `try_acquire` needs). `f64::INFINITY` is accepted for both
|
||||
/// fields — that is the explicit "unlimited" sentinel.
|
||||
pub fn try_limited(signals_per_second: f64, burst_capacity: f64) -> Result<Self, String> {
|
||||
let config = Self {
|
||||
signals_per_second,
|
||||
burst_capacity,
|
||||
};
|
||||
config.validate()?;
|
||||
Ok(config)
|
||||
}
|
||||
|
||||
/// Validate the configuration's invariants.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `Err` describing the first violated invariant. See
|
||||
/// [`RateLimiterConfig::try_limited`] for the rules. The all-`INFINITY`
|
||||
/// default (unlimited) always validates.
|
||||
pub fn validate(&self) -> Result<(), String> {
|
||||
let rate = self.signals_per_second;
|
||||
// INFINITY is the explicit unlimited sentinel and is always valid.
|
||||
if !rate.is_infinite() && (rate.is_nan() || rate <= 0.0) {
|
||||
return Err(format!(
|
||||
"signals_per_second must be > 0.0 (or INFINITY for unlimited); got {rate}"
|
||||
));
|
||||
}
|
||||
let burst = self.burst_capacity;
|
||||
// A bucket that can never hold one whole token can never grant a token,
|
||||
// so anything below 1.0 (or NaN) locks the agent out. INFINITY is fine.
|
||||
if !burst.is_infinite() && (burst.is_nan() || burst < 1.0) {
|
||||
return Err(format!(
|
||||
"burst_capacity must be >= 1.0 (or INFINITY for unlimited); got {burst}"
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for RateLimiterConfig {
|
||||
@ -141,15 +196,56 @@ pub struct RateLimiter {
|
||||
/// Per-(agent, session) token buckets.
|
||||
buckets: DashMap<(String, u64), TokenBucket>,
|
||||
config: RateLimiterConfig,
|
||||
/// Whether limiting is actually enforced. Computed once at construction:
|
||||
/// `false` iff both `signals_per_second` and `burst_capacity` are infinite
|
||||
/// (the unlimited default). When `false`, `check` short-circuits to `Ok(())`
|
||||
/// without allocating a key or touching the `DashMap` — keeping the common
|
||||
/// unlimited path off the per-`session_signal` allocation hot path.
|
||||
enabled: bool,
|
||||
}
|
||||
|
||||
impl RateLimiter {
|
||||
/// Create a new rate limiter with the given configuration.
|
||||
///
|
||||
/// If `config` is degenerate (see [`RateLimiterConfig::validate`]), the
|
||||
/// limiter falls back to the unlimited configuration and logs an error
|
||||
/// rather than constructing a bucket that locks every agent out forever.
|
||||
/// Use [`RateLimiter::try_new`] to surface the error to the caller instead.
|
||||
#[must_use]
|
||||
pub fn new(config: RateLimiterConfig) -> Self {
|
||||
match Self::try_new(config) {
|
||||
Ok(rl) => rl,
|
||||
Err(err) => {
|
||||
tracing::error!(
|
||||
%err,
|
||||
"RateLimiter: rejecting degenerate config; falling back to unlimited \
|
||||
(no per-agent rate limiting will be enforced)"
|
||||
);
|
||||
Self::build(RateLimiterConfig::default())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a new rate limiter, rejecting a degenerate configuration.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `Err` if `config` fails [`RateLimiterConfig::validate`] — i.e. a
|
||||
/// zero/NaN/negative `signals_per_second` or a `burst_capacity < 1.0`,
|
||||
/// either of which would make every bucket permanently empty.
|
||||
pub fn try_new(config: RateLimiterConfig) -> Result<Self, String> {
|
||||
config.validate()?;
|
||||
Ok(Self::build(config))
|
||||
}
|
||||
|
||||
/// Construct without validating. Caller guarantees `config` is valid.
|
||||
fn build(config: RateLimiterConfig) -> Self {
|
||||
let enabled =
|
||||
!(config.signals_per_second.is_infinite() && config.burst_capacity.is_infinite());
|
||||
Self {
|
||||
buckets: DashMap::new(),
|
||||
config,
|
||||
enabled,
|
||||
}
|
||||
}
|
||||
|
||||
@ -169,6 +265,13 @@ impl RateLimiter {
|
||||
/// `retry_after_ms` is the estimated wait before one token refills;
|
||||
/// `limit` is the configured signals-per-second rate.
|
||||
pub fn check(&self, agent_id: &str, session_id: u64) -> Result<(), (u64, f64)> {
|
||||
// Fast path for the unlimited (default) config: an infinite bucket can
|
||||
// never deny, so skip the `agent_id.to_owned()` allocation and the
|
||||
// `DashMap` shard lock entirely. `enabled` is computed once at
|
||||
// construction (see `RateLimiter::build`).
|
||||
if !self.enabled {
|
||||
return Ok(());
|
||||
}
|
||||
let key = (agent_id.to_owned(), session_id);
|
||||
// Compute the result inside a block so the DashMap entry guard drops
|
||||
// before we return, releasing the shard lock as early as possible.
|
||||
@ -245,12 +348,62 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn rate_limiter_creates_bucket_lazily() {
|
||||
let rl = RateLimiter::new(RateLimiterConfig::default());
|
||||
// A *limited* config creates the bucket on first access.
|
||||
let rl = RateLimiter::new(RateLimiterConfig::limited(10.0, 10.0));
|
||||
assert_eq!(rl.active_buckets(), 0);
|
||||
let _ = rl.check("agent-a", 1);
|
||||
assert_eq!(rl.active_buckets(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unlimited_config_takes_fast_path_without_allocating() {
|
||||
// The default (unlimited) config must never touch the DashMap: the
|
||||
// fast path returns Ok(()) before constructing the key, so no bucket
|
||||
// is ever created no matter how many checks run.
|
||||
let rl = RateLimiter::new(RateLimiterConfig::default());
|
||||
assert!(!rl.enabled, "default config must be disabled (unlimited)");
|
||||
for i in 0..1000 {
|
||||
assert!(rl.check("agent-a", i).is_ok());
|
||||
}
|
||||
assert_eq!(
|
||||
rl.active_buckets(),
|
||||
0,
|
||||
"unlimited mode must not allocate buckets"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn zero_rate_is_rejected_at_construction() {
|
||||
// A 0.0 rate would lock the agent out forever and emit a u64::MAX
|
||||
// retry hint — it must be rejected eagerly, not at first write.
|
||||
assert!(RateLimiterConfig::try_limited(0.0, 10.0).is_err());
|
||||
assert!(RateLimiter::try_new(RateLimiterConfig::limited(0.0, 10.0)).is_err());
|
||||
assert!(RateLimiterConfig::limited(0.0, 10.0).validate().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn degenerate_rates_rejected_but_infinity_accepted() {
|
||||
// Non-positive / NaN sustained rates are rejected.
|
||||
assert!(RateLimiterConfig::try_limited(-1.0, 10.0).is_err());
|
||||
assert!(RateLimiterConfig::try_limited(f64::NAN, 10.0).is_err());
|
||||
// Sub-unit / NaN burst is rejected (can never hold the one token needed).
|
||||
assert!(RateLimiterConfig::try_limited(10.0, 0.5).is_err());
|
||||
assert!(RateLimiterConfig::try_limited(10.0, f64::NAN).is_err());
|
||||
// INFINITY is the explicit unlimited sentinel and must validate.
|
||||
assert!(RateLimiterConfig::try_limited(f64::INFINITY, f64::INFINITY).is_ok());
|
||||
assert!(RateLimiterConfig::default().validate().is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_falls_back_to_unlimited_on_degenerate_config() {
|
||||
// The infallible `new` must never produce a lock-everyone-out limiter:
|
||||
// a degenerate config degrades to unlimited (enabled == false).
|
||||
let rl = RateLimiter::new(RateLimiterConfig::limited(0.0, 0.0));
|
||||
assert!(!rl.enabled, "degenerate config must fall back to unlimited");
|
||||
assert!(rl.check("agent-a", 1).is_ok());
|
||||
assert_eq!(rl.active_buckets(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rate_limiter_separate_buckets_per_session() {
|
||||
let config = RateLimiterConfig {
|
||||
@ -269,7 +422,12 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn rate_limiter_remove_cleans_up() {
|
||||
let rl = RateLimiter::new(RateLimiterConfig::default());
|
||||
// A LIMITED config so a bucket is actually created (the default unlimited
|
||||
// config now short-circuits in `check` without touching the map).
|
||||
let rl = RateLimiter::new(RateLimiterConfig {
|
||||
signals_per_second: 1.0,
|
||||
burst_capacity: 1.0,
|
||||
});
|
||||
let _ = rl.check("agent-a", 1);
|
||||
assert_eq!(rl.active_buckets(), 1);
|
||||
rl.remove("agent-a", 1);
|
||||
|
||||
@ -43,6 +43,23 @@ pub(crate) fn scan_candidates(
|
||||
}
|
||||
|
||||
/// Rank candidates by a specific signal's decay score.
|
||||
///
|
||||
/// # Cost
|
||||
///
|
||||
/// This performs an `O(N)` scan over the **entire** signal ledger (every
|
||||
/// `(entity, signal_type)` cell), filtering to the requested signal type. There
|
||||
/// is no secondary per-signal-type index that would let us avoid the scan, so
|
||||
/// the `O(N)` pass is inherent to the current ledger layout — `N` is the number
|
||||
/// of live entity/signal cells, not the catalog size, and the per-cell work is a
|
||||
/// single `current_score` evaluation. To keep this off the heap, the working
|
||||
/// set is **bounded to `O(cap)`** where `cap` is the top-K candidate budget:
|
||||
/// rather than collecting every matching cell into an unbounded `Vec` and
|
||||
/// partitioning once at the end, we cap the buffer at `2 * cap` and
|
||||
/// partition-then-truncate to `cap` whenever it fills. That is amortized `O(N)`
|
||||
/// time with `O(cap)` memory instead of `O(N)` memory — important when one
|
||||
/// signal type dominates a large ledger. If this scan ever shows up in a
|
||||
/// profile, add a per-signal-type inverted index (entity IDs sorted by decay
|
||||
/// score) updated on the write path.
|
||||
pub(crate) fn signal_ranked_candidates(
|
||||
ledger: &SignalLedger,
|
||||
signal_name: &str,
|
||||
@ -54,7 +71,19 @@ pub(crate) fn signal_ranked_candidates(
|
||||
};
|
||||
|
||||
let now_ns = Timestamp::now().as_nanos();
|
||||
let mut scored: Vec<(EntityId, f64)> = Vec::new();
|
||||
|
||||
// Descending by score; NaN (degenerate decay math) falls back to descending
|
||||
// entity-ID order so output is deterministic and NaNs sink to the bottom.
|
||||
let cmp = |a: &(EntityId, f64), b: &(EntityId, f64)| {
|
||||
b.1.partial_cmp(&a.1)
|
||||
.unwrap_or_else(|| b.0.as_u64().cmp(&a.0.as_u64()))
|
||||
};
|
||||
|
||||
// Bounded buffer: never exceeds 2 * max_candidates entries. When it fills, an
|
||||
// O(buffer-len) select-and-truncate keeps the top max_candidates and drops the
|
||||
// rest, so peak memory is O(max_candidates) regardless of ledger size.
|
||||
let buffer_cap = max_candidates.saturating_mul(2).max(max_candidates + 1);
|
||||
let mut scored: Vec<(EntityId, f64)> = Vec::with_capacity(buffer_cap);
|
||||
|
||||
for entry in ledger.entries() {
|
||||
let (entity_id, signal_type_id) = entry.key();
|
||||
@ -65,17 +94,17 @@ pub(crate) fn signal_ranked_candidates(
|
||||
// happens in Stage 3 via ProfileExecutor.
|
||||
let score = entry.value().hot.current_score(0, now_ns, 0.0);
|
||||
scored.push((*entity_id, score));
|
||||
|
||||
if scored.len() >= buffer_cap {
|
||||
// Partition so the top `max_candidates` occupy [0, max_candidates)
|
||||
// then discard the tail, keeping the buffer bounded.
|
||||
scored.select_nth_unstable_by(max_candidates - 1, cmp);
|
||||
scored.truncate(max_candidates);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// O(N) partition to top-K via `select_nth_unstable_by`, then O(K log K)
|
||||
// sort of just the top K. Better than O(N log N) full sort when N >> limit.
|
||||
// NaN scores (from degenerate decay math) fall back to entity-ID order for
|
||||
// deterministic output.
|
||||
let cmp = |a: &(EntityId, f64), b: &(EntityId, f64)| {
|
||||
b.1.partial_cmp(&a.1)
|
||||
.unwrap_or_else(|| b.0.as_u64().cmp(&a.0.as_u64()))
|
||||
};
|
||||
// Final top-K: partition (O(N)) then sort just the survivors (O(K log K)).
|
||||
if scored.len() > max_candidates {
|
||||
scored.select_nth_unstable_by(max_candidates - 1, cmp);
|
||||
scored.truncate(max_candidates);
|
||||
|
||||
@ -134,6 +134,19 @@ impl RetrieveExecutor<'_> {
|
||||
}
|
||||
|
||||
/// Re-score candidates using cohort signal values and re-sort.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns [`QueryError::InvalidFilter`] when a boost references an
|
||||
/// aggregation that has no executor implementation (`Ratio` /
|
||||
/// `RelativeVelocity`). This mirrors the main scoring path
|
||||
/// (`ranking::executor::helpers::read_agg`), which surfaces the same
|
||||
/// aggregations as an internal invariant violation rather than silently
|
||||
/// scoring `0.0`. The profile registry rejects these aggregations up front
|
||||
/// (`ProfileError::UnsupportedAggregation`), so reaching the error arm means
|
||||
/// a profile bypassed validation — failing loud beats a silently-wrong
|
||||
/// cohort ranking. Keeping this consistent with the main path is the point:
|
||||
/// the cohort override must not score `0.0` where direct scoring would error.
|
||||
#[allow(clippy::unused_self)]
|
||||
pub(super) fn rescore_with_cohort(
|
||||
&self,
|
||||
@ -141,12 +154,14 @@ impl RetrieveExecutor<'_> {
|
||||
cohort_name: &str,
|
||||
cohort_ledger: &crate::cohort::CohortSignalLedger,
|
||||
boosts: &[crate::ranking::profile::Boost],
|
||||
) {
|
||||
) -> Result<(), QueryError> {
|
||||
use crate::ranking::profile::SignalAgg;
|
||||
|
||||
for candidate in scored.iter_mut() {
|
||||
let mut cohort_score = 0.0;
|
||||
for boost in boosts {
|
||||
let value = match &boost.agg {
|
||||
crate::ranking::profile::SignalAgg::Value => {
|
||||
SignalAgg::Value => {
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
let count = cohort_ledger
|
||||
.read_windowed_count(
|
||||
@ -158,7 +173,7 @@ impl RetrieveExecutor<'_> {
|
||||
.unwrap_or(0) as f64;
|
||||
count
|
||||
}
|
||||
crate::ranking::profile::SignalAgg::Velocity => cohort_ledger
|
||||
SignalAgg::Velocity => cohort_ledger
|
||||
.read_velocity(
|
||||
cohort_name,
|
||||
candidate.entity_id,
|
||||
@ -166,11 +181,20 @@ impl RetrieveExecutor<'_> {
|
||||
boost.window,
|
||||
)
|
||||
.unwrap_or(0.0),
|
||||
crate::ranking::profile::SignalAgg::DecayScore => cohort_ledger
|
||||
SignalAgg::DecayScore => cohort_ledger
|
||||
.read_decay_score(cohort_name, candidate.entity_id, &boost.signal, 0)
|
||||
.unwrap_or(None)
|
||||
.unwrap_or(0.0),
|
||||
_ => 0.0,
|
||||
// `Ratio` / `RelativeVelocity` have no executor implementation.
|
||||
// The registry rejects them at profile registration, so this is
|
||||
// an invariant violation if reached — surface it the same way
|
||||
// the main scoring path does instead of silently scoring 0.0.
|
||||
agg @ (SignalAgg::Ratio | SignalAgg::RelativeVelocity) => {
|
||||
return Err(QueryError::InvalidFilter {
|
||||
field: format!("cohort boost '{}'", boost.signal),
|
||||
reason: format!("aggregation '{}' is not implemented", agg.label()),
|
||||
});
|
||||
}
|
||||
};
|
||||
cohort_score += value * boost.weight;
|
||||
}
|
||||
@ -190,6 +214,8 @@ impl RetrieveExecutor<'_> {
|
||||
.partial_cmp(&a.score)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Apply notification caps as a post-diversity pass.
|
||||
|
||||
@ -22,6 +22,9 @@ pub mod user_filter;
|
||||
mod helpers;
|
||||
mod pipeline;
|
||||
|
||||
#[cfg(test)]
|
||||
mod test_support;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
@ -460,7 +463,7 @@ impl<'a> RetrieveExecutor<'a> {
|
||||
});
|
||||
}
|
||||
if cohort_exists {
|
||||
self.rescore_with_cohort(&mut scored, cohort_name, cohort_ledger, &profile.boosts);
|
||||
self.rescore_with_cohort(&mut scored, cohort_name, cohort_ledger, &profile.boosts)?;
|
||||
}
|
||||
}
|
||||
|
||||
@ -494,12 +497,23 @@ impl<'a> RetrieveExecutor<'a> {
|
||||
|
||||
/// Extract `SocialGraph { user_id, depth }` parameters from a filter expression.
|
||||
///
|
||||
/// Walks the top-level expression (including `And` lists) and returns the first
|
||||
/// Walks `And` AND `Or` lists (deliberately NOT `Not`) and returns the first
|
||||
/// `SocialGraph` variant found. Returns `None` if no social graph filter is present.
|
||||
///
|
||||
/// This MUST share the exact traversal discipline of the Stage 2.5 filter
|
||||
/// extractor [`user_filter::extract_user_state_filters`], which recurses into
|
||||
/// both `And` and `Or` (skipping `Not`). If this resolver only walked `And`, an
|
||||
/// `Or`-nested `SocialGraph` would fire the inclusion FILTER (Stage 2.5) but not
|
||||
/// the trending-SCOPE resolution (`with_social_graph_users`) — the filtering and
|
||||
/// scoring scopes would disagree about which social graph is active. `Not` is
|
||||
/// skipped because a negated deferred filter has no positive scope and is
|
||||
/// rejected up front by `reject_negated_deferred_filters`.
|
||||
fn extract_social_graph_params(expr: &FilterExpr) -> Option<(u64, u8)> {
|
||||
match expr {
|
||||
FilterExpr::SocialGraph { user_id, depth } => Some((*user_id, *depth)),
|
||||
FilterExpr::And(exprs) => exprs.iter().find_map(extract_social_graph_params),
|
||||
FilterExpr::And(exprs) | FilterExpr::Or(exprs) => {
|
||||
exprs.iter().find_map(extract_social_graph_params)
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
@ -50,10 +50,12 @@ impl RetrieveExecutor<'_> {
|
||||
|
||||
// Reject filter shapes with no sound semantics BEFORE any stage runs, so
|
||||
// a bad query fails loudly instead of silently returning the inverse set
|
||||
// (negated deferred post-filters) or AND-instead-of-OR (OR-of-thresholds).
|
||||
// (negated deferred post-filters) or AND-instead-of-OR (any deferred
|
||||
// post-filter under an OR — they resolve to the full universe in Stage 2
|
||||
// and are then applied conjunctively, so the OR would collapse to AND).
|
||||
if let Some(ref filter_expr) = combined_filter {
|
||||
user_filter::reject_negated_deferred_filters(filter_expr)?;
|
||||
user_filter::reject_or_of_signal_thresholds(filter_expr)?;
|
||||
user_filter::reject_or_of_deferred_filters(filter_expr)?;
|
||||
}
|
||||
|
||||
tracing::trace!(
|
||||
@ -84,17 +86,38 @@ impl RetrieveExecutor<'_> {
|
||||
// items using the CreatorItemsBitmap. This is O(k) where k = the
|
||||
// creator's item count, avoiding the O(N) scan cap that would miss
|
||||
// items at high IDs in a large catalog.
|
||||
//
|
||||
// The two `for_creator` arms below are kept SEPARATE on purpose. A single
|
||||
// `if let creator_id && let creator_items` chain would make the whole
|
||||
// condition false when `for_creator` is set but `creator_items` is absent,
|
||||
// falling through to the full-universe scan and SILENTLY dropping the
|
||||
// creator restriction (returning every creator's items). Instead, when the
|
||||
// scope is requested but cannot be satisfied, we warn and return an EMPTY
|
||||
// candidate set — the requested scope yields no items rather than the
|
||||
// inverse (all items). Mirrors the warn-on-missing-index pattern in Stage 2.
|
||||
let has_user_context = query.for_user.is_some();
|
||||
let mut candidates = if let Some(creator_id) = query.for_creator
|
||||
&& let Some(creator_items) = self.creator_items
|
||||
{
|
||||
creator_items
|
||||
.get(creator_id.as_u64())
|
||||
.map_or_else(Vec::new, |bm| {
|
||||
bm.iter()
|
||||
.map(|id_u32| EntityId::new(u64::from(id_u32)))
|
||||
.collect()
|
||||
})
|
||||
let mut candidates = if let Some(creator_id) = query.for_creator {
|
||||
if let Some(creator_items) = self.creator_items {
|
||||
creator_items
|
||||
.get(creator_id.as_u64())
|
||||
.map_or_else(Vec::new, |bm| {
|
||||
bm.iter()
|
||||
.map(|id_u32| EntityId::new(u64::from(id_u32)))
|
||||
.collect()
|
||||
})
|
||||
} else {
|
||||
warnings.push(
|
||||
"FOR CREATOR scope requested but the creator-items index is unavailable; \
|
||||
returning no candidates (refusing to fall back to a full-universe scan \
|
||||
that would ignore the creator restriction)"
|
||||
.to_string(),
|
||||
);
|
||||
tracing::warn!(
|
||||
creator_id = creator_id.as_u64(),
|
||||
"for_creator scope unsatisfiable: creator_items index absent; returning empty"
|
||||
);
|
||||
Vec::new()
|
||||
}
|
||||
} else {
|
||||
match &profile.candidate_strategy {
|
||||
CandidateStrategy::Scan { .. } => {
|
||||
@ -162,9 +185,28 @@ impl RetrieveExecutor<'_> {
|
||||
}
|
||||
|
||||
// M7p2: ReducedCandidates — cap scan set to reduce Stage 3 scoring work.
|
||||
//
|
||||
// A blind `truncate(cap)` keeps the FIRST `cap` candidates in scan order.
|
||||
// `scan_candidates` iterates the universe bitmap in ASCENDING entity-ID
|
||||
// order, so a blind truncate keeps the LOWEST IDs. For an ID-ordered
|
||||
// profile that ranks by DESCENDING entity ID (e.g. `New`), the highest IDs
|
||||
// are exactly the items that should rank at the top — so a blind truncate
|
||||
// drops the top-ranked items before scoring and corrupts the "new"
|
||||
// ranking. We therefore truncate by the profile's primary ordering key:
|
||||
// for `New`, keep the `cap` HIGHEST IDs via an O(N) `select_nth_unstable`
|
||||
// partition (no full sort; Stage 3 still orders the survivors). Other
|
||||
// profiles score on signal state rather than scan position, so their
|
||||
// truncation order is not load-bearing and a plain truncate is retained.
|
||||
if self.degradation_level.reduces_candidates() {
|
||||
let cap = (query.limit * 4).max(100);
|
||||
candidates.truncate(cap);
|
||||
if candidates.len() > cap {
|
||||
if matches!(profile.sort, Some(Sort::New)) {
|
||||
// Partition so the `cap` highest IDs occupy [0, cap): the items
|
||||
// that rank highest under descending-ID order survive.
|
||||
candidates.select_nth_unstable_by(cap - 1, |a, b| b.as_u64().cmp(&a.as_u64()));
|
||||
}
|
||||
candidates.truncate(cap);
|
||||
}
|
||||
}
|
||||
|
||||
stats.candidates_considered = candidates.len();
|
||||
@ -434,7 +476,10 @@ impl RetrieveExecutor<'_> {
|
||||
);
|
||||
|
||||
// ── Stage 5: Result Assembly ────────────────────────────────────
|
||||
let offset = query.cursor.as_ref().map_or(0, Cursor::offset);
|
||||
// `Cursor::offset` is fallible: a keyset cursor used as an offset (or a
|
||||
// 32-bit-overflowing value) is a loud `InvalidCursor` rather than a silent
|
||||
// misread.
|
||||
let offset = query.cursor.as_ref().map_or(Ok(0), Cursor::offset)?;
|
||||
// Guard against cursor-offset overflow: an attacker-supplied cursor can
|
||||
// carry an `offset` near `usize::MAX`, and `offset + limit` would wrap.
|
||||
// `saturating_add` clamps to `usize::MAX`; the `.min(len)` then keeps the
|
||||
|
||||
98
tidal/src/query/executor/test_support.rs
Normal file
98
tidal/src/query/executor/test_support.rs
Normal file
@ -0,0 +1,98 @@
|
||||
//! Shared test fixtures for the RETRIEVE executor unit tests.
|
||||
//!
|
||||
//! `tests.rs` and `tests_part2.rs` were split solely to keep each file under the
|
||||
//! ≤ 600-line cap (`CODING_GUIDELINES` §9), but both need the SAME scaffolding:
|
||||
//! a minimal schema, the built-in profile registry, an item-insertion helper,
|
||||
//! and an executor constructor. Rather than copy-paste these verbatim across the
|
||||
//! two modules (a DRY hazard — a fix to one fixture silently skips the other),
|
||||
//! they live here once and are imported via `use super::test_support::*`.
|
||||
|
||||
#![cfg(test)]
|
||||
#![allow(clippy::unwrap_used)]
|
||||
|
||||
use std::{sync::RwLock, time::Duration};
|
||||
|
||||
use roaring::RoaringBitmap;
|
||||
|
||||
use super::RetrieveExecutor;
|
||||
use crate::{
|
||||
ranking::{builtins::register_builtins, registry::ProfileRegistry},
|
||||
schema::{DecaySpec, EntityKind, SchemaBuilder, Timestamp, Window},
|
||||
signals::SignalLedger,
|
||||
storage::indexes::{bitmap::BitmapIndex, range::RangeIndex},
|
||||
};
|
||||
|
||||
/// Minimal item schema: the generic signal vocabulary the built-in profiles read.
|
||||
pub(super) fn test_schema() -> crate::schema::Schema {
|
||||
let mut builder = SchemaBuilder::new();
|
||||
for sig in &["view", "like", "share", "skip", "completion"] {
|
||||
let _ = builder
|
||||
.signal(
|
||||
sig,
|
||||
EntityKind::Item,
|
||||
DecaySpec::Exponential {
|
||||
half_life: Duration::from_secs(7 * 24 * 3600),
|
||||
},
|
||||
)
|
||||
.windows(&[Window::OneHour, Window::TwentyFourHours, Window::SevenDays])
|
||||
.velocity(true)
|
||||
.add();
|
||||
}
|
||||
builder.build().unwrap()
|
||||
}
|
||||
|
||||
/// Profile registry pre-loaded with the built-in profiles.
|
||||
pub(super) fn setup_registry() -> ProfileRegistry {
|
||||
let mut reg = ProfileRegistry::new();
|
||||
register_builtins(&mut reg).unwrap();
|
||||
reg
|
||||
}
|
||||
|
||||
/// Add an item to the in-memory indexes and universe bitmap.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) fn add_item(
|
||||
category_idx: &BitmapIndex,
|
||||
format_idx: &BitmapIndex,
|
||||
creator_idx: &BitmapIndex,
|
||||
duration_idx: &RangeIndex<u32>,
|
||||
created_at_idx: &RangeIndex<u64>,
|
||||
universe: &mut RoaringBitmap,
|
||||
id: u64,
|
||||
category: &str,
|
||||
format: &str,
|
||||
creator: u64,
|
||||
) {
|
||||
let id_u32 = id as u32;
|
||||
category_idx.insert(id_u32, category);
|
||||
format_idx.insert(id_u32, format);
|
||||
creator_idx.insert(id_u32, creator.to_string());
|
||||
duration_idx.insert(id_u32, 0u32);
|
||||
created_at_idx.insert(id_u32, Timestamp::now().as_nanos());
|
||||
universe.insert(id_u32);
|
||||
}
|
||||
|
||||
/// Build an executor from test indexes.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) fn make_executor<'a>(
|
||||
ledger: &'a SignalLedger,
|
||||
profile_reg: &'a ProfileRegistry,
|
||||
cat: &'a BitmapIndex,
|
||||
fmt: &'a BitmapIndex,
|
||||
creator: &'a BitmapIndex,
|
||||
tag: &'a BitmapIndex,
|
||||
dur: &'a RangeIndex<u32>,
|
||||
ts: &'a RangeIndex<u64>,
|
||||
universe: &'a RwLock<RoaringBitmap>,
|
||||
) -> RetrieveExecutor<'a> {
|
||||
RetrieveExecutor::new(
|
||||
ledger,
|
||||
profile_reg,
|
||||
Some(cat),
|
||||
Some(fmt),
|
||||
Some(creator),
|
||||
Some(tag),
|
||||
Some(dur),
|
||||
Some(ts),
|
||||
Some(universe),
|
||||
)
|
||||
}
|
||||
@ -3,93 +3,19 @@
|
||||
#![cfg(test)]
|
||||
#![allow(clippy::unwrap_used)]
|
||||
|
||||
use std::{sync::RwLock, time::Duration};
|
||||
use std::sync::RwLock;
|
||||
|
||||
use roaring::RoaringBitmap;
|
||||
|
||||
use super::*;
|
||||
use super::test_support::{add_item, make_executor, setup_registry, test_schema};
|
||||
use crate::{
|
||||
entities::{CreatorItemsBitmap, HardNegIndex, InteractionLedger, UserStateIndex},
|
||||
query::retrieve::{QueryError, Retrieve},
|
||||
ranking::{builtins::register_builtins, registry::ProfileRegistry},
|
||||
schema::{DecaySpec, EntityId, EntityKind, SchemaBuilder, Timestamp, Window},
|
||||
schema::EntityId,
|
||||
signals::{NoopWalWriter, SignalLedger},
|
||||
storage::indexes::{bitmap::BitmapIndex, filter::FilterExpr, range::RangeIndex},
|
||||
};
|
||||
|
||||
fn test_schema() -> crate::schema::Schema {
|
||||
let mut builder = SchemaBuilder::new();
|
||||
for sig in &["view", "like", "share", "skip", "completion"] {
|
||||
let _ = builder
|
||||
.signal(
|
||||
sig,
|
||||
EntityKind::Item,
|
||||
DecaySpec::Exponential {
|
||||
half_life: Duration::from_secs(7 * 24 * 3600),
|
||||
},
|
||||
)
|
||||
.windows(&[Window::OneHour, Window::TwentyFourHours, Window::SevenDays])
|
||||
.velocity(true)
|
||||
.add();
|
||||
}
|
||||
builder.build().unwrap()
|
||||
}
|
||||
|
||||
fn setup_registry() -> ProfileRegistry {
|
||||
let mut reg = ProfileRegistry::new();
|
||||
register_builtins(&mut reg).unwrap();
|
||||
reg
|
||||
}
|
||||
|
||||
/// Add an item to the in-memory indexes and universe bitmap.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn add_item(
|
||||
category_idx: &BitmapIndex,
|
||||
format_idx: &BitmapIndex,
|
||||
creator_idx: &BitmapIndex,
|
||||
duration_idx: &RangeIndex<u32>,
|
||||
created_at_idx: &RangeIndex<u64>,
|
||||
universe: &mut RoaringBitmap,
|
||||
id: u64,
|
||||
category: &str,
|
||||
format: &str,
|
||||
creator: u64,
|
||||
) {
|
||||
let id_u32 = id as u32;
|
||||
category_idx.insert(id_u32, category);
|
||||
format_idx.insert(id_u32, format);
|
||||
creator_idx.insert(id_u32, creator.to_string());
|
||||
duration_idx.insert(id_u32, 0u32);
|
||||
created_at_idx.insert(id_u32, Timestamp::now().as_nanos());
|
||||
universe.insert(id_u32);
|
||||
}
|
||||
|
||||
/// Helper: build executor from test indexes.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn make_executor<'a>(
|
||||
ledger: &'a SignalLedger,
|
||||
profile_reg: &'a ProfileRegistry,
|
||||
cat: &'a BitmapIndex,
|
||||
fmt: &'a BitmapIndex,
|
||||
creator: &'a BitmapIndex,
|
||||
tag: &'a BitmapIndex,
|
||||
dur: &'a RangeIndex<u32>,
|
||||
ts: &'a RangeIndex<u64>,
|
||||
universe: &'a RwLock<RoaringBitmap>,
|
||||
) -> RetrieveExecutor<'a> {
|
||||
RetrieveExecutor::new(
|
||||
ledger,
|
||||
profile_reg,
|
||||
Some(cat),
|
||||
Some(fmt),
|
||||
Some(creator),
|
||||
Some(tag),
|
||||
Some(dur),
|
||||
Some(ts),
|
||||
Some(universe),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_database_returns_empty_results() {
|
||||
let schema = test_schema();
|
||||
@ -441,3 +367,141 @@ fn saved_filter_retains_only_saved_items() {
|
||||
assert!(ids.contains(&3));
|
||||
assert!(ids.contains(&7));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn for_creator_without_creator_items_returns_empty_not_full_scan() {
|
||||
// Regression: when FOR CREATOR is requested but the creator-items index is
|
||||
// absent, the executor must return NO candidates (the scope cannot be
|
||||
// satisfied) and warn — never fall through to a full-universe scan that
|
||||
// silently drops the creator restriction and returns every creator's items.
|
||||
let schema = test_schema();
|
||||
let ledger = SignalLedger::new(schema, Box::new(NoopWalWriter));
|
||||
let profile_reg = setup_registry();
|
||||
let cat = BitmapIndex::new("category");
|
||||
let fmt = BitmapIndex::new("format");
|
||||
let creator_idx = BitmapIndex::new("creator");
|
||||
let tag = BitmapIndex::new("tags");
|
||||
let dur: RangeIndex<u32> = RangeIndex::new("duration");
|
||||
let ts: RangeIndex<u64> = RangeIndex::new("created_at");
|
||||
let mut universe_bm = RoaringBitmap::new();
|
||||
|
||||
// 10 items, all owned by creator 1 in the creator bitmap index.
|
||||
for i in 1..=10u64 {
|
||||
add_item(
|
||||
&cat,
|
||||
&fmt,
|
||||
&creator_idx,
|
||||
&dur,
|
||||
&ts,
|
||||
&mut universe_bm,
|
||||
i,
|
||||
"jazz",
|
||||
"video",
|
||||
1,
|
||||
);
|
||||
}
|
||||
|
||||
let universe = RwLock::new(universe_bm);
|
||||
|
||||
// `make_executor` wires every bitmap/range index but NOT `creator_items`
|
||||
// (that is only set via `with_user_context`), so `self.creator_items` is
|
||||
// None here — exactly the condition the fix guards.
|
||||
let exec = make_executor(
|
||||
&ledger,
|
||||
&profile_reg,
|
||||
&cat,
|
||||
&fmt,
|
||||
&creator_idx,
|
||||
&tag,
|
||||
&dur,
|
||||
&ts,
|
||||
&universe,
|
||||
);
|
||||
|
||||
let query = Retrieve::builder()
|
||||
.profile("new")
|
||||
.for_creator(EntityId::new(1))
|
||||
.limit(20)
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
let results = exec.execute(&query).unwrap();
|
||||
assert!(
|
||||
results.items.is_empty(),
|
||||
"for_creator with no creator-items index must yield no candidates, got {} items",
|
||||
results.items.len()
|
||||
);
|
||||
assert!(
|
||||
results
|
||||
.warnings
|
||||
.iter()
|
||||
.any(|w| w.contains("FOR CREATOR scope requested")),
|
||||
"expected a warning explaining the unsatisfiable for_creator scope, got {:?}",
|
||||
results.warnings
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_degradation_truncation_keeps_top_ranked_new_items() {
|
||||
// Regression: under ReducedCandidates load, the candidate cap is
|
||||
// `(limit*4).max(100)`. `scan_candidates` iterates the universe in ASCENDING
|
||||
// ID order, so a blind `truncate(cap)` would keep the LOWEST IDs. The "new"
|
||||
// profile sorts by DESCENDING ID, so the highest IDs are the top-ranked
|
||||
// items — a blind truncate would drop exactly the items that should rank #1.
|
||||
// The fix truncates by descending ID for `New`, so the highest IDs survive.
|
||||
let schema = test_schema();
|
||||
let ledger = SignalLedger::new(schema, Box::new(NoopWalWriter));
|
||||
let profile_reg = setup_registry();
|
||||
let cat = BitmapIndex::new("category");
|
||||
let fmt = BitmapIndex::new("format");
|
||||
let creator_idx = BitmapIndex::new("creator");
|
||||
let tag = BitmapIndex::new("tags");
|
||||
let dur: RangeIndex<u32> = RangeIndex::new("duration");
|
||||
let ts: RangeIndex<u64> = RangeIndex::new("created_at");
|
||||
let mut universe_bm = RoaringBitmap::new();
|
||||
|
||||
// 150 items: more than the degradation cap of (1*4).max(100) = 100, so the
|
||||
// truncation path is exercised. The highest IDs (101..=150) would be dropped
|
||||
// by a blind ascending-order truncate.
|
||||
for i in 1..=150u64 {
|
||||
add_item(
|
||||
&cat,
|
||||
&fmt,
|
||||
&creator_idx,
|
||||
&dur,
|
||||
&ts,
|
||||
&mut universe_bm,
|
||||
i,
|
||||
"jazz",
|
||||
"video",
|
||||
1,
|
||||
);
|
||||
}
|
||||
|
||||
let universe = RwLock::new(universe_bm);
|
||||
let exec = make_executor(
|
||||
&ledger,
|
||||
&profile_reg,
|
||||
&cat,
|
||||
&fmt,
|
||||
&creator_idx,
|
||||
&tag,
|
||||
&dur,
|
||||
&ts,
|
||||
&universe,
|
||||
)
|
||||
.with_degradation_level(crate::load::DegradationLevel::ReducedCandidates);
|
||||
|
||||
let query = Retrieve::builder().profile("new").limit(1).build().unwrap();
|
||||
let results = exec.execute(&query).unwrap();
|
||||
|
||||
// The #1 "new" item is the highest ID (150). A blind truncate would have
|
||||
// dropped 101..=150 before scoring and returned 100 instead.
|
||||
assert_eq!(results.items.len(), 1);
|
||||
assert_eq!(
|
||||
results.items[0].entity_id,
|
||||
EntityId::new(150),
|
||||
"top 'new' item must be the highest ID, not the highest surviving \
|
||||
ascending-order id after truncation"
|
||||
);
|
||||
}
|
||||
|
||||
@ -1,20 +1,21 @@
|
||||
//! Executor unit tests — part 2.
|
||||
//!
|
||||
//! Sibling to `tests.rs`; test helpers are re-declared here.
|
||||
//! This keeps each file ≤ 600 lines.
|
||||
//! Sibling to `tests.rs`; both files share fixtures from `super::test_support`.
|
||||
//! The split exists only to keep each file ≤ 600 lines (`CODING_GUIDELINES` §9).
|
||||
|
||||
#![cfg(test)]
|
||||
#![allow(clippy::unwrap_used)]
|
||||
|
||||
use std::{sync::RwLock, time::Duration};
|
||||
use std::sync::RwLock;
|
||||
|
||||
use roaring::RoaringBitmap;
|
||||
|
||||
use super::*;
|
||||
use super::test_support::{add_item, make_executor, setup_registry, test_schema};
|
||||
use crate::{
|
||||
entities::{CreatorItemsBitmap, HardNegIndex, InteractionLedger, UserStateIndex},
|
||||
query::retrieve::Retrieve,
|
||||
ranking::{builtins::register_builtins, executor::ScoredCandidate, registry::ProfileRegistry},
|
||||
schema::{DecaySpec, EntityKind, SchemaBuilder, Timestamp, Window},
|
||||
ranking::executor::ScoredCandidate,
|
||||
schema::{EntityId, EntityKind, Timestamp},
|
||||
signals::{NoopWalWriter, SignalLedger},
|
||||
storage::{
|
||||
Tag, encode_key,
|
||||
@ -22,79 +23,6 @@ use crate::{
|
||||
},
|
||||
};
|
||||
|
||||
fn test_schema() -> crate::schema::Schema {
|
||||
let mut builder = SchemaBuilder::new();
|
||||
for sig in &["view", "like", "share", "skip", "completion"] {
|
||||
let _ = builder
|
||||
.signal(
|
||||
sig,
|
||||
EntityKind::Item,
|
||||
DecaySpec::Exponential {
|
||||
half_life: Duration::from_secs(7 * 24 * 3600),
|
||||
},
|
||||
)
|
||||
.windows(&[Window::OneHour, Window::TwentyFourHours, Window::SevenDays])
|
||||
.velocity(true)
|
||||
.add();
|
||||
}
|
||||
builder.build().unwrap()
|
||||
}
|
||||
|
||||
fn setup_registry() -> ProfileRegistry {
|
||||
let mut reg = ProfileRegistry::new();
|
||||
register_builtins(&mut reg).unwrap();
|
||||
reg
|
||||
}
|
||||
|
||||
/// Add an item to the in-memory indexes and universe bitmap.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn add_item(
|
||||
category_idx: &BitmapIndex,
|
||||
format_idx: &BitmapIndex,
|
||||
creator_idx: &BitmapIndex,
|
||||
duration_idx: &RangeIndex<u32>,
|
||||
created_at_idx: &RangeIndex<u64>,
|
||||
universe: &mut RoaringBitmap,
|
||||
id: u64,
|
||||
category: &str,
|
||||
format: &str,
|
||||
creator: u64,
|
||||
) {
|
||||
let id_u32 = id as u32;
|
||||
category_idx.insert(id_u32, category);
|
||||
format_idx.insert(id_u32, format);
|
||||
creator_idx.insert(id_u32, creator.to_string());
|
||||
duration_idx.insert(id_u32, 0u32);
|
||||
created_at_idx.insert(id_u32, Timestamp::now().as_nanos());
|
||||
universe.insert(id_u32);
|
||||
}
|
||||
|
||||
/// Helper: build executor from test indexes.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn make_executor<'a>(
|
||||
ledger: &'a SignalLedger,
|
||||
profile_reg: &'a ProfileRegistry,
|
||||
cat: &'a BitmapIndex,
|
||||
fmt: &'a BitmapIndex,
|
||||
creator: &'a BitmapIndex,
|
||||
tag: &'a BitmapIndex,
|
||||
dur: &'a RangeIndex<u32>,
|
||||
ts: &'a RangeIndex<u64>,
|
||||
universe: &'a RwLock<RoaringBitmap>,
|
||||
) -> RetrieveExecutor<'a> {
|
||||
RetrieveExecutor::new(
|
||||
ledger,
|
||||
profile_reg,
|
||||
Some(cat),
|
||||
Some(fmt),
|
||||
Some(creator),
|
||||
Some(tag),
|
||||
Some(dur),
|
||||
Some(ts),
|
||||
Some(universe),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn liked_filter_retains_only_liked_items() {
|
||||
let schema = test_schema();
|
||||
@ -360,97 +288,8 @@ fn signal_ranked_strategy_returns_candidates() {
|
||||
assert!(!results.items.is_empty());
|
||||
}
|
||||
|
||||
// ── Exploration injection tests ─────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn exploration_injects_random_candidates() {
|
||||
// Use inject_exploration directly to verify that exploration candidates
|
||||
// from the full candidate list appear in the scored output.
|
||||
|
||||
// Scored list: entities 1-5 (the "top ranked" items).
|
||||
let mut scored: Vec<ScoredCandidate> = (1..=5)
|
||||
.map(|i| ScoredCandidate {
|
||||
entity_id: EntityId::new(i),
|
||||
// i ∈ 1..=5 — exactly representable in f64; fixture score only.
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
score: (i as f64).mul_add(-0.1, 1.0),
|
||||
signal_snapshot: vec![],
|
||||
creator_id: None,
|
||||
format: None,
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Full candidate list: entities 1-10 (so 6-10 are available for exploration).
|
||||
let all_candidates: Vec<EntityId> = (1..=10).map(EntityId::new).collect();
|
||||
|
||||
// 40% exploration -> ceil(0.4 * 5) = 2 exploration slots.
|
||||
candidate_gen::inject_exploration(&mut scored, &all_candidates, 0.4);
|
||||
|
||||
assert_eq!(scored.len(), 5, "total output should remain 5");
|
||||
|
||||
// At least one of the exploration candidates (6-10) should be present.
|
||||
let exploration_ids: Vec<u64> = scored
|
||||
.iter()
|
||||
.filter(|c| c.entity_id.as_u64() > 5)
|
||||
.map(|c| c.entity_id.as_u64())
|
||||
.collect();
|
||||
assert_eq!(
|
||||
exploration_ids.len(),
|
||||
2,
|
||||
"expected 2 exploration candidates, got {}",
|
||||
exploration_ids.len()
|
||||
);
|
||||
|
||||
// Exploration candidates should have score 0.0.
|
||||
for c in &scored {
|
||||
if c.entity_id.as_u64() > 5 {
|
||||
assert!(
|
||||
c.score.abs() < 1e-12,
|
||||
"exploration candidates should have score 0.0"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exploration_zero_fraction_is_noop() {
|
||||
let mut scored: Vec<ScoredCandidate> = (1..=5)
|
||||
.map(|i| ScoredCandidate {
|
||||
entity_id: EntityId::new(i),
|
||||
score: 1.0,
|
||||
signal_snapshot: vec![],
|
||||
creator_id: None,
|
||||
format: None,
|
||||
})
|
||||
.collect();
|
||||
|
||||
let all: Vec<EntityId> = (1..=10).map(EntityId::new).collect();
|
||||
candidate_gen::inject_exploration(&mut scored, &all, 0.0);
|
||||
|
||||
// No change when exploration fraction is 0.
|
||||
assert_eq!(scored.len(), 5);
|
||||
assert!(scored.iter().all(|c| c.entity_id.as_u64() <= 5));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exploration_no_extra_candidates_is_noop() {
|
||||
// All candidates are already scored -- no exploration pool.
|
||||
let mut scored: Vec<ScoredCandidate> = (1..=5)
|
||||
.map(|i| ScoredCandidate {
|
||||
entity_id: EntityId::new(i),
|
||||
score: 1.0,
|
||||
signal_snapshot: vec![],
|
||||
creator_id: None,
|
||||
format: None,
|
||||
})
|
||||
.collect();
|
||||
|
||||
let all: Vec<EntityId> = (1..=5).map(EntityId::new).collect();
|
||||
candidate_gen::inject_exploration(&mut scored, &all, 0.5);
|
||||
|
||||
// No change when there are no extra candidates for exploration.
|
||||
assert_eq!(scored.len(), 5);
|
||||
}
|
||||
// Exploration-injection tests live in `candidate_gen::tests` (the module that
|
||||
// owns `inject_exploration`); they are intentionally NOT duplicated here.
|
||||
|
||||
// ── Cohort rescore coverage (WARNING: re-normalize after cohort rescore) ──────
|
||||
|
||||
@ -471,8 +310,9 @@ fn cohort_rescore_normalizes_score_to_unit_range_and_reorders() {
|
||||
// "now" so the decay score read back in rescore (which uses Timestamp::now)
|
||||
// is essentially undecayed and the magnitudes are comparable.
|
||||
let now_ns = Timestamp::now().as_nanos();
|
||||
cohort_ledger.record("tech", EntityId::new(1), view_id, 1.0, now_ns);
|
||||
cohort_ledger.record("tech", EntityId::new(2), view_id, 100.0, now_ns);
|
||||
let tech: std::sync::Arc<str> = std::sync::Arc::from("tech");
|
||||
cohort_ledger.record(&tech, EntityId::new(1), view_id, 1.0, now_ns);
|
||||
cohort_ledger.record(&tech, EntityId::new(2), view_id, 100.0, now_ns);
|
||||
|
||||
let registry = CohortRegistry::new();
|
||||
registry
|
||||
@ -532,7 +372,8 @@ fn cohort_rescore_normalizes_score_to_unit_range_and_reorders() {
|
||||
weight: 1.0,
|
||||
}];
|
||||
|
||||
exec.rescore_with_cohort(&mut scored, "tech", &cohort_ledger, &boosts);
|
||||
exec.rescore_with_cohort(&mut scored, "tech", &cohort_ledger, &boosts)
|
||||
.unwrap();
|
||||
|
||||
// After rescore: item 2 (high cohort view) must rank first, and EVERY score
|
||||
// must be back in [0, 1] (the WARNING was that rescore left un-normalized
|
||||
@ -767,3 +608,48 @@ fn notification_caps_limit_total_and_per_creator() {
|
||||
let ids: Vec<u64> = retained.iter().map(|c| c.entity_id.as_u64()).collect();
|
||||
assert_eq!(ids, vec![1, 3]);
|
||||
}
|
||||
|
||||
// ── Social-graph scope resolution traversal (mod.rs) ──────────────────────────
|
||||
|
||||
#[test]
|
||||
fn extract_social_graph_params_walks_and_or_but_not_not() {
|
||||
use super::extract_social_graph_params;
|
||||
|
||||
// Bare SocialGraph.
|
||||
assert_eq!(
|
||||
extract_social_graph_params(&FilterExpr::social_graph(42, 2)),
|
||||
Some((42, 2))
|
||||
);
|
||||
|
||||
// Nested under And (already worked before the fix).
|
||||
let and = FilterExpr::And(vec![
|
||||
FilterExpr::CategoryEq("jazz".into()),
|
||||
FilterExpr::social_graph(7, 1),
|
||||
]);
|
||||
assert_eq!(extract_social_graph_params(&and), Some((7, 1)));
|
||||
|
||||
// Nested under Or: the Stage 2.5 inclusion FILTER (`extract_user_state_filters`)
|
||||
// recurses into Or, so the scope resolver MUST too — otherwise filtering and
|
||||
// scoring scopes disagree. This is the case the fix added.
|
||||
let or = FilterExpr::Or(vec![
|
||||
FilterExpr::CategoryEq("jazz".into()),
|
||||
FilterExpr::social_graph(9, 2),
|
||||
]);
|
||||
assert_eq!(extract_social_graph_params(&or), Some((9, 2)));
|
||||
|
||||
// Deeply nested Or-within-And.
|
||||
let nested = FilterExpr::And(vec![
|
||||
FilterExpr::FormatEq("video".into()),
|
||||
FilterExpr::Or(vec![
|
||||
FilterExpr::CategoryEq("a".into()),
|
||||
FilterExpr::social_graph(3, 1),
|
||||
]),
|
||||
]);
|
||||
assert_eq!(extract_social_graph_params(&nested), Some((3, 1)));
|
||||
|
||||
// Under Not: deliberately NOT extracted (a negated deferred filter has no
|
||||
// positive scope; `reject_negated_deferred_filters` rejects such queries up
|
||||
// front, mirroring `extract_user_state_filters`).
|
||||
let negated = FilterExpr::Not(Box::new(FilterExpr::social_graph(5, 1)));
|
||||
assert_eq!(extract_social_graph_params(&negated), None);
|
||||
}
|
||||
|
||||
@ -127,7 +127,7 @@ fn collect_user_state_filters<'a>(expr: &'a FilterExpr, out: &mut Vec<&'a Filter
|
||||
/// These variants pass through the `FilterEvaluator` as "return full universe"
|
||||
/// and are applied as post-filters in Stage 2.3 of the executor pipeline, where
|
||||
/// the signal ledger is available. Thresholds are applied conjunctively, so
|
||||
/// OR-of-thresholds is rejected up front by [`reject_or_of_signal_thresholds`].
|
||||
/// OR-of-thresholds is rejected up front by [`reject_or_of_deferred_filters`].
|
||||
pub(crate) fn extract_signal_threshold_filters(expr: &FilterExpr) -> Vec<&FilterExpr> {
|
||||
let mut result = Vec::new();
|
||||
collect_signal_threshold_filters(expr, &mut result);
|
||||
@ -142,7 +142,7 @@ fn collect_signal_threshold_filters<'a>(expr: &'a FilterExpr, out: &mut Vec<&'a
|
||||
FilterExpr::And(children) | FilterExpr::Or(children) => {
|
||||
// All MinSignal/MaxSignal variants are extracted and applied
|
||||
// CONJUNCTIVELY in Stage 2.3. OR-of-threshold is therefore wrong
|
||||
// semantics if it reaches here, so `reject_or_of_signal_thresholds`
|
||||
// semantics if it reaches here, so `reject_or_of_deferred_filters`
|
||||
// (called before extraction) rejects it up front rather than letting
|
||||
// an OR be silently evaluated as AND. AND is walked normally.
|
||||
for child in children {
|
||||
@ -157,53 +157,73 @@ fn collect_signal_threshold_filters<'a>(expr: &'a FilterExpr, out: &mut Vec<&'a
|
||||
}
|
||||
}
|
||||
|
||||
/// Reject any `OR` node that contains a `MinSignal`/`MaxSignal` threshold filter.
|
||||
/// Reject any `OR` node that contains *any* deferred post-filter variant
|
||||
/// (`Saved`, `Liked`, `InProgress`, `SocialGraph`, `MinSignal`, `MaxSignal`,
|
||||
/// `NearLocation`, `InCollection` — see [`is_deferred_post_filter`]).
|
||||
///
|
||||
/// Stage 2.3 applies extracted thresholds conjunctively (every threshold must
|
||||
/// pass). An `OR` of thresholds therefore cannot be honored: it would be
|
||||
/// silently evaluated as an `AND`, returning fewer results than requested with
|
||||
/// no error. Until composable OR-over-thresholds is implemented, such a query is
|
||||
/// rejected so the caller never silently gets `AND` semantics.
|
||||
/// Deferred variants resolve to the full universe in Stage 2 (the
|
||||
/// `FilterEvaluator` cannot apply them at the bitmap layer), so an
|
||||
/// `Or([Saved(u), CategoryEq("jazz")])` evaluates to `universe ∪ jazz =
|
||||
/// universe` — every candidate retained. The Stage 2.x extractors then recurse
|
||||
/// INTO the `Or`, pull the deferred variant out, and apply it as an
|
||||
/// unconditional `candidates.retain(...)` — an *intersection*. The net effect is
|
||||
/// that `saved OR jazz` silently collapses to `saved` (the inverse of the
|
||||
/// requested union). Thresholds are the same trap: Stage 2.3 applies them
|
||||
/// conjunctively, so `OR` of thresholds is silently evaluated as `AND`.
|
||||
///
|
||||
/// Until a polarity/union-aware deferred evaluator exists, every such query is
|
||||
/// rejected so the caller never silently gets `AND` (intersection) semantics in
|
||||
/// place of the requested `OR`. Failing loudly beats returning the inverse set
|
||||
/// (see the module-level contract).
|
||||
///
|
||||
/// Index-backed `Or` (e.g. `Or([CategoryEq, FormatEq])`) is fine: those variants
|
||||
/// union correctly at the bitmap layer and are never deferred.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns [`QueryError::InvalidFilter`] if any `OR` subtree contains a signal
|
||||
/// threshold filter.
|
||||
pub(crate) fn reject_or_of_signal_thresholds(expr: &FilterExpr) -> Result<(), QueryError> {
|
||||
fn subtree_has_threshold(expr: &FilterExpr) -> bool {
|
||||
/// Returns [`QueryError::InvalidFilter`] naming the first deferred variant found
|
||||
/// inside any `OR` subtree.
|
||||
pub(crate) fn reject_or_of_deferred_filters(expr: &FilterExpr) -> Result<(), QueryError> {
|
||||
/// First deferred variant anywhere in `expr`'s subtree, or `None`.
|
||||
fn first_deferred_in_subtree(expr: &FilterExpr) -> Option<&FilterExpr> {
|
||||
if is_deferred_post_filter(expr) {
|
||||
return Some(expr);
|
||||
}
|
||||
match expr {
|
||||
FilterExpr::MinSignal { .. } | FilterExpr::MaxSignal { .. } => true,
|
||||
FilterExpr::And(children) | FilterExpr::Or(children) => {
|
||||
children.iter().any(subtree_has_threshold)
|
||||
children.iter().find_map(first_deferred_in_subtree)
|
||||
}
|
||||
FilterExpr::Not(inner) => subtree_has_threshold(inner),
|
||||
_ => false,
|
||||
FilterExpr::Not(inner) => first_deferred_in_subtree(inner),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
match expr {
|
||||
FilterExpr::Or(children) => {
|
||||
if children.iter().any(subtree_has_threshold) {
|
||||
if let Some(offender) = children.iter().find_map(first_deferred_in_subtree) {
|
||||
return Err(QueryError::InvalidFilter {
|
||||
field: "OR(MinSignal/MaxSignal)".to_string(),
|
||||
reason: "OR-composed signal threshold filters are not supported: thresholds \
|
||||
are applied conjunctively (AND), so an OR would be silently \
|
||||
evaluated as AND. Use separate AND-combined thresholds, or split \
|
||||
the query."
|
||||
field: format!("OR({offender:?})"),
|
||||
reason: "OR over a deferred post-filter (Saved/Liked/InProgress/\
|
||||
SocialGraph/MinSignal/MaxSignal/NearLocation/InCollection) is not \
|
||||
supported: these filters resolve to the full universe in Stage 2 and \
|
||||
are then applied conjunctively (AND) in a later stage, so an OR would \
|
||||
silently collapse to AND (an intersection, the inverse of the \
|
||||
requested union). Use AND-combined deferred filters, replace it with \
|
||||
an index-backed filter, or split the query."
|
||||
.to_string(),
|
||||
});
|
||||
}
|
||||
for child in children {
|
||||
reject_or_of_signal_thresholds(child)?;
|
||||
reject_or_of_deferred_filters(child)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
FilterExpr::And(children) => {
|
||||
for child in children {
|
||||
reject_or_of_signal_thresholds(child)?;
|
||||
reject_or_of_deferred_filters(child)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
FilterExpr::Not(inner) => reject_or_of_signal_thresholds(inner),
|
||||
FilterExpr::Not(inner) => reject_or_of_deferred_filters(inner),
|
||||
_ => Ok(()),
|
||||
}
|
||||
}
|
||||
@ -511,7 +531,12 @@ mod tests {
|
||||
assert!(reject_negated_deferred_filters(&expr).is_ok());
|
||||
}
|
||||
|
||||
// ── OR-of-threshold rejection ───────────────────────────────────────────
|
||||
// ── OR-of-deferred rejection ────────────────────────────────────────────
|
||||
//
|
||||
// Thresholds are deferred filters, so the threshold cases below are now a
|
||||
// special case of the general `reject_or_of_deferred_filters` guard. The
|
||||
// `*_variant_*` tests extend the same coverage to every other deferred
|
||||
// variant (BLOCKER 5: OR over them silently collapsed to AND).
|
||||
|
||||
#[test]
|
||||
fn reject_or_of_thresholds_errors() {
|
||||
@ -520,7 +545,7 @@ mod tests {
|
||||
FilterExpr::min_signal("like", 5.0),
|
||||
]);
|
||||
let err =
|
||||
reject_or_of_signal_thresholds(&expr).expect_err("OR-of-threshold must be rejected");
|
||||
reject_or_of_deferred_filters(&expr).expect_err("OR-of-threshold must be rejected");
|
||||
assert!(matches!(err, QueryError::InvalidFilter { .. }));
|
||||
}
|
||||
|
||||
@ -530,15 +555,77 @@ mod tests {
|
||||
FilterExpr::min_signal("view", 10.0),
|
||||
FilterExpr::max_signal("hide", 5.0),
|
||||
]);
|
||||
assert!(reject_or_of_signal_thresholds(&expr).is_ok());
|
||||
assert!(reject_or_of_deferred_filters(&expr).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reject_or_of_thresholds_allows_or_without_thresholds() {
|
||||
fn reject_or_of_deferred_allows_or_without_deferred() {
|
||||
// A legitimate index-backed `Or` (no deferred variant) unions correctly
|
||||
// at the bitmap layer and must still succeed.
|
||||
let expr = FilterExpr::Or(vec![
|
||||
FilterExpr::CategoryEq("jazz".into()),
|
||||
FilterExpr::CategoryEq("blues".into()),
|
||||
]);
|
||||
assert!(reject_or_of_signal_thresholds(&expr).is_ok());
|
||||
assert!(reject_or_of_deferred_filters(&expr).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reject_or_of_deferred_covers_all_variants() {
|
||||
// BLOCKER 5: every deferred variant placed under an `Or` must be
|
||||
// rejected, not silently degraded to AND. Each is OR'd with a harmless
|
||||
// index-backed sibling (`CategoryEq`) — exactly the `saved OR jazz`
|
||||
// shape that previously collapsed to `saved`.
|
||||
let cid = crate::entities::CollectionId::new(1);
|
||||
let deferred = [
|
||||
FilterExpr::Saved(1),
|
||||
FilterExpr::Liked(1),
|
||||
FilterExpr::InProgress {
|
||||
user_id: 1,
|
||||
threshold: 0.5,
|
||||
},
|
||||
FilterExpr::SocialGraph {
|
||||
user_id: 1,
|
||||
depth: 1,
|
||||
},
|
||||
FilterExpr::min_signal("view", 1.0),
|
||||
FilterExpr::max_signal("hide", 1.0),
|
||||
FilterExpr::near_location(0.0, 0.0, 1.0),
|
||||
FilterExpr::InCollection(cid),
|
||||
];
|
||||
for v in deferred {
|
||||
let expr = FilterExpr::Or(vec![v.clone(), FilterExpr::CategoryEq("jazz".into())]);
|
||||
assert!(
|
||||
matches!(
|
||||
reject_or_of_deferred_filters(&expr),
|
||||
Err(QueryError::InvalidFilter { .. })
|
||||
),
|
||||
"Or([{v:?}, CategoryEq]) must be rejected"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reject_or_of_deferred_finds_nested_or() {
|
||||
// A deferred variant buried inside an `Or` nested under an `And` must
|
||||
// still be caught — the walk recurses through both combinators.
|
||||
let expr = FilterExpr::And(vec![
|
||||
FilterExpr::FormatEq("video".into()),
|
||||
FilterExpr::Or(vec![
|
||||
FilterExpr::Saved(7),
|
||||
FilterExpr::CategoryEq("jazz".into()),
|
||||
]),
|
||||
]);
|
||||
assert!(reject_or_of_deferred_filters(&expr).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reject_or_of_deferred_allows_and_of_deferred() {
|
||||
// Deferred filters combined with AND are honored (applied
|
||||
// conjunctively in Stage 2.x), so an `And` of them must succeed.
|
||||
let expr = FilterExpr::And(vec![
|
||||
FilterExpr::Saved(1),
|
||||
FilterExpr::InCollection(crate::entities::CollectionId::new(3)),
|
||||
]);
|
||||
assert!(reject_or_of_deferred_filters(&expr).is_ok());
|
||||
}
|
||||
}
|
||||
|
||||
@ -7,54 +7,148 @@ use base64::Engine as _;
|
||||
use super::errors::QueryError;
|
||||
use crate::schema::EntityId;
|
||||
|
||||
// ── Cursor kind discriminant ──────────────────────────────────────────────────
|
||||
|
||||
/// Wire discriminant distinguishing a keyset cursor from a bare offset cursor.
|
||||
///
|
||||
/// Encoded as the first byte of the on-wire format so a keyset cursor handed to
|
||||
/// an offset-only consumer (or vice versa) is a loud `InvalidCursor` error rather
|
||||
/// than a silent misread of the payload as the wrong shape.
|
||||
const KIND_KEYSET: u8 = 0;
|
||||
const KIND_OFFSET: u8 = 1;
|
||||
|
||||
/// The two semantics a [`Cursor`] can carry.
|
||||
///
|
||||
/// A `Keyset` cursor anchors the next page on the `(score, entity_id)` of the
|
||||
/// last row of the previous page. An `Offset` cursor anchors on a bare list
|
||||
/// position. The RETRIEVE/SEARCH executors currently interpret cursors as
|
||||
/// offsets; handing them a keyset cursor is rejected by [`Cursor::offset`]
|
||||
/// rather than silently reinterpreting the keyset's entity ID as a position.
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
enum CursorKind {
|
||||
/// Keyset anchor: `(last_score, last_id)` of the previous page's last row.
|
||||
Keyset { score: f64, id: EntityId },
|
||||
/// Bare list-position offset for offset-based pagination.
|
||||
Offset(u64),
|
||||
}
|
||||
|
||||
// ── Cursor ──────────────────────────────────────────────────────────────────
|
||||
|
||||
/// Opaque pagination cursor for keyset pagination.
|
||||
///
|
||||
/// Encodes the last-seen entity ID and score as a base64 string. The cursor
|
||||
/// is opaque to callers -- they receive it from `Results::next_cursor` and
|
||||
/// pass it back to `RetrieveBuilder::cursor()`.
|
||||
/// Encodes either a keyset anchor `(score, entity_id)` or a bare offset as a
|
||||
/// base64url string. The cursor is opaque to callers -- they receive it from
|
||||
/// `Results::next_cursor` and pass it back to `RetrieveBuilder::cursor()`.
|
||||
///
|
||||
/// Internal format: `{score_bits_be}{entity_id_be}` (16 bytes), base64url-encoded
|
||||
/// without padding.
|
||||
/// The two shapes are tagged with a discriminant byte so a keyset cursor can
|
||||
/// never be silently misread as an offset (or vice versa): [`Cursor::offset`]
|
||||
/// returns an error on a keyset cursor, and the keyset accessors return an error
|
||||
/// on an offset cursor.
|
||||
///
|
||||
/// Internal wire format: `{kind:1}{score_bits_be:8}{id_be:8}` (17 bytes),
|
||||
/// base64url-encoded without padding. For an offset cursor the score field is
|
||||
/// zero and the id field carries the offset.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct Cursor {
|
||||
/// The score of the last result on the previous page.
|
||||
pub last_score: f64,
|
||||
/// The entity ID of the last result on the previous page.
|
||||
pub last_id: EntityId,
|
||||
kind: CursorKind,
|
||||
}
|
||||
|
||||
impl Cursor {
|
||||
/// Create a new cursor from a score and entity ID.
|
||||
/// Create a keyset cursor from the score and entity ID of the previous
|
||||
/// page's last row.
|
||||
///
|
||||
/// Non-finite scores (NaN/Inf) are clamped to `0.0`: a keyset anchor must be
|
||||
/// a real, totally-ordered value, and `encode`/`decode` reject non-finite
|
||||
/// scores on the wire, so admitting one here would only produce a cursor that
|
||||
/// cannot round-trip.
|
||||
#[must_use]
|
||||
pub const fn new(last_score: f64, last_id: EntityId) -> Self {
|
||||
let score = if last_score.is_finite() {
|
||||
last_score
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
Self {
|
||||
last_score,
|
||||
last_id,
|
||||
kind: CursorKind::Keyset { score, id: last_id },
|
||||
}
|
||||
}
|
||||
|
||||
/// The score of the last result on the previous page.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `QueryError::InvalidCursor` if this is an offset cursor, which
|
||||
/// carries no keyset score.
|
||||
pub fn last_score(&self) -> Result<f64, QueryError> {
|
||||
match self.kind {
|
||||
CursorKind::Keyset { score, .. } => Ok(score),
|
||||
CursorKind::Offset(_) => Err(QueryError::InvalidCursor(
|
||||
"offset cursor has no keyset score".to_owned(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
/// The entity ID of the last result on the previous page.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `QueryError::InvalidCursor` if this is an offset cursor, which
|
||||
/// carries no keyset entity ID.
|
||||
pub fn last_id(&self) -> Result<EntityId, QueryError> {
|
||||
match self.kind {
|
||||
CursorKind::Keyset { id, .. } => Ok(id),
|
||||
CursorKind::Offset(_) => Err(QueryError::InvalidCursor(
|
||||
"offset cursor has no keyset entity id".to_owned(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a cursor from a simple offset (for offset-based pagination).
|
||||
///
|
||||
/// The score is set to `0.0` and the entity ID encodes the offset.
|
||||
#[must_use]
|
||||
pub const fn from_offset(offset: usize) -> Self {
|
||||
Self::new(0.0, EntityId::new(offset as u64))
|
||||
// `usize as u64` is lossless on all targets tidalDB supports (<= 64-bit);
|
||||
// the inverse cast in `offset()` is the one that needs a guard.
|
||||
Self {
|
||||
kind: CursorKind::Offset(offset as u64),
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract the offset from a cursor created via `from_offset`.
|
||||
#[must_use]
|
||||
pub const fn offset(&self) -> usize {
|
||||
self.last_id.as_u64() as usize
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `QueryError::InvalidCursor` if:
|
||||
/// - this is a keyset cursor (its entity ID is not a page position), or
|
||||
/// - the stored offset does not fit in `usize` (e.g. a value above
|
||||
/// `u32::MAX` decoded on a 32-bit target, or a crafted hostile cursor).
|
||||
///
|
||||
/// Both cases previously truncated silently via `as usize`, changing the
|
||||
/// page position without any signal to the caller.
|
||||
pub fn offset(&self) -> Result<usize, QueryError> {
|
||||
let raw = match self.kind {
|
||||
CursorKind::Offset(raw) => raw,
|
||||
CursorKind::Keyset { .. } => {
|
||||
return Err(QueryError::InvalidCursor(
|
||||
"keyset cursor cannot be used as a page offset".to_owned(),
|
||||
));
|
||||
}
|
||||
};
|
||||
usize::try_from(raw).map_err(|_| {
|
||||
QueryError::InvalidCursor(format!("offset {raw} exceeds usize::MAX on this target"))
|
||||
})
|
||||
}
|
||||
|
||||
/// Encode the cursor as an opaque base64url string.
|
||||
#[must_use]
|
||||
pub fn encode(&self) -> String {
|
||||
let mut buf = [0u8; 16];
|
||||
buf[..8].copy_from_slice(&self.last_score.to_be_bytes());
|
||||
buf[8..].copy_from_slice(&self.last_id.as_u64().to_be_bytes());
|
||||
let mut buf = [0u8; 17];
|
||||
let (kind, score, id) = match self.kind {
|
||||
CursorKind::Keyset { score, id } => (KIND_KEYSET, score, id.as_u64()),
|
||||
CursorKind::Offset(raw) => (KIND_OFFSET, 0.0_f64, raw),
|
||||
};
|
||||
buf[0] = kind;
|
||||
buf[1..9].copy_from_slice(&score.to_be_bytes());
|
||||
buf[9..].copy_from_slice(&id.to_be_bytes());
|
||||
base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(buf)
|
||||
}
|
||||
|
||||
@ -62,35 +156,58 @@ impl Cursor {
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `QueryError::InvalidCursor` if the string cannot be decoded or
|
||||
/// has the wrong length.
|
||||
/// Returns `QueryError::InvalidCursor` if the string cannot be decoded, has
|
||||
/// the wrong length, carries an unknown kind discriminant, or carries a
|
||||
/// non-finite (NaN/Inf) keyset score. Rejecting non-finite scores keeps a
|
||||
/// corrupt or hostile cursor from injecting a value that breaks the total
|
||||
/// ordering keyset pagination relies on.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Does not panic in practice. The `unwrap` calls on byte-slice conversion
|
||||
/// are unreachable because the length is verified to be exactly 16 bytes.
|
||||
/// are unreachable because the length is verified to be exactly 17 bytes.
|
||||
pub fn decode(s: &str) -> Result<Self, QueryError> {
|
||||
let bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD
|
||||
.decode(s)
|
||||
.map_err(|e| QueryError::InvalidCursor(format!("base64 decode: {e}")))?;
|
||||
|
||||
if bytes.len() != 16 {
|
||||
if bytes.len() != 17 {
|
||||
return Err(QueryError::InvalidCursor(format!(
|
||||
"expected 16 bytes, got {}",
|
||||
"expected 17 bytes, got {}",
|
||||
bytes.len()
|
||||
)));
|
||||
}
|
||||
|
||||
// Safety: len == 16 verified above; slices are exactly 8 bytes.
|
||||
// Safety: len == 17 verified above; slices are exactly 8 bytes.
|
||||
#[allow(clippy::unwrap_used)]
|
||||
let score_bytes = <[u8; 8]>::try_from(&bytes[..8]).unwrap();
|
||||
let score_bytes = <[u8; 8]>::try_from(&bytes[1..9]).unwrap();
|
||||
#[allow(clippy::unwrap_used)]
|
||||
let id_bytes = <[u8; 8]>::try_from(&bytes[8..16]).unwrap();
|
||||
let id_bytes = <[u8; 8]>::try_from(&bytes[9..17]).unwrap();
|
||||
|
||||
Ok(Self {
|
||||
last_score: f64::from_be_bytes(score_bytes),
|
||||
last_id: EntityId::new(u64::from_be_bytes(id_bytes)),
|
||||
})
|
||||
let score = f64::from_be_bytes(score_bytes);
|
||||
let id = u64::from_be_bytes(id_bytes);
|
||||
|
||||
let kind = match bytes[0] {
|
||||
KIND_KEYSET => {
|
||||
if !score.is_finite() {
|
||||
return Err(QueryError::InvalidCursor(format!(
|
||||
"non-finite keyset score: {score}"
|
||||
)));
|
||||
}
|
||||
CursorKind::Keyset {
|
||||
score,
|
||||
id: EntityId::new(id),
|
||||
}
|
||||
}
|
||||
KIND_OFFSET => CursorKind::Offset(id),
|
||||
other => {
|
||||
return Err(QueryError::InvalidCursor(format!(
|
||||
"unknown cursor kind byte: {other}"
|
||||
)));
|
||||
}
|
||||
};
|
||||
|
||||
Ok(Self { kind })
|
||||
}
|
||||
}
|
||||
|
||||
@ -99,3 +216,158 @@ impl fmt::Display for Cursor {
|
||||
f.write_str(&self.encode())
|
||||
}
|
||||
}
|
||||
|
||||
// ── Tests ───────────────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
#[allow(clippy::unwrap_used)]
|
||||
mod tests {
|
||||
use base64::Engine as _;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn keyset_roundtrip() {
|
||||
let cursor = Cursor::new(0.95, EntityId::new(42));
|
||||
let decoded = Cursor::decode(&cursor.encode()).unwrap();
|
||||
assert_eq!(decoded, cursor);
|
||||
assert!((decoded.last_score().unwrap() - 0.95).abs() < f64::EPSILON);
|
||||
assert_eq!(decoded.last_id().unwrap(), EntityId::new(42));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn offset_roundtrip() {
|
||||
let cursor = Cursor::from_offset(12_345);
|
||||
let decoded = Cursor::decode(&cursor.encode()).unwrap();
|
||||
assert_eq!(decoded, cursor);
|
||||
assert_eq!(decoded.offset().unwrap(), 12_345);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keyset_cursor_used_as_offset_is_a_loud_error() {
|
||||
// Finding 4: a keyset cursor must never be silently reinterpreted as an
|
||||
// offset (its entity ID is not a page position).
|
||||
let cursor = Cursor::new(0.5, EntityId::new(99));
|
||||
assert!(matches!(cursor.offset(), Err(QueryError::InvalidCursor(_))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn offset_cursor_has_no_keyset_accessors() {
|
||||
let cursor = Cursor::from_offset(7);
|
||||
assert!(matches!(
|
||||
cursor.last_score(),
|
||||
Err(QueryError::InvalidCursor(_))
|
||||
));
|
||||
assert!(matches!(
|
||||
cursor.last_id(),
|
||||
Err(QueryError::InvalidCursor(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_invalid_base64() {
|
||||
assert!(matches!(
|
||||
Cursor::decode("!!!not-base64!!!"),
|
||||
Err(QueryError::InvalidCursor(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_wrong_length() {
|
||||
let short = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0u8; 8]);
|
||||
assert!(matches!(
|
||||
Cursor::decode(&short),
|
||||
Err(QueryError::InvalidCursor(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_unknown_kind_byte_rejected() {
|
||||
let mut buf = [0u8; 17];
|
||||
buf[0] = 99; // neither KIND_KEYSET nor KIND_OFFSET
|
||||
let encoded = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(buf);
|
||||
assert!(matches!(
|
||||
Cursor::decode(&encoded),
|
||||
Err(QueryError::InvalidCursor(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_rejects_non_finite_keyset_score() {
|
||||
// Finding 5: a hostile/corrupt cursor must not inject NaN/Inf scores.
|
||||
for bad in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
|
||||
let mut buf = [0u8; 17];
|
||||
buf[0] = KIND_KEYSET;
|
||||
buf[1..9].copy_from_slice(&bad.to_be_bytes());
|
||||
buf[9..].copy_from_slice(&1u64.to_be_bytes());
|
||||
let encoded = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(buf);
|
||||
assert!(
|
||||
matches!(Cursor::decode(&encoded), Err(QueryError::InvalidCursor(_))),
|
||||
"non-finite score {bad} should be rejected"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_clamps_non_finite_score_to_zero() {
|
||||
let cursor = Cursor::new(f64::NAN, EntityId::new(3));
|
||||
assert!(cursor.last_score().unwrap().abs() < f64::EPSILON);
|
||||
// And the clamped cursor still round-trips.
|
||||
assert_eq!(Cursor::decode(&cursor.encode()).unwrap(), cursor);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn offset_rejects_value_above_u32_on_decode() {
|
||||
// Finding 2: a crafted offset cursor carrying a value > u32::MAX must not
|
||||
// truncate silently. On a 64-bit target it round-trips; on a 32-bit
|
||||
// target `offset()` errors. We assert the value survives decode here and
|
||||
// that `usize::try_from` is the only narrowing gate.
|
||||
let big: u64 = u64::from(u32::MAX) + 1;
|
||||
let mut buf = [0u8; 17];
|
||||
buf[0] = KIND_OFFSET;
|
||||
buf[9..].copy_from_slice(&big.to_be_bytes());
|
||||
let encoded = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(buf);
|
||||
let cursor = Cursor::decode(&encoded).unwrap();
|
||||
// Re-encode preserves the full 64-bit value (no truncation in the codec).
|
||||
assert_eq!(Cursor::decode(&cursor.encode()).unwrap(), cursor);
|
||||
// On 64-bit, offset() returns the exact value; on 32-bit it would error.
|
||||
match cursor.offset() {
|
||||
Ok(v) => assert_eq!(u64::try_from(v).unwrap(), big),
|
||||
Err(QueryError::InvalidCursor(_)) => {}
|
||||
Err(other) => panic!("unexpected error: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn display_matches_encode() {
|
||||
let cursor = Cursor::new(0.5, EntityId::new(100));
|
||||
assert_eq!(cursor.to_string(), cursor.encode());
|
||||
}
|
||||
|
||||
mod proptests {
|
||||
use proptest::prelude::*;
|
||||
|
||||
use super::*;
|
||||
|
||||
proptest! {
|
||||
#[test]
|
||||
fn keyset_roundtrip_prop(
|
||||
score in proptest::num::f64::NORMAL,
|
||||
id in 0u64..u64::MAX,
|
||||
) {
|
||||
let cursor = Cursor::new(score, EntityId::new(id));
|
||||
let decoded = Cursor::decode(&cursor.encode()).unwrap();
|
||||
prop_assert_eq!(decoded.last_id().unwrap(), EntityId::new(id));
|
||||
// f64 roundtrip through be_bytes is exact for finite values.
|
||||
prop_assert!((decoded.last_score().unwrap() - score).abs() < f64::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn offset_roundtrip_prop(offset in 0usize..=usize::MAX) {
|
||||
let cursor = Cursor::from_offset(offset);
|
||||
let decoded = Cursor::decode(&cursor.encode()).unwrap();
|
||||
prop_assert_eq!(decoded.offset().unwrap(), offset);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -13,6 +13,31 @@ use crate::{
|
||||
storage::indexes::filter::FilterExpr,
|
||||
};
|
||||
|
||||
/// Reject a `for_creator` entity ID that does not fit the bitmap index's u32
|
||||
/// key space.
|
||||
///
|
||||
/// `CreatorEq` filters live in the `RoaringBitmap` item-ID space (u32). A
|
||||
/// creator ID above `u32::MAX` would be silently narrowed by the `as u32` cast
|
||||
/// in [`RetrieveBuilder::for_creator`], producing a wrong filter while
|
||||
/// `for_creator` retains the full 64-bit value. Until the bitmap index is
|
||||
/// upgraded to 64-bit keys, such an ID is rejected with `InvalidFilter` rather
|
||||
/// than narrowed.
|
||||
fn check_for_creator_fits_u32(for_creator: Option<EntityId>) -> Result<(), QueryError> {
|
||||
if let Some(creator_id) = for_creator
|
||||
&& u32::try_from(creator_id.as_u64()).is_err()
|
||||
{
|
||||
return Err(QueryError::InvalidFilter {
|
||||
field: "for_creator".to_owned(),
|
||||
reason: format!(
|
||||
"creator id {} exceeds the u32 bitmap-index key space (max {})",
|
||||
creator_id.as_u64(),
|
||||
u32::MAX
|
||||
),
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ── NotificationCaps ────────────────────────────────────────────────────────
|
||||
|
||||
/// Notification frequency caps for the `notification` profile.
|
||||
@ -138,14 +163,15 @@ impl Retrieve {
|
||||
/// Checks:
|
||||
/// 1. Profile exists (by name, optionally by version)
|
||||
/// 2. Limit is in `[1, 500]`
|
||||
/// 3. `for_user` and `similar_to` are supported (M3)
|
||||
/// 4. Candidate strategy is supported
|
||||
/// 3. `for_creator` fits the u32 bitmap key space
|
||||
/// 4. `for_user` and `similar_to` are supported (M3)
|
||||
/// 5. Candidate strategy is supported
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `QueryError` describing the first validation failure.
|
||||
// Used by the query executor in Task 02. Until then, only tests reference it.
|
||||
#[allow(dead_code)]
|
||||
// Called by `RetrieveExecutor::execute` (query/executor/pipeline.rs) before
|
||||
// planning, in addition to the unit tests below.
|
||||
pub(crate) fn validate(&self, registry: &ProfileRegistry) -> Result<(), QueryError> {
|
||||
// 1. Profile existence.
|
||||
let profile_name = &self.profile.name;
|
||||
@ -167,6 +193,12 @@ impl Retrieve {
|
||||
});
|
||||
}
|
||||
|
||||
// 2b. for_creator must fit the bitmap index's u32 key space. Reject
|
||||
// oversized IDs loudly rather than narrowing them with a silent `as u32`
|
||||
// (which would apply a wrong CreatorEq filter while `for_creator` keeps
|
||||
// the full 64-bit value, so the two could disagree for one query).
|
||||
check_for_creator_fits_u32(self.for_creator)?;
|
||||
|
||||
// 3. M3: for_user and similar_to are now supported.
|
||||
// No validation rejection -- the executor handles user context.
|
||||
|
||||
@ -362,15 +394,20 @@ impl RetrieveBuilder {
|
||||
/// Adds a `CreatorEq` filter for the creator's ID and sets `for_creator`
|
||||
/// on the query for creator-bitmap-based candidate restriction.
|
||||
///
|
||||
/// # Known Limitation
|
||||
/// # Bitmap key-space limit
|
||||
///
|
||||
/// `CreatorEq` stores the creator ID as `u32` (matching the `RoaringBitmap`
|
||||
/// item-ID space). Creator IDs > `u32::MAX` (4,294,967,295) will be silently
|
||||
/// truncated, producing incorrect filter results. This is a systemic constraint
|
||||
/// of the bitmap index design and will be addressed when the bitmap index is
|
||||
/// upgraded to 64-bit keys (planned post-M7).
|
||||
/// `CreatorEq` lives in the `RoaringBitmap` item-ID space (`u32`). A creator
|
||||
/// ID above `u32::MAX` (4,294,967,295) cannot be represented there. Rather
|
||||
/// than silently truncate, the builder defers the range check: an oversized
|
||||
/// ID is rejected with `QueryError::InvalidFilter` at [`RetrieveBuilder::build`]
|
||||
/// (and again in `Retrieve::validate`), never producing a wrong filter. The
|
||||
/// `u32` ceiling is a systemic constraint of the bitmap index design and will
|
||||
/// lift when the index is upgraded to 64-bit keys (planned post-M7).
|
||||
#[must_use]
|
||||
pub fn for_creator(mut self, creator_id: EntityId) -> Self {
|
||||
// The narrowing cast is sound only for IDs <= u32::MAX; `build()` /
|
||||
// `validate()` reject anything larger via `check_for_creator_fits_u32`,
|
||||
// so this cast never silently changes a query that survives validation.
|
||||
self.filters
|
||||
.push(FilterExpr::CreatorEq(creator_id.as_u64() as u32));
|
||||
self.for_creator = Some(creator_id);
|
||||
@ -391,11 +428,13 @@ impl RetrieveBuilder {
|
||||
self
|
||||
}
|
||||
|
||||
/// Build the `Retrieve` query, validating the limit range.
|
||||
/// Build the `Retrieve` query, validating the limit range and `for_creator`.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `QueryError::InvalidLimit` if `limit` is 0 or greater than 500.
|
||||
/// - `QueryError::InvalidLimit` if `limit` is 0 or greater than 500.
|
||||
/// - `QueryError::InvalidFilter` if a `for_creator` ID exceeds `u32::MAX`
|
||||
/// (the bitmap index key-space ceiling).
|
||||
pub fn build(self) -> Result<Retrieve, QueryError> {
|
||||
if self.limit == 0 || self.limit > 500 {
|
||||
return Err(QueryError::InvalidLimit {
|
||||
@ -405,6 +444,10 @@ impl RetrieveBuilder {
|
||||
});
|
||||
}
|
||||
|
||||
// Reject a `for_creator` ID that would be silently narrowed into the
|
||||
// u32 bitmap key space (see `check_for_creator_fits_u32`).
|
||||
check_for_creator_fits_u32(self.for_creator)?;
|
||||
|
||||
Ok(Retrieve {
|
||||
entity_kind: self.entity_kind,
|
||||
profile: self.profile,
|
||||
|
||||
@ -110,6 +110,60 @@ fn builder_with_cursor() {
|
||||
assert!(q.cursor.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builder_for_creator_within_u32_accepted() {
|
||||
let q = RetrieveBuilder::new(EntityKind::Item, ProfileRef::new("trending"))
|
||||
.for_creator(EntityId::new(u64::from(u32::MAX)))
|
||||
.build()
|
||||
.unwrap();
|
||||
assert_eq!(q.for_creator, Some(EntityId::new(u64::from(u32::MAX))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builder_for_creator_above_u32_rejected_at_build() {
|
||||
// Finding 3: an oversized creator ID must not be silently narrowed.
|
||||
let big = u64::from(u32::MAX) + 1;
|
||||
let result = RetrieveBuilder::new(EntityKind::Item, ProfileRef::new("trending"))
|
||||
.for_creator(EntityId::new(big))
|
||||
.build();
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(QueryError::InvalidFilter { ref field, .. }) if field == "for_creator"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_for_creator_above_u32_rejected() {
|
||||
let registry = make_registry_with_profile(
|
||||
"creator_test",
|
||||
CandidateStrategy::Scan {
|
||||
sort_field: "created_at".into(),
|
||||
},
|
||||
);
|
||||
// Construct directly to bypass build()'s check and exercise validate().
|
||||
let q = Retrieve {
|
||||
entity_kind: EntityKind::Item,
|
||||
profile: ProfileRef::new("creator_test"),
|
||||
filters: vec![],
|
||||
diversity: None,
|
||||
limit: 20,
|
||||
exclude: vec![],
|
||||
cursor: None,
|
||||
for_user: None,
|
||||
similar_to: None,
|
||||
context: None,
|
||||
for_session: None,
|
||||
cohort: None,
|
||||
cohort_predicate: None,
|
||||
notification_caps: None,
|
||||
for_creator: Some(EntityId::new(u64::from(u32::MAX) + 1)),
|
||||
};
|
||||
assert!(matches!(
|
||||
q.validate(®istry),
|
||||
Err(QueryError::InvalidFilter { ref field, .. }) if field == "for_creator"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builder_with_for_user() {
|
||||
let q = RetrieveBuilder::new(EntityKind::Item, ProfileRef::new("trending"))
|
||||
@ -176,6 +230,7 @@ fn cursor_roundtrip() {
|
||||
let encoded = cursor.encode();
|
||||
let decoded = Cursor::decode(&encoded).unwrap();
|
||||
assert_eq!(decoded, cursor);
|
||||
assert_eq!(decoded.last_id().unwrap(), EntityId::new(42));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@ -453,10 +508,9 @@ mod proptests {
|
||||
let cursor = Cursor::new(score, EntityId::new(id));
|
||||
let encoded = cursor.encode();
|
||||
let decoded = Cursor::decode(&encoded).unwrap();
|
||||
prop_assert_eq!(decoded.last_id, cursor.last_id);
|
||||
// f64 roundtrip through be_bytes is exact.
|
||||
prop_assert!((decoded.last_score - cursor.last_score).abs() < f64::EPSILON
|
||||
|| (decoded.last_score.is_nan() && cursor.last_score.is_nan()));
|
||||
prop_assert_eq!(decoded.last_id().unwrap(), EntityId::new(id));
|
||||
// f64 roundtrip through be_bytes is exact for finite values.
|
||||
prop_assert!((decoded.last_score().unwrap() - score).abs() < f64::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@ -53,6 +53,13 @@ use crate::{
|
||||
text::collectors::AllScoresCollector,
|
||||
};
|
||||
|
||||
/// Canonical default embedding-slot name.
|
||||
///
|
||||
/// Used when neither the schema nor the executor threads an explicit item slot
|
||||
/// name. Centralizes the literal so the `"content"` default is declared once
|
||||
/// instead of being duplicated at each `unwrap_or("content")` call site.
|
||||
const DEFAULT_EMBEDDING_SLOT: &str = "content";
|
||||
|
||||
/// Convert a 64-bit entity id to the 32-bit key space used by every
|
||||
/// `RoaringBitmap` in the suppression/inclusion path.
|
||||
///
|
||||
@ -74,6 +81,7 @@ impl SearchExecutor<'_> {
|
||||
///
|
||||
/// Returns `QueryError` on validation failure, missing index, profile not
|
||||
/// found, or storage error during BM25/ANN retrieval.
|
||||
#[allow(clippy::too_many_lines)] // the 8-stage SEARCH pipeline reads best linearly
|
||||
pub fn execute(&self, query: &Search) -> Result<SearchResults, QueryError> {
|
||||
let query_start = Instant::now();
|
||||
let mut stats = QueryStats::new(query.profile.name.clone());
|
||||
@ -88,19 +96,29 @@ impl SearchExecutor<'_> {
|
||||
|
||||
let profile = self.resolve_profile(query)?;
|
||||
|
||||
// Compute the combined filter tree ONCE for the whole pipeline. Each
|
||||
// stage helper borrows it (`Option<&FilterExpr>`) instead of calling
|
||||
// `query.combined_filter()` again — that method deep-clones the filter
|
||||
// Vec (and allocates a fresh `FilterExpr::And` for >1 filter), so the
|
||||
// previous five independent calls did up to five full tree clones per
|
||||
// query on the hot path.
|
||||
let combined_filter = query.combined_filter();
|
||||
let combined_filter = combined_filter.as_ref();
|
||||
|
||||
// Reject filter shapes with no sound semantics before any stage runs:
|
||||
// negated deferred post-filters (would return the inverse set) and
|
||||
// OR-of-thresholds (would be silently evaluated as AND). Shares the
|
||||
// RETRIEVE-path validators so both query surfaces fail identically.
|
||||
if let Some(ref filter_expr) = query.combined_filter() {
|
||||
// negated deferred post-filters (would return the inverse set) and any
|
||||
// deferred post-filter under an OR (would be silently evaluated as AND —
|
||||
// see `reject_or_of_deferred_filters`). Shares the RETRIEVE-path
|
||||
// validators so both query surfaces fail identically.
|
||||
if let Some(filter_expr) = combined_filter {
|
||||
crate::query::executor::user_filter::reject_negated_deferred_filters(filter_expr)?;
|
||||
crate::query::executor::user_filter::reject_or_of_signal_thresholds(filter_expr)?;
|
||||
crate::query::executor::user_filter::reject_or_of_deferred_filters(filter_expr)?;
|
||||
}
|
||||
|
||||
let mut warnings: Vec<String> = Vec::new();
|
||||
|
||||
// ── Stage 0: Scope Pre-Filter ─────────────────────────────────────
|
||||
let scope_bitmap = self.resolve_scope(query)?;
|
||||
let scope_bitmap = self.resolve_scope(query, &mut warnings)?;
|
||||
|
||||
// ── Stage 1a: Text Retrieval (BM25) ────────────────────────────────
|
||||
let mut bm25_results = self.retrieve_bm25(query, &mut warnings)?;
|
||||
@ -158,14 +176,14 @@ impl SearchExecutor<'_> {
|
||||
);
|
||||
|
||||
// ── Stage 2 + 2b: Metadata Filters ──────────────────────────────────
|
||||
self.apply_metadata_filter(query, &mut candidates);
|
||||
self.apply_creator_metadata_filter(query, &mut candidates);
|
||||
self.apply_metadata_filter(combined_filter, &mut candidates, &mut warnings);
|
||||
self.apply_creator_metadata_filter(query, combined_filter, &mut candidates);
|
||||
|
||||
// ── Stage 2.2/2.3/2.4: Deferred Post-Filters ────────────────────────
|
||||
self.apply_deferred_post_filters(query, &mut candidates)?;
|
||||
self.apply_deferred_post_filters(combined_filter, &mut candidates)?;
|
||||
|
||||
// ── Stage 2.5: User-Context Filtering ────────────────────────────────
|
||||
self.apply_user_context_filter(query, &mut candidates);
|
||||
self.apply_user_context_filter(query, combined_filter, &mut candidates);
|
||||
|
||||
stats.candidates_after_filter = candidates.len();
|
||||
stats.filters_applied = query.filters.len();
|
||||
@ -242,12 +260,31 @@ impl SearchExecutor<'_> {
|
||||
}
|
||||
|
||||
/// Stage 0: resolve the optional scope pre-filter into a bitmap.
|
||||
fn resolve_scope(&self, query: &Search) -> Result<Option<RoaringBitmap>, QueryError> {
|
||||
///
|
||||
/// On a poisoned universe lock we recover the last-known bitmap via
|
||||
/// [`std::sync::PoisonError::into_inner`] (the same poison-recovery pattern
|
||||
/// `suggest.rs` uses) and push a warning, rather than silently degrading to
|
||||
/// an EMPTY universe — which would make `WithinScope::Trending` resolve to
|
||||
/// nothing and return zero results with no diagnostic.
|
||||
fn resolve_scope(
|
||||
&self,
|
||||
query: &Search,
|
||||
warnings: &mut Vec<String>,
|
||||
) -> Result<Option<RoaringBitmap>, QueryError> {
|
||||
let Some(ref scope) = query.within_scope else {
|
||||
return Ok(None);
|
||||
};
|
||||
let universe = self.universe.map_or_else(RoaringBitmap::new, |u| {
|
||||
u.read().map(|g| g.clone()).unwrap_or_default()
|
||||
u.read().map_or_else(
|
||||
|poisoned| {
|
||||
warnings.push(
|
||||
"universe lock poisoned; scope pre-filter recovered last-known universe"
|
||||
.to_string(),
|
||||
);
|
||||
poisoned.into_inner().clone()
|
||||
},
|
||||
|g| g.clone(),
|
||||
)
|
||||
});
|
||||
|
||||
let mut resolver = ScopeResolver::new(self.ledger, universe);
|
||||
@ -297,7 +334,15 @@ impl SearchExecutor<'_> {
|
||||
.search(tq.as_ref(), &collector)
|
||||
.map_err(|e| QueryError::StorageError(format!("BM25 search: {e}")))?;
|
||||
// Sort descending by BM25 score (AllScoresCollector does not sort).
|
||||
raw.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||
//
|
||||
// `partial_cmp(...).unwrap_or(Equal)` treats every comparison against a
|
||||
// `NaN` score as "equal", which scrambles the order non-deterministically
|
||||
// when a NaN slips through (BM25 over degenerate term statistics can emit
|
||||
// one). Neutralize NaN to 0.0 — the same neutral value the ranking module
|
||||
// uses (`finite_score`) — then sort with `total_cmp`, a true total order
|
||||
// over all finite scores, so the result is deterministic.
|
||||
let finite = |s: f32| if s.is_nan() { 0.0 } else { s };
|
||||
raw.sort_by(|a, b| finite(b.1).total_cmp(&finite(a.1)));
|
||||
Ok(raw)
|
||||
}
|
||||
|
||||
@ -360,14 +405,50 @@ impl SearchExecutor<'_> {
|
||||
}
|
||||
|
||||
/// Stage 2: index/predicate metadata filter.
|
||||
fn apply_metadata_filter(&self, query: &Search, candidates: &mut Vec<EntityId>) {
|
||||
let Some(filter_expr) = query.combined_filter() else {
|
||||
///
|
||||
/// The `FilterEvaluator` resolves deferred variants
|
||||
/// (`MinSignal`/`MaxSignal`/`NearLocation`/`InCollection`/user-state) to the
|
||||
/// *universe* and `evaluate_and` intersects every child against it. That is
|
||||
/// only sound when `self.universe` is a populated, trustworthy bitmap. When
|
||||
/// the universe is unavailable — not wired (`None`) or the read lock is
|
||||
/// poisoned — its clone is empty, so the intersection would silently drop
|
||||
/// EVERY candidate (including ones that should pass) before the Stage 2.2/2.5
|
||||
/// post-filters ever run. The actual deferred filtering happens in
|
||||
/// [`Self::apply_deferred_post_filters`], so the correct degraded behavior is
|
||||
/// to skip the index intersection entirely (retain all) and surface a warning
|
||||
/// rather than collapse to an empty result set.
|
||||
#[allow(clippy::significant_drop_tightening)] // universe read guard held across the evaluator
|
||||
fn apply_metadata_filter(
|
||||
&self,
|
||||
combined_filter: Option<&FilterExpr>,
|
||||
candidates: &mut Vec<EntityId>,
|
||||
warnings: &mut Vec<String>,
|
||||
) {
|
||||
let Some(filter_expr) = combined_filter else {
|
||||
return;
|
||||
};
|
||||
|
||||
// Resolve the universe up front, decoupled from the deferred-filter
|
||||
// sentinel: if it is missing or poisoned we must NOT run the bitmap
|
||||
// intersection, because every deferred variant resolves to that (empty)
|
||||
// universe and would drop all candidates. Fall through to the post-filter
|
||||
// stage instead.
|
||||
let Some(universe_lock) = self.universe else {
|
||||
return;
|
||||
};
|
||||
let Ok(universe_guard) = universe_lock.read() else {
|
||||
warnings.push(
|
||||
"universe lock poisoned; metadata index pre-filter degraded — relying on \
|
||||
post-filters"
|
||||
.to_string(),
|
||||
);
|
||||
return;
|
||||
};
|
||||
let universe_ref: &RoaringBitmap = &universe_guard;
|
||||
|
||||
let empty_bitmap = BitmapIndex::new("_empty");
|
||||
let empty_dur = RangeIndex::<u32>::new("_empty");
|
||||
let empty_ts = RangeIndex::<u64>::new("_empty");
|
||||
let empty_universe = RoaringBitmap::new();
|
||||
|
||||
let cat = self.category_index.unwrap_or(&empty_bitmap);
|
||||
let fmt = self.format_index.unwrap_or(&empty_bitmap);
|
||||
@ -376,21 +457,8 @@ impl SearchExecutor<'_> {
|
||||
let dur = self.duration_index.unwrap_or(&empty_dur);
|
||||
let ts = self.created_at_index.unwrap_or(&empty_ts);
|
||||
|
||||
let universe_guard;
|
||||
#[allow(clippy::option_if_let_else)]
|
||||
let universe_ref = match self.universe {
|
||||
Some(u) => match u.read() {
|
||||
Ok(guard) => {
|
||||
universe_guard = guard;
|
||||
&*universe_guard
|
||||
}
|
||||
Err(_) => &empty_universe,
|
||||
},
|
||||
None => &empty_universe,
|
||||
};
|
||||
|
||||
let evaluator = FilterEvaluator::new(cat, fmt, cre, tag, dur, ts, universe_ref);
|
||||
match evaluator.evaluate(&filter_expr) {
|
||||
match evaluator.evaluate(filter_expr) {
|
||||
FilterResult::Bitmap(bitmap) => {
|
||||
candidates.retain(|id| entity_as_u32(*id).is_some_and(|i| bitmap.contains(i)));
|
||||
}
|
||||
@ -406,14 +474,19 @@ impl SearchExecutor<'_> {
|
||||
}
|
||||
|
||||
/// Stage 2b: creator metadata post-filter against stored creator metadata.
|
||||
fn apply_creator_metadata_filter(&self, query: &Search, candidates: &mut Vec<EntityId>) {
|
||||
fn apply_creator_metadata_filter(
|
||||
&self,
|
||||
query: &Search,
|
||||
combined_filter: Option<&FilterExpr>,
|
||||
candidates: &mut Vec<EntityId>,
|
||||
) {
|
||||
if query.entity_kind != EntityKind::Creator || candidates.is_empty() {
|
||||
return;
|
||||
}
|
||||
let Some(combined) = query.combined_filter() else {
|
||||
let Some(combined) = combined_filter else {
|
||||
return;
|
||||
};
|
||||
if !has_creator_metadata_filter(&combined) {
|
||||
if !has_creator_metadata_filter(combined) {
|
||||
return;
|
||||
}
|
||||
let Some(storage) = self.creators_storage else {
|
||||
@ -430,7 +503,7 @@ impl SearchExecutor<'_> {
|
||||
meta
|
||||
})
|
||||
.unwrap_or_default();
|
||||
evaluate_metadata_filter(&combined, &meta)
|
||||
evaluate_metadata_filter(combined, &meta)
|
||||
});
|
||||
|
||||
tracing::trace!(
|
||||
@ -453,10 +526,10 @@ impl SearchExecutor<'_> {
|
||||
/// name referenced by a Min/MaxSignal threshold).
|
||||
fn apply_deferred_post_filters(
|
||||
&self,
|
||||
query: &Search,
|
||||
combined_filter: Option<&FilterExpr>,
|
||||
candidates: &mut Vec<EntityId>,
|
||||
) -> Result<(), QueryError> {
|
||||
if let Some(ref filter_expr) = query.combined_filter() {
|
||||
if let Some(filter_expr) = combined_filter {
|
||||
let ctx = crate::query::executor::post_filter::PostFilterCtx {
|
||||
ledger: self.ledger,
|
||||
collection_index: self.collection_index,
|
||||
@ -472,7 +545,12 @@ impl SearchExecutor<'_> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn apply_user_context_filter(&self, query: &Search, candidates: &mut Vec<EntityId>) {
|
||||
fn apply_user_context_filter(
|
||||
&self,
|
||||
query: &Search,
|
||||
combined_filter: Option<&FilterExpr>,
|
||||
candidates: &mut Vec<EntityId>,
|
||||
) {
|
||||
let Some(user_id) = query.for_user else {
|
||||
return;
|
||||
};
|
||||
@ -502,9 +580,9 @@ impl SearchExecutor<'_> {
|
||||
// User-state inclusion filters (Saved, Liked, InProgress, SocialGraph).
|
||||
// Routed through the canonical RETRIEVE extractor (which includes
|
||||
// SocialGraph) so both surfaces share one collector and cannot drift.
|
||||
if let Some(filter_expr) = query.combined_filter() {
|
||||
if let Some(filter_expr) = combined_filter {
|
||||
for uf in
|
||||
crate::query::executor::user_filter::extract_user_state_filters(&filter_expr)
|
||||
crate::query::executor::user_filter::extract_user_state_filters(filter_expr)
|
||||
{
|
||||
match uf {
|
||||
FilterExpr::Saved(uid) => {
|
||||
@ -619,7 +697,7 @@ impl SearchExecutor<'_> {
|
||||
candidates,
|
||||
self.preference_vectors,
|
||||
self.items_storage,
|
||||
self.item_embedding_slot.unwrap_or("content"),
|
||||
self.item_embedding_slot.unwrap_or(DEFAULT_EMBEDDING_SLOT),
|
||||
);
|
||||
executor.score_personalized(
|
||||
candidates,
|
||||
@ -664,7 +742,9 @@ impl SearchExecutor<'_> {
|
||||
query_start: Instant,
|
||||
mut stats: QueryStats,
|
||||
) -> Result<SearchResults, QueryError> {
|
||||
let offset = query.cursor.as_ref().map_or(0, Cursor::offset);
|
||||
// `Cursor::offset` is fallible: a keyset cursor used as an offset (or a
|
||||
// 32-bit-overflowing value) errors loudly rather than misreading silently.
|
||||
let offset = query.cursor.as_ref().map_or(Ok(0), Cursor::offset)?;
|
||||
let limit = query.limit as usize;
|
||||
// Guard against cursor-offset overflow from an attacker-supplied cursor:
|
||||
// `offset + limit` could wrap. `saturating_add` clamps to usize::MAX,
|
||||
@ -1026,7 +1106,7 @@ mod tests {
|
||||
};
|
||||
|
||||
/// Items-only length-prefixed metadata encoding that `db::deserialize_metadata`
|
||||
/// (used by the NearLocation post-filter) reads back. Mirrors
|
||||
/// (used by the `NearLocation` post-filter) reads back. Mirrors
|
||||
/// `db::metadata::serialize_metadata`, which is not visible outside `db`.
|
||||
fn encode_meta(pairs: &[(&str, &str)]) -> Vec<u8> {
|
||||
let mut buf = Vec::new();
|
||||
@ -1218,6 +1298,82 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
/// BLOCKER 5: SEARCH must reject an `OR` containing a deferred post-filter
|
||||
/// rather than silently collapsing `saved OR jazz` to `saved`. The guard is
|
||||
/// shared verbatim with the RETRIEVE pipeline
|
||||
/// (`user_filter::reject_or_of_deferred_filters`), so the two surfaces cannot
|
||||
/// diverge; this test proves it fires through the public `execute()`.
|
||||
#[test]
|
||||
fn search_or_of_deferred_filter_rejected() {
|
||||
let schema = test_schema();
|
||||
let ledger = SignalLedger::new(schema, Box::new(NoopWalWriter));
|
||||
let profile_reg = setup_registry();
|
||||
let registry = two_item_registry();
|
||||
let universe = two_item_universe();
|
||||
let exec = filtered_executor(&ledger, &profile_reg, ®istry, &universe);
|
||||
|
||||
// `saved OR jazz` — previously degraded to `saved` (an intersection).
|
||||
let saved_or_jazz = Search::builder()
|
||||
.vector(vec![1.0f32, 0.0])
|
||||
.using_profile("search")
|
||||
.for_user(42)
|
||||
.filter(FilterExpr::Or(vec![
|
||||
FilterExpr::Saved(42),
|
||||
FilterExpr::CategoryEq("jazz".into()),
|
||||
]))
|
||||
.limit(10)
|
||||
.build()
|
||||
.unwrap();
|
||||
let err = exec
|
||||
.execute(&saved_or_jazz)
|
||||
.expect_err("OR over a deferred Saved filter must be rejected");
|
||||
assert!(matches!(err, QueryError::InvalidFilter { .. }));
|
||||
|
||||
// OR-of-thresholds is the same trap and must error identically.
|
||||
let or_thresholds = Search::builder()
|
||||
.vector(vec![1.0f32, 0.0])
|
||||
.using_profile("search")
|
||||
.filter(FilterExpr::Or(vec![
|
||||
FilterExpr::min_signal("view", 1.0),
|
||||
FilterExpr::min_signal("like", 1.0),
|
||||
]))
|
||||
.limit(10)
|
||||
.build()
|
||||
.unwrap();
|
||||
let err = exec
|
||||
.execute(&or_thresholds)
|
||||
.expect_err("OR-of-thresholds must be rejected");
|
||||
assert!(matches!(err, QueryError::InvalidFilter { .. }));
|
||||
}
|
||||
|
||||
/// A legitimate index-backed `Or` (no deferred variant) must still succeed
|
||||
/// through SEARCH — the guard rejects only deferred variants under `Or`.
|
||||
#[test]
|
||||
fn search_or_of_index_filters_allowed() {
|
||||
let schema = test_schema();
|
||||
let ledger = SignalLedger::new(schema, Box::new(NoopWalWriter));
|
||||
let profile_reg = setup_registry();
|
||||
let registry = two_item_registry();
|
||||
let universe = two_item_universe();
|
||||
let exec = filtered_executor(&ledger, &profile_reg, ®istry, &universe);
|
||||
|
||||
let query = Search::builder()
|
||||
.vector(vec![1.0f32, 0.0])
|
||||
.using_profile("search")
|
||||
.filter(FilterExpr::Or(vec![
|
||||
FilterExpr::CategoryEq("jazz".into()),
|
||||
FilterExpr::CategoryEq("blues".into()),
|
||||
]))
|
||||
.limit(10)
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
// No category index is wired, so the filter matches nothing, but the
|
||||
// query must SUCCEED (no rejection) — proving the guard does not fire on
|
||||
// a pure index-backed OR.
|
||||
assert!(exec.execute(&query).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn search_social_graph_scopes_to_followed_creator_items() {
|
||||
let schema = test_schema();
|
||||
@ -1254,4 +1410,106 @@ mod tests {
|
||||
the SocialGraph arm was swallowed and both were returned"
|
||||
);
|
||||
}
|
||||
|
||||
/// Accuracy-W regression: a deferred `MinSignal` filter with NO universe
|
||||
/// wired must still return candidates from retrieval and let the Stage 2.2
|
||||
/// post-filter trim them — not collapse to an empty set.
|
||||
///
|
||||
/// Before the fix, `apply_metadata_filter` ran the `FilterEvaluator`, which
|
||||
/// resolves `MinSignal` to `self.universe`; with the universe unwired the
|
||||
/// clone was empty, so `evaluate_and` intersected against nothing and
|
||||
/// `candidates.retain` dropped EVERY candidate before the post-filter could
|
||||
/// run. Now the index pre-filter is skipped (retain all) when the universe is
|
||||
/// `None`, and the deferred post-filter does the real trimming.
|
||||
#[test]
|
||||
fn search_min_signal_with_no_universe_still_postfilters() {
|
||||
let schema = test_schema();
|
||||
let ledger = SignalLedger::new(schema, Box::new(NoopWalWriter));
|
||||
let profile_reg = setup_registry();
|
||||
// Two equidistant items so the ANN path retrieves both; the deferred
|
||||
// MinSignal filter (not relevance) decides which survives.
|
||||
let registry = RwLock::new(registry_with_slot(
|
||||
"content",
|
||||
&[(1, [1.0, 0.0]), (2, [1.0, 0.0])],
|
||||
));
|
||||
|
||||
// Item 1 has 5 likes (>= 3), item 2 has 1 like (< 3).
|
||||
let now = Timestamp::now();
|
||||
for _ in 0..5 {
|
||||
ledger
|
||||
.record_signal("like", EntityId::new(1), 1.0, now)
|
||||
.unwrap();
|
||||
}
|
||||
ledger
|
||||
.record_signal("like", EntityId::new(2), 1.0, now)
|
||||
.unwrap();
|
||||
|
||||
// NOTE: `vector_executor` wires NO universe — the exact condition that
|
||||
// previously dropped everything.
|
||||
let exec = vector_executor(&ledger, &profile_reg, ®istry);
|
||||
let query = Search::builder()
|
||||
.vector(vec![1.0f32, 0.0])
|
||||
.using_profile("search")
|
||||
.filter(FilterExpr::min_signal("like", 3.0))
|
||||
.limit(10)
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
let results = exec.execute(&query).unwrap();
|
||||
assert_eq!(
|
||||
ids_of(&results),
|
||||
vec![1],
|
||||
"with no universe wired, min_signal(like,3) must keep item 1 and drop item 2 via the \
|
||||
post-filter; before the fix the empty-universe intersection dropped BOTH"
|
||||
);
|
||||
}
|
||||
|
||||
/// Completeness-W regression: a poisoned universe lock must not silently
|
||||
/// return an empty result set. `resolve_scope` recovers the last-known
|
||||
/// universe via `PoisonError::into_inner` and surfaces a warning, and
|
||||
/// `apply_metadata_filter` skips the index intersection (retain all) with its
|
||||
/// own warning rather than dropping every candidate.
|
||||
#[test]
|
||||
fn search_poisoned_universe_surfaces_warning_not_empty() {
|
||||
use std::{panic::catch_unwind, sync::Arc};
|
||||
|
||||
let schema = test_schema();
|
||||
let ledger = SignalLedger::new(schema, Box::new(NoopWalWriter));
|
||||
let profile_reg = setup_registry();
|
||||
let registry = RwLock::new(registry_with_slot(
|
||||
"content",
|
||||
&[(1, [1.0, 0.0]), (2, [1.0, 0.0])],
|
||||
));
|
||||
|
||||
// Poison the universe lock by panicking while holding the write guard.
|
||||
let universe = Arc::new(two_item_universe());
|
||||
let poison_target = Arc::clone(&universe);
|
||||
let _ = catch_unwind(move || {
|
||||
let _guard = poison_target.write().unwrap();
|
||||
panic!("intentional panic to poison the universe lock");
|
||||
});
|
||||
assert!(universe.is_poisoned(), "universe lock should be poisoned");
|
||||
|
||||
let exec = filtered_executor(&ledger, &profile_reg, ®istry, &universe);
|
||||
// `&universe` (`&Arc<RwLock<_>>`) deref-coerces to the `&RwLock<_>` the
|
||||
// executor expects.
|
||||
let query = Search::builder()
|
||||
.vector(vec![1.0f32, 0.0])
|
||||
.using_profile("search")
|
||||
.within(crate::query::search::WithinScope::Trending { window_hours: 24 })
|
||||
.filter(FilterExpr::CategoryEq("jazz".into()))
|
||||
.limit(10)
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
let results = exec.execute(&query).unwrap();
|
||||
assert!(
|
||||
results
|
||||
.warnings
|
||||
.iter()
|
||||
.any(|w| w.contains("universe lock poisoned")),
|
||||
"a poisoned universe must surface a warning, got {:?}",
|
||||
results.warnings
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@ -11,7 +11,7 @@ use crate::{
|
||||
cohort::{CohortRegistry, CohortSignalLedger},
|
||||
entities::{CollectionIndex, CreatorItemsBitmap, UserStateIndex},
|
||||
query::{retrieve::QueryError, search::types::WithinScope},
|
||||
schema::Window,
|
||||
schema::{Timestamp, Window},
|
||||
signals::SignalLedger,
|
||||
storage::indexes::bitmap::BitmapIndex,
|
||||
};
|
||||
@ -128,7 +128,11 @@ impl<'a> ScopeResolver<'a> {
|
||||
return self.universe.clone();
|
||||
}
|
||||
|
||||
// Accumulate per-entity velocity from the ledger entries.
|
||||
// Accumulate per-entity velocity from the ledger entries. Read the clock
|
||||
// once for the whole scan and pass it to every windowed read so a viral
|
||||
// item that has gone quiet ages out of the trending set instead of
|
||||
// reporting its peak velocity forever (CRITICAL #6).
|
||||
let now_ns = Timestamp::now().as_nanos();
|
||||
let mut entity_velocity: HashMap<u64, u64> = HashMap::new();
|
||||
|
||||
for entry in self.ledger.entries() {
|
||||
@ -138,7 +142,7 @@ impl<'a> ScopeResolver<'a> {
|
||||
if !is_view && !is_share {
|
||||
continue;
|
||||
}
|
||||
let count = entry.value().warm.windowed_count(window);
|
||||
let count = entry.value().warm.windowed_count(window, now_ns);
|
||||
*entity_velocity.entry(entity_id.as_u64()).or_insert(0) += count;
|
||||
}
|
||||
|
||||
@ -185,11 +189,14 @@ impl<'a> ScopeResolver<'a> {
|
||||
return Ok(self.universe.clone());
|
||||
}
|
||||
|
||||
// Read the clock once for the whole scan and pass it to every windowed
|
||||
// read so quiet cohort items age out of CohortTrending (CRITICAL #6).
|
||||
let now_ns = Timestamp::now().as_nanos();
|
||||
let mut entity_velocity: HashMap<u64, u64> = HashMap::new();
|
||||
|
||||
for entry in cohort_ledger.iter_entries() {
|
||||
let (ref entry_cohort, entity_id, type_id) = *entry.key();
|
||||
if entry_cohort != cohort {
|
||||
if entry_cohort.as_ref() != cohort {
|
||||
continue;
|
||||
}
|
||||
let is_view = view_id.is_some_and(|id| type_id == id);
|
||||
@ -197,7 +204,7 @@ impl<'a> ScopeResolver<'a> {
|
||||
if !is_view && !is_share {
|
||||
continue;
|
||||
}
|
||||
let count = entry.value().warm.windowed_count(window);
|
||||
let count = entry.value().warm.windowed_count(window, now_ns);
|
||||
*entity_velocity.entry(entity_id.as_u64()).or_insert(0) += count;
|
||||
}
|
||||
|
||||
|
||||
@ -295,10 +295,11 @@ fn resolve_cohort_trending_returns_top_items() {
|
||||
let ts = Timestamp::now().as_nanos();
|
||||
|
||||
// Record signals in cohort.
|
||||
let tech: std::sync::Arc<str> = std::sync::Arc::from("tech");
|
||||
for _ in 0..10 {
|
||||
cohort_ledger.record("tech", crate::schema::EntityId::new(1), view_id, 1.0, ts);
|
||||
cohort_ledger.record(&tech, crate::schema::EntityId::new(1), view_id, 1.0, ts);
|
||||
}
|
||||
cohort_ledger.record("tech", crate::schema::EntityId::new(2), view_id, 1.0, ts);
|
||||
cohort_ledger.record(&tech, crate::schema::EntityId::new(2), view_id, 1.0, ts);
|
||||
|
||||
let resolver =
|
||||
ScopeResolver::new(&ledger, build_universe(3)).with_cohort(&cohort_ledger, ®istry);
|
||||
|
||||
@ -10,10 +10,10 @@
|
||||
|
||||
use std::sync::{
|
||||
RwLock,
|
||||
atomic::{AtomicU64, Ordering},
|
||||
atomic::{AtomicU64, AtomicUsize, Ordering},
|
||||
};
|
||||
|
||||
use dashmap::DashMap;
|
||||
use dashmap::{DashMap, mapref::entry::Entry};
|
||||
|
||||
use crate::schema::EntityId;
|
||||
|
||||
@ -80,6 +80,17 @@ pub struct SuggestionIndex {
|
||||
terms: RwLock<Vec<String>>,
|
||||
/// Query frequency counter for trending searches.
|
||||
query_freq: DashMap<String, AtomicU64>,
|
||||
/// Distinct-key count for `query_freq`, used to enforce the
|
||||
/// `MAX_TRENDING_KEYS` cap atomically.
|
||||
///
|
||||
/// `DashMap::len()` is racy against concurrent inserts on other shards, so a
|
||||
/// plain `len() < cap` check-then-insert can let several threads each pass
|
||||
/// the check and overshoot the cap. This counter is reserved via a CAS
|
||||
/// *before* the insert is committed, making "is there room?" and "insert"
|
||||
/// a single atomic decision. It is kept in lockstep with the map: the only
|
||||
/// inserts go through [`SuggestionIndex::record_query`] and the only removals
|
||||
/// through [`SuggestionIndex::evict_low_frequency`], both of which adjust it.
|
||||
key_count: AtomicUsize,
|
||||
}
|
||||
|
||||
impl SuggestionIndex {
|
||||
@ -89,6 +100,7 @@ impl SuggestionIndex {
|
||||
Self {
|
||||
terms: RwLock::new(Vec::new()),
|
||||
query_freq: DashMap::new(),
|
||||
key_count: AtomicUsize::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
@ -140,17 +152,47 @@ impl SuggestionIndex {
|
||||
if normalized.is_empty() {
|
||||
return;
|
||||
}
|
||||
// Fast path: entry already exists — just increment.
|
||||
if let Some(entry) = self.query_freq.get(&normalized) {
|
||||
entry.fetch_add(1, Ordering::Relaxed); // Approximate: Relaxed is correct, no sync needed
|
||||
return;
|
||||
// `entry` holds the shard lock for `normalized`, so the Occupied/Vacant
|
||||
// decision and the insert are atomic *for this key*. To make the GLOBAL
|
||||
// cap atomic too, a vacant key reserves a slot via CAS on `key_count`
|
||||
// BEFORE committing the insert: if the reservation would exceed
|
||||
// `MAX_TRENDING_KEYS`, the key is dropped without inserting. This closes
|
||||
// the check-then-insert race where two threads both saw `len() < cap`
|
||||
// and both inserted, overshooting the cap.
|
||||
match self.query_freq.entry(normalized) {
|
||||
Entry::Occupied(occupied) => {
|
||||
// Approximate: Relaxed is correct, no cross-thread sync needed.
|
||||
occupied.get().fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
Entry::Vacant(vacant) => {
|
||||
if Self::reserve_key_slot(&self.key_count) {
|
||||
vacant.insert(AtomicU64::new(1));
|
||||
}
|
||||
// Cap reached: drop the vacant entry without inserting.
|
||||
}
|
||||
}
|
||||
// Slow path: insert new entry only if below the cap.
|
||||
if self.query_freq.len() < MAX_TRENDING_KEYS {
|
||||
self.query_freq
|
||||
.entry(normalized)
|
||||
.or_insert_with(|| AtomicU64::new(0))
|
||||
.fetch_add(1, Ordering::Relaxed); // Approximate: Relaxed is correct, no sync needed
|
||||
}
|
||||
|
||||
/// Atomically reserve one slot under the `MAX_TRENDING_KEYS` cap.
|
||||
///
|
||||
/// Returns `true` and increments `key_count` iff there was room; returns
|
||||
/// `false` without mutating on a full map. The CAS loop guarantees no two
|
||||
/// threads can both reserve the last slot.
|
||||
fn reserve_key_slot(key_count: &AtomicUsize) -> bool {
|
||||
let mut current = key_count.load(Ordering::Relaxed);
|
||||
loop {
|
||||
if current >= MAX_TRENDING_KEYS {
|
||||
return false;
|
||||
}
|
||||
match key_count.compare_exchange_weak(
|
||||
current,
|
||||
current + 1,
|
||||
Ordering::Relaxed,
|
||||
Ordering::Relaxed,
|
||||
) {
|
||||
Ok(_) => return true,
|
||||
Err(observed) => current = observed,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -227,7 +269,12 @@ impl SuggestionIndex {
|
||||
pairs.sort_by_key(|&(_, f)| f);
|
||||
let remove_count = pairs.len() / 2;
|
||||
for (key, _) in pairs.into_iter().take(remove_count) {
|
||||
self.query_freq.remove(&key);
|
||||
// Keep `key_count` in lockstep with the map: decrement only when a
|
||||
// key was actually present and removed, so a concurrent
|
||||
// `record_query` removing/inserting the same key cannot double-count.
|
||||
if self.query_freq.remove(&key).is_some() {
|
||||
self.key_count.fetch_sub(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -364,9 +411,55 @@ mod tests {
|
||||
idx.record_query(&format!("query_{i}"));
|
||||
}
|
||||
assert_eq!(idx.query_freq.len(), MAX_TRENDING_KEYS);
|
||||
// The reserved-slot counter must track the map exactly.
|
||||
assert_eq!(idx.key_count.load(Ordering::Relaxed), MAX_TRENDING_KEYS);
|
||||
// New unique key should be silently dropped.
|
||||
idx.record_query("brand_new_unique_query");
|
||||
assert_eq!(idx.query_freq.len(), MAX_TRENDING_KEYS);
|
||||
assert_eq!(idx.key_count.load(Ordering::Relaxed), MAX_TRENDING_KEYS);
|
||||
}
|
||||
|
||||
/// Accuracy-S regression: the cap must hold under concurrent inserts. The
|
||||
/// previous `len() < cap` check-then-insert let multiple threads each pass
|
||||
/// the check and overshoot `MAX_TRENDING_KEYS`. With the atomic
|
||||
/// reserve-then-insert path, the distinct-key count can NEVER exceed the cap,
|
||||
/// no matter how many threads race.
|
||||
#[test]
|
||||
fn record_query_cap_holds_under_concurrency() {
|
||||
use std::{sync::Arc, thread};
|
||||
|
||||
let idx = Arc::new(SuggestionIndex::new());
|
||||
// Pre-fill to one slot below the cap so every racing insert competes for
|
||||
// the SAME remaining headroom — the worst case for the old race.
|
||||
for i in 0..(MAX_TRENDING_KEYS - 1) {
|
||||
idx.record_query(&format!("seed_{i}"));
|
||||
}
|
||||
assert_eq!(idx.query_freq.len(), MAX_TRENDING_KEYS - 1);
|
||||
|
||||
// Many threads each try to insert a DISTINCT brand-new key. At most ONE
|
||||
// may succeed (one slot left); all others must be rejected.
|
||||
let threads: Vec<_> = (0..32)
|
||||
.map(|t| {
|
||||
let idx = Arc::clone(&idx);
|
||||
thread::spawn(move || {
|
||||
idx.record_query(&format!("racer_{t}"));
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
for h in threads {
|
||||
h.join().unwrap();
|
||||
}
|
||||
|
||||
assert_eq!(
|
||||
idx.query_freq.len(),
|
||||
MAX_TRENDING_KEYS,
|
||||
"concurrent inserts must not exceed the cap"
|
||||
);
|
||||
assert_eq!(
|
||||
idx.key_count.load(Ordering::Relaxed),
|
||||
MAX_TRENDING_KEYS,
|
||||
"reserved-slot counter must equal the map size at the cap"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@ -31,14 +31,25 @@ pub struct DiversitySelector;
|
||||
impl DiversitySelector {
|
||||
/// Select up to `target_count` candidates satisfying diversity constraints.
|
||||
///
|
||||
/// Uses a three-stage relaxation strategy:
|
||||
/// Uses a four-stage relaxation strategy:
|
||||
/// - Stage 0: full constraints
|
||||
/// - Stage 1: `max_per_creator` doubled (if used)
|
||||
/// - Stage 2: `format_mix` ignored
|
||||
/// - Stage 3: accept anything (fill to target)
|
||||
///
|
||||
/// INV-RANK-5: `result.selected.len() == min(target_count, candidates.len())` always.
|
||||
/// INV-RANK-6: score order preserved within same-creator groups.
|
||||
/// # Guaranteed invariants
|
||||
///
|
||||
/// - **INV-RANK-5:** `result.selected.len() == min(target_count, candidates.len())`
|
||||
/// always -- the relaxation stages fill to the target even when constraints
|
||||
/// are unsatisfiable.
|
||||
/// - **INV-RANK-6:** *global* descending score order is preserved across the
|
||||
/// entire result, not merely within same-creator groups. The relaxation
|
||||
/// stages only decide *which* candidates are accepted; the final result is
|
||||
/// emitted by a single pass over the already-score-sorted `candidates`
|
||||
/// (retaining accepted IDs), so every emitted item precedes every
|
||||
/// lower-scored emitted item regardless of creator or format. A
|
||||
/// `debug_assert!` below enforces the descending-input precondition this
|
||||
/// guarantee relies on.
|
||||
#[must_use]
|
||||
pub fn select(
|
||||
candidates: &[ScoredCandidate],
|
||||
|
||||
@ -12,6 +12,24 @@ use crate::{
|
||||
signals::SignalLedger,
|
||||
};
|
||||
|
||||
// -- Degradation window substitution ------------------------------------------
|
||||
|
||||
/// Coarsest velocity bucket used when `DegradationLevel::coarsens_aggregates()`.
|
||||
///
|
||||
/// Under coarse-aggregate degradation, a `SignalAgg::Velocity` read is forced to
|
||||
/// this window (the coarsest velocity bucket) to cut fine-grained ledger scan
|
||||
/// cost. Exposed at module scope so the social-graph trending path in
|
||||
/// [`super::scoring`] -- which reads velocity directly off the per-user signal
|
||||
/// index rather than through [`read_agg`] -- can read the *same* window and stay
|
||||
/// in lockstep with the substitution rule here.
|
||||
pub(super) const COARSE_VELOCITY_WINDOW: Window = Window::TwentyFourHours;
|
||||
|
||||
/// Coarsest count bucket used when `DegradationLevel::coarsens_aggregates()`.
|
||||
///
|
||||
/// Under coarse-aggregate degradation, a `SignalAgg::Value` read is forced to the
|
||||
/// warm-tier all-time count (O(1)) instead of a fine-grained windowed scan.
|
||||
pub(super) const COARSE_VALUE_WINDOW: Window = Window::AllTime;
|
||||
|
||||
// -- Signal reading -----------------------------------------------------------
|
||||
|
||||
/// Read a signal aggregation for a candidate.
|
||||
@ -43,8 +61,8 @@ pub(crate) fn read_agg(
|
||||
) -> crate::Result<f64> {
|
||||
let window = if degradation.coarsens_aggregates() {
|
||||
match agg {
|
||||
SignalAgg::Value => Window::AllTime,
|
||||
SignalAgg::Velocity => Window::TwentyFourHours,
|
||||
SignalAgg::Value => COARSE_VALUE_WINDOW,
|
||||
SignalAgg::Velocity => COARSE_VELOCITY_WINDOW,
|
||||
_ => window,
|
||||
}
|
||||
} else {
|
||||
@ -191,28 +209,68 @@ pub(super) fn passes_excludes(
|
||||
|
||||
/// Min-max normalize candidate scores to `[0.0, 1.0]`.
|
||||
///
|
||||
/// If all candidates have the same score, they are all set to 1.0.
|
||||
/// If all (finite) candidates have the same score, they are all set to 1.0.
|
||||
///
|
||||
/// # Infinite "sort last / first" sentinels
|
||||
///
|
||||
/// Several sort modes emit `f64::NEG_INFINITY` (or `f64::INFINITY`) as a
|
||||
/// deliberate "sort last / first" sentinel for missing metadata (`Shortest`,
|
||||
/// `Longest`, `DateSaved`). `f64::total_cmp` (used by `finalize`'s sort) already
|
||||
/// ranks those infinities at the correct extreme, so the *order* is right by the
|
||||
/// time we get here. Normalization must keep `score` **monotonic with that rank**.
|
||||
///
|
||||
/// The min/max range is therefore computed over the **finite** scores only, and
|
||||
/// the infinite sentinels are mapped directly to the matching end of the output
|
||||
/// range afterwards: `NEG_INFINITY -> 0.0` (rank-last bottom), `INFINITY -> 1.0`
|
||||
/// (rank-first top). This fixes the anomaly where clamping `NEG_INFINITY -> 0.0`
|
||||
/// *before* taking the min let a missing item become the maximum on a negated
|
||||
/// scale: for `Shortest`, every present item has a *negative* score (`-duration`),
|
||||
/// so a pre-clamped `0.0` floated above them and the last-ranked (missing) item
|
||||
/// reported the highest score `1.0`. Computing the range over finite scores and
|
||||
/// folding sentinels to the correct end keeps the reported score consistent with
|
||||
/// the rank for the whole class of `±Inf`-sentinel sorts.
|
||||
///
|
||||
/// `NaN` has no defined order and is neutralized to `0.0` at score *construction*
|
||||
/// (`finite_score`) before it can reach the sort, so it does not appear here; any
|
||||
/// `NaN` that slipped through is treated as a finite `0.0` for robustness.
|
||||
pub(super) fn normalize(candidates: &mut [ScoredCandidate]) {
|
||||
if candidates.is_empty() {
|
||||
return;
|
||||
}
|
||||
// Clamp non-finite scores (NaN, +/-Inf) to 0.0 before normalization.
|
||||
// Defensive: a stray NaN (which has no order) is mapped to the neutral 0.0,
|
||||
// matching `finite_score`. `±Inf` sentinels are intentionally preserved --
|
||||
// they are folded to the correct end of the output range below.
|
||||
for c in candidates.iter_mut() {
|
||||
if !c.score.is_finite() {
|
||||
if c.score.is_nan() {
|
||||
c.score = 0.0;
|
||||
}
|
||||
}
|
||||
let min = candidates
|
||||
.iter()
|
||||
.map(|c| c.score)
|
||||
.fold(f64::INFINITY, f64::min);
|
||||
let max = candidates
|
||||
.iter()
|
||||
.map(|c| c.score)
|
||||
.fold(f64::NEG_INFINITY, f64::max);
|
||||
// Range is measured over FINITE scores only, so an infinite sentinel never
|
||||
// skews min/max (and so never becomes the artificial extreme of the range).
|
||||
let mut min = f64::INFINITY;
|
||||
let mut max = f64::NEG_INFINITY;
|
||||
for c in candidates.iter() {
|
||||
if c.score.is_finite() {
|
||||
min = min.min(c.score);
|
||||
max = max.max(c.score);
|
||||
}
|
||||
}
|
||||
// No finite scores at all (every candidate is a ±Inf sentinel): fold each to
|
||||
// its end of the range directly. With no finite reference, NEG_INFINITY -> 0,
|
||||
// INFINITY -> 1. (`is_sign_positive()` distinguishes the two infinities.)
|
||||
if !min.is_finite() {
|
||||
for c in candidates.iter_mut() {
|
||||
c.score = if c.score.is_sign_positive() { 1.0 } else { 0.0 };
|
||||
}
|
||||
return;
|
||||
}
|
||||
let range = max - min;
|
||||
for c in candidates.iter_mut() {
|
||||
c.score = if range < f64::EPSILON {
|
||||
c.score = if c.score.is_infinite() {
|
||||
// Sentinel folds to the matching end of the output range:
|
||||
// +Inf "sort first" -> top (1.0); -Inf "sort last" -> bottom (0.0).
|
||||
if c.score.is_sign_positive() { 1.0 } else { 0.0 }
|
||||
} else if range < f64::EPSILON {
|
||||
1.0
|
||||
} else {
|
||||
(c.score - min) / range
|
||||
@ -225,47 +283,16 @@ pub(super) fn normalize(candidates: &mut [ScoredCandidate]) {
|
||||
#[cfg(test)]
|
||||
#[allow(clippy::unwrap_used, clippy::float_cmp, clippy::cast_precision_loss)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use super::*;
|
||||
use crate::{
|
||||
schema::{DecaySpec, EntityId, EntityKind, SchemaBuilder, Timestamp},
|
||||
signals::NoopWalWriter,
|
||||
// Shared ledger fixtures (DRY-S, M0-M10 review).
|
||||
ranking::test_fixtures::{ledger_for, seeded_ledger},
|
||||
schema::Timestamp,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn read_agg_unimplemented_aggregation_errors() {
|
||||
let mut builder = SchemaBuilder::new();
|
||||
for sig in &["view", "share", "like"] {
|
||||
let _ = builder
|
||||
.signal(
|
||||
sig,
|
||||
EntityKind::Item,
|
||||
DecaySpec::Exponential {
|
||||
half_life: Duration::from_secs(7 * 24 * 3600),
|
||||
},
|
||||
)
|
||||
.windows(&[Window::OneHour, Window::TwentyFourHours, Window::SevenDays])
|
||||
.velocity(true)
|
||||
.add();
|
||||
}
|
||||
let schema = builder.build().unwrap();
|
||||
let ledger = SignalLedger::new(schema, Box::new(NoopWalWriter));
|
||||
|
||||
let base_ns = 1_708_000_000_000_000_000u64;
|
||||
for i in 0u64..5 {
|
||||
let entity_id = EntityId::new(i + 1);
|
||||
let ts = Timestamp::from_nanos(base_ns - i * 3_600_000_000_000);
|
||||
ledger
|
||||
.record_signal("view", entity_id, (5 - i) as f64, ts)
|
||||
.unwrap();
|
||||
ledger
|
||||
.record_signal("share", entity_id, (i % 3) as f64, ts)
|
||||
.unwrap();
|
||||
ledger
|
||||
.record_signal("like", entity_id, (i % 2) as f64, ts)
|
||||
.unwrap();
|
||||
}
|
||||
let ledger = seeded_ledger();
|
||||
|
||||
let entity_id = EntityId::new(1);
|
||||
// SignalAgg::Ratio is not implemented; reaching read_agg with it is an
|
||||
@ -284,20 +311,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn read_agg_unknown_signal_errors() {
|
||||
let mut builder = SchemaBuilder::new();
|
||||
let _ = builder
|
||||
.signal(
|
||||
"view",
|
||||
EntityKind::Item,
|
||||
DecaySpec::Exponential {
|
||||
half_life: Duration::from_secs(3600),
|
||||
},
|
||||
)
|
||||
.windows(&[Window::OneHour])
|
||||
.velocity(false)
|
||||
.add();
|
||||
let schema = builder.build().unwrap();
|
||||
let ledger = SignalLedger::new(schema, Box::new(NoopWalWriter));
|
||||
let ledger = ledger_for(&["view"]);
|
||||
// "nonexistent" is not in the schema -> schema error propagates.
|
||||
let result = read_agg(
|
||||
EntityId::new(1),
|
||||
@ -312,20 +326,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn passes_gates_below_threshold_excluded() {
|
||||
let mut builder = SchemaBuilder::new();
|
||||
let _ = builder
|
||||
.signal(
|
||||
"view",
|
||||
EntityKind::Item,
|
||||
DecaySpec::Exponential {
|
||||
half_life: Duration::from_secs(3600),
|
||||
},
|
||||
)
|
||||
.windows(&[Window::OneHour])
|
||||
.velocity(false)
|
||||
.add();
|
||||
let schema = builder.build().unwrap();
|
||||
let ledger = SignalLedger::new(schema, Box::new(NoopWalWriter));
|
||||
let ledger = ledger_for(&["view"]);
|
||||
let entity_id = EntityId::new(1);
|
||||
// Record 3 signals; windowed count = 3, below threshold of 5.
|
||||
for _ in 0..3 {
|
||||
@ -345,20 +346,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn passes_gates_at_threshold_included() {
|
||||
let mut builder = SchemaBuilder::new();
|
||||
let _ = builder
|
||||
.signal(
|
||||
"view",
|
||||
EntityKind::Item,
|
||||
DecaySpec::Exponential {
|
||||
half_life: Duration::from_secs(3600),
|
||||
},
|
||||
)
|
||||
.windows(&[Window::OneHour])
|
||||
.velocity(false)
|
||||
.add();
|
||||
let schema = builder.build().unwrap();
|
||||
let ledger = SignalLedger::new(schema, Box::new(NoopWalWriter));
|
||||
let ledger = ledger_for(&["view"]);
|
||||
let entity_id = EntityId::new(1);
|
||||
// Record exactly 5 signals; windowed count = 5, meets threshold exactly.
|
||||
for _ in 0..5 {
|
||||
@ -389,6 +377,67 @@ mod tests {
|
||||
assert_eq!(candidates[0].score, 1.0);
|
||||
}
|
||||
|
||||
fn candidate(id: u64, score: f64) -> ScoredCandidate {
|
||||
ScoredCandidate {
|
||||
entity_id: EntityId::new(id),
|
||||
score,
|
||||
signal_snapshot: vec![],
|
||||
creator_id: None,
|
||||
format: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_neg_inf_sentinel_on_negated_scale_folds_to_bottom() {
|
||||
// Regression (Accuracy-W): `Shortest` scores present items as `-duration`
|
||||
// (all negative) and a missing item as NEG_INFINITY. The candidates arrive
|
||||
// already sorted descending: -60 (60s, shortest), -120 (120s), -Inf
|
||||
// (missing). Before the fix, NEG_INFINITY was clamped to 0.0 *before*
|
||||
// taking the min, so it became the maximum on the negated scale and
|
||||
// normalized to 1.0 -- the last-ranked item reporting the highest score.
|
||||
let mut candidates = vec![
|
||||
candidate(3, -60.0),
|
||||
candidate(1, -120.0),
|
||||
candidate(2, f64::NEG_INFINITY),
|
||||
];
|
||||
normalize(&mut candidates);
|
||||
|
||||
// Finite range is [-120, -60]; the shortest (−60) maps to 1.0, the 120s
|
||||
// (−120) to 0.0, and the missing sentinel folds to the bottom (0.0).
|
||||
assert_eq!(candidates[0].score, 1.0);
|
||||
assert_eq!(candidates[1].score, 0.0);
|
||||
assert_eq!(candidates[2].score, 0.0);
|
||||
// Score stays monotonic with rank -- the tail never outscores the head.
|
||||
assert!(candidates.last().unwrap().score <= candidates.first().unwrap().score);
|
||||
for w in candidates.windows(2) {
|
||||
assert!(w[0].score >= w[1].score);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_pos_inf_sentinel_folds_to_top() {
|
||||
// The `+Inf` "sort first" sentinel folds to the top of the output range,
|
||||
// while finite scores are min-max normalized over their own range.
|
||||
let mut candidates = vec![
|
||||
candidate(1, f64::INFINITY),
|
||||
candidate(2, 10.0),
|
||||
candidate(3, 2.0),
|
||||
];
|
||||
normalize(&mut candidates);
|
||||
assert_eq!(candidates[0].score, 1.0); // +Inf -> top
|
||||
assert_eq!(candidates[1].score, 1.0); // max finite -> top of finite range
|
||||
assert_eq!(candidates[2].score, 0.0); // min finite -> bottom
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_all_infinite_sentinels_fold_to_ends() {
|
||||
// No finite scores at all: each sentinel folds directly to its end.
|
||||
let mut candidates = vec![candidate(1, f64::INFINITY), candidate(2, f64::NEG_INFINITY)];
|
||||
normalize(&mut candidates);
|
||||
assert_eq!(candidates[0].score, 1.0);
|
||||
assert_eq!(candidates[1].score, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_all_nan_clamps_to_one() {
|
||||
let mut candidates: Vec<ScoredCandidate> = (1u64..=3)
|
||||
|
||||
@ -49,7 +49,10 @@ use crate::{
|
||||
/// metadata (`Shortest`, `Longest`, `DateSaved`), and `f64::total_cmp` (used
|
||||
/// for the sort below) is a true total order over `[-Inf, +Inf]`, so an
|
||||
/// infinite score ranks correctly without any clamp. [`normalize`] then folds
|
||||
/// the surviving infinities to `0.0` for the `[0, 1]` output range.
|
||||
/// the surviving sentinels to the matching *end* of the `[0, 1]` output range
|
||||
/// (`-Inf -> 0.0`, `+Inf -> 1.0`), computing the min/max range over the finite
|
||||
/// scores only so a missing-metadata item stays monotonic with its rank rather
|
||||
/// than reporting the maximum score (see [`normalize`]).
|
||||
#[inline]
|
||||
const fn finite_score(raw: f64) -> f64 {
|
||||
if raw.is_nan() { 0.0 } else { raw }
|
||||
@ -169,7 +172,16 @@ impl<'a> ProfileExecutor<'a> {
|
||||
///
|
||||
/// - `candidates`: entity IDs to consider; excluded/gate-failing candidates are removed.
|
||||
/// - `profile`: the ranking profile controlling sort mode, boosts, penalties, gates, excludes, and decay.
|
||||
/// - `now`: the current timestamp, used for decay calculations.
|
||||
/// - `now`: the query's logical timestamp. **Clock contract:** decay is
|
||||
/// currently evaluated at *call time*, not at `now`. Every time-dependent
|
||||
/// read -- `SignalAgg::DecayScore` and the `ProfileDecay` recency factor --
|
||||
/// decays forward to wall-clock `Timestamp::now()` *inside* the
|
||||
/// [`SignalLedger`], so `now` does not yet influence those reads. The
|
||||
/// parameter is retained because it is part of the public scoring API and is
|
||||
/// threaded by the RETRIEVE/SEARCH callers and benchmarks; making scoring a
|
||||
/// pure function of `(ledger snapshot, now)` requires a `read_decay_score_at`
|
||||
/// on the ledger (out of scope for the ranking layer). See the M0-M10 review
|
||||
/// (scoring.rs Maintainability-W).
|
||||
///
|
||||
/// Returns candidates sorted by descending normalized score in `[0.0, 1.0]`.
|
||||
///
|
||||
|
||||
@ -26,6 +26,11 @@ impl ProfileExecutor<'_> {
|
||||
/// Returns `(score, snapshot)` where `snapshot` lists the raw signal values
|
||||
/// that contributed to the score, for explain-ability in API responses.
|
||||
///
|
||||
/// **Clock contract:** `now` is threaded for API symmetry but does not
|
||||
/// currently steer any read -- every time-dependent sort (`DecayScore`-backed
|
||||
/// modes, velocity modes) decays/ages forward to wall-clock time *inside* the
|
||||
/// [`SignalLedger`] at call time. See [`ProfileExecutor::score`]'s `now` doc.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Sort modes read a fixed generic signal vocabulary via
|
||||
@ -114,6 +119,13 @@ impl ProfileExecutor<'_> {
|
||||
}
|
||||
}
|
||||
|
||||
/// `_now` is accepted for signature symmetry with [`Self::score_by_sort`] but
|
||||
/// is intentionally unused: Hot scoring reads only the all-time view count
|
||||
/// (decayed forward to wall-clock time inside the ledger) and applies a
|
||||
/// *uniform* 24h age term, because no per-entity `created_at` is plumbed into
|
||||
/// the executor (see the `age_hours` comment below). When age-aware Hot
|
||||
/// scoring lands, this `now` becomes the age reference. See the M0-M10 review
|
||||
/// clock-contract note (scoring.rs Maintainability-W).
|
||||
fn score_hot(
|
||||
&self,
|
||||
entity_id: EntityId,
|
||||
@ -172,11 +184,25 @@ impl ProfileExecutor<'_> {
|
||||
let view_type_id = self.ledger.resolve_signal_type("view");
|
||||
let share_type_id = self.ledger.resolve_signal_type("share");
|
||||
|
||||
// Degradation exemption (Tech Debt-S, M0-M10 review): this social-graph
|
||||
// path does NOT route through `read_agg`'s degradation window
|
||||
// substitution, and it does not need to. `read_agg` substitutes
|
||||
// `SignalAgg::Velocity` reads to `helpers::COARSE_VELOCITY_WINDOW`
|
||||
// (`Window::TwentyFourHours`) under `coarsens_aggregates()`; the
|
||||
// trending velocity already reads that exact window (the coarsest
|
||||
// velocity bucket, matching the global-ledger fallback below). So the
|
||||
// substitution would be a no-op here -- the path is already at the
|
||||
// coarse-aggregate window in every degradation level. Reusing the
|
||||
// shared `COARSE_VELOCITY_WINDOW` constant keeps this site in lockstep
|
||||
// with `read_agg` so the two can never silently diverge if the
|
||||
// velocity coarsening target ever changes.
|
||||
let vel_window = super::helpers::COARSE_VELOCITY_WINDOW;
|
||||
|
||||
let view_vel = view_type_id.map_or(0.0, |tid| {
|
||||
user_signal_idx.aggregate_velocity(entity_id, users, tid, Window::TwentyFourHours)
|
||||
user_signal_idx.aggregate_velocity(entity_id, users, tid, vel_window)
|
||||
});
|
||||
let share_vel = share_type_id.map_or(0.0, |tid| {
|
||||
user_signal_idx.aggregate_velocity(entity_id, users, tid, Window::TwentyFourHours)
|
||||
user_signal_idx.aggregate_velocity(entity_id, users, tid, vel_window)
|
||||
});
|
||||
return Ok((
|
||||
trending_score(view_vel, share_vel),
|
||||
@ -399,15 +425,21 @@ impl ProfileExecutor<'_> {
|
||||
}
|
||||
|
||||
/// Score for `Shortest`: read "duration" metadata, return `-duration` so
|
||||
/// shorter items score higher (more negative = lower after negation).
|
||||
/// Missing duration returns `f64::NEG_INFINITY` (sorted last).
|
||||
/// shorter items score higher (less negative = higher after negation).
|
||||
/// Missing duration returns `f64::NEG_INFINITY`, a "sort last" sentinel that
|
||||
/// [`super::helpers::normalize`] folds to the bottom of the `[0, 1]` output
|
||||
/// range (computing the min/max range over the finite scores only), so the
|
||||
/// missing item ranks last *and* reports the lowest score -- monotonic with
|
||||
/// rank, never the spurious maximum `1.0`.
|
||||
fn score_shortest(&self, entity_id: EntityId) -> f64 {
|
||||
self.read_duration(entity_id)
|
||||
.map_or(f64::NEG_INFINITY, |dur| -dur)
|
||||
}
|
||||
|
||||
/// Score for `Longest`: read "duration" metadata, return `duration` directly.
|
||||
/// Missing duration returns `f64::NEG_INFINITY` (sorted last).
|
||||
/// Missing duration returns `f64::NEG_INFINITY`, a "sort last" sentinel that
|
||||
/// [`super::helpers::normalize`] folds to the bottom of the `[0, 1]` output
|
||||
/// range (see [`Self::score_shortest`]).
|
||||
fn score_longest(&self, entity_id: EntityId) -> f64 {
|
||||
self.read_duration(entity_id).unwrap_or(f64::NEG_INFINITY)
|
||||
}
|
||||
@ -422,7 +454,9 @@ impl ProfileExecutor<'_> {
|
||||
|
||||
/// Score for `DateSaved`: look up the save timestamp from user state.
|
||||
/// Returns timestamp as f64 (later saves = higher score = first in descending).
|
||||
/// Missing save returns `f64::NEG_INFINITY` (sorted last).
|
||||
/// Missing save returns `f64::NEG_INFINITY`, a "sort last" sentinel that
|
||||
/// [`super::helpers::normalize`] folds to the bottom of the `[0, 1]` output
|
||||
/// range (see [`Self::score_shortest`]).
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
fn score_date_saved(&self, entity_id: EntityId) -> f64 {
|
||||
self.user_state_for_date_saved
|
||||
|
||||
@ -1,74 +1,28 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use super::*;
|
||||
use crate::{
|
||||
ranking::{
|
||||
builtins::register_builtins,
|
||||
profile::{
|
||||
Boost, CandidateStrategy, DiversitySpec, Exclude, Gate, Penalty, ProfileDecay,
|
||||
SignalAgg, Sort,
|
||||
},
|
||||
profile::{Boost, Exclude, Gate, Penalty, ProfileDecay, SignalAgg, Sort},
|
||||
registry::ProfileRegistry,
|
||||
// Shared fixtures (extracted to de-duplicate the per-module ledger/profile
|
||||
// builders across the five ranking test modules; DRY-S, M0-M10 review).
|
||||
test_fixtures::{ledger_for, profile_with, seeded_ledger},
|
||||
},
|
||||
schema::{DecaySpec, EntityKind, SchemaBuilder, Window},
|
||||
signals::{NoopWalWriter, SignalLedger},
|
||||
schema::EntityId,
|
||||
signals::SignalLedger,
|
||||
};
|
||||
|
||||
/// Re-export the canonical seeded ledger under the historical `test_ledger` name
|
||||
/// so `sort_tests::use super::test_ledger` keeps resolving without churn.
|
||||
fn test_ledger() -> SignalLedger {
|
||||
let mut builder = SchemaBuilder::new();
|
||||
for sig in &["view", "share", "like"] {
|
||||
let _ = builder
|
||||
.signal(
|
||||
sig,
|
||||
EntityKind::Item,
|
||||
DecaySpec::Exponential {
|
||||
half_life: Duration::from_secs(7 * 24 * 3600),
|
||||
},
|
||||
)
|
||||
.windows(&[Window::OneHour, Window::TwentyFourHours, Window::SevenDays])
|
||||
.velocity(true)
|
||||
.add();
|
||||
}
|
||||
let schema = builder.build().unwrap();
|
||||
let ledger = SignalLedger::new(schema, Box::new(NoopWalWriter));
|
||||
|
||||
let base_ns = 1_708_000_000_000_000_000u64;
|
||||
for i in 0u64..5 {
|
||||
let entity_id = EntityId::new(i + 1);
|
||||
let ts = Timestamp::from_nanos(base_ns - i * 3_600_000_000_000);
|
||||
ledger
|
||||
.record_signal("view", entity_id, (5 - i) as f64, ts)
|
||||
.unwrap();
|
||||
ledger
|
||||
.record_signal("share", entity_id, (i % 3) as f64, ts)
|
||||
.unwrap();
|
||||
ledger
|
||||
.record_signal("like", entity_id, (i % 2) as f64, ts)
|
||||
.unwrap();
|
||||
}
|
||||
ledger
|
||||
seeded_ledger()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn score_empty_candidates_returns_empty() {
|
||||
let ledger = test_ledger();
|
||||
let executor = ProfileExecutor::new(&ledger);
|
||||
let profile = RankingProfile {
|
||||
name: "test".into(),
|
||||
version: 1,
|
||||
candidate_strategy: CandidateStrategy::Scan {
|
||||
sort_field: "created_at".into(),
|
||||
},
|
||||
boosts: vec![],
|
||||
decay: None,
|
||||
gates: vec![],
|
||||
penalties: vec![],
|
||||
excludes: vec![],
|
||||
diversity: DiversitySpec::default(),
|
||||
exploration: 0.0,
|
||||
sort: Some(Sort::Trending),
|
||||
is_builtin: false,
|
||||
};
|
||||
let profile = profile_with("test", |p| p.sort = Some(Sort::Trending));
|
||||
let result = executor
|
||||
.score(
|
||||
&[],
|
||||
@ -121,22 +75,7 @@ fn score_normalization_all_equal() {
|
||||
let executor = ProfileExecutor::new(&ledger);
|
||||
// Use a profile with no sort and no boosts -- all candidates get 0.0
|
||||
// raw score, which normalizes to 1.0.
|
||||
let profile = RankingProfile {
|
||||
name: "test_equal".into(),
|
||||
version: 1,
|
||||
candidate_strategy: CandidateStrategy::Scan {
|
||||
sort_field: "created_at".into(),
|
||||
},
|
||||
boosts: vec![],
|
||||
decay: None,
|
||||
gates: vec![],
|
||||
penalties: vec![],
|
||||
excludes: vec![],
|
||||
diversity: DiversitySpec::default(),
|
||||
exploration: 0.0,
|
||||
sort: None,
|
||||
is_builtin: false,
|
||||
};
|
||||
let profile = profile_with("test_equal", |_| {});
|
||||
let candidates: Vec<EntityId> = (1..=3).map(EntityId::new).collect();
|
||||
let now = Timestamp::from_nanos(1_708_000_000_000_000_000);
|
||||
|
||||
@ -151,28 +90,16 @@ fn score_normalization_all_equal() {
|
||||
fn score_gate_filters_candidates() {
|
||||
let ledger = test_ledger();
|
||||
let executor = ProfileExecutor::new(&ledger);
|
||||
let profile = RankingProfile {
|
||||
name: "test_gated".into(),
|
||||
version: 1,
|
||||
candidate_strategy: CandidateStrategy::Scan {
|
||||
sort_field: "created_at".into(),
|
||||
},
|
||||
boosts: vec![],
|
||||
decay: None,
|
||||
gates: vec![Gate {
|
||||
let profile = profile_with("test_gated", |p| {
|
||||
p.gates = vec![Gate {
|
||||
signal: "view".into(),
|
||||
agg: SignalAgg::Value,
|
||||
window: Window::AllTime,
|
||||
// Threshold higher than any recorded count -- all filtered.
|
||||
min_threshold: f64::MAX,
|
||||
}],
|
||||
penalties: vec![],
|
||||
excludes: vec![],
|
||||
diversity: DiversitySpec::default(),
|
||||
exploration: 0.0,
|
||||
sort: Some(Sort::Trending),
|
||||
is_builtin: false,
|
||||
};
|
||||
}];
|
||||
p.sort = Some(Sort::Trending);
|
||||
});
|
||||
let candidates: Vec<EntityId> = (1..=5).map(EntityId::new).collect();
|
||||
let now = Timestamp::from_nanos(1_708_000_000_000_000_000);
|
||||
|
||||
@ -184,45 +111,21 @@ fn score_gate_filters_candidates() {
|
||||
fn score_gate_on_missing_signal_errors_loudly() {
|
||||
// A gate referencing a signal the schema does not define is a
|
||||
// misconfigured profile: `score` must fail loudly with a schema error
|
||||
// rather than silently scoring 0.0 and filtering every candidate.
|
||||
let mut builder = SchemaBuilder::new();
|
||||
let _ = builder
|
||||
.signal(
|
||||
"view",
|
||||
EntityKind::Item,
|
||||
DecaySpec::Exponential {
|
||||
half_life: Duration::from_secs(3600),
|
||||
},
|
||||
)
|
||||
.windows(&[Window::OneHour])
|
||||
.velocity(false)
|
||||
.add();
|
||||
let schema = builder.build().unwrap();
|
||||
let ledger = SignalLedger::new(schema, Box::new(NoopWalWriter));
|
||||
// rather than silently scoring 0.0 and filtering every candidate. The schema
|
||||
// defines only "view", so "nonexistent_signal" is genuinely absent.
|
||||
let ledger = ledger_for(&["view"]);
|
||||
let executor = ProfileExecutor::new(&ledger);
|
||||
|
||||
let candidates: Vec<EntityId> = (1..=5).map(EntityId::new).collect();
|
||||
let profile = RankingProfile {
|
||||
name: "test_gate".into(),
|
||||
version: 1,
|
||||
candidate_strategy: CandidateStrategy::Scan {
|
||||
sort_field: "created_at".into(),
|
||||
},
|
||||
boosts: vec![],
|
||||
decay: None,
|
||||
gates: vec![Gate {
|
||||
let profile = profile_with("test_gate", |p| {
|
||||
p.gates = vec![Gate {
|
||||
signal: "nonexistent_signal".into(),
|
||||
agg: SignalAgg::Value,
|
||||
window: Window::AllTime,
|
||||
min_threshold: 1.0,
|
||||
}],
|
||||
penalties: vec![],
|
||||
excludes: vec![],
|
||||
diversity: DiversitySpec::default(),
|
||||
exploration: 0.0,
|
||||
sort: Some(Sort::New),
|
||||
is_builtin: false,
|
||||
};
|
||||
}];
|
||||
p.sort = Some(Sort::New);
|
||||
});
|
||||
let now = Timestamp::from_nanos(1_708_000_000_000_000_000u64);
|
||||
let result = executor.score(&candidates, &profile, now);
|
||||
assert!(
|
||||
@ -242,22 +145,7 @@ fn score_nan_does_not_propagate() {
|
||||
let now = Timestamp::from_nanos(1_708_000_000_000_000_000u64);
|
||||
|
||||
// Test with Trending sort (exercises velocity reads).
|
||||
let profile = RankingProfile {
|
||||
name: "test_nan".into(),
|
||||
version: 1,
|
||||
candidate_strategy: CandidateStrategy::Scan {
|
||||
sort_field: "created_at".into(),
|
||||
},
|
||||
boosts: vec![],
|
||||
decay: None,
|
||||
gates: vec![],
|
||||
penalties: vec![],
|
||||
excludes: vec![],
|
||||
diversity: DiversitySpec::default(),
|
||||
exploration: 0.0,
|
||||
sort: Some(Sort::Trending),
|
||||
is_builtin: false,
|
||||
};
|
||||
let profile = profile_with("test_nan", |p| p.sort = Some(Sort::Trending));
|
||||
let result = executor.score(&candidates, &profile, now).unwrap();
|
||||
for c in &result {
|
||||
assert!(
|
||||
@ -295,22 +183,7 @@ fn score_personalized_boosts_interacted_creators() {
|
||||
let executor = ProfileExecutor::new(&ledger);
|
||||
|
||||
// All candidates get the same base score (no sort, no boosts).
|
||||
let profile = RankingProfile {
|
||||
name: "test_personalized".into(),
|
||||
version: 1,
|
||||
candidate_strategy: CandidateStrategy::Scan {
|
||||
sort_field: "created_at".into(),
|
||||
},
|
||||
boosts: vec![],
|
||||
decay: None,
|
||||
gates: vec![],
|
||||
penalties: vec![],
|
||||
excludes: vec![],
|
||||
diversity: DiversitySpec::default(),
|
||||
exploration: 0.0,
|
||||
sort: None,
|
||||
is_builtin: false,
|
||||
};
|
||||
let profile = profile_with("test_personalized", |_| {});
|
||||
|
||||
let candidates: Vec<EntityId> = (1..=5).map(EntityId::new).collect();
|
||||
let now = Timestamp::from_nanos(1_708_000_000_000_000_000);
|
||||
@ -359,22 +232,7 @@ fn score_personalized_applies_preference_boost() {
|
||||
let ledger = test_ledger();
|
||||
let executor = ProfileExecutor::new(&ledger);
|
||||
|
||||
let profile = RankingProfile {
|
||||
name: "test_preference".into(),
|
||||
version: 1,
|
||||
candidate_strategy: CandidateStrategy::Scan {
|
||||
sort_field: "created_at".into(),
|
||||
},
|
||||
boosts: vec![],
|
||||
decay: None,
|
||||
gates: vec![],
|
||||
penalties: vec![],
|
||||
excludes: vec![],
|
||||
diversity: DiversitySpec::default(),
|
||||
exploration: 0.0,
|
||||
sort: None,
|
||||
is_builtin: false,
|
||||
};
|
||||
let profile = profile_with("test_preference", |_| {});
|
||||
|
||||
let candidates: Vec<EntityId> = (1..=3).map(EntityId::new).collect();
|
||||
let now = Timestamp::from_nanos(1_708_000_000_000_000_000);
|
||||
@ -434,22 +292,7 @@ fn score_personalized_without_boosts_matches_base() {
|
||||
let ledger = test_ledger();
|
||||
let executor = ProfileExecutor::new(&ledger);
|
||||
|
||||
let profile = RankingProfile {
|
||||
name: "test_no_boost".into(),
|
||||
version: 1,
|
||||
candidate_strategy: CandidateStrategy::Scan {
|
||||
sort_field: "created_at".into(),
|
||||
},
|
||||
boosts: vec![],
|
||||
decay: None,
|
||||
gates: vec![],
|
||||
penalties: vec![],
|
||||
excludes: vec![],
|
||||
diversity: DiversitySpec::default(),
|
||||
exploration: 0.0,
|
||||
sort: Some(Sort::New),
|
||||
is_builtin: false,
|
||||
};
|
||||
let profile = profile_with("test_no_boost", |p| p.sort = Some(Sort::New));
|
||||
|
||||
let candidates: Vec<EntityId> = (1..=5).map(EntityId::new).collect();
|
||||
let now = Timestamp::from_nanos(1_708_000_000_000_000_000);
|
||||
@ -518,23 +361,8 @@ fn nan_interaction_boost_does_not_scramble_ranking() {
|
||||
let ledger = test_ledger();
|
||||
let executor = ProfileExecutor::new(&ledger);
|
||||
|
||||
let profile = RankingProfile {
|
||||
name: "nan_boost".into(),
|
||||
version: 1,
|
||||
candidate_strategy: CandidateStrategy::Scan {
|
||||
sort_field: "created_at".into(),
|
||||
},
|
||||
boosts: vec![],
|
||||
decay: None,
|
||||
gates: vec![],
|
||||
penalties: vec![],
|
||||
excludes: vec![],
|
||||
diversity: DiversitySpec::default(),
|
||||
exploration: 0.0,
|
||||
// `New` scores entity-id descending, so without any boost the order is 5,4,3,2,1.
|
||||
sort: Some(Sort::New),
|
||||
is_builtin: false,
|
||||
};
|
||||
// `New` scores entity-id descending, so without any boost the order is 5,4,3,2,1.
|
||||
let profile = profile_with("nan_boost", |p| p.sort = Some(Sort::New));
|
||||
|
||||
let candidates: Vec<EntityId> = (1..=5).map(EntityId::new).collect();
|
||||
let now = Timestamp::from_nanos(1_708_000_000_000_000_000);
|
||||
@ -582,52 +410,7 @@ fn nan_interaction_boost_does_not_scramble_ranking() {
|
||||
}
|
||||
|
||||
// -- Penalty / exclude / decay tests (BLOCKER: negative signals are first-class) --
|
||||
|
||||
/// Build a ledger over `signals` (all windows + velocity) with a 7-day half-life.
|
||||
fn ledger_for(signals: &[&str]) -> SignalLedger {
|
||||
let mut builder = SchemaBuilder::new();
|
||||
for sig in signals {
|
||||
let _ = builder
|
||||
.signal(
|
||||
sig,
|
||||
EntityKind::Item,
|
||||
DecaySpec::Exponential {
|
||||
half_life: Duration::from_secs(7 * 24 * 3600),
|
||||
},
|
||||
)
|
||||
.windows(&[
|
||||
Window::OneHour,
|
||||
Window::TwentyFourHours,
|
||||
Window::SevenDays,
|
||||
Window::AllTime,
|
||||
])
|
||||
.velocity(true)
|
||||
.add();
|
||||
}
|
||||
let schema = builder.build().unwrap();
|
||||
SignalLedger::new(schema, Box::new(NoopWalWriter))
|
||||
}
|
||||
|
||||
fn profile_with(name: &str, mut mutate: impl FnMut(&mut RankingProfile)) -> RankingProfile {
|
||||
let mut p = RankingProfile {
|
||||
name: name.into(),
|
||||
version: 1,
|
||||
candidate_strategy: CandidateStrategy::Scan {
|
||||
sort_field: "created_at".into(),
|
||||
},
|
||||
boosts: vec![],
|
||||
decay: None,
|
||||
gates: vec![],
|
||||
penalties: vec![],
|
||||
excludes: vec![],
|
||||
diversity: DiversitySpec::default(),
|
||||
exploration: 0.0,
|
||||
sort: None,
|
||||
is_builtin: false,
|
||||
};
|
||||
mutate(&mut p);
|
||||
p
|
||||
}
|
||||
// (`ledger_for` and `profile_with` now come from `ranking::test_fixtures`.)
|
||||
|
||||
#[test]
|
||||
fn penalty_reduces_score_for_skipped_item() {
|
||||
@ -784,6 +567,17 @@ fn profile_decay_demotes_aged_item() {
|
||||
// `read_decay_score` decays forward to wall-clock now, so we record entity 1's
|
||||
// freshness signal far in the past and entity 2's at now. A short half-life
|
||||
// makes the gap decisive and deterministic.
|
||||
//
|
||||
// This test needs two signals with *different* half-lives (7-day `view`,
|
||||
// 60-second `freshness`), which the shared `ledger_for` fixture (uniform 7-day
|
||||
// half-life) cannot express -- so it builds its schema inline.
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::{
|
||||
schema::{DecaySpec, EntityKind, SchemaBuilder},
|
||||
signals::NoopWalWriter,
|
||||
};
|
||||
|
||||
let mut builder = SchemaBuilder::new();
|
||||
let _ = builder
|
||||
.signal(
|
||||
|
||||
@ -1,36 +1,20 @@
|
||||
use std::{collections::HashMap, time::Duration};
|
||||
use std::collections::HashMap;
|
||||
|
||||
use super::test_ledger;
|
||||
use crate::{
|
||||
ranking::{
|
||||
executor::ProfileExecutor,
|
||||
profile::{CandidateStrategy, DiversitySpec, RankingProfile, Sort},
|
||||
profile::Sort,
|
||||
// Shared fixtures (DRY-S, M0-M10 review): `make_profile` and `ledger_with`
|
||||
// were per-module duplicates of `profile_with_sort` and `ledger_for`.
|
||||
test_fixtures::{ledger_for, profile_with_sort as make_profile},
|
||||
},
|
||||
schema::{DecaySpec, EntityId, EntityKind, SchemaBuilder, Timestamp, Window},
|
||||
signals::{NoopWalWriter, SignalLedger},
|
||||
schema::{EntityId, Timestamp, Window},
|
||||
signals::SignalLedger,
|
||||
};
|
||||
|
||||
// -- M6p3: New sort variant unit tests ────────────────────────────────
|
||||
|
||||
fn make_profile(sort: Sort) -> RankingProfile {
|
||||
RankingProfile {
|
||||
name: "test".into(),
|
||||
version: 1,
|
||||
candidate_strategy: CandidateStrategy::Scan {
|
||||
sort_field: "created_at".into(),
|
||||
},
|
||||
boosts: vec![],
|
||||
decay: None,
|
||||
gates: vec![],
|
||||
penalties: vec![],
|
||||
excludes: vec![],
|
||||
diversity: DiversitySpec::default(),
|
||||
exploration: 0.0,
|
||||
sort: Some(sort),
|
||||
is_builtin: false,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sort_alphabetical_asc_a_before_z() {
|
||||
let ledger = test_ledger();
|
||||
@ -182,33 +166,74 @@ fn sort_missing_duration_last() {
|
||||
|
||||
// Entity 2 (no duration) should be last.
|
||||
assert_eq!(result[2].entity_id, EntityId::new(2));
|
||||
|
||||
// Regression (Accuracy-W): the normalized score must stay monotonic with
|
||||
// rank. `Shortest` scores present items as `-duration` (negative), and the
|
||||
// missing item used to be clamped to `0.0` *before* the min/max range was
|
||||
// taken -- which made it the maximum on the negated scale and let the
|
||||
// last-ranked item report the highest score `1.0`. With the fix, the
|
||||
// missing-metadata sentinel folds to the bottom of the output range, so the
|
||||
// last result never outscores the first.
|
||||
assert!(
|
||||
result.last().unwrap().score <= result.first().unwrap().score,
|
||||
"last-ranked score {} must not exceed first-ranked score {}",
|
||||
result.last().unwrap().score,
|
||||
result.first().unwrap().score
|
||||
);
|
||||
// Scores are monotonically non-increasing across the whole result, including
|
||||
// the missing item at the tail.
|
||||
for w in result.windows(2) {
|
||||
assert!(
|
||||
w[0].score >= w[1].score,
|
||||
"scores not monotonic with rank: {} then {}",
|
||||
w[0].score,
|
||||
w[1].score
|
||||
);
|
||||
}
|
||||
// The present items keep their relative order: the 60s item (shortest) is
|
||||
// first and outscores the 120s item.
|
||||
assert_eq!(result[0].entity_id, EntityId::new(3)); // 60s
|
||||
assert_eq!(result[1].entity_id, EntityId::new(1)); // 120s
|
||||
assert!(result[0].score >= result[1].score);
|
||||
// The missing item folds to the bottom of the range.
|
||||
assert_eq!(
|
||||
result[2].score, 0.0,
|
||||
"missing-duration item must report the bottom score 0.0, not 1.0"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sort_longest_missing_duration_score_monotonic_with_rank() {
|
||||
// Sibling of `sort_missing_duration_last` for `Longest`: present items score
|
||||
// `+duration` (positive), the missing item is `NEG_INFINITY`. Verify the
|
||||
// missing item ranks last AND reports the bottom normalized score, so the
|
||||
// score stays monotonic with rank across the whole class of ±Inf-sentinel
|
||||
// sorts (not just `Shortest`).
|
||||
let ledger = test_ledger();
|
||||
let mut meta: HashMap<u64, HashMap<String, String>> = HashMap::new();
|
||||
meta.insert(1, [("duration".into(), "120".into())].into());
|
||||
meta.insert(2, HashMap::new()); // no duration
|
||||
meta.insert(3, [("duration".into(), "300".into())].into());
|
||||
|
||||
let executor = ProfileExecutor::new(&ledger).with_item_metadata(&meta);
|
||||
let candidates: Vec<EntityId> = (1..=3).map(EntityId::new).collect();
|
||||
let now = Timestamp::from_nanos(1_708_000_000_000_000_000);
|
||||
let profile = make_profile(Sort::Longest);
|
||||
let result = executor.score(&candidates, &profile, now).unwrap();
|
||||
|
||||
assert_eq!(result[0].entity_id, EntityId::new(3)); // 300s (longest)
|
||||
assert_eq!(result[1].entity_id, EntityId::new(1)); // 120s
|
||||
assert_eq!(result[2].entity_id, EntityId::new(2)); // missing -> last
|
||||
assert!(result.last().unwrap().score <= result.first().unwrap().score);
|
||||
for w in result.windows(2) {
|
||||
assert!(w[0].score >= w[1].score);
|
||||
}
|
||||
assert_eq!(result[2].score, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sort_most_commented() {
|
||||
// Build a ledger with "comment" signal.
|
||||
let mut builder = SchemaBuilder::new();
|
||||
for sig in &["view", "share", "like", "comment"] {
|
||||
let _ = builder
|
||||
.signal(
|
||||
sig,
|
||||
EntityKind::Item,
|
||||
DecaySpec::Exponential {
|
||||
half_life: Duration::from_secs(7 * 24 * 3600),
|
||||
},
|
||||
)
|
||||
.windows(&[
|
||||
Window::OneHour,
|
||||
Window::TwentyFourHours,
|
||||
Window::SevenDays,
|
||||
Window::AllTime,
|
||||
])
|
||||
.velocity(true)
|
||||
.add();
|
||||
}
|
||||
let schema = builder.build().unwrap();
|
||||
let ledger = SignalLedger::new(schema, Box::new(NoopWalWriter));
|
||||
|
||||
let ledger = ledger_for(&["view", "share", "like", "comment"]);
|
||||
let now = Timestamp::now();
|
||||
// Entity 1: 3 comments, Entity 2: 1 comment, Entity 3: 5 comments
|
||||
for _ in 0..3 {
|
||||
@ -240,24 +265,7 @@ fn sort_most_commented() {
|
||||
|
||||
#[test]
|
||||
fn sort_most_shared() {
|
||||
// Build a dedicated ledger with "share" signal and known event counts.
|
||||
let mut builder = SchemaBuilder::new();
|
||||
for sig in &["view", "share", "like"] {
|
||||
let _ = builder
|
||||
.signal(
|
||||
sig,
|
||||
EntityKind::Item,
|
||||
DecaySpec::Exponential {
|
||||
half_life: Duration::from_secs(7 * 24 * 3600),
|
||||
},
|
||||
)
|
||||
.windows(&[Window::OneHour, Window::AllTime])
|
||||
.velocity(false)
|
||||
.add();
|
||||
}
|
||||
let schema = builder.build().unwrap();
|
||||
let ledger = SignalLedger::new(schema, Box::new(NoopWalWriter));
|
||||
|
||||
let ledger = ledger_for(&["view", "share", "like"]);
|
||||
let now = Timestamp::now();
|
||||
// Entity 1: 1 share, Entity 2: 4 shares, Entity 3: 7 shares
|
||||
for _ in 0..1 {
|
||||
@ -292,6 +300,15 @@ fn sort_most_shared() {
|
||||
|
||||
#[test]
|
||||
fn sort_live_viewer_count() {
|
||||
// Needs a short (300s) half-life on `viewer_count` so decay scores stay live,
|
||||
// which the uniform 7-day `ledger_for` fixture cannot express -- built inline.
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::{
|
||||
schema::{DecaySpec, EntityKind, SchemaBuilder},
|
||||
signals::NoopWalWriter,
|
||||
};
|
||||
|
||||
let mut builder = SchemaBuilder::new();
|
||||
let _ = builder
|
||||
.signal(
|
||||
@ -367,35 +384,10 @@ fn sort_date_saved_latest_first() {
|
||||
assert_eq!(result[2].entity_id, EntityId::new(1)); // saved at 1000
|
||||
}
|
||||
|
||||
/// Build a ledger over `signals` with all common windows + velocity enabled.
|
||||
fn ledger_with(signals: &[&str]) -> SignalLedger {
|
||||
let mut builder = SchemaBuilder::new();
|
||||
for sig in signals {
|
||||
let _ = builder
|
||||
.signal(
|
||||
sig,
|
||||
EntityKind::Item,
|
||||
DecaySpec::Exponential {
|
||||
half_life: Duration::from_secs(7 * 24 * 3600),
|
||||
},
|
||||
)
|
||||
.windows(&[
|
||||
Window::OneHour,
|
||||
Window::TwentyFourHours,
|
||||
Window::SevenDays,
|
||||
Window::AllTime,
|
||||
])
|
||||
.velocity(true)
|
||||
.add();
|
||||
}
|
||||
let schema = builder.build().unwrap();
|
||||
SignalLedger::new(schema, Box::new(NoopWalWriter))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sort_controversial_balanced_outranks_one_sided() {
|
||||
// controversial_score = (like*dislike)/(like+dislike)^2 peaks at like==dislike.
|
||||
let ledger = ledger_with(&["like", "dislike", "view"]);
|
||||
let ledger = ledger_for(&["like", "dislike", "view"]);
|
||||
let now = Timestamp::now();
|
||||
// Entity 1: 50/50 split (most controversial). Entity 2: heavily one-sided.
|
||||
// Entity 3: mild split.
|
||||
@ -443,7 +435,7 @@ fn sort_controversial_balanced_outranks_one_sided() {
|
||||
#[test]
|
||||
fn sort_hidden_gems_low_views_high_quality_first() {
|
||||
// hidden_gems_score = completion_decay / log10(view_count + 10).
|
||||
let ledger = ledger_with(&["completion", "view"]);
|
||||
let ledger = ledger_for(&["completion", "view"]);
|
||||
let now = Timestamp::now();
|
||||
// Entity 1: high quality, few views (the gem).
|
||||
// Entity 2: high quality, many views (already discovered).
|
||||
@ -506,7 +498,7 @@ fn sort_rising_accelerating_item_first() {
|
||||
// windowed counts (no wall-clock decay), so fixed timestamps make it fully
|
||||
// deterministic. The 12h gap rotates the old tranche out of the 1h (minute)
|
||||
// window while keeping it inside the 24h window.
|
||||
let ledger = ledger_with(&["view"]);
|
||||
let ledger = ledger_for(&["view"]);
|
||||
let now = Timestamp::from_nanos(1_708_000_000_000_000_000);
|
||||
let twelve_hours_ago = Timestamp::from_nanos(now.as_nanos() - 12 * 3_600_000_000_000);
|
||||
|
||||
@ -546,7 +538,7 @@ fn sort_rising_accelerating_item_first() {
|
||||
fn sort_top_window_orders_by_weighted_engagement() {
|
||||
// TopWindow = view*0.3 + like*0.3 + share*0.2 + completion*view*0.1, all
|
||||
// counts within the requested window. Higher total engagement -> higher rank.
|
||||
let ledger = ledger_with(&["view", "like", "share", "completion"]);
|
||||
let ledger = ledger_for(&["view", "like", "share", "completion"]);
|
||||
let now = Timestamp::now();
|
||||
// Entity 1: high across the board. Entity 2: low. Entity 3: medium.
|
||||
for _ in 0..100 {
|
||||
@ -599,4 +591,12 @@ fn sort_date_saved_no_save_sorted_last() {
|
||||
|
||||
// Entity 2 (no save timestamp) should be last.
|
||||
assert_eq!(result[2].entity_id, EntityId::new(2));
|
||||
// Sibling of `sort_missing_duration_last` for the `DateSaved` ±Inf sentinel:
|
||||
// the missing item folds to the bottom of the range and stays monotonic with
|
||||
// rank rather than reporting the spurious maximum `1.0`.
|
||||
assert!(result.last().unwrap().score <= result.first().unwrap().score);
|
||||
for w in result.windows(2) {
|
||||
assert!(w[0].score >= w[1].score);
|
||||
}
|
||||
assert_eq!(result[2].score, 0.0);
|
||||
}
|
||||
|
||||
@ -4,6 +4,9 @@ pub mod executor;
|
||||
pub mod profile;
|
||||
pub mod registry;
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) mod test_fixtures;
|
||||
|
||||
pub use builtins::register_builtins;
|
||||
pub use diversity::{
|
||||
ConstraintViolation, DiversityConstraints, DiversityResult, DiversitySelector,
|
||||
|
||||
@ -240,25 +240,9 @@ pub struct DiversitySpec {
|
||||
#[allow(clippy::unwrap_used)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn minimal_profile(name: &str) -> RankingProfile {
|
||||
RankingProfile {
|
||||
name: name.to_owned(),
|
||||
version: 1,
|
||||
candidate_strategy: CandidateStrategy::Scan {
|
||||
sort_field: "created_at".into(),
|
||||
},
|
||||
boosts: vec![],
|
||||
decay: None,
|
||||
gates: vec![],
|
||||
penalties: vec![],
|
||||
excludes: vec![],
|
||||
diversity: DiversitySpec::default(),
|
||||
exploration: 0.0,
|
||||
sort: None,
|
||||
is_builtin: false,
|
||||
}
|
||||
}
|
||||
// Shared profile skeleton (DRY-S, M0-M10 review): the per-module
|
||||
// `minimal_profile` 14-field literal now lives in `ranking::test_fixtures`.
|
||||
use crate::ranking::test_fixtures::minimal_profile;
|
||||
|
||||
#[test]
|
||||
fn profile_serializes_to_json() {
|
||||
|
||||
@ -287,8 +287,12 @@ mod tests {
|
||||
|
||||
use super::*;
|
||||
use crate::{
|
||||
ranking::profile::{
|
||||
Boost, CandidateStrategy, DiversitySpec, Gate, RankingProfile, SignalAgg,
|
||||
ranking::{
|
||||
profile::{Boost, Gate, SignalAgg},
|
||||
// Shared profile skeleton (DRY-S, M0-M10 review): the per-module
|
||||
// `minimal_profile(name, version)` literal now lives in
|
||||
// `ranking::test_fixtures` as `minimal_profile_versioned`.
|
||||
test_fixtures::minimal_profile_versioned as minimal_profile,
|
||||
},
|
||||
schema::{DecaySpec, EntityKind, SchemaBuilder, Window},
|
||||
};
|
||||
@ -310,25 +314,6 @@ mod tests {
|
||||
builder.build().unwrap()
|
||||
}
|
||||
|
||||
fn minimal_profile(name: &str, version: u32) -> RankingProfile {
|
||||
RankingProfile {
|
||||
name: name.to_owned(),
|
||||
version,
|
||||
candidate_strategy: CandidateStrategy::Scan {
|
||||
sort_field: "created_at".into(),
|
||||
},
|
||||
boosts: vec![],
|
||||
decay: None,
|
||||
gates: vec![],
|
||||
penalties: vec![],
|
||||
excludes: vec![],
|
||||
diversity: DiversitySpec::default(),
|
||||
exploration: 0.0,
|
||||
sort: None,
|
||||
is_builtin: false,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn register_and_get() {
|
||||
let mut registry = ProfileRegistry::new();
|
||||
|
||||
132
tidal/src/ranking/test_fixtures.rs
Normal file
132
tidal/src/ranking/test_fixtures.rs
Normal file
@ -0,0 +1,132 @@
|
||||
//! Shared test fixtures for the `ranking` subsystem.
|
||||
//!
|
||||
//! Before this module the same `SignalLedger` and `RankingProfile` builders were
|
||||
//! re-declared verbatim across five ranking test modules (`executor::tests`,
|
||||
//! `executor::tests::sort_tests`, `executor::helpers::tests`, `profile::tests`,
|
||||
//! `registry::tests`). The duplication drifted (different window sets, different
|
||||
//! seed timestamps) and every new sort test re-pasted the 14-field profile
|
||||
//! literal. These builders are the single home for that boilerplate.
|
||||
//!
|
||||
//! Compiled only under `cfg(test)`; never linked into a release build.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::{
|
||||
ranking::profile::{CandidateStrategy, DiversitySpec, RankingProfile, Sort},
|
||||
schema::{DecaySpec, EntityId, EntityKind, SchemaBuilder, Timestamp, Window},
|
||||
signals::{NoopWalWriter, SignalLedger},
|
||||
};
|
||||
|
||||
/// Fixed base timestamp shared by ledger seeds so windowed reads are
|
||||
/// deterministic across modules (not wall-clock-relative).
|
||||
pub const BASE_NS: u64 = 1_708_000_000_000_000_000;
|
||||
|
||||
/// A 7-day exponential half-life, the default for ranking-test signals.
|
||||
const WEEK: Duration = Duration::from_secs(7 * 24 * 3600);
|
||||
|
||||
/// The standard window set used by ranking tests: every coarse-grained bucket
|
||||
/// plus `AllTime`, with velocity enabled.
|
||||
const STD_WINDOWS: &[Window] = &[
|
||||
Window::OneHour,
|
||||
Window::TwentyFourHours,
|
||||
Window::SevenDays,
|
||||
Window::AllTime,
|
||||
];
|
||||
|
||||
/// Build a `SignalLedger` defining each name in `signals` with a 7-day half-life,
|
||||
/// the standard window set, and velocity enabled. No events are recorded.
|
||||
///
|
||||
/// This is the union of the previously-duplicated `ledger_for` / `ledger_with`
|
||||
/// builders; both used identical schemas.
|
||||
#[must_use]
|
||||
pub fn ledger_for(signals: &[&str]) -> SignalLedger {
|
||||
let mut builder = SchemaBuilder::new();
|
||||
for sig in signals {
|
||||
let _ = builder
|
||||
.signal(
|
||||
sig,
|
||||
EntityKind::Item,
|
||||
DecaySpec::Exponential { half_life: WEEK },
|
||||
)
|
||||
.windows(STD_WINDOWS)
|
||||
.velocity(true)
|
||||
.add();
|
||||
}
|
||||
let schema = builder.build().expect("ranking test schema must build");
|
||||
SignalLedger::new(schema, Box::new(NoopWalWriter))
|
||||
}
|
||||
|
||||
/// The canonical seeded ledger used by the bulk of executor tests: `view`,
|
||||
/// `share`, and `like` signals, seeded with a small deterministic event spread
|
||||
/// across five entities (ids 1..=5) anchored at [`BASE_NS`].
|
||||
///
|
||||
/// Entity `i` (1-indexed) gets `5 - (i-1)` views, `(i-1) % 3` shares, and
|
||||
/// `(i-1) % 2` likes, each timestamped one hour earlier than the previous
|
||||
/// entity. This is the exact seed the old per-module `test_ledger` produced.
|
||||
#[must_use]
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
pub fn seeded_ledger() -> SignalLedger {
|
||||
let ledger = ledger_for(&["view", "share", "like"]);
|
||||
for i in 0u64..5 {
|
||||
let entity_id = EntityId::new(i + 1);
|
||||
let ts = Timestamp::from_nanos(BASE_NS - i * 3_600_000_000_000);
|
||||
ledger
|
||||
.record_signal("view", entity_id, (5 - i) as f64, ts)
|
||||
.expect("seed view");
|
||||
ledger
|
||||
.record_signal("share", entity_id, (i % 3) as f64, ts)
|
||||
.expect("seed share");
|
||||
ledger
|
||||
.record_signal("like", entity_id, (i % 2) as f64, ts)
|
||||
.expect("seed like");
|
||||
}
|
||||
ledger
|
||||
}
|
||||
|
||||
/// Build a minimal `RankingProfile` skeleton: empty boosts/penalties/excludes/
|
||||
/// gates, no decay, no sort, version 1, not built-in. This is the 14-field
|
||||
/// literal every ranking test re-pasted.
|
||||
#[must_use]
|
||||
pub fn minimal_profile(name: &str) -> RankingProfile {
|
||||
RankingProfile {
|
||||
name: name.to_owned(),
|
||||
version: 1,
|
||||
candidate_strategy: CandidateStrategy::Scan {
|
||||
sort_field: "created_at".into(),
|
||||
},
|
||||
boosts: vec![],
|
||||
decay: None,
|
||||
gates: vec![],
|
||||
penalties: vec![],
|
||||
excludes: vec![],
|
||||
diversity: DiversitySpec::default(),
|
||||
exploration: 0.0,
|
||||
sort: None,
|
||||
is_builtin: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// A minimal profile at an explicit `version` (used by registry version tests).
|
||||
#[must_use]
|
||||
pub fn minimal_profile_versioned(name: &str, version: u32) -> RankingProfile {
|
||||
let mut p = minimal_profile(name);
|
||||
p.version = version;
|
||||
p
|
||||
}
|
||||
|
||||
/// A minimal profile with the given [`Sort`] mode set.
|
||||
#[must_use]
|
||||
pub fn profile_with_sort(sort: Sort) -> RankingProfile {
|
||||
let mut p = minimal_profile("test");
|
||||
p.sort = Some(sort);
|
||||
p
|
||||
}
|
||||
|
||||
/// A minimal profile mutated by `mutate` -- the ergonomic builder for tests that
|
||||
/// tweak boosts/penalties/decay/excludes on top of the skeleton.
|
||||
#[must_use]
|
||||
pub fn profile_with(name: &str, mut mutate: impl FnMut(&mut RankingProfile)) -> RankingProfile {
|
||||
let mut p = minimal_profile(name);
|
||||
mutate(&mut p);
|
||||
p
|
||||
}
|
||||
@ -146,6 +146,7 @@ mod tests {
|
||||
id: WalSegmentId::new(RegionId::SINGLE, shard, seqno),
|
||||
bytes: vec![0xAB; 100],
|
||||
event_count: 5,
|
||||
leader_last_seq: seqno,
|
||||
}
|
||||
}
|
||||
|
||||
@ -204,6 +205,7 @@ mod tests {
|
||||
id: WalSegmentId::single_node(1),
|
||||
bytes: vec![0u8; MAX_PAYLOAD_BYTES + 1],
|
||||
event_count: 0,
|
||||
leader_last_seq: 0,
|
||||
};
|
||||
let result = t0.send_segment(ShardId(1), payload);
|
||||
assert!(result.is_err());
|
||||
|
||||
@ -27,6 +27,12 @@ pub struct TenantMigration {
|
||||
state: Mutex<MigrationState>,
|
||||
_control_plane: Arc<ControlPlane>,
|
||||
tenant_router: Arc<TenantRouter>,
|
||||
/// Durable sink for the tenant router's migration routing state. When set,
|
||||
/// every routing transition (enter dual-write, finalize) is persisted so a
|
||||
/// crash mid-migration cannot silently revert a finalized tenant to
|
||||
/// jump-hash routing (Accuracy-W). `None` (the default / ephemeral) keeps
|
||||
/// routing in memory only, identical to the previous behaviour.
|
||||
persistence: Option<Arc<dyn crate::storage::StorageEngine>>,
|
||||
}
|
||||
|
||||
impl TenantMigration {
|
||||
@ -46,6 +52,42 @@ impl TenantMigration {
|
||||
state: Mutex::new(MigrationState::Idle),
|
||||
_control_plane: control_plane,
|
||||
tenant_router,
|
||||
persistence: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Attach a durable storage sink so routing transitions are checkpointed.
|
||||
///
|
||||
/// Call this in persistent (non-ephemeral) deployments with the items
|
||||
/// storage engine. Each subsequent `enter_dual_write` / `finalize` persists
|
||||
/// the tenant router's routing state via
|
||||
/// [`checkpoint_migration_routing`](crate::replication::tenant::checkpoint_migration_routing),
|
||||
/// so the migration survives a crash. Returns `self` for chaining.
|
||||
#[must_use]
|
||||
pub fn with_persistence(mut self, storage: Arc<dyn crate::storage::StorageEngine>) -> Self {
|
||||
self.persistence = Some(storage);
|
||||
self
|
||||
}
|
||||
|
||||
/// Persist the router's routing state if a durable sink is attached.
|
||||
///
|
||||
/// A checkpoint failure is logged (the migration is already applied in
|
||||
/// memory) but not propagated, so a transient disk error does not strand the
|
||||
/// in-memory transition — the next transition (or open-time restore) reaches
|
||||
/// a consistent point. A persistent failure is loud in the logs.
|
||||
fn persist_routing(&self) {
|
||||
if let Some(storage) = &self.persistence
|
||||
&& let Err(e) = crate::replication::tenant::checkpoint_migration_routing(
|
||||
storage.as_ref(),
|
||||
&self.tenant_router,
|
||||
)
|
||||
{
|
||||
tracing::error!(
|
||||
tenant_id = self.tenant_id.0,
|
||||
error = %e,
|
||||
"failed to checkpoint migration routing state; a crash before the next \
|
||||
checkpoint would revert this tenant's routing"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@ -88,6 +130,9 @@ impl TenantMigration {
|
||||
.set_dual_write(self.tenant_id, self.source_shard, self.target_shard);
|
||||
*state = MigrationState::DualWrite { cutover_seqno };
|
||||
drop(state);
|
||||
// Persist the new dual-write routing so a crash during dual-write does
|
||||
// not lose the fact that this tenant is mid-migration.
|
||||
self.persist_routing();
|
||||
Ok(cutover_seqno)
|
||||
}
|
||||
|
||||
@ -118,6 +163,10 @@ impl TenantMigration {
|
||||
let switched_at_ns = crate::replication::now_ns();
|
||||
*state = MigrationState::Finalizing { switched_at_ns };
|
||||
drop(state);
|
||||
// Persist the finalized pin (and the cleared dual-write) so a crash after
|
||||
// finalization cannot revert the tenant to jump-hash routing — the exact
|
||||
// bug this closes (Accuracy-W).
|
||||
self.persist_routing();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@ -261,4 +310,72 @@ mod tests {
|
||||
let err = m.gc_source(600_000_000_000).unwrap_err();
|
||||
assert!(matches!(err, TidalError::InvalidState(_)));
|
||||
}
|
||||
|
||||
/// Accuracy-W: a finalized migration's routing must survive a crash. The
|
||||
/// coordinator persists routing on every transition; a fresh router restored
|
||||
/// from the durable checkpoint keeps the pin instead of reverting to
|
||||
/// jump-hash routing.
|
||||
#[test]
|
||||
fn migration_routing_survives_simulated_crash() {
|
||||
use std::sync::RwLock;
|
||||
|
||||
use crate::{
|
||||
replication::tenant::{ShardAssignment, restore_migration_routing},
|
||||
storage::InMemoryBackend,
|
||||
};
|
||||
|
||||
let topo = Arc::new(RwLock::new(ClusterTopology {
|
||||
shards: vec![
|
||||
ShardAssignment {
|
||||
shard_id: ShardId(0),
|
||||
region_id: crate::replication::RegionId(0),
|
||||
},
|
||||
ShardAssignment {
|
||||
shard_id: ShardId(1),
|
||||
region_id: crate::replication::RegionId(1),
|
||||
},
|
||||
],
|
||||
}));
|
||||
let router = Arc::new(TenantRouter::new(Arc::clone(&topo)));
|
||||
let state = Arc::new(ReplicationState::single());
|
||||
let lag = Arc::new(ReplicationLagGauge::new(ShardId::SINGLE, state));
|
||||
let cp = Arc::new(ControlPlane::new(
|
||||
Arc::clone(&topo),
|
||||
Arc::clone(&router),
|
||||
lag,
|
||||
));
|
||||
|
||||
// Durable sink shared with the migration coordinator.
|
||||
let storage = Arc::new(InMemoryBackend::new());
|
||||
|
||||
let migration = TenantMigration::new(
|
||||
TenantId(42),
|
||||
ShardId(0),
|
||||
ShardId(1),
|
||||
cp,
|
||||
Arc::clone(&router),
|
||||
)
|
||||
.with_persistence(Arc::clone(&storage) as Arc<dyn crate::storage::StorageEngine>);
|
||||
|
||||
// Drive the migration to finalization. Each transition persists routing.
|
||||
migration.prepare_target(10).unwrap();
|
||||
migration.enter_dual_write(20).unwrap();
|
||||
migration.finalize(25).unwrap();
|
||||
|
||||
// "Crash": discard the in-memory router and rebuild routing from disk.
|
||||
let recovered = TenantRouter::new(Arc::clone(&topo));
|
||||
let n = restore_migration_routing(storage.as_ref(), &recovered).unwrap();
|
||||
assert_eq!(
|
||||
n, 1,
|
||||
"the finalized pin must be restored from the checkpoint"
|
||||
);
|
||||
assert_eq!(
|
||||
recovered.pinned_shard(TenantId(42)),
|
||||
Some(ShardId(1)),
|
||||
"finalized tenant must stay pinned to its target after a crash, not \
|
||||
revert to jump-hash routing"
|
||||
);
|
||||
// It must NOT still be in dual-write (finalize cleared it).
|
||||
assert!(!recovered.is_dual_write(TenantId(42)));
|
||||
}
|
||||
}
|
||||
|
||||
@ -26,7 +26,9 @@ pub use idempotency::{IdempotencyKey, IdempotencyStore};
|
||||
pub use in_process::{InProcessTransport, InProcessTransportFactory};
|
||||
pub use lag::ReplicationLagGauge;
|
||||
pub use migration::{MigrationState, TenantMigration};
|
||||
pub use receiver::{SegmentReceiverHandle, apply_payload, spawn_receiver};
|
||||
pub use receiver::{
|
||||
SegmentReceiverHandle, apply_payload, apply_payload_with_leader_seq, spawn_receiver,
|
||||
};
|
||||
pub use reconcile::{HardNegAction, MergePlan, ReconciliationEngine, StateSnapshot};
|
||||
pub use segment_id::WalSegmentId;
|
||||
pub use session_bridge::{
|
||||
@ -34,7 +36,7 @@ pub use session_bridge::{
|
||||
SessionShardTransport,
|
||||
};
|
||||
pub use shard::{EntityIdRange, RegionId, RouterError, RoutingStrategy, ShardId, ShardRouter};
|
||||
pub use shipper::{ShipperConfig, WalShipperHandle, spawn_shipper};
|
||||
pub use shipper::{ShipperConfig, ShipperState, WalShipperHandle, spawn_shipper};
|
||||
pub use state::ReplicationState;
|
||||
pub use tenant::{
|
||||
ClusterTopology, ShardAssignment, TenantConfig, TenantId, TenantRateLimiter, TenantRouter,
|
||||
|
||||
@ -3,15 +3,30 @@
|
||||
//!
|
||||
//! The receiver runs in a background thread, blocking on
|
||||
//! [`Transport::recv_segment`] and replaying each batch into the shared
|
||||
//! [`SignalLedger`] via `apply_wal_event`. Idempotent replay is ensured by
|
||||
//! checking the per-shard high-water-mark in [`ReplicationState`].
|
||||
//! [`SignalLedger`] via `apply_replicated_event`. Idempotent replay is ensured
|
||||
//! by checking the per-shard high-water-mark in [`ReplicationState`].
|
||||
//!
|
||||
//! # Durability (BLOCKER 6)
|
||||
//!
|
||||
//! Each accepted event is applied **WAL-first** on the follower:
|
||||
//! `apply_replicated_event` appends the event to the follower's *own* WAL and
|
||||
//! fsyncs (via the group-commit writer) before mutating in-memory state. The
|
||||
//! follower's normal WAL replay on open therefore reconstructs all replicated
|
||||
//! state — without this, every replicated event between the last checkpoint and
|
||||
//! a follower crash was silently lost. The leader-seqno high-water-mark is
|
||||
//! persisted separately by the periodic checkpoint thread (a re-shipped segment
|
||||
//! whose last seq is at or below the restored high-water-mark is an idempotent
|
||||
//! no-op).
|
||||
//!
|
||||
//! # Error handling
|
||||
//!
|
||||
//! If a batch fails BLAKE3 verification or structural decode, `apply_payload`
|
||||
//! returns [`WalError::Corruption`]. The receiver thread propagates this
|
||||
//! error immediately (it does **not** skip the payload) so that
|
||||
//! [`SegmentReceiverHandle::join`] can surface it to the operator.
|
||||
//! returns [`WalError::Corruption`]. If the follower's own WAL append fails (a
|
||||
//! disk fault), it returns [`WalError::Io`]. In both cases the receiver thread
|
||||
//! propagates the error immediately (it does **not** skip the payload) so that
|
||||
//! [`SegmentReceiverHandle::join`] can surface it to the operator — the
|
||||
//! follower must never silently acknowledge an event it failed to durably
|
||||
//! record.
|
||||
|
||||
use std::{sync::Arc, thread::JoinHandle};
|
||||
|
||||
@ -85,12 +100,19 @@ pub fn spawn_receiver<T: Transport + ?Sized>(
|
||||
};
|
||||
|
||||
let shard_id = payload.id.shard_id;
|
||||
if let Err(e) = apply_payload(
|
||||
// Carry the segment's authoritative leader high-water-mark (known
|
||||
// before community-overlay filtering) so the lag gauge tracks the
|
||||
// leader even across an all-local (empty) segment that decodes to
|
||||
// zero batches. `0` means "unknown" and falls back to per-batch
|
||||
// boundaries inside apply_payload_with_leader_seq.
|
||||
let leader_seqno = (payload.leader_last_seq > 0).then_some(payload.leader_last_seq);
|
||||
if let Err(e) = apply_payload_with_leader_seq(
|
||||
&payload.bytes,
|
||||
shard_id,
|
||||
&ledger,
|
||||
&replication_state,
|
||||
lag_gauge.as_deref(),
|
||||
leader_seqno,
|
||||
) {
|
||||
// A corrupt segment halts the receiver thread. Log at the
|
||||
// failure site so a stalled replica is observable rather
|
||||
@ -119,17 +141,42 @@ pub fn spawn_receiver<T: Transport + ?Sized>(
|
||||
/// Idempotent: batches whose last sequence number is at or below the
|
||||
/// replication state's high-water-mark for the source shard are skipped.
|
||||
///
|
||||
/// When `lag_gauge` is `Some`, the leader high-water-mark is advanced to the
|
||||
/// last WAL sequence of every decoded batch — including idempotently skipped
|
||||
/// ones — *before* application. The gauge's applied seqno tracks
|
||||
/// `state.applied_seqno`, so `lag_segments()` is non-zero whenever a batch has
|
||||
/// # Atomic payload application (WARNING fix)
|
||||
///
|
||||
/// The payload is applied **all-or-nothing**. We first decode and validate
|
||||
/// *every* batch (BLAKE3, structural decode, and the `first_seq + event_count`
|
||||
/// overflow check) into a staged list, advancing only the lag gauge's *leader*
|
||||
/// high-water-mark as boundaries are observed. Only after the whole payload
|
||||
/// validates do we apply staged events to the follower ledger (WAL-first) and
|
||||
/// advance the *applied* high-water-mark to the payload's final boundary in one
|
||||
/// step.
|
||||
///
|
||||
/// This means a corrupt batch *anywhere* in the payload causes the function to
|
||||
/// return `Err` **before** any event is applied or any applied-seqno is
|
||||
/// advanced, so the follower's ledger and idempotency boundary never reflect a
|
||||
/// half-applied payload. The leader can cleanly re-ship the same un-acked
|
||||
/// segment and it replays from a single consistent resume point. (Without this,
|
||||
/// a later corrupt batch left earlier batches applied with the applied-seqno
|
||||
/// stuck mid-payload, with no clean point to resume from.)
|
||||
///
|
||||
/// `leader_seqno` is the segment's authoritative last WAL seqno as known by the
|
||||
/// leader *before* community-overlay filtering (see [`WalSegmentPayload`]). It
|
||||
/// feeds the lag gauge's leader high-water-mark regardless of how many batches
|
||||
/// survived filtering — an all-local segment ships zero batches yet still
|
||||
/// advances the leader HWM so the gauge does not under-report (obs-REPL-1,
|
||||
/// WARNING fix). `None` falls back to the per-batch boundaries (legacy callers
|
||||
/// / tests with no authoritative value).
|
||||
///
|
||||
/// When `lag_gauge` is `Some`, the gauge's applied seqno tracks
|
||||
/// `state.applied_seqno`, so `lag_segments()` is non-zero whenever a payload has
|
||||
/// been received from the leader but not yet applied, and returns to 0 once the
|
||||
/// follower catches up (obs-REPL-1).
|
||||
/// follower catches up.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `WalError::Corruption` on any BLAKE3 or structural decode failure.
|
||||
/// The offset of the first corrupt batch is included in the error message.
|
||||
/// Returns `WalError::Corruption` on any BLAKE3 or structural decode failure
|
||||
/// (or a wrapped seqno boundary). The offset of the first corrupt batch is
|
||||
/// included in the error message. On error, no state has been mutated.
|
||||
pub fn apply_payload(
|
||||
bytes: &[u8],
|
||||
from_shard: ShardId,
|
||||
@ -137,76 +184,140 @@ pub fn apply_payload(
|
||||
state: &ReplicationState,
|
||||
lag_gauge: Option<&ReplicationLagGauge>,
|
||||
) -> Result<(), WalError> {
|
||||
apply_payload_with_leader_seq(bytes, from_shard, ledger, state, lag_gauge, None)
|
||||
}
|
||||
|
||||
/// Like [`apply_payload`] but carrying the segment's authoritative leader
|
||||
/// high-water-mark (the true last WAL seqno before community-overlay filtering).
|
||||
///
|
||||
/// See [`apply_payload`] for the atomicity and leader-seqno contract. The
|
||||
/// `leader_seqno` is fed to the lag gauge independently of how many batches
|
||||
/// survive filtering, so an all-local (empty) segment still advances the leader
|
||||
/// HWM.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `WalError::Corruption` on any decode/checksum/overflow failure,
|
||||
/// before any state is mutated.
|
||||
pub fn apply_payload_with_leader_seq(
|
||||
bytes: &[u8],
|
||||
from_shard: ShardId,
|
||||
ledger: &SignalLedger,
|
||||
state: &ReplicationState,
|
||||
lag_gauge: Option<&ReplicationLagGauge>,
|
||||
leader_seqno: Option<u64>,
|
||||
) -> Result<(), WalError> {
|
||||
// ── Phase 1: decode + validate the WHOLE payload before mutating state ──
|
||||
//
|
||||
// Stage every batch's events and track the highest last-seq across the
|
||||
// payload. If any batch fails decode/checksum/overflow we return here, with
|
||||
// neither the follower ledger nor the applied high-water-mark touched.
|
||||
let applied_floor = state.applied_seqno(from_shard);
|
||||
let mut staged: Vec<(SignalTypeId, EntityId, f64, Timestamp)> = Vec::new();
|
||||
// The highest last-seq seen across all decoded batches. WAL batches in a
|
||||
// segment are ascending, but taking the max (rather than the trailing
|
||||
// batch's boundary) keeps the advance order-independent and monotonic.
|
||||
let mut payload_last_seq: Option<u64> = None;
|
||||
let mut max_batch_last_seq: u64 = 0;
|
||||
|
||||
let mut offset = 0;
|
||||
while offset < bytes.len() {
|
||||
let remaining = &bytes[offset..];
|
||||
if remaining.len() < HEADER_SIZE {
|
||||
break;
|
||||
}
|
||||
match decode_batch(remaining) {
|
||||
Ok((header, events)) => {
|
||||
// The last WAL seq is `first_seq + event_count - 1`. Both
|
||||
// operands originate from peer- or disk-controlled header bytes,
|
||||
// so a corrupt/hostile `first_seq` near `u64::MAX` could wrap. A
|
||||
// wrapped seqno would corrupt the idempotency boundary and the
|
||||
// lag gauge, so we reject the batch as corruption rather than
|
||||
// applying it. (`event_count == 0` is a legal empty batch whose
|
||||
// boundary is just `first_seq`.)
|
||||
let batch_last_seq = if header.event_count > 0 {
|
||||
header
|
||||
.first_seq
|
||||
.checked_add(u64::from(header.event_count) - 1)
|
||||
.ok_or_else(|| WalError::Corruption {
|
||||
message: format!(
|
||||
"WAL seq overflow at payload offset {offset}: \
|
||||
first_seq={} + event_count={} would wrap u64",
|
||||
header.first_seq, header.event_count
|
||||
),
|
||||
})?
|
||||
} else {
|
||||
header.first_seq
|
||||
};
|
||||
let (header, events) = decode_batch(remaining).map_err(|e| WalError::Corruption {
|
||||
message: format!("corrupt batch at payload offset {offset}: {e}"),
|
||||
})?;
|
||||
|
||||
// Record the leader high-water-mark from every batch the leader
|
||||
// has shipped us, before the idempotency check: the leader's
|
||||
// seqno is known the moment a segment carrying it arrives, even
|
||||
// if we then skip the batch as already-applied. `lag_segments()`
|
||||
// = leader_seqno − applied_seqno is the real follower lag.
|
||||
if let Some(gauge) = lag_gauge {
|
||||
gauge.update_leader_seqno(batch_last_seq);
|
||||
}
|
||||
// The last WAL seq is `first_seq + event_count - 1`. Both operands
|
||||
// originate from peer- or disk-controlled header bytes, so a
|
||||
// corrupt/hostile `first_seq` near `u64::MAX` could wrap. A wrapped
|
||||
// seqno would corrupt the idempotency boundary and the lag gauge, so we
|
||||
// reject the batch as corruption rather than applying it.
|
||||
// (`event_count == 0` is a legal empty batch whose boundary is just
|
||||
// `first_seq`.)
|
||||
let batch_last_seq = if header.event_count > 0 {
|
||||
header
|
||||
.first_seq
|
||||
.checked_add(u64::from(header.event_count) - 1)
|
||||
.ok_or_else(|| WalError::Corruption {
|
||||
message: format!(
|
||||
"WAL seq overflow at payload offset {offset}: \
|
||||
first_seq={} + event_count={} would wrap u64",
|
||||
header.first_seq, header.event_count
|
||||
),
|
||||
})?
|
||||
} else {
|
||||
header.first_seq
|
||||
};
|
||||
max_batch_last_seq = max_batch_last_seq.max(batch_last_seq);
|
||||
payload_last_seq = Some(max_batch_last_seq);
|
||||
|
||||
// Idempotency: skip if already applied.
|
||||
if let Some(applied) = state.applied_seqno(from_shard)
|
||||
&& batch_last_seq <= applied
|
||||
{
|
||||
let batch_size = HEADER_SIZE + header.payload_len as usize;
|
||||
offset += batch_size;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Apply each event to the ledger.
|
||||
for event in &events {
|
||||
let signal_type_id = SignalTypeId::new(u16::from(event.signal_type));
|
||||
let entity_id = EntityId::new(event.entity_id);
|
||||
let weight = f64::from(event.weight);
|
||||
let timestamp = Timestamp::from_nanos(event.timestamp_nanos);
|
||||
ledger.apply_wal_event(signal_type_id, entity_id, weight, timestamp);
|
||||
}
|
||||
|
||||
// Advance replication high-water-mark.
|
||||
state.advance(from_shard, batch_last_seq);
|
||||
|
||||
let batch_size = HEADER_SIZE + header.payload_len as usize;
|
||||
offset += batch_size;
|
||||
}
|
||||
Err(e) => {
|
||||
return Err(WalError::Corruption {
|
||||
message: format!("corrupt batch at payload offset {offset}: {e}"),
|
||||
});
|
||||
// Idempotency: a batch whose boundary is at or below the already-applied
|
||||
// floor is a re-ship; stage nothing for it but still let its boundary
|
||||
// count toward the payload's final advance (the advance is monotonic, so
|
||||
// re-applying the same boundary is a no-op).
|
||||
let already_applied = applied_floor.is_some_and(|applied| batch_last_seq <= applied);
|
||||
if !already_applied {
|
||||
for event in &events {
|
||||
staged.push((
|
||||
SignalTypeId::new(u16::from(event.signal_type)),
|
||||
EntityId::new(event.entity_id),
|
||||
f64::from(event.weight),
|
||||
Timestamp::from_nanos(event.timestamp_nanos),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
let batch_size = HEADER_SIZE + header.payload_len as usize;
|
||||
offset += batch_size;
|
||||
}
|
||||
|
||||
// Record the leader high-water-mark. Prefer the authoritative segment-level
|
||||
// value (known before filtering, so an all-local segment that decoded zero
|
||||
// batches still advances it); otherwise fall back to the highest per-batch
|
||||
// boundary we observed. The leader's seqno is known the moment the segment
|
||||
// arrives, independent of how many batches survived filtering or the
|
||||
// idempotency check, so `lag_segments() = leader_seqno − applied_seqno` is
|
||||
// the real follower lag.
|
||||
if let Some(gauge) = lag_gauge {
|
||||
let observed = leader_seqno.unwrap_or(max_batch_last_seq);
|
||||
if observed > 0 {
|
||||
gauge.update_leader_seqno(observed);
|
||||
}
|
||||
}
|
||||
|
||||
// ── Phase 2: apply staged events, then advance the applied HWM once ──
|
||||
//
|
||||
// Every staged event is appended to THIS follower's own WAL (and fsynced by
|
||||
// the group-commit writer) before its in-memory aggregate is mutated, so the
|
||||
// follower's normal WAL replay on open reconstructs the applied state. A WAL
|
||||
// append failure (disk fault) returns before the high-water-mark advances,
|
||||
// so the follower never acknowledges an event it could not durably record —
|
||||
// the re-shipped segment is replayed in full from one consistent point.
|
||||
for (signal_type_id, entity_id, weight, timestamp) in &staged {
|
||||
ledger
|
||||
.apply_replicated_event(*signal_type_id, *entity_id, *weight, *timestamp)
|
||||
.map_err(|e| {
|
||||
WalError::Io(std::io::Error::other(format!(
|
||||
"follower WAL append failed applying replicated event \
|
||||
(entity={}, signal_type={}): {e}",
|
||||
entity_id.as_u64(),
|
||||
signal_type_id.as_u16()
|
||||
)))
|
||||
})?;
|
||||
}
|
||||
|
||||
// Advance the replication high-water-mark to the payload's final boundary in
|
||||
// a single step, only AFTER every staged event is durably appended above.
|
||||
// This is the leader-seqno idempotency boundary the periodic checkpoint
|
||||
// persists; a re-shipped payload at or below it is skipped before
|
||||
// re-application.
|
||||
if let Some(last) = payload_last_seq {
|
||||
state.advance(from_shard, last);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@ -509,6 +620,105 @@ mod tests {
|
||||
assert!(!ledger.entries().contains_key(&(EntityId::new(7), type_id)));
|
||||
}
|
||||
|
||||
/// WARNING (REPL): a payload whose LATER batch is corrupt must apply
|
||||
/// **none** of its batches and must NOT advance the applied high-water-mark.
|
||||
/// The earlier (valid) batch must not be half-applied — the leader can then
|
||||
/// cleanly re-ship the same un-acked segment and replay it from one
|
||||
/// consistent resume point.
|
||||
#[test]
|
||||
fn apply_payload_atomic_on_mid_payload_corruption() {
|
||||
let schema = make_schema();
|
||||
let ledger = Arc::new(SignalLedger::new(schema, Box::new(NoopWalWriter)));
|
||||
let state = ReplicationState::new(&[ShardId::SINGLE]);
|
||||
let type_id = ledger.resolve_signal_type("view").unwrap();
|
||||
|
||||
// First batch (seqs 1..=1): valid, entity 100. Second batch (seqs
|
||||
// 2..=2): valid then corrupted so decode fails AFTER the first batch.
|
||||
let good = vec![make_event(100, type_id.as_u16() as u8, 100)];
|
||||
let bad = vec![make_event(200, type_id.as_u16() as u8, 200)];
|
||||
let mut bytes = encode_batch(&good, 1, 100).unwrap();
|
||||
let good_len = bytes.len();
|
||||
bytes.extend(encode_batch(&bad, 2, 200).unwrap());
|
||||
// Corrupt the SECOND batch's checksum region (bytes 32..64 of it).
|
||||
for b in &mut bytes[good_len + 32..good_len + 64] {
|
||||
*b = b.wrapping_add(1);
|
||||
}
|
||||
|
||||
let result = apply_payload(&bytes, ShardId::SINGLE, &ledger, &state, None);
|
||||
assert!(
|
||||
matches!(result, Err(WalError::Corruption { .. })),
|
||||
"mid-payload corruption must return Corruption, got {result:?}"
|
||||
);
|
||||
// NEITHER batch may be applied: the first (valid) batch must not have
|
||||
// been committed before the second batch failed to decode.
|
||||
assert!(
|
||||
!ledger
|
||||
.entries()
|
||||
.contains_key(&(EntityId::new(100), type_id)),
|
||||
"earlier valid batch must NOT be applied when a later batch is corrupt"
|
||||
);
|
||||
assert!(
|
||||
!ledger
|
||||
.entries()
|
||||
.contains_key(&(EntityId::new(200), type_id)),
|
||||
"corrupt batch must not be applied"
|
||||
);
|
||||
// The applied high-water-mark must NOT have advanced past 0.
|
||||
assert_eq!(
|
||||
state.applied_seqno(ShardId::SINGLE),
|
||||
Some(0),
|
||||
"half-applied payload must not advance the applied high-water-mark"
|
||||
);
|
||||
}
|
||||
|
||||
/// WARNING (REPL): in `community_share_only` mode a segment whose events are
|
||||
/// all local filters to empty bytes that decode to ZERO batches. The
|
||||
/// receiver must still advance the lag gauge's leader high-water-mark from
|
||||
/// the segment's authoritative `leader_last_seq` (known before filtering),
|
||||
/// so the gauge does not under-report the leader's progress across the
|
||||
/// all-local segment.
|
||||
#[test]
|
||||
fn lag_gauge_tracks_leader_hwm_across_all_local_empty_segment() {
|
||||
use crate::replication::lag::ReplicationLagGauge;
|
||||
|
||||
let schema = make_schema();
|
||||
let ledger = Arc::new(SignalLedger::new(schema, Box::new(NoopWalWriter)));
|
||||
let state = Arc::new(ReplicationState::new(&[ShardId::SINGLE]));
|
||||
let gauge = Arc::new(ReplicationLagGauge::new(
|
||||
ShardId::SINGLE,
|
||||
Arc::clone(&state),
|
||||
));
|
||||
|
||||
// An all-local segment filters to empty bytes: the leader's TRUE last
|
||||
// WAL seqno for it is 5 (e.g. a batch first_seq=4 with 2 events), but the
|
||||
// shipped bytes are empty.
|
||||
let empty_bytes: Vec<u8> = Vec::new();
|
||||
apply_payload_with_leader_seq(
|
||||
&empty_bytes,
|
||||
ShardId::SINGLE,
|
||||
&ledger,
|
||||
&state,
|
||||
Some(&gauge),
|
||||
Some(5),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
gauge.leader_seqno(),
|
||||
5,
|
||||
"leader HWM must advance from the authoritative leader_last_seq even \
|
||||
when the filtered segment decodes zero batches"
|
||||
);
|
||||
// No events were applied (empty segment), so the applied seqno stays 0.
|
||||
assert_eq!(state.applied_seqno(ShardId::SINGLE), Some(0));
|
||||
// The gauge therefore reports the leader being ahead of the follower.
|
||||
assert_eq!(
|
||||
gauge.lag_segments(),
|
||||
5,
|
||||
"leader at 5, applied 0 → 5 segments behind across the empty segment"
|
||||
);
|
||||
}
|
||||
|
||||
/// A minimal transport that returns one payload then signals shutdown.
|
||||
struct OneShot {
|
||||
rx: crossbeam::channel::Receiver<crate::replication::WalSegmentPayload>,
|
||||
@ -554,6 +764,7 @@ mod tests {
|
||||
id: WalSegmentId::new(RegionId::SINGLE, ShardId(0), 1),
|
||||
bytes: payload_bytes,
|
||||
event_count: 1,
|
||||
leader_last_seq: 1,
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
@ -599,6 +810,7 @@ mod tests {
|
||||
id: WalSegmentId::new(RegionId::SINGLE, ShardId(0), 1),
|
||||
bytes: corrupt_bytes,
|
||||
event_count: 1,
|
||||
leader_last_seq: 1,
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
|
||||
@ -4,15 +4,10 @@
|
||||
//! sealed segments (those with a newer segment after them) and sending each
|
||||
//! to all configured peer shards via the [`Transport`] trait.
|
||||
|
||||
use std::{
|
||||
collections::{HashMap, HashSet},
|
||||
path::PathBuf,
|
||||
sync::Arc,
|
||||
thread::JoinHandle,
|
||||
time::Duration,
|
||||
};
|
||||
use std::{collections::HashMap, path::PathBuf, sync::Arc, thread::JoinHandle, time::Duration};
|
||||
|
||||
use crossbeam::channel::bounded;
|
||||
use dashmap::DashMap;
|
||||
|
||||
use crate::{
|
||||
replication::{
|
||||
@ -60,12 +55,114 @@ impl Default for ShipperConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// Shared, operator-visible shipper state: per-peer high-water-mark and the set
|
||||
/// of quarantined peers.
|
||||
///
|
||||
/// Previously these lived as local variables on the shipper thread's stack
|
||||
/// (`peer_hwm`, `quarantined`), so an operator could neither observe which peers
|
||||
/// were quarantined nor clear a quarantine without dropping and re-spawning the
|
||||
/// whole shipper. Hoisting them into an `Arc<ShipperState>` shared with
|
||||
/// [`WalShipperHandle`] gives the control plane read accessors
|
||||
/// ([`quarantined_peers`](Self::quarantined_peers), [`peer_hwm`](Self::peer_hwm))
|
||||
/// and a [`clear_quarantine`](Self::clear_quarantine) method an operator
|
||||
/// endpoint can call after fixing a peer's config — no shipper restart required.
|
||||
///
|
||||
/// Both maps are [`DashMap`]s so the shipper thread and an operator thread can
|
||||
/// touch them concurrently without a coarse lock. The shipper advances a peer's
|
||||
/// HWM only on its own successful send and re-reads the quarantine set at the
|
||||
/// top of each peer's run, so a concurrent `clear_quarantine` takes effect on
|
||||
/// the next poll.
|
||||
#[derive(Debug, Default)]
|
||||
pub struct ShipperState {
|
||||
/// Per-peer high-water-mark: the highest seqno each peer has acknowledged.
|
||||
peer_hwm: DashMap<ShardId, u64>,
|
||||
/// Peers quarantined after a PERMANENT send failure. Value is the seqno the
|
||||
/// peer was stuck on when quarantined (for operator diagnostics).
|
||||
quarantined: DashMap<ShardId, u64>,
|
||||
}
|
||||
|
||||
impl ShipperState {
|
||||
/// Build a fresh state with every configured peer at HWM 0 and no
|
||||
/// quarantines.
|
||||
#[must_use]
|
||||
fn new(peers: &[ShardId]) -> Self {
|
||||
let peer_hwm = DashMap::new();
|
||||
for &p in peers {
|
||||
peer_hwm.insert(p, 0);
|
||||
}
|
||||
Self {
|
||||
peer_hwm,
|
||||
quarantined: DashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// The high-water-mark (last acknowledged seqno) for `peer`, or 0 if unknown.
|
||||
#[must_use]
|
||||
pub fn peer_hwm(&self, peer: ShardId) -> u64 {
|
||||
self.peer_hwm.get(&peer).map_or(0, |r| *r)
|
||||
}
|
||||
|
||||
/// Snapshot of every tracked peer's high-water-mark.
|
||||
#[must_use]
|
||||
pub fn peer_hwms(&self) -> HashMap<ShardId, u64> {
|
||||
self.peer_hwm
|
||||
.iter()
|
||||
.map(|r| (*r.key(), *r.value()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// `true` if `peer` is currently quarantined (permanently failed).
|
||||
#[must_use]
|
||||
pub fn is_quarantined(&self, peer: ShardId) -> bool {
|
||||
self.quarantined.contains_key(&peer)
|
||||
}
|
||||
|
||||
/// The set of currently-quarantined peers, each with the seqno it was stuck
|
||||
/// on when quarantined.
|
||||
#[must_use]
|
||||
pub fn quarantined_peers(&self) -> HashMap<ShardId, u64> {
|
||||
self.quarantined
|
||||
.iter()
|
||||
.map(|r| (*r.key(), *r.value()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Clear `peer`'s quarantine so the shipper retries it on the next poll.
|
||||
///
|
||||
/// Intended for the operator endpoint after a peer's config/TLS/CA has been
|
||||
/// fixed. Returns `true` if the peer was quarantined (and is now cleared),
|
||||
/// `false` if it was not quarantined. The peer resumes from its existing
|
||||
/// high-water-mark, so no segment is re-shipped that the peer already acked.
|
||||
pub fn clear_quarantine(&self, peer: ShardId) -> bool {
|
||||
let was = self.quarantined.remove(&peer).is_some();
|
||||
if was {
|
||||
tracing::info!(
|
||||
peer = %peer,
|
||||
"shipper: quarantine cleared by operator; peer will be retried next poll"
|
||||
);
|
||||
}
|
||||
was
|
||||
}
|
||||
|
||||
// ── Internal mutators used by the shipper thread ────────────────────────
|
||||
|
||||
fn advance_hwm(&self, peer: ShardId, seqno: u64) {
|
||||
self.peer_hwm.insert(peer, seqno);
|
||||
}
|
||||
|
||||
fn quarantine(&self, peer: ShardId, stuck_seqno: u64) {
|
||||
self.quarantined.insert(peer, stuck_seqno);
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle to a running WAL shipper thread.
|
||||
///
|
||||
/// Call [`stop`](Self::stop) to signal shutdown and join the thread.
|
||||
pub struct WalShipperHandle {
|
||||
shutdown_tx: crossbeam::channel::Sender<()>,
|
||||
thread: Option<JoinHandle<()>>,
|
||||
/// Shared, operator-visible shipper state (per-peer HWM + quarantine set).
|
||||
state: Arc<ShipperState>,
|
||||
}
|
||||
|
||||
impl WalShipperHandle {
|
||||
@ -76,6 +173,22 @@ impl WalShipperHandle {
|
||||
let _ = handle.join();
|
||||
}
|
||||
}
|
||||
|
||||
/// The shared shipper state for operator/control-plane inspection and
|
||||
/// quarantine clearing. The returned `Arc` stays valid after [`stop`](Self::stop).
|
||||
#[must_use]
|
||||
pub fn state(&self) -> Arc<ShipperState> {
|
||||
Arc::clone(&self.state)
|
||||
}
|
||||
|
||||
/// Clear a peer's quarantine without restarting the shipper.
|
||||
///
|
||||
/// Returns `true` if the peer was quarantined and is now cleared. See
|
||||
/// [`ShipperState::clear_quarantine`].
|
||||
#[must_use]
|
||||
pub fn clear_quarantine(&self, peer: ShardId) -> bool {
|
||||
self.state.clear_quarantine(peer)
|
||||
}
|
||||
}
|
||||
|
||||
/// Spawn a background thread that polls for sealed WAL segments and ships
|
||||
@ -94,25 +207,27 @@ pub fn spawn_shipper<T: Transport + ?Sized>(
|
||||
transport: Arc<T>,
|
||||
) -> WalShipperHandle {
|
||||
let (shutdown_tx, shutdown_rx) = bounded::<()>(1);
|
||||
// Shared, operator-visible per-peer HWM + quarantine set. Owned by both the
|
||||
// handle (for inspection / clear_quarantine) and the shipper thread.
|
||||
let state = Arc::new(ShipperState::new(&config.peer_shards));
|
||||
let thread_state = Arc::clone(&state);
|
||||
let thread = std::thread::Builder::new()
|
||||
.name("tidaldb-wal-shipper".into())
|
||||
.spawn(move || {
|
||||
// Per-peer high-water-mark: the highest seqno each peer has
|
||||
// acknowledged. A peer's cursor advances only when ITS send
|
||||
// succeeds, so a peer that fails (or applies backpressure) keeps
|
||||
// re-receiving the segment on subsequent polls instead of being
|
||||
// permanently skipped.
|
||||
let mut peer_hwm: HashMap<ShardId, u64> =
|
||||
config.peer_shards.iter().map(|&p| (p, 0)).collect();
|
||||
// Peers that hit a PERMANENT send failure (bad TLS/CA, auth
|
||||
// rejection, malformed payload, unimplemented RPC). Re-shipping the
|
||||
// same seqno to these peers will fail identically forever, so we
|
||||
// quarantine them: they are skipped entirely on every subsequent
|
||||
// poll and their cursor never advances. This trades a silent
|
||||
// infinite retry stall for a loud, actionable, bounded failure —
|
||||
// an operator must fix the peer/config and restart the shipper to
|
||||
// clear the quarantine. Healthy peers are unaffected.
|
||||
let mut quarantined: HashSet<ShardId> = HashSet::new();
|
||||
// Per-peer high-water-mark and the quarantine set now live in the
|
||||
// shared `thread_state` (an `Arc<ShipperState>`) instead of on this
|
||||
// thread's stack, so an operator can observe quarantined peers / HWMs
|
||||
// and clear a quarantine via the handle without restarting the
|
||||
// shipper. The shipper still advances a peer's cursor ONLY when ITS
|
||||
// own send succeeds, and re-reads the quarantine set at the top of
|
||||
// each peer's run so a concurrent clear takes effect next poll.
|
||||
//
|
||||
// A peer hits PERMANENT failure (bad TLS/CA, auth rejection,
|
||||
// malformed payload, unimplemented RPC) and is quarantined: skipped
|
||||
// on every subsequent poll, cursor frozen, until an operator clears
|
||||
// it. This trades a silent infinite retry stall for a loud,
|
||||
// actionable, bounded failure. Healthy peers are unaffected.
|
||||
let state = thread_state;
|
||||
loop {
|
||||
// Sleep for poll_interval, exit early on shutdown signal.
|
||||
match shutdown_rx.recv_timeout(config.poll_interval) {
|
||||
@ -144,8 +259,13 @@ pub fn spawn_shipper<T: Transport + ?Sized>(
|
||||
let sealed = &segments[..sealed_count];
|
||||
|
||||
// Cache segment bytes read this poll so peers at the same lag
|
||||
// don't each re-read the same file. Keyed by seqno.
|
||||
let mut bytes_cache: HashMap<u64, Vec<u8>> = HashMap::new();
|
||||
// don't each re-read the same file. Keyed by seqno. Each entry
|
||||
// carries the (possibly filtered) shippable bytes AND the
|
||||
// segment's authoritative last WAL seqno computed from the
|
||||
// ORIGINAL (pre-filter) bytes — the latter is what the receiver's
|
||||
// lag gauge needs so an all-local empty filtration still advances
|
||||
// the leader high-water-mark.
|
||||
let mut bytes_cache: HashMap<u64, (Vec<u8>, u64)> = HashMap::new();
|
||||
|
||||
// Ship each peer its own contiguous run of segments, in seqno
|
||||
// order, starting just past its high-water-mark. A peer is
|
||||
@ -159,28 +279,35 @@ pub fn spawn_shipper<T: Transport + ?Sized>(
|
||||
// re-attempt it (re-shipping would fail identically). It
|
||||
// stays skipped until an operator fixes the peer/config and
|
||||
// restarts the shipper.
|
||||
if quarantined.contains(&peer) {
|
||||
if state.is_quarantined(peer) {
|
||||
continue;
|
||||
}
|
||||
for (seqno, path) in sealed {
|
||||
if *seqno <= peer_hwm.get(&peer).copied().unwrap_or(0) {
|
||||
if *seqno <= state.peer_hwm(peer) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Read (and cache) the segment bytes lazily. In
|
||||
// community-overlay mode, strip local-scope events from
|
||||
// the segment before caching so no peer ever receives a
|
||||
// local signal.
|
||||
let segment_bytes = match bytes_cache.get(seqno) {
|
||||
// local signal. We compute the authoritative last WAL
|
||||
// seqno from the ORIGINAL bytes BEFORE filtering, so an
|
||||
// all-local segment (which filters to empty) still carries
|
||||
// the leader's true high-water-mark for the receiver's lag
|
||||
// gauge.
|
||||
let entry = match bytes_cache.get(seqno) {
|
||||
Some(b) => b,
|
||||
None => match std::fs::read(path) {
|
||||
Ok(b) => {
|
||||
let b = if config.community_share_only {
|
||||
filter_segment_drop_local(&b)
|
||||
Ok(original) => {
|
||||
let leader_last_seq = last_seq_in_segment(&original);
|
||||
let shippable = if config.community_share_only {
|
||||
filter_segment_drop_local(&original)
|
||||
} else {
|
||||
b
|
||||
original
|
||||
};
|
||||
bytes_cache.entry(*seqno).or_insert(b)
|
||||
bytes_cache
|
||||
.entry(*seqno)
|
||||
.or_insert((shippable, leader_last_seq))
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
@ -195,11 +322,13 @@ pub fn spawn_shipper<T: Transport + ?Sized>(
|
||||
}
|
||||
},
|
||||
};
|
||||
let (segment_bytes, leader_last_seq) = entry;
|
||||
let event_count = count_events_in_segment(segment_bytes);
|
||||
let payload = WalSegmentPayload {
|
||||
id: WalSegmentId::new(RegionId::SINGLE, config.shard_id, *seqno),
|
||||
bytes: segment_bytes.clone(),
|
||||
event_count,
|
||||
leader_last_seq: *leader_last_seq,
|
||||
};
|
||||
match transport.send_segment(peer, payload) {
|
||||
Ok(()) => {
|
||||
@ -210,7 +339,7 @@ pub fn spawn_shipper<T: Transport + ?Sized>(
|
||||
);
|
||||
// Advance ONLY this peer's cursor on its own
|
||||
// successful send.
|
||||
peer_hwm.insert(peer, *seqno);
|
||||
state.advance_hwm(peer, *seqno);
|
||||
}
|
||||
Err(TransportError::Closed) => {
|
||||
// Transient (peer down, circuit open, or
|
||||
@ -231,18 +360,21 @@ pub fn spawn_shipper<T: Transport + ?Sized>(
|
||||
// RPC). Quarantine the peer instead of stalling
|
||||
// on an unbounded silent retry: stop advancing
|
||||
// it, skip it on every future poll, and surface
|
||||
// a LOUD, actionable error. An operator must fix
|
||||
// the peer/config and restart the shipper to
|
||||
// clear the quarantine. The peer's HWM is left
|
||||
// untouched so no segment is silently skipped.
|
||||
quarantined.insert(peer);
|
||||
// a LOUD, actionable error. The peer's HWM is
|
||||
// left untouched so no segment is silently
|
||||
// skipped. The quarantine is now operator-clearable
|
||||
// via WalShipperHandle::clear_quarantine — no
|
||||
// shipper restart required once the peer/config is
|
||||
// fixed.
|
||||
state.quarantine(peer, *seqno);
|
||||
tracing::error!(
|
||||
peer = %peer,
|
||||
seqno = *seqno,
|
||||
reason = %reason,
|
||||
"shipper: PERMANENT transport failure — peer quarantined, \
|
||||
replication to this peer is STOPPED until operator fixes \
|
||||
peer/config and restarts the shipper; retrying would not help"
|
||||
replication to this peer is STOPPED until an operator clears \
|
||||
the quarantine (clear_quarantine) after fixing peer/config; \
|
||||
retrying without a fix would not help"
|
||||
);
|
||||
break;
|
||||
}
|
||||
@ -268,6 +400,7 @@ pub fn spawn_shipper<T: Transport + ?Sized>(
|
||||
WalShipperHandle {
|
||||
shutdown_tx,
|
||||
thread: Some(thread),
|
||||
state,
|
||||
}
|
||||
}
|
||||
|
||||
@ -376,6 +509,45 @@ pub fn filter_segment_drop_local(bytes: &[u8]) -> Vec<u8> {
|
||||
out
|
||||
}
|
||||
|
||||
/// Compute a WAL segment's authoritative last sequence number from its
|
||||
/// **original** (pre-filter) bytes: the maximum `first_seq + event_count - 1`
|
||||
/// across all decodable batches.
|
||||
///
|
||||
/// This is the value the leader knows before any community-overlay filtering,
|
||||
/// and the value the receiver feeds to the replication lag gauge's leader
|
||||
/// high-water-mark. Computing it from the original bytes — not the filtered
|
||||
/// ones — is load-bearing: an all-local segment filters to empty and would
|
||||
/// otherwise report a leader seqno of 0, making the receiver's gauge
|
||||
/// under-report the leader's progress across that segment (obs-REPL-1).
|
||||
///
|
||||
/// Returns `0` for an empty or fully-undecodable segment (the receiver treats
|
||||
/// `0` as "unknown" and falls back to the per-batch boundaries it decodes).
|
||||
/// Corrupt or trailing bytes terminate the scan (matching
|
||||
/// [`count_events_in_segment`]); the last seq of the already-decoded batches is
|
||||
/// returned.
|
||||
fn last_seq_in_segment(bytes: &[u8]) -> u64 {
|
||||
let mut offset = 0;
|
||||
let mut last_seq = 0u64;
|
||||
while offset < bytes.len() {
|
||||
let remaining = &bytes[offset..];
|
||||
if remaining.len() < HEADER_SIZE {
|
||||
break;
|
||||
}
|
||||
match decode_batch(remaining) {
|
||||
Ok((header, _events)) => {
|
||||
let batch_last_seq = header
|
||||
.first_seq
|
||||
.saturating_add(u64::from(header.event_count).saturating_sub(1));
|
||||
last_seq = last_seq.max(batch_last_seq);
|
||||
let batch_size = HEADER_SIZE + header.payload_len as usize;
|
||||
offset += batch_size;
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
last_seq
|
||||
}
|
||||
|
||||
/// Count the total number of events in a WAL segment by parsing all batches.
|
||||
fn count_events_in_segment(bytes: &[u8]) -> u64 {
|
||||
let mut offset = 0;
|
||||
@ -602,6 +774,13 @@ mod tests {
|
||||
fn heal(&self, peer: ShardId) {
|
||||
self.failing.lock().unwrap().remove(&peer);
|
||||
}
|
||||
|
||||
/// Clear a peer from the PERMANENT-failure set (the operator fixed the
|
||||
/// peer's config/TLS). Used together with `clear_quarantine` to prove a
|
||||
/// quarantined peer can recover without a shipper restart.
|
||||
fn heal_permanent(&self, peer: ShardId) {
|
||||
self.permanent.lock().unwrap().remove(&peer);
|
||||
}
|
||||
}
|
||||
|
||||
impl Transport for RecordingTransport {
|
||||
@ -831,4 +1010,185 @@ mod tests {
|
||||
|
||||
handle.stop();
|
||||
}
|
||||
|
||||
// ── Operator-visible quarantine/HWM state (Extensibility-W) ──────────────
|
||||
|
||||
/// `last_seq_in_segment` returns the maximum `first_seq + event_count - 1`
|
||||
/// across all batches, computed from the ORIGINAL (pre-filter) bytes.
|
||||
#[test]
|
||||
fn last_seq_in_segment_is_max_batch_boundary() {
|
||||
// Batch A: first_seq=10, 1 event → last 10. Batch B: first_seq=11, 3
|
||||
// events → last 13. Segment last seq = 13.
|
||||
let a = encode_batch(
|
||||
&[EventRecord::signal(1, RECORD_TYPE_SIGNAL, 1.0, 100)],
|
||||
10,
|
||||
10,
|
||||
)
|
||||
.unwrap();
|
||||
let b = encode_batch(
|
||||
&(0..3)
|
||||
.map(|i| EventRecord::signal(i, RECORD_TYPE_SIGNAL, 1.0, 200))
|
||||
.collect::<Vec<_>>(),
|
||||
11,
|
||||
11,
|
||||
)
|
||||
.unwrap();
|
||||
let mut seg = a;
|
||||
seg.extend(b);
|
||||
assert_eq!(last_seq_in_segment(&seg), 13);
|
||||
// Empty segment → 0 (unknown).
|
||||
assert_eq!(last_seq_in_segment(&[]), 0);
|
||||
}
|
||||
|
||||
/// The shipper's quarantine and per-peer HWM are visible to an operator via
|
||||
/// the handle's shared `ShipperState`, and a quarantine can be cleared
|
||||
/// WITHOUT restarting the shipper — after which the (now-healed) peer is
|
||||
/// retried and catches up.
|
||||
#[test]
|
||||
fn quarantine_and_hwm_visible_and_clearable_without_restart() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
seed_sealed_segments(dir.path(), 3);
|
||||
|
||||
let healthy = ShardId(1);
|
||||
let permanent_peer = ShardId(2);
|
||||
let transport = Arc::new(RecordingTransport::with_permanent(&[permanent_peer]));
|
||||
|
||||
let config = ShipperConfig {
|
||||
wal_dir: dir.path().to_path_buf(),
|
||||
shard_id: ShardId::SINGLE,
|
||||
peer_shards: vec![healthy, permanent_peer],
|
||||
poll_interval: Duration::from_millis(5),
|
||||
community_share_only: false,
|
||||
};
|
||||
let handle = spawn_shipper(config, Arc::clone(&transport));
|
||||
let state = handle.state();
|
||||
|
||||
// The healthy peer advances; its HWM becomes observable as 3.
|
||||
wait_until(Duration::from_secs(2), || state.peer_hwm(healthy) == 3);
|
||||
assert_eq!(
|
||||
state.peer_hwm(healthy),
|
||||
3,
|
||||
"operator can read the healthy peer's high-water-mark"
|
||||
);
|
||||
|
||||
// The permanent peer becomes observably quarantined (stuck at seqno 1).
|
||||
wait_until(Duration::from_secs(2), || {
|
||||
state.is_quarantined(permanent_peer)
|
||||
});
|
||||
let quarantined = state.quarantined_peers();
|
||||
assert_eq!(
|
||||
quarantined.get(&permanent_peer).copied(),
|
||||
Some(1),
|
||||
"quarantine set must expose the peer and the seqno it stuck on, got: {quarantined:?}"
|
||||
);
|
||||
|
||||
// Operator fixes the peer's config, then clears the quarantine via the
|
||||
// handle — no shipper restart.
|
||||
transport.heal_permanent(permanent_peer);
|
||||
assert!(
|
||||
handle.clear_quarantine(permanent_peer),
|
||||
"clear_quarantine must report it cleared a live quarantine"
|
||||
);
|
||||
assert!(
|
||||
!state.is_quarantined(permanent_peer),
|
||||
"peer must no longer be quarantined after clear"
|
||||
);
|
||||
|
||||
// The peer is now retried and catches up to all sealed segments.
|
||||
wait_until(Duration::from_secs(2), || {
|
||||
state.peer_hwm(permanent_peer) == 3
|
||||
});
|
||||
assert_eq!(
|
||||
state.peer_hwm(permanent_peer),
|
||||
3,
|
||||
"cleared peer resumes shipping and catches up without a restart"
|
||||
);
|
||||
|
||||
// Clearing a non-quarantined peer reports false.
|
||||
assert!(!handle.clear_quarantine(ShardId(99)));
|
||||
|
||||
handle.stop();
|
||||
}
|
||||
|
||||
/// In community-overlay mode an all-local segment ships empty bytes but the
|
||||
/// payload must still carry the authoritative `leader_last_seq` so the
|
||||
/// receiver's lag gauge does not under-report (obs-REPL-1).
|
||||
#[test]
|
||||
fn all_local_segment_ships_empty_with_authoritative_leader_last_seq() {
|
||||
use std::sync::Mutex as StdMutex;
|
||||
|
||||
/// Captures the last payload's (`bytes_len`, `leader_last_seq`) per send.
|
||||
struct CapturingTransport {
|
||||
captured: StdMutex<Vec<(usize, u64)>>,
|
||||
}
|
||||
impl Transport for CapturingTransport {
|
||||
fn send_segment(
|
||||
&self,
|
||||
_to: ShardId,
|
||||
payload: WalSegmentPayload,
|
||||
) -> Result<(), TransportError> {
|
||||
self.captured
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push((payload.bytes.len(), payload.leader_last_seq));
|
||||
Ok(())
|
||||
}
|
||||
fn recv_segment(&self) -> Option<WalSegmentPayload> {
|
||||
None
|
||||
}
|
||||
fn local_shard(&self) -> ShardId {
|
||||
ShardId::SINGLE
|
||||
}
|
||||
}
|
||||
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
// Write one sealed all-local segment (file seqno 1) whose single batch
|
||||
// has first_seq=4 and 2 events → true last WAL seq = 5; plus a trailing
|
||||
// active segment so the first is considered sealed.
|
||||
let all_local = encode_batch(
|
||||
&[
|
||||
EventRecord::signal(1, RECORD_TYPE_SIGNAL, 1.0, 100),
|
||||
EventRecord::signal(2, RECORD_TYPE_SIGNAL, 1.0, 200),
|
||||
],
|
||||
4,
|
||||
4,
|
||||
)
|
||||
.unwrap();
|
||||
std::fs::write(
|
||||
dir.path().join(segment_filename(ShardId::SINGLE, 1)),
|
||||
all_local,
|
||||
)
|
||||
.unwrap();
|
||||
std::fs::write(
|
||||
dir.path().join(segment_filename(ShardId::SINGLE, 2)),
|
||||
encode_batch(&[community_event(9)], 6, 6).unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let peer = ShardId(1);
|
||||
let transport = Arc::new(CapturingTransport {
|
||||
captured: StdMutex::new(Vec::new()),
|
||||
});
|
||||
let config = ShipperConfig {
|
||||
wal_dir: dir.path().to_path_buf(),
|
||||
shard_id: ShardId::SINGLE,
|
||||
peer_shards: vec![peer],
|
||||
poll_interval: Duration::from_millis(5),
|
||||
community_share_only: true,
|
||||
};
|
||||
let handle = spawn_shipper(config, Arc::clone(&transport));
|
||||
|
||||
wait_until(Duration::from_secs(2), || {
|
||||
!transport.captured.lock().unwrap().is_empty()
|
||||
});
|
||||
let captured = transport.captured.lock().unwrap().clone();
|
||||
handle.stop();
|
||||
|
||||
let (bytes_len, leader_last_seq) = captured[0];
|
||||
assert_eq!(bytes_len, 0, "all-local segment filters to empty bytes");
|
||||
assert_eq!(
|
||||
leader_last_seq, 5,
|
||||
"empty segment must still carry the leader's true last WAL seq (4 + 2 - 1)"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@ -196,6 +196,18 @@ impl TenantRateLimiter {
|
||||
}
|
||||
}
|
||||
|
||||
/// On-disk shape of the migration routing checkpoint (dual-write pairs + pins).
|
||||
///
|
||||
/// String keys are decimal [`TenantId`]s; dual-write values are
|
||||
/// `(source_shard, target_shard)` raw `u16`s and pin values are the target
|
||||
/// `ShardId` `u16`. `Default` (both maps empty) is the safe restore for missing
|
||||
/// or malformed bytes.
|
||||
#[derive(Debug, Default, serde::Serialize, serde::Deserialize)]
|
||||
struct MigrationRoutingCheckpoint {
|
||||
dual: std::collections::HashMap<String, (u16, u16)>,
|
||||
pins: std::collections::HashMap<String, u16>,
|
||||
}
|
||||
|
||||
/// Describes the shard topology for the cluster.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ClusterTopology {
|
||||
@ -371,6 +383,69 @@ impl TenantRouter {
|
||||
self.shard_pin_map.get(&tenant_id).map(|r| *r)
|
||||
}
|
||||
|
||||
// ── Migration routing durability (Accuracy-W) ───────────────────────────
|
||||
//
|
||||
// The dual-write and shard-pin maps decide where a tenant's writes go during
|
||||
// and after a migration. They used to live ONLY in memory, so a crash
|
||||
// mid-migration silently reverted a finalized tenant to jump-hash routing —
|
||||
// re-scattering its data to the wrong shard. These maps are now serialized to
|
||||
// a single durable row and restored on open.
|
||||
|
||||
/// Serialize the migration routing state (dual-write pairs + shard pins) for
|
||||
/// the durable checkpoint.
|
||||
///
|
||||
/// JSON with string keys (the [`TenantId`] as a decimal string) because JSON
|
||||
/// object keys must be strings; values are the raw `ShardId` `u16`s. The two
|
||||
/// maps are written as one object so a single row captures the whole routing
|
||||
/// state atomically.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Panics only if `serde_json` fails to serialize a `HashMap` of plain
|
||||
/// integers, which is unreachable (both key and value types are always
|
||||
/// representable).
|
||||
#[must_use]
|
||||
pub fn to_checkpoint_bytes(&self) -> Vec<u8> {
|
||||
let dual: std::collections::HashMap<String, (u16, u16)> = self
|
||||
.dual_write_map
|
||||
.iter()
|
||||
.map(|r| (r.key().0.to_string(), (r.value().0.0, r.value().1.0)))
|
||||
.collect();
|
||||
let pins: std::collections::HashMap<String, u16> = self
|
||||
.shard_pin_map
|
||||
.iter()
|
||||
.map(|r| (r.key().0.to_string(), r.value().0))
|
||||
.collect();
|
||||
let snap = MigrationRoutingCheckpoint { dual, pins };
|
||||
serde_json::to_vec(&snap).expect("migration routing serialization is infallible")
|
||||
}
|
||||
|
||||
/// Restore the migration routing state from checkpoint bytes, replacing any
|
||||
/// current in-memory routing.
|
||||
///
|
||||
/// Malformed bytes restore nothing (the maps are left empty), which is the
|
||||
/// safe degradation: a tenant with no pin/dual-write falls back to jump-hash
|
||||
/// routing, identical to a never-migrated tenant. Returns the number of
|
||||
/// tenants whose routing was restored (dual-write + pinned).
|
||||
#[must_use]
|
||||
pub fn restore_from_checkpoint_bytes(&self, bytes: &[u8]) -> usize {
|
||||
let snap: MigrationRoutingCheckpoint = serde_json::from_slice(bytes).unwrap_or_default();
|
||||
self.dual_write_map.clear();
|
||||
self.shard_pin_map.clear();
|
||||
for (tid, (src, tgt)) in snap.dual {
|
||||
if let Ok(id) = tid.parse::<u64>() {
|
||||
self.dual_write_map
|
||||
.insert(TenantId(id), (ShardId(src), ShardId(tgt)));
|
||||
}
|
||||
}
|
||||
for (tid, shard) in snap.pins {
|
||||
if let Ok(id) = tid.parse::<u64>() {
|
||||
self.shard_pin_map.insert(TenantId(id), ShardId(shard));
|
||||
}
|
||||
}
|
||||
self.dual_write_map.len() + self.shard_pin_map.len()
|
||||
}
|
||||
|
||||
fn eligible_shards_for(
|
||||
&self,
|
||||
tenant_id: TenantId,
|
||||
@ -473,6 +548,72 @@ pub fn tenant_wal_dir(data_dir: &std::path::Path, tenant_id: TenantId) -> std::p
|
||||
}
|
||||
}
|
||||
|
||||
// ── Migration routing checkpoint persistence ────────────────────────────────
|
||||
|
||||
/// Storage key for the single migration-routing checkpoint row (entity id 0,
|
||||
/// [`Tag::MigrationRouting`](crate::storage::Tag::MigrationRouting)).
|
||||
#[must_use]
|
||||
fn migration_routing_key() -> Vec<u8> {
|
||||
crate::storage::encode_key(
|
||||
EntityId::new(0),
|
||||
crate::storage::Tag::MigrationRouting,
|
||||
b"route",
|
||||
)
|
||||
}
|
||||
|
||||
/// Durably persist the tenant router's migration routing state (dual-write pairs
|
||||
/// + shard pins).
|
||||
///
|
||||
/// Called from the migration coordinator on every routing transition (enter
|
||||
/// dual-write, finalize) so a crash mid-migration cannot silently revert a
|
||||
/// finalized tenant to jump-hash routing. Writes one row and fsyncs.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TidalError::Storage` if the write or flush fails.
|
||||
pub(crate) fn checkpoint_migration_routing(
|
||||
storage: &dyn crate::storage::StorageEngine,
|
||||
router: &TenantRouter,
|
||||
) -> crate::Result<()> {
|
||||
let bytes = router.to_checkpoint_bytes();
|
||||
let mut batch = crate::storage::WriteBatch::new();
|
||||
batch.put(migration_routing_key(), bytes);
|
||||
storage.write_batch(batch)?;
|
||||
storage.flush()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Restore the tenant router's migration routing state from its durable
|
||||
/// checkpoint on open.
|
||||
///
|
||||
/// On first boot / no checkpoint / ephemeral mode / corrupt bytes this restores
|
||||
/// nothing (every tenant uses jump-hash routing), which is correct — identical
|
||||
/// to a cluster that never migrated a tenant. Returns the number of tenants
|
||||
/// whose routing was restored.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TidalError::Storage` if the checkpoint read fails (the caller may
|
||||
/// treat this as non-fatal and start with empty routing).
|
||||
pub(crate) fn restore_migration_routing(
|
||||
storage: &dyn crate::storage::StorageEngine,
|
||||
router: &TenantRouter,
|
||||
) -> crate::Result<usize> {
|
||||
match storage.get(&migration_routing_key())? {
|
||||
Some(bytes) => {
|
||||
let n = router.restore_from_checkpoint_bytes(&bytes);
|
||||
if n > 0 {
|
||||
tracing::info!(
|
||||
tenants = n,
|
||||
"migration routing state restored from checkpoint"
|
||||
);
|
||||
}
|
||||
Ok(n)
|
||||
}
|
||||
None => Ok(0),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[allow(clippy::unwrap_used)]
|
||||
mod tests {
|
||||
@ -707,4 +848,107 @@ mod tests {
|
||||
std::path::PathBuf::from("/data/tenants/7/wal")
|
||||
);
|
||||
}
|
||||
|
||||
// ── Migration routing durability (Accuracy-W) ───────────────────────────
|
||||
|
||||
fn two_shard_router() -> TenantRouter {
|
||||
let topo = Arc::new(RwLock::new(ClusterTopology {
|
||||
shards: vec![
|
||||
ShardAssignment {
|
||||
shard_id: ShardId(0),
|
||||
region_id: RegionId(0),
|
||||
},
|
||||
ShardAssignment {
|
||||
shard_id: ShardId(1),
|
||||
region_id: RegionId(1),
|
||||
},
|
||||
],
|
||||
}));
|
||||
TenantRouter::new(topo)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn migration_routing_checkpoint_roundtrips_dual_write_and_pins() {
|
||||
let router = two_shard_router();
|
||||
router.set_dual_write(TenantId(5), ShardId(0), ShardId(1));
|
||||
router.finalize_migration(TenantId(9), ShardId(1)); // pins 9 -> shard 1
|
||||
|
||||
let bytes = router.to_checkpoint_bytes();
|
||||
|
||||
// Restore into a fresh router (different identity, same topology).
|
||||
let restored = two_shard_router();
|
||||
let n = restored.restore_from_checkpoint_bytes(&bytes);
|
||||
assert_eq!(n, 2, "one dual-write tenant + one pinned tenant restored");
|
||||
|
||||
assert!(
|
||||
restored.is_dual_write(TenantId(5)),
|
||||
"dual-write routing must survive the checkpoint"
|
||||
);
|
||||
assert_eq!(
|
||||
restored.pinned_shard(TenantId(9)),
|
||||
Some(ShardId(1)),
|
||||
"the finalized pin must survive the checkpoint (no revert to jump-hash)"
|
||||
);
|
||||
// The pinned tenant routes to its target shard, NOT a jump-hash result.
|
||||
let assignment = restored.route(TenantId(9), EntityId::new(123)).unwrap();
|
||||
assert_eq!(assignment.shard_id, ShardId(1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn migration_routing_restore_replaces_current_state() {
|
||||
let router = two_shard_router();
|
||||
// Pre-existing (stale) routing that the restore must overwrite.
|
||||
router.set_dual_write(TenantId(1), ShardId(1), ShardId(0));
|
||||
|
||||
// Checkpoint from a DIFFERENT router with a single pinned tenant.
|
||||
let source = two_shard_router();
|
||||
source.finalize_migration(TenantId(2), ShardId(1));
|
||||
let bytes = source.to_checkpoint_bytes();
|
||||
|
||||
let _ = router.restore_from_checkpoint_bytes(&bytes);
|
||||
assert!(
|
||||
!router.is_dual_write(TenantId(1)),
|
||||
"restore must clear stale in-memory routing"
|
||||
);
|
||||
assert_eq!(router.pinned_shard(TenantId(2)), Some(ShardId(1)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn migration_routing_malformed_bytes_restore_empty() {
|
||||
let router = two_shard_router();
|
||||
let n = router.restore_from_checkpoint_bytes(b"not json");
|
||||
assert_eq!(
|
||||
n, 0,
|
||||
"malformed bytes restore nothing (safe jump-hash fallback)"
|
||||
);
|
||||
assert!(!router.is_dual_write(TenantId(1)));
|
||||
assert_eq!(router.pinned_shard(TenantId(1)), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn migration_routing_persist_and_restore_through_storage() {
|
||||
use crate::storage::InMemoryBackend;
|
||||
|
||||
let storage = InMemoryBackend::new();
|
||||
let router = two_shard_router();
|
||||
router.finalize_migration(TenantId(7), ShardId(1));
|
||||
|
||||
super::checkpoint_migration_routing(&storage, &router).unwrap();
|
||||
|
||||
// A fresh router restores the pin from the durable row.
|
||||
let restored = two_shard_router();
|
||||
let n = super::restore_migration_routing(&storage, &restored).unwrap();
|
||||
assert_eq!(n, 1);
|
||||
assert_eq!(restored.pinned_shard(TenantId(7)), Some(ShardId(1)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn migration_routing_restore_missing_row_is_empty() {
|
||||
use crate::storage::InMemoryBackend;
|
||||
|
||||
let storage = InMemoryBackend::new();
|
||||
let router = two_shard_router();
|
||||
let n = super::restore_migration_routing(&storage, &router).unwrap();
|
||||
assert_eq!(n, 0, "no checkpoint row → nothing restored");
|
||||
}
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in New Issue
Block a user