diff --git a/.env.example b/.env.example index 8e0f345..ed47929 100644 --- a/.env.example +++ b/.env.example @@ -59,6 +59,27 @@ RUST_LOG=compass=info # AZURE_STORAGE_CONNECTION_STRING= # AZURE_STORAGE_SAS_KEY= +# ── Warm serverless (cloud mode only) ─────────────────────────────────────── +# Node role: `full` (default — serve reads + writes with local indexes) or +# `writer` (durable-append-only: no local indexes, no read serving, instant +# boot). Writers validate against the bucket's collection config. +# COMPASS_ROLE=full + +# Seconds between manifest refreshes (convergence with other nodes' writes). +# Default 5; 0 disables the background refresher. +# COMPASS_REFRESH_INTERVAL=5 + +# Lazy attach: register bucket collections at boot and attach (rebuild local +# indexes) on first request instead of eagerly. Default false. +# COMPASS_LAZY_ATTACH=false + +# Max simultaneously-attached collections when lazy attach is on (LRU detach +# past the budget; detached collections re-attach on demand). 0 = unbounded. +# COMPASS_MAX_ATTACHED=0 + # ── Telemetry (anonymous; opt out) ────────────────────────────────────────── +# Global in-flight request cap (backpressure). Unset = effectively unlimited. +# COMPASS_MAX_CONCURRENCY=1024 + # COMPASS_TELEMETRY=off # DO_NOT_TRACK=1 diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 779ea3a..68c27cb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -57,6 +57,71 @@ jobs: sudo apt-get install -y cmake pkg-config libssl-dev - run: cargo test --workspace --exclude compass-vector-gpu + # Cloud-feature coverage: the object-storage build + the real-S3 integration + # tests against MinIO. Not in the required-checks list (new job), but a + # failure here still blocks review attention. + test-cloud: + runs-on: ubuntu-24.04 + env: + COMPASS_TEST_S3_BUCKET: compass-data + COMPASS_S3_ENDPOINT: http://localhost:9000 + COMPASS_S3_ALLOW_HTTP: "true" + AWS_ACCESS_KEY_ID: minioadmin + AWS_SECRET_ACCESS_KEY: minioadmin + AWS_DEFAULT_REGION: us-east-1 + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@stable + - uses: Swatinem/rust-cache@v2 + - run: | + sudo apt-get update + sudo apt-get install -y cmake pkg-config libssl-dev + # MinIO as a plain container (service containers can't override the + # image command, and minio/minio needs `server /data`). Same images as + # docker-compose.minio.yml. + - run: | + docker run -d --name minio -p 9000:9000 \ + -e MINIO_ROOT_USER=minioadmin -e MINIO_ROOT_PASSWORD=minioadmin \ + minio/minio:latest server /data + for i in $(seq 1 30); do + curl -sf http://localhost:9000/minio/health/live && break + sleep 1 + done + docker run --rm --network host --entrypoint sh minio/mc:latest -c \ + "mc alias set local http://localhost:9000 minioadmin minioadmin && mc mb -p local/compass-data" + # NON-TEST compile of the cloud feature: no other job builds this + # combination (clippy/msrv build without the feature; tests build with + # cfg(test)), so a cfg gate that hides an item from the release cloud + # build otherwise sails through green checks and breaks docker builds. + - run: cargo check -p compass --features object-storage + # The s3_integration tests skip silently without the env; guard against + # env-name drift turning this job into a green no-op. + - run: | + cargo test -p compass --features object-storage -- --nocapture 2>&1 | tee /tmp/cloud-tests.log + if grep -q '^skipped: COMPASS_TEST_S3_BUCKET' /tmp/cloud-tests.log; then + echo '::error::s3_integration tests were skipped — MinIO env wiring is broken' + exit 1 + fi + + # Developer Certificate of Origin: every PR commit carries a Signed-off-by + # trailer. Dependency-free check over the PR range. + dco: + if: github.event_name == 'pull_request' + runs-on: ubuntu-24.04 + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + - run: | + missing=0 + for sha in $(git rev-list --no-merges ${{ github.event.pull_request.base.sha }}..${{ github.event.pull_request.head.sha }}); do + if ! git log -1 --format=%B "$sha" | grep -q '^Signed-off-by: '; then + echo "::error::commit $sha is missing a Signed-off-by trailer (git commit -s)" + missing=1 + fi + done + exit $missing + msrv: runs-on: ubuntu-24.04 steps: diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 00db435..ff7cb09 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -67,18 +67,22 @@ Per-collection state lives under `$DATA_DIR//`: ``` data// - meta.json CollectionMetadata (name, default vector space, vector_spaces map) - chunks.bin Append-only log of Chunk records - metadata.bin Per-chunk metadata (typed values, bitset-faceted) - fts/ Tantivy directory - vectors// - index.usearch USearch HNSW (CPU) — mmap-backed - index.cuvs cuVS HNSW (GPU build) — when COMPASS_BACKEND=gpu - index.keymap Internal HNSW key -> external chunk id mapping - vectors.bin Raw float buffer (used for brute-force fallback + rebuilds) + collection.json Collection metadata (name, config, vector_spaces map, applied_seq) + chunks.redb Chunk bodies + metadata (redb; disk source of truth) + relations.redb Typed many-to-many chunk relations (redb) relationships.bin Parent-child + sibling edges + tantivy/ Tantivy FTS index directory + vectors/ + .index USearch HNSW graph — mmap-backed + .keymap Internal HNSW key -> external chunk id mapping + .bin CMV2 mmap vector file (torn-append-safe, per-batch durable) ``` +In cloud mode the object-storage bucket additionally holds, per collection: +`collection.json` (bucket config), `manifest` (LSM manifest, CAS-committed), +`wal/{uuid}.frag` (WAL fragments), `segments/{uuid}` (CSEG0002 sectioned +segments), and `id-alloc` (CAS-leased chunk-id blocks). + The disk format is the contract. Bumping it requires a migration path documented in CHANGELOG.md. ## Rebuild flow (model upgrades) diff --git a/CHANGELOG.md b/CHANGELOG.md index 83e591a..0d87856 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,33 @@ All notable changes to this project are documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### Added — "warm serverless" + +- **Stateless writer role** (`COMPASS_ROLE=writer`): durable-append-only nodes with no local indexes and instant boot. Writes validate against the bucket's collection config, mint ids from CAS-leased blocks, append one WAL fragment, and return its `seq`. Reads and delete-by-filter are refused with clear errors. Consistency contract: durable immediately, searchable on serving nodes within the refresh interval. +- **Id-block allocator** (`{ns}/id-alloc`): in cloud mode every ingest path claims id blocks via CAS, so attached nodes and stateless writers can never mint colliding ids. Pre-v0.4 namespaces migrate automatically (seeded from the bucket-derived high-water mark). Do not run v0.3 and v0.4 writers against one bucket during a rolling upgrade. +- **Bucket collection config** (`{ns}/collection.json`): vector-space specs, default space, `created_at`, and `CollectionConfig` are durable in the bucket and survive cold rebuilds (previously specs were re-inferred as `model:"recovered"` and `embed_model` was silently lost). Vector-space CRUD is bucket-first CAS; zero-ingest collections are discoverable from a fresh disk. +- **Manifest refresh + read-your-writes**: serving nodes converge with other nodes' writes via a background refresher (`COMPASS_REFRESH_INTERVAL`, default 5s) using a per-collection seq tracker that never double-applies a node's own fragments. Config changes sync on refresh; a deleted collection detaches; a recreated one re-attaches. Write responses carry `seq`; `SearchRequest.min_seq` refreshes-then-serves with a bounded wait. +- **Lazy attach + LRU detach** (`COMPASS_LAZY_ATTACH`, `COMPASS_MAX_ATTACHED`): boot registers bucket namespaces and attaches on first request (stampede-safe, one rebuild); past the budget the least-recently-used collection detaches and re-attaches on demand — the bucket is the source of truth. Default off; local mode unchanged. +- **CI**: object-storage build + real-S3 integration tests run against MinIO on every PR (with a silent-skip guard); DCO sign-off enforced on PR commits (merge commits exempt). + +### Fixed + +- A completed vector-space rebuild (`POST .../rebuild`) never activated: the space stayed `status="building"` and the rebuilt index was not served until restart. Rebuild completion now flips the persisted status (CAS in cloud mode) and hot-loads the index; activation failure is reported as a failed rebuild. +- Searching or ingesting into a missing collection returned HTTP 500/400; typed not-found errors now map to 404 across all endpoints (replacing three copies of substring-based status sniffing). +- Facet counts were wiped by every ingest after the first (each batch replaced the accumulated facet state; latent since v0.2), came back empty after any restart (nothing rebuilt them from disk), and counted deleted chunks until a full FTS rebuild. Facets are now roaring treemaps keyed by chunk id: batches accumulate, the load/rebuild scan reconstructs them, and counts intersect the live-id universe so tombstoned chunks are excluded. Found by the new live-stack E2E harness (`scripts/e2e.sh`, 44 checks across every endpoint and both node roles). +- Warm restarts of an actively-written collection were O(collection size): batched HNSW persistence legitimately leaves the index file behind the mmap, and the load path treated that as corruption and re-inserted every vector (20.2s vs v0.3.0's 1.1s at 100k chunks in the comparison bench). Load now heals incrementally — append only the missing tail rows from the mmap, save, and serve mmap-backed. Warm restart at 100k: 1.6s. +- Sub-1000-vector collections never persisted the vector keymap, silently relying on identity key→id mapping that returned wrong chunk ids once ids were non-dense (exposed by block allocation; latent since v0.2). The keymap is now saved on every build and synthesized as identity for pre-fix directories. + +### Changed + +- Pork audit (three independent review passes): −1,200 lines of dead weight removed — the unwired VectorIndex/GPU backend plumbing (`COMPASS_BACKEND` did nothing), a third never-called filter evaluator, never-wired filter-index persistence codecs, the legacy vector writer, the `rayon` dependency, and assorted dead fields/params. `delete_by_filter` now resolves ids through the same roaring filter-index pushdown as search (one filter semantics, not three). The `dead_code` lint is enabled again crate-wide. `collections/mod.rs` shrank from 7,100 to 3,700 lines (test modules extracted to files). + +### Scope & limitations (honest) + +- Warm, not cold: attach cost is proportional to collection size until the sectioned segment format + serve-from-storage indexes land (roadmap Phases 5–6). Cross-node convergence is periodic (refresh interval), not synchronous — use `min_seq` when you need read-your-writes. + ## [0.3.0] - 2026-07-03 ### Added diff --git a/CLAUDE.md b/CLAUDE.md index 69c67bc..e52add0 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -26,7 +26,7 @@ docker run -p 4001:4001 -v ./data:/app/data compass crates/ compass/ Main engine binary (Axum API, search, scoring, embed) compass-index-api/ VectorIndex trait (no I/O, no async) - compass-vector-gpu/ Optional cuVS GPU backend (--features gpu, Linux + CUDA) + compass-vector-gpu/ cuVS GPU backend crate (standalone; not yet wired into the engine) ``` ## Architecture @@ -100,7 +100,10 @@ POST /collections/:name/vector-spaces/:space/rebuild Trigger re-embedding GET /collections/:name/vector-spaces/:space/status Rebuild progress PUT /collections/:name/default-vector-space Switch default space +GET /collections/:name/segments/at Temporal segment lookup (TAMS) + GET /health Health check +GET /metrics Prometheus-text metrics ``` ## Embedding Models diff --git a/Cargo.lock b/Cargo.lock index 1a302a3..2a4f3f9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -473,14 +473,11 @@ dependencies = [ "candle-nn", "candle-transformers", "chrono", - "compass-index-api", - "compass-vector-gpu", "futures", "half", "lru", "memmap2", "object_store", - "rayon", "redb", "reqwest", "roaring", @@ -492,6 +489,7 @@ dependencies = [ "thiserror 1.0.69", "tokenizers", "tokio", + "tower", "tower-http", "tracing", "tracing-subscriber", @@ -3562,6 +3560,7 @@ dependencies = [ "pin-project-lite", "sync_wrapper", "tokio", + "tokio-util", "tower-layer", "tower-service", "tracing", diff --git a/Cargo.toml b/Cargo.toml index e9c22f1..3e43044 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,6 +18,7 @@ rust-version = "1.88" [workspace.dependencies] # Async runtime + web framework axum = { version = "0.8", features = ["json"] } +tower = { version = "0.5", features = ["limit"] } tokio = { version = "1", features = ["full"] } tower-http = { version = "0.6", features = ["cors"] } diff --git a/README.md b/README.md index 2a4f8ab..0318bc5 100644 --- a/README.md +++ b/README.md @@ -566,7 +566,10 @@ POST /collections/:name/vector-spaces/:space/rebuild Trigger re-embedding GET /collections/:name/vector-spaces/:space/status Rebuild progress PUT /collections/:name/default-vector-space Switch default space +GET /collections/:name/segments/at Temporal segment lookup (TAMS) + GET /health Health check +GET /metrics Prometheus-text metrics ``` ## Contributing diff --git a/crates/compass/Cargo.toml b/crates/compass/Cargo.toml index 19379ba..088c101 100644 --- a/crates/compass/Cargo.toml +++ b/crates/compass/Cargo.toml @@ -14,9 +14,7 @@ path = "src/main.rs" [features] default = [] -# Opt-in GPU vector backend via the compass-vector-gpu crate. # Requires CUDA 12+ and a Linux host. See ARCHITECTURE.md for build details. -gpu = ["dep:compass-vector-gpu"] # Opt-in object-storage backend (S3/GCS/Azure/MinIO) via the object_store crate. # Off by default — local-first deployments pull in zero extra dependencies. object-storage = ["dep:object_store", "dep:futures"] @@ -24,16 +22,15 @@ object-storage = ["dep:object_store", "dep:futures"] [dependencies] # Internal trait crate — defines VectorIndex, IndexParams, IndexError. # Stable surface that pluggable backends bind to. -compass-index-api = { workspace = true } axum = { workspace = true } +tower = { workspace = true } tokio = { workspace = true } tower-http = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } chrono = { workspace = true } uuid = { workspace = true } -rayon = { workspace = true } half = { workspace = true } tantivy = { workspace = true } usearch = { workspace = true } @@ -58,5 +55,3 @@ reqwest = { workspace = true } tracing = { workspace = true } tracing-subscriber = { workspace = true } -# Optional GPU backend, enabled with --features gpu -compass-vector-gpu = { path = "../compass-vector-gpu", optional = true } diff --git a/crates/compass/src/api/collections.rs b/crates/compass/src/api/collections.rs index f32d23e..6cc6124 100644 --- a/crates/compass/src/api/collections.rs +++ b/crates/compass/src/api/collections.rs @@ -16,7 +16,7 @@ pub async fn create_collection( .manager .create_collection(&req.name, req.vector_spaces, req.embedding_dims, req.config) .await - .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?; + .map_err(|e| crate::api::error_response(e, StatusCode::BAD_REQUEST))?; Ok((StatusCode::CREATED, Json(collection_to_info(&collection)))) } @@ -66,7 +66,7 @@ pub async fn add_vector_space( .manager .add_vector_space(&name, &req.name, req.dims, &req.model) .await - .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?; + .map_err(|e| crate::api::error_response(e, StatusCode::BAD_REQUEST))?; Ok(( StatusCode::CREATED, @@ -114,7 +114,7 @@ pub async fn delete_vector_space( .manager .delete_vector_space(&name, &space) .await - .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?; + .map_err(|e| crate::api::error_response(e, StatusCode::BAD_REQUEST))?; Ok(StatusCode::NO_CONTENT) } @@ -128,7 +128,7 @@ pub async fn set_default_vector_space( .manager .set_default_vector_space(&name, &req.name) .await - .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?; + .map_err(|e| crate::api::error_response(e, StatusCode::BAD_REQUEST))?; Ok(StatusCode::OK) } @@ -162,7 +162,7 @@ pub async fn trigger_rebuild( .manager .get_all_chunk_data(&name) .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + .map_err(|e| crate::api::error_response(e, StatusCode::INTERNAL_SERVER_ERROR))?; let vectors_dir = state.manager.vectors_dir(&name); @@ -177,6 +177,7 @@ pub async fn trigger_rebuild( req.batch_size, state.manager.rebuild_tracker.clone(), name, + state.manager.clone(), ) .await .map_err(|e| (StatusCode::CONFLICT, e))?; diff --git a/crates/compass/src/api/delete.rs b/crates/compass/src/api/delete.rs index 053e15c..cf96d4c 100644 --- a/crates/compass/src/api/delete.rs +++ b/crates/compass/src/api/delete.rs @@ -16,18 +16,7 @@ use axum::Json; use std::sync::Arc; fn map_err(e: Box) -> (StatusCode, String) { - let msg = e.to_string(); - if msg.contains("not found") { - (StatusCode::NOT_FOUND, msg) - } else { - // Log the detail server-side; internal errors (paths, backends, redb - // internals) don't belong in response bodies. - tracing::error!("delete handler error: {msg}"); - ( - StatusCode::INTERNAL_SERVER_ERROR, - "internal error (see server logs)".to_string(), - ) - } + crate::api::error_response(e, StatusCode::INTERNAL_SERVER_ERROR) } /// DELETE /collections/:name/chunks/:id @@ -35,7 +24,7 @@ pub async fn delete_chunk( State(state): State>, Path((name, id)): Path<(String, u64)>, ) -> Result, (StatusCode, String)> { - let deleted = state + let (deleted, seq) = state .manager .delete_chunks(&name, &[id]) .await @@ -47,7 +36,7 @@ pub async fn delete_chunk( format!("chunk {id} not found or already deleted"), )); } - Ok(Json(DeleteResponse { deleted })) + Ok(Json(DeleteResponse { deleted, seq })) } /// POST /collections/:name/compact — fold S3 segments + WAL into one segment, @@ -78,18 +67,24 @@ pub async fn delete_by_query( } let mut deleted = 0usize; + let mut seq: Option = None; if !req.ids.is_empty() { - deleted += state + let (n, s) = state .manager .delete_chunks(&name, &req.ids) .await .map_err(map_err)?; + deleted += n; + seq = s.or(seq); } if !req.filters.is_empty() { // If the filter-delete fails after an ids-delete succeeded, report the // partial progress — deletes already applied are not undone. match state.manager.delete_by_filter(&name, &req.filters).await { - Ok(n) => deleted += n, + Ok((n, s)) => { + deleted += n; + seq = s.or(seq); + } Err(e) => { let (code, msg) = map_err(e); return Err(( @@ -101,5 +96,5 @@ pub async fn delete_by_query( } } } - Ok(Json(DeleteResponse { deleted })) + Ok(Json(DeleteResponse { deleted, seq })) } diff --git a/crates/compass/src/api/ingest.rs b/crates/compass/src/api/ingest.rs index 11f2b03..6c651da 100644 --- a/crates/compass/src/api/ingest.rs +++ b/crates/compass/src/api/ingest.rs @@ -22,11 +22,11 @@ pub async fn ingest_chunks( ) -> Result, (StatusCode, String)> { let start = std::time::Instant::now(); - let (count, id_map) = state + let (count, id_map, seq) = state .manager .ingest(&name, req.chunks, &state.embed_state) .await - .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?; + .map_err(|e| crate::api::error_response(e, StatusCode::BAD_REQUEST))?; let took_ms = start.elapsed().as_millis() as u64; @@ -34,5 +34,6 @@ pub async fn ingest_chunks( indexed: count, id_map, took_ms, + seq, })) } diff --git a/crates/compass/src/api/mod.rs b/crates/compass/src/api/mod.rs index 896ee75..6eb0096 100644 --- a/crates/compass/src/api/mod.rs +++ b/crates/compass/src/api/mod.rs @@ -25,6 +25,23 @@ use axum::{Json, Router}; use std::sync::Arc; /// Shared application state passed to every request handler. +/// Map an engine error to an HTTP response. Typed `NotFound` becomes 404 +/// regardless of the handler's default; a 500 default logs the detail and +/// returns a generic body (backend/path internals don't belong in responses). +pub(crate) fn error_response( + e: Box, + default: StatusCode, +) -> (StatusCode, String) { + if e.downcast_ref::().is_some() { + return (StatusCode::NOT_FOUND, e.to_string()); + } + if default == StatusCode::INTERNAL_SERVER_ERROR { + tracing::error!("handler error: {e}"); + return (default, "internal error (see server logs)".to_string()); + } + (default, e.to_string()) +} + pub struct AppState { pub manager: Arc, pub embed_state: Arc, @@ -148,12 +165,45 @@ pub fn build_router(state: Arc, auth: Arc) -> Router { Router::new() // ── Health (unauthenticated) ───────────────────────────────────── .route("/health", get(health_check)) + .route("/metrics", get(metrics_endpoint)) .merge(protected) // 64 MB body limit. Default 2 MB is too small for batched ingest with embeddings. .layer(axum::extract::DefaultBodyLimit::max(64 * 1024 * 1024)) + // Backpressure: bound in-flight requests instead of queueing without + // limit (COMPASS_MAX_CONCURRENCY; unset = unlimited). The cap must stay + // under tokio's Semaphore::MAX_PERMITS (usize::MAX >> 3) — a larger + // value PANICS at startup. + .layer(tower::limit::GlobalConcurrencyLimitLayer::new( + std::env::var("COMPASS_MAX_CONCURRENCY") + .ok() + .and_then(|v| v.parse().ok()) + .filter(|n: &usize| *n > 0) + .unwrap_or(usize::MAX >> 4) + .min(usize::MAX >> 4), + )) .with_state(state) } +/// GET /metrics — Prometheus text. Unauthenticated (like /health); carries +/// operational counters plus per-collection gauges. +async fn metrics_endpoint(State(state): State>) -> String { + let mut gauges = String::new(); + // Attached collections only: list_collections in lazy mode does one S3 + // GET per registered namespace — an unauthenticated request-amplifier if + // exposed to a scraper. + for c in state.manager.attached_collections().await { + gauges.push_str(&format!( + "compass_collection_chunks{{collection=\"{}\"}} {}\n", + c.name, c.chunk_count + )); + gauges.push_str(&format!( + "compass_collection_applied_seq{{collection=\"{}\"}} {}\n", + c.name, c.applied_seq + )); + } + crate::metrics::render(&gauges) +} + /// GET /health async fn health_check(State(state): State>) -> Json { let collections = state.manager.list_collections().await; diff --git a/crates/compass/src/api/relations.rs b/crates/compass/src/api/relations.rs index 611f27d..2a54ff5 100644 --- a/crates/compass/src/api/relations.rs +++ b/crates/compass/src/api/relations.rs @@ -25,19 +25,10 @@ const MAX_RELATIONS_PER_REQUEST: usize = 10_000; fn map_err(e: Box) -> (StatusCode, String) { let msg = e.to_string(); - if msg.contains("not found") { - (StatusCode::NOT_FOUND, msg) - } else if msg.contains("must differ") { - (StatusCode::BAD_REQUEST, msg) - } else { - // Log the detail server-side; internal errors (paths, backends, redb - // internals) don't belong in response bodies. - tracing::error!("relations handler error: {msg}"); - ( - StatusCode::INTERNAL_SERVER_ERROR, - "internal error (see server logs)".to_string(), - ) + if msg.contains("must differ") { + return (StatusCode::BAD_REQUEST, msg); } + crate::api::error_response(e, StatusCode::INTERNAL_SERVER_ERROR) } /// POST /collections/:name/relations diff --git a/crates/compass/src/api/search.rs b/crates/compass/src/api/search.rs index cb9a068..377c6cd 100644 --- a/crates/compass/src/api/search.rs +++ b/crates/compass/src/api/search.rs @@ -34,7 +34,7 @@ pub async fn search_collection( .manager .search(&name, &req, &state.embed_state) .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + .map_err(|e| crate::api::error_response(e, StatusCode::INTERNAL_SERVER_ERROR))?; let hits: Vec = results .into_iter() @@ -70,7 +70,7 @@ pub async fn get_facets( .manager .get_facets(&name, query_str, &req.fields) .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + .map_err(|e| crate::api::error_response(e, StatusCode::INTERNAL_SERVER_ERROR))?; Ok(Json(FacetResponse { facets, took_us })) } diff --git a/crates/compass/src/api/segments.rs b/crates/compass/src/api/segments.rs index 1abcf96..175da9c 100644 --- a/crates/compass/src/api/segments.rs +++ b/crates/compass/src/api/segments.rs @@ -67,14 +67,7 @@ pub async fn segments_at( params.time_end_ms, ) .await - .map_err(|e| { - let msg = e.to_string(); - if msg.contains("not found") { - (StatusCode::NOT_FOUND, msg) - } else { - (StatusCode::INTERNAL_SERVER_ERROR, msg) - } - })?; + .map_err(|e| crate::api::error_response(e, StatusCode::INTERNAL_SERVER_ERROR))?; let took_ms = t0.elapsed().as_secs_f64() * 1_000.0; Ok(Json(SegmentsAtResponse { results, took_ms })) diff --git a/crates/compass/src/collections/cloud.rs b/crates/compass/src/collections/cloud.rs index 2264e81..b9b0c19 100644 --- a/crates/compass/src/collections/cloud.rs +++ b/crates/compass/src/collections/cloud.rs @@ -18,12 +18,116 @@ //! A segment payload is a versioned JSON object `{chunks, relations}` — the full //! live set at compaction time. -use crate::models::{ChunkRelation, DocumentChunk}; +use crate::models::{ + ChunkRelation, Collection, CollectionConfig, DocumentChunk, VectorSpaceConfig, +}; use crate::storage::lsm::{self, FragmentKind, Manifest}; use crate::storage::{Storage, StorageError}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; +/// Bucket-resident collection config — `{ns}/collection.json` in object +/// storage. The durable source of truth for everything in [`Collection`] +/// EXCEPT the node-local counters (`chunk_count`, `next_id`). Without it, a +/// cold rebuild has to fabricate metadata (inferring vector-space specs from +/// recovered embeddings and silently losing `CollectionConfig.embed_model`). +/// +/// Distinct from the LOCAL file `data/{ns}/collection.json` (node cache); +/// bucket writes are strictly gated on cloud mode so a local-disk Storage +/// backend can never clobber the real local metadata file. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct BucketConfig { + #[serde(default)] + pub version: u8, + pub name: String, + pub created_at: chrono::DateTime, + pub vector_spaces: HashMap, + pub default_vector_space: Option, + pub embedding_dims: usize, + #[serde(default)] + pub config: CollectionConfig, +} + +const BUCKET_CONFIG_VERSION: u8 = 1; + +impl BucketConfig { + pub fn from_collection(c: &Collection) -> Self { + Self { + version: BUCKET_CONFIG_VERSION, + name: c.name.clone(), + created_at: c.created_at, + vector_spaces: c.vector_spaces.clone(), + default_vector_space: c.default_vector_space.clone(), + embedding_dims: c.embedding_dims, + config: c.config.clone(), + } + } +} + +pub fn config_key(ns: &str) -> String { + format!("{ns}/collection.json") +} + +/// Read the bucket config, or None when absent (pre-v0.4 collection). +pub async fn read_bucket_config( + storage: &dyn Storage, + ns: &str, +) -> Result, StorageError> { + match storage.get(&config_key(ns)).await { + Ok(bytes) => Ok(Some(serde_json::from_slice(&bytes).map_err(|e| { + StorageError::Io(format!("bucket config decode for '{ns}': {e}")) + })?)), + Err(StorageError::NotFound(_)) => Ok(None), + Err(e) => Err(e), + } +} + +/// Create-only write of the bucket config. `AlreadyExists` bubbles up so the +/// caller can distinguish "fresh create" from "collection already in bucket". +pub async fn write_bucket_config_if_absent( + storage: &dyn Storage, + ns: &str, + cfg: &BucketConfig, +) -> Result<(), StorageError> { + let bytes = serde_json::to_vec(cfg) + .map_err(|e| StorageError::Io(format!("bucket config encode: {e}")))?; + storage + .put_if_not_exists(&config_key(ns), bytes::Bytes::from(bytes)) + .await + .map(|_| ()) +} + +/// CAS read-modify-write on the bucket config. `mutate` sees the LATEST doc +/// each attempt and may fail validation (e.g. "space already exists") — that +/// error aborts the loop. Retries only on version conflicts. +pub async fn cas_update_bucket_config( + storage: &dyn Storage, + ns: &str, + mut mutate: F, +) -> Result> +where + F: FnMut(&mut BucketConfig) -> Result<(), Box>, +{ + const MAX_RETRIES: u32 = 10; + for _ in 0..MAX_RETRIES { + let (bytes, version) = storage.get_versioned(&config_key(ns)).await?; + let mut cfg: BucketConfig = serde_json::from_slice(&bytes) + .map_err(|e| StorageError::Io(format!("bucket config decode for '{ns}': {e}")))?; + mutate(&mut cfg)?; + let encoded = serde_json::to_vec(&cfg) + .map_err(|e| StorageError::Io(format!("bucket config encode: {e}")))?; + match storage + .put_if_match(&config_key(ns), bytes::Bytes::from(encoded), &version) + .await + { + Ok(_) => return Ok(cfg), + Err(StorageError::VersionConflict { .. }) => continue, + Err(e) => return Err(e.into()), + } + } + Err("bucket config CAS failed after max retries (persistent contention)".into()) +} + /// The live materialized state of a collection reconstructed from object storage. pub struct Materialized { /// Live chunks, keyed by id (deletes already applied). @@ -49,30 +153,179 @@ pub struct Segment { /// only and REUSE deleted ids, breaking the monotonic-id invariant. #[serde(default)] pub max_id: u64, + /// Chunk ids deleted in the folded range that may still exist in OLDER + /// segments (partitioned compaction folds only the WAL tail, so deletes + /// must carry across segment boundaries until a full merge drops them). + #[serde(default)] + pub tombstones: Vec, + /// Relation ids deleted in the folded range (same cross-segment rule). + #[serde(default)] + pub relation_tombstones: Vec, +} + +/// v2 binary segment magic. v1 segments are JSON (decoded via fallback). +const SEG_MAGIC_V2: [u8; 8] = *b"CSEG0002"; + +/// Encode a segment in the v2 sectioned binary layout: +/// `[magic][u64 max_id][u32 toc_len][toc JSON][sections...]` +/// Sections: `meta` (JSON chunks with embeddings STRIPPED), `emb:` +/// (`[u32 dims][u64 n][n × (u64 id + dims×f32 LE)]`), `rels` (JSON), +/// `tombs` (u64 LE array), `rtombs` (JSON ids). Embeddings dominate segment +/// size; storing them as raw f32 instead of JSON decimals is ~10× smaller and +/// range-readable by section. +pub fn encode_segment_v2(seg: &Segment) -> Result, StorageError> { + let err = |e: String| StorageError::Io(format!("segment v2 encode: {e}")); + let mut sections: Vec<(String, Vec)> = Vec::new(); + + let mut meta_chunks: Vec = Vec::with_capacity(seg.chunks.len()); + let mut by_space: std::collections::BTreeMap)>> = + std::collections::BTreeMap::new(); + for c in &seg.chunks { + let mut m = c.clone(); + for (space, emb) in std::mem::take(&mut m.embeddings) { + by_space.entry(space).or_default().push((c.id, emb)); + } + meta_chunks.push(m); + } + sections.push(( + "meta".into(), + serde_json::to_vec(&meta_chunks).map_err(|e| err(e.to_string()))?, + )); + for (space, rows) in by_space { + let dims = rows.first().map(|(_, v)| v.len()).unwrap_or(0) as u32; + let mut buf = Vec::with_capacity(12 + rows.len() * (8 + dims as usize * 4)); + buf.extend_from_slice(&dims.to_le_bytes()); + buf.extend_from_slice(&(rows.len() as u64).to_le_bytes()); + for (id, v) in &rows { + if v.len() as u32 != dims { + return Err(err(format!("ragged dims in space '{space}'"))); + } + buf.extend_from_slice(&id.to_le_bytes()); + for x in v { + buf.extend_from_slice(&x.to_le_bytes()); + } + } + sections.push((format!("emb:{space}"), buf)); + } + sections.push(( + "rels".into(), + serde_json::to_vec(&seg.relations).map_err(|e| err(e.to_string()))?, + )); + let mut tombs = Vec::with_capacity(seg.tombstones.len() * 8); + for id in &seg.tombstones { + tombs.extend_from_slice(&id.to_le_bytes()); + } + sections.push(("tombs".into(), tombs)); + sections.push(( + "rtombs".into(), + serde_json::to_vec(&seg.relation_tombstones).map_err(|e| err(e.to_string()))?, + )); + + let toc: Vec<(String, u64)> = sections + .iter() + .map(|(n, b)| (n.clone(), b.len() as u64)) + .collect(); + let toc_bytes = serde_json::to_vec(&toc).map_err(|e| err(e.to_string()))?; + let mut out = Vec::new(); + out.extend_from_slice(&SEG_MAGIC_V2); + out.extend_from_slice(&seg.max_id.to_le_bytes()); + out.extend_from_slice(&(toc_bytes.len() as u32).to_le_bytes()); + out.extend_from_slice(&toc_bytes); + for (_, b) in sections { + out.extend_from_slice(&b); + } + Ok(out) } -const SEGMENT_VERSION: u8 = 1; +fn decode_segment_v2(bytes: &[u8]) -> Result { + let err = |e: String| StorageError::Io(format!("segment v2 decode: {e}")); + let need = |n: usize, have: usize| -> Result<(), StorageError> { + if have < n { + Err(err("truncated".into())) + } else { + Ok(()) + } + }; + need(20, bytes.len())?; + let max_id = u64::from_le_bytes(bytes[8..16].try_into().unwrap()); + let toc_len = u32::from_le_bytes(bytes[16..20].try_into().unwrap()) as usize; + need(20 + toc_len, bytes.len())?; + let toc: Vec<(String, u64)> = + serde_json::from_slice(&bytes[20..20 + toc_len]).map_err(|e| err(e.to_string()))?; + let mut pos = 20 + toc_len; + let mut seg = Segment { + version: 2, + max_id, + ..Default::default() + }; + let mut embs: HashMap>> = HashMap::new(); + for (name, len) in toc { + let len = len as usize; + need(pos + len, bytes.len())?; + let body = &bytes[pos..pos + len]; + pos += len; + if name == "meta" { + seg.chunks = serde_json::from_slice(body).map_err(|e| err(e.to_string()))?; + } else if let Some(space) = name.strip_prefix("emb:") { + need(12, body.len())?; + let dims = u32::from_le_bytes(body[0..4].try_into().unwrap()) as usize; + let n = u64::from_le_bytes(body[4..12].try_into().unwrap()) as usize; + let row = 8 + dims * 4; + need(12 + n * row, body.len())?; + for i in 0..n { + let off = 12 + i * row; + let id = u64::from_le_bytes(body[off..off + 8].try_into().unwrap()); + let mut v = Vec::with_capacity(dims); + for d in 0..dims { + let o = off + 8 + d * 4; + v.push(f32::from_le_bytes(body[o..o + 4].try_into().unwrap())); + } + embs.entry(id).or_default().insert(space.to_string(), v); + } + } else if name == "rels" { + seg.relations = serde_json::from_slice(body).map_err(|e| err(e.to_string()))?; + } else if name == "tombs" { + seg.tombstones = body + .chunks_exact(8) + .map(|c| u64::from_le_bytes(c.try_into().unwrap())) + .collect(); + } else if name == "rtombs" { + seg.relation_tombstones = + serde_json::from_slice(body).map_err(|e| err(e.to_string()))?; + } + // Unknown sections are skipped (forward compat). + } + for c in &mut seg.chunks { + if let Some(e) = embs.remove(&c.id) { + c.embeddings = e; + } + } + Ok(seg) +} -/// Serialize a live set as a segment payload. `max_id` must be the id -/// high-water mark INCLUDING tombstoned ids (pass `Materialized::max_id`, not -/// the max of the live set). +/// Serialize a live set as a segment payload (v2 binary). `max_id` must be +/// the id high-water mark INCLUDING tombstoned ids. pub fn encode_segment( chunks: &[DocumentChunk], relations: &[ChunkRelation], max_id: u64, ) -> Result, StorageError> { - let seg = Segment { - version: SEGMENT_VERSION, + encode_segment_v2(&Segment { + version: 2, chunks: chunks.to_vec(), relations: relations.to_vec(), max_id, - }; - serde_json::to_vec(&seg).map_err(|e| StorageError::Io(format!("segment encode: {e}"))) + tombstones: Vec::new(), + relation_tombstones: Vec::new(), + }) } fn decode_segment(bytes: &[u8]) -> Result { - // Back-compat: an older segment was a bare JSON array of chunks. Try the - // versioned object first, then fall back to a plain chunk array. + // v2 binary (magic-tagged) first; then v1 JSON object; then the oldest + // bare-JSON-array form. + if bytes.len() >= 8 && bytes[0..8] == SEG_MAGIC_V2 { + return decode_segment_v2(bytes); + } if let Ok(seg) = serde_json::from_slice::(bytes) { return Ok(seg); } @@ -81,8 +334,7 @@ fn decode_segment(bytes: &[u8]) -> Result { Ok(Segment { version: 0, chunks, - relations: Vec::new(), - max_id: 0, + ..Default::default() }) } @@ -103,6 +355,57 @@ fn decode_relation_ids(bytes: &[u8]) -> Result, StorageError> { .map_err(|e| StorageError::Io(format!("relation-delete decode: {e}"))) } +/// Fold ONLY a WAL tail (uncompacted fragments, in seq order) into a Segment +/// — the bounded-work unit of partitioned compaction. Deletes that don't hit +/// a chunk/relation within the tail are carried as segment tombstones so they +/// still apply to OLDER segments at materialize time. +pub fn fold_tail(frags: &[(lsm::FragmentRef, bytes::Bytes)]) -> Result { + let mut chunks: HashMap = HashMap::new(); + let mut relations: HashMap = HashMap::new(); + let mut tombs: std::collections::BTreeSet = Default::default(); + let mut rtombs: std::collections::BTreeSet = Default::default(); + let mut max_id = 0u64; + for (fref, bytes) in frags { + match fref.kind { + FragmentKind::Data => { + for chunk in decode_chunks(bytes)? { + max_id = max_id.max(chunk.id); + tombs.remove(&chunk.id); // re-created after an earlier delete + chunks.insert(chunk.id, chunk); + } + } + FragmentKind::Tombstone => { + for id in decode_ids(bytes)? { + max_id = max_id.max(id); + chunks.remove(&id); + relations.retain(|_, r| r.source_chunk_id != id && r.target_chunk_id != id); + tombs.insert(id); // must ALSO apply to older segments + } + } + FragmentKind::RelationUpsert => { + for rel in decode_relations(bytes)? { + rtombs.remove(&rel.relation_id); + relations.insert(rel.relation_id.clone(), rel); + } + } + FragmentKind::RelationDelete => { + for rid in decode_relation_ids(bytes)? { + relations.remove(&rid); + rtombs.insert(rid); + } + } + } + } + Ok(Segment { + version: 2, + chunks: chunks.into_values().collect(), + relations: relations.into_values().collect(), + max_id, + tombstones: tombs.into_iter().collect(), + relation_tombstones: rtombs.into_iter().collect(), + }) +} + /// Materialize the full live state (chunks + relations) from a manifest: read /// all segments, then replay uncompacted fragments in seq order (latest-wins, /// deletes applied). A chunk delete (tombstone) also drops any relation incident @@ -123,6 +426,17 @@ pub async fn materialize( // The stored high-water mark covers tombstoned ids that compaction // physically dropped — required so next_id never regresses/reuses. max_id = max_id.max(segment.max_id); + // Cross-segment deletes first: a tail-fold segment's tombstones apply + // to everything OLDER than it (already accumulated), never to its own + // surviving chunks (compaction removed those before encoding). + for id in &segment.tombstones { + max_id = max_id.max(*id); + chunks.remove(id); + relations.retain(|_, r| r.source_chunk_id != *id && r.target_chunk_id != *id); + } + for rid in &segment.relation_tombstones { + relations.remove(rid); + } for chunk in segment.chunks { max_id = max_id.max(chunk.id); chunks.insert(chunk.id, chunk); @@ -344,4 +658,94 @@ mod tests { r.chunks.keys().collect::>() ); } + + // ── Segment v2 / partitioned compaction ─────────────────────────────── + + #[test] + fn segment_v2_roundtrip_with_embeddings_and_tombstones() { + let mut c1 = chunk(1, "one"); + c1.embeddings + .insert("default".into(), vec![0.1, 0.2, 0.3, 0.4]); + let mut c2 = chunk(2, "two"); + c2.embeddings + .insert("default".into(), vec![0.5, 0.6, 0.7, 0.8]); + c2.embeddings.insert("wide".into(), vec![1.0; 8]); + let seg = Segment { + version: 2, + chunks: vec![c1, c2], + relations: vec![relation("r1", 1, 2)], + max_id: 42, + tombstones: vec![7, 9], + relation_tombstones: vec!["dead".into()], + }; + let bytes = encode_segment_v2(&seg).unwrap(); + assert_eq!(&bytes[0..8], b"CSEG0002"); + let back = decode_segment(&bytes).unwrap(); + assert_eq!(back.max_id, 42); + assert_eq!(back.tombstones, vec![7, 9]); + assert_eq!(back.relation_tombstones, vec!["dead".to_string()]); + assert_eq!(back.chunks.len(), 2); + let c2b = back.chunks.iter().find(|c| c.id == 2).unwrap(); + assert_eq!(c2b.embeddings["default"], vec![0.5, 0.6, 0.7, 0.8]); + assert_eq!(c2b.embeddings["wide"].len(), 8); + assert_eq!(back.relations.len(), 1); + } + + // A delete folded into a NEWER tail segment must erase a chunk living in + // an OLDER segment at materialize time (cross-segment tombstones). + #[tokio::test] + async fn tail_segment_tombstones_apply_to_older_segments() { + let s = store("xseg"); + // Older state via a REAL fold: data + relation fragments -> segment A. + let mut c1 = chunk(1, "old"); + c1.embeddings + .insert("default".into(), vec![0.1, 0.2, 0.3, 0.4]); + let data = serde_json::to_vec(&vec![c1, chunk(2, "keep")]).unwrap(); + lsm::append_fragment(s.as_ref(), "ns", Bytes::from(data), 2) + .await + .unwrap(); + let rels = serde_json::to_vec(&vec![relation("r1", 1, 2)]).unwrap(); + lsm::append_relation_upsert(s.as_ref(), "ns", Bytes::from(rels), 1) + .await + .unwrap(); + let fold_once = |sref: Arc| async move { + let (m1, v1) = lsm::read_manifest(sref.as_ref(), "ns").await.unwrap(); + let frags = lsm::read_uncompacted_fragments(sref.as_ref(), "ns", &m1) + .await + .unwrap(); + let tail = fold_tail(&frags).unwrap(); + let folded_through = m1.uncompacted().map(|f| f.seq).max().unwrap(); + let records = tail.chunks.len() as u64; + lsm::append_segment( + sref.as_ref(), + "ns", + &v1, + &m1, + Bytes::from(encode_segment_v2(&tail).unwrap()), + records, + folded_through, + ) + .await + .unwrap(); + tail + }; + let seg_a = fold_once(s.clone()).await; + assert!(seg_a.tombstones.is_empty()); + + // Newer tail: delete chunk 1; the delete finds nothing IN the tail so + // it must be carried as a cross-segment tombstone. + lsm::append_tombstone(s.as_ref(), "ns", &[1]).await.unwrap(); + let seg_b = fold_once(s.clone()).await; + assert_eq!( + seg_b.tombstones, + vec![1], + "unmatched delete carried forward" + ); + + let r = mat(s.as_ref(), "ns").await; + assert!(!r.chunks.contains_key(&1), "older-segment chunk deleted"); + assert!(r.chunks.contains_key(&2)); + assert!(r.relations.is_empty(), "incident relation pruned"); + assert_eq!(r.max_id, 2); + } } diff --git a/crates/compass/src/collections/cloud_ingest_tests.rs b/crates/compass/src/collections/cloud_ingest_tests.rs new file mode 100644 index 0000000..d77c59b --- /dev/null +++ b/crates/compass/src/collections/cloud_ingest_tests.rs @@ -0,0 +1,2209 @@ +// collections/cloud_ingest_tests.rs — extracted test module (was inline in mod.rs). +// Child module of `collections`: `super::*` sees the parent's private items. + +//! Verifies that in object-storage (cloud) mode, ingest mirrors the batch +//! into the LSM as a WAL fragment + CAS-committed manifest — the S3-native +//! path. Uses the in-memory object_store backend, which +//! exercises the identical `Storage`/`ObjectStoreBackend` code an S3 bucket +//! would, without needing real credentials. + +use super::*; +use crate::embed::EmbedState; +use crate::storage::object_store_backend::ObjectStoreBackend; +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; + +fn unique_data_dir() -> std::path::PathBuf { + static N: AtomicU64 = AtomicU64::new(0); + std::env::temp_dir().join(format!( + "compass-cloud-ingest-{}-{}", + std::process::id(), + N.fetch_add(1, Ordering::Relaxed) + )) +} + +fn embed_state() -> EmbedState { + // No models needed: chunks carry precomputed embeddings. + EmbedState { + bge: None, + distilled: None, + } +} + +fn ingest_chunk(idx: u32) -> IngestChunk { + let mut metadata = HashMap::new(); + metadata.insert( + "org_id".to_string(), + MetadataValue::String("acme".to_string()), + ); + let mut embeddings = HashMap::new(); + embeddings.insert("default".to_string(), vec![0.1, 0.2, 0.3, 0.4]); + IngestChunk { + client_id: None, + file_id: format!("f{idx}"), + chunk_index: 0, + page: None, + text: format!("chunk-{idx}"), + metadata, + doc_type: "chunk".to_string(), + parent_id: None, + parent_ref: None, + group_id: None, + embeddings, + embedding: None, + } +} + +#[tokio::test] +async fn ingest_writes_wal_fragment_and_manifest_to_object_storage() { + let data_dir = unique_data_dir(); + std::fs::create_dir_all(&data_dir).unwrap(); + let embed = embed_state(); + + // In-memory object storage backend (same code path as s3://). + let storage: Arc = Arc::new(ObjectStoreBackend::from_store( + std::sync::Arc::new(object_store::memory::InMemory::new()), + "object-store:memory", + )); + let manager = CollectionManager::new_with_storage(&data_dir, storage.clone()) + .await + .unwrap(); + manager + .create_collection("cloudcoll", None, Some(4), None) + .await + .unwrap(); + + // Ingest two batches. + manager + .ingest("cloudcoll", vec![ingest_chunk(0), ingest_chunk(1)], &embed) + .await + .unwrap(); + manager + .ingest("cloudcoll", vec![ingest_chunk(2)], &embed) + .await + .unwrap(); + + // The manifest exists and records two WAL fragments. + let (manifest, version) = crate::storage::lsm::read_manifest(storage.as_ref(), "cloudcoll") + .await + .unwrap(); + assert!(version.is_some(), "manifest must exist in object storage"); + assert_eq!(manifest.fragments.len(), 2, "one fragment per ingest batch"); + assert_eq!(manifest.next_seq, 2); + + // The WAL fragment objects exist and decode back to the ingested chunks. + let frags = + crate::storage::lsm::read_uncompacted_fragments(storage.as_ref(), "cloudcoll", &manifest) + .await + .unwrap(); + assert_eq!(frags.len(), 2); + + let batch0: Vec = serde_json::from_slice(&frags[0].1).unwrap(); + assert_eq!(batch0.len(), 2); + assert_eq!(batch0[0].text, "chunk-0"); + let batch1: Vec = serde_json::from_slice(&frags[1].1).unwrap(); + assert_eq!(batch1.len(), 1); + assert_eq!(batch1[0].text, "chunk-2"); + + // Total records across fragments == total chunks ingested. + let total: u64 = manifest.fragments.iter().map(|f| f.records).sum(); + assert_eq!(total, 3); + + let _ = std::fs::remove_dir_all(&data_dir); +} + +#[tokio::test] +async fn local_mode_writes_no_wal() { + // Sanity: a local-disk manager must NOT create any WAL/manifest objects. + let data_dir = unique_data_dir(); + std::fs::create_dir_all(&data_dir).unwrap(); + let embed = embed_state(); + let manager = CollectionManager::new(&data_dir).await.unwrap(); + manager + .create_collection("localcoll", None, Some(4), None) + .await + .unwrap(); + manager + .ingest("localcoll", vec![ingest_chunk(0)], &embed) + .await + .unwrap(); + + // No manifest object should exist under the collection prefix. + let manifest_path = data_dir.join("localcoll").join("manifest"); + assert!( + !manifest_path.exists(), + "local mode must not write an LSM manifest" + ); + // Nor an id-block allocator: local mode allocates from next_id. + assert!( + !data_dir.join("localcoll").join("id-alloc").exists(), + "local mode must not seed the id-block allocator" + ); + // And ids stay dense from 0 (block allocation would start at 0 too, + // but a second ingest would jump; assert both batches are contiguous). + manager + .ingest("localcoll", vec![ingest_chunk(1)], &embed) + .await + .unwrap(); + let (_, mut ids) = manager.get_all_chunk_data("localcoll").await.unwrap(); + ids.sort_unstable(); + assert_eq!(ids, vec![0, 1], "local ids must be dense next_id values"); + let _ = std::fs::remove_dir_all(&data_dir); +} + +// A stray COMPASS_ROLE=writer on a local-disk deployment must be +// neutralized: cloud_mode is false, so the constructor forces Full and +// the node keeps serving reads and creating collections normally. +#[tokio::test] +async fn writer_role_is_neutralized_in_local_mode() { + let data_dir = unique_data_dir(); + std::fs::create_dir_all(&data_dir).unwrap(); + let embed = embed_state(); + let storage: Arc = + Arc::new(crate::storage::local::LocalDiskStorage::new(&data_dir).unwrap()); + let manager = CollectionManager::new_with_storage_opts( + &data_dir, + storage, + NodeRole::Writer, + false, + usize::MAX, + 0, + ) + .await + .unwrap(); + manager + .create_collection("localwriter", None, Some(4), None) + .await + .expect("local node must create collections despite COMPASS_ROLE=writer"); + manager + .ingest("localwriter", vec![ingest_chunk(0)], &embed) + .await + .unwrap(); + let (_, ids) = manager + .get_all_chunk_data("localwriter") + .await + .expect("local node must serve reads despite COMPASS_ROLE=writer"); + assert_eq!(ids.len(), 1); + let _ = std::fs::remove_dir_all(&data_dir); +} + +#[tokio::test] +async fn delete_writes_tombstone_wal_fragment() { + use crate::storage::lsm::{read_manifest, read_uncompacted_fragments, FragmentKind}; + + let data_dir = unique_data_dir(); + std::fs::create_dir_all(&data_dir).unwrap(); + let embed = embed_state(); + let storage: Arc = Arc::new(ObjectStoreBackend::from_store( + std::sync::Arc::new(object_store::memory::InMemory::new()), + "object-store:memory", + )); + let manager = CollectionManager::new_with_storage(&data_dir, storage.clone()) + .await + .unwrap(); + manager + .create_collection("delcloud", None, Some(4), None) + .await + .unwrap(); + manager + .ingest( + "delcloud", + vec![ingest_chunk(0), ingest_chunk(1), ingest_chunk(2)], + &embed, + ) + .await + .unwrap(); + + // Delete chunk 1 -> a tombstone WAL fragment lands in object storage. + let (n, _) = manager.delete_chunks("delcloud", &[1]).await.unwrap(); + assert_eq!(n, 1); + + let (manifest, _) = read_manifest(storage.as_ref(), "delcloud").await.unwrap(); + // seq 0 = data fragment (the ingest), seq 1 = tombstone fragment. + assert_eq!(manifest.fragments.len(), 2); + assert_eq!(manifest.fragments[0].kind, FragmentKind::Data); + assert_eq!(manifest.fragments[1].kind, FragmentKind::Tombstone); + + // The tombstone fragment decodes to the deleted id [1]. + let frags = read_uncompacted_fragments(storage.as_ref(), "delcloud", &manifest) + .await + .unwrap(); + let deleted_ids: Vec = serde_json::from_slice(&frags[1].1).unwrap(); + assert_eq!(deleted_ids, vec![1]); + + let _ = std::fs::remove_dir_all(&data_dir); +} + +// THE structural fix: a cloud collection must survive a restart on a FRESH +// local disk by rebuilding from S3. Ingest, delete one, then drop the manager +// AND wipe the local data dir, then reload from the SAME object store — the +// data (minus the deleted chunk) must come back. +#[tokio::test] +async fn cloud_restart_rehydrates_from_object_storage() { + let embed = embed_state(); + // Shared object store persists across the "restart". + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + + let data_dir_a = unique_data_dir(); + std::fs::create_dir_all(&data_dir_a).unwrap(); + { + let storage: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m = CollectionManager::new_with_storage(&data_dir_a, storage) + .await + .unwrap(); + m.create_collection("survive", None, Some(4), None) + .await + .unwrap(); + m.ingest( + "survive", + vec![ingest_chunk(0), ingest_chunk(1), ingest_chunk(2)], + &embed, + ) + .await + .unwrap(); + m.delete_chunks("survive", &[1]).await.unwrap(); + } + // Simulate node loss: wipe the local disk entirely. + std::fs::remove_dir_all(&data_dir_a).unwrap(); + + // Restart on a BRAND-NEW empty local dir, same object store. + let data_dir_b = unique_data_dir(); + std::fs::create_dir_all(&data_dir_b).unwrap(); + let storage: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m2 = CollectionManager::new_with_storage(&data_dir_b, storage) + .await + .unwrap(); + + // The collection is back, recovered from S3. + let info = m2.get_collection("survive").await; + assert!(info.is_some(), "collection must be recovered from S3"); + + // Search finds the surviving chunks (0 and 2), not the deleted one (1). + let req = SearchRequest { + query: "chunk".to_string(), + mode: "semantic".to_string(), + vector_space: None, + top_k: 10, + query_vector: Some(vec![0.1, 0.2, 0.3, 0.4]), + filters: HashMap::new(), + score_weights: None, + recency: None, + recency_preset: None, + recency_field: None, + boosts: Vec::new(), + relationship_boost: None, + explain: false, + include_relations: false, + relation_types: None, + relation_direction: RelationDirection::Outgoing, + min_seq: None, + }; + let (hits, _, _, _) = m2.search("survive", &req, &embed).await.unwrap(); + let ids: std::collections::HashSet = hits.iter().map(|(c, _, _, _, _)| c.id).collect(); + assert!(ids.contains(&0), "chunk 0 recovered"); + assert!(ids.contains(&2), "chunk 2 recovered"); + assert!(!ids.contains(&1), "deleted chunk 1 must NOT reappear"); + + let _ = std::fs::remove_dir_all(&data_dir_b); +} + +// Compaction folds segments+fragments into one segment, dropping tombstoned +// records so they can never resurrect. +#[tokio::test] +async fn compaction_reclaims_tombstoned_data() { + use crate::storage::lsm::read_manifest; + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let data_dir = unique_data_dir(); + std::fs::create_dir_all(&data_dir).unwrap(); + let storage: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m = CollectionManager::new_with_storage(&data_dir, storage.clone()) + .await + .unwrap(); + m.create_collection("comp", None, Some(4), None) + .await + .unwrap(); + m.ingest( + "comp", + vec![ingest_chunk(0), ingest_chunk(1), ingest_chunk(2)], + &embed, + ) + .await + .unwrap(); + m.delete_chunks("comp", &[1]).await.unwrap(); + + // Before: manifest has data + tombstone fragments, no segment. + let (before, _) = read_manifest(storage.as_ref(), "comp").await.unwrap(); + assert!(before.segments.is_empty()); + assert_eq!(before.fragments.len(), 2); + + // Compact. + let live = m.compact_collection("comp").await.unwrap(); + assert_eq!(live, 2, "2 live records (0 and 2) after dropping deleted 1"); + + // After: the WAL tail folded into an appended segment, no live fragments. + let (after, _) = read_manifest(storage.as_ref(), "comp").await.unwrap(); + assert_eq!(after.segments.len(), 1); + assert!(after.uncompacted().count() == 0); + + // Durable truth via materialize (exercises the v2 binary codec): + // live chunks 0 and 2 survive, deleted 1 is gone. + let mat = cloud::materialize(storage.as_ref(), "comp", &after) + .await + .unwrap(); + let ids: std::collections::HashSet = mat.chunks.keys().copied().collect(); + assert!(ids.contains(&0) && ids.contains(&2)); + assert!( + !ids.contains(&1), + "compaction must drop the tombstoned chunk" + ); + + let _ = std::fs::remove_dir_all(&data_dir); +} + +// #1 regression: typed RELATIONS must survive a cold restart from S3 (the bug +// where relation_store was local-redb-only and vanished on rebuild). Create +// relations, wipe the local disk, restart on a fresh dir, relations return. +#[tokio::test] +async fn cloud_restart_recovers_relations() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + + let data_dir_a = unique_data_dir(); + std::fs::create_dir_all(&data_dir_a).unwrap(); + { + let storage: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m = CollectionManager::new_with_storage(&data_dir_a, storage) + .await + .unwrap(); + m.create_collection("relsurv", None, Some(4), None) + .await + .unwrap(); + m.ingest( + "relsurv", + vec![ingest_chunk(0), ingest_chunk(1), ingest_chunk(2)], + &embed, + ) + .await + .unwrap(); + // Create two relations, then delete one — only the survivor should + // come back. + let created = m + .create_relations( + "relsurv", + vec![ + CreateRelation { + source_chunk_id: 0, + target_chunk_id: 1, + target_document_id: None, + relation_type: "cites".into(), + metadata: HashMap::new(), + }, + CreateRelation { + source_chunk_id: 0, + target_chunk_id: 2, + target_document_id: None, + relation_type: "supersedes".into(), + metadata: HashMap::new(), + }, + ], + ) + .await + .unwrap(); + m.delete_relation("relsurv", &created[1].relation_id) + .await + .unwrap(); + } + // Node loss: wipe local disk. + std::fs::remove_dir_all(&data_dir_a).unwrap(); + + // Restart on a fresh local dir, same object store. + let data_dir_b = unique_data_dir(); + std::fs::create_dir_all(&data_dir_b).unwrap(); + let storage: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m2 = CollectionManager::new_with_storage(&data_dir_b, storage) + .await + .unwrap(); + + // The surviving relation (0 --cites--> 1) must be recovered from S3; + // the deleted one (0 --supersedes--> 2) must NOT reappear. + let out = m2 + .get_chunk_relations("relsurv", 0, RelationDirection::Outgoing, None) + .await + .unwrap(); + assert_eq!(out.len(), 1, "exactly one relation should survive restart"); + assert_eq!(out[0].relation_type, "cites"); + assert_eq!(out[0].target_chunk_id, 1); + + let _ = std::fs::remove_dir_all(&data_dir_b); +} + +// #3: auto-compaction. Ingest enough batches to cross the fragment threshold; +// the background trigger should fold them into a segment. We poll briefly for +// the detached task to run, then assert the WAL is bounded. +#[tokio::test] +async fn auto_compaction_bounds_the_wal() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let data_dir = unique_data_dir(); + std::fs::create_dir_all(&data_dir).unwrap(); + let storage: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m = CollectionManager::new_with_storage(&data_dir, storage.clone()) + .await + .unwrap(); + m.create_collection("auto", None, Some(4), None) + .await + .unwrap(); + + // One chunk per ingest = one fragment per ingest. Cross the threshold. + let batches = AUTO_COMPACT_FRAGMENT_THRESHOLD + 2; + for i in 0..batches { + m.ingest("auto", vec![ingest_chunk(i as u32)], &embed) + .await + .unwrap(); + } + + // Poll up to ~3s for the detached auto-compaction to land a segment and + // shrink the uncompacted fragment set. + let mut compacted = false; + for _ in 0..30 { + let (man, _) = crate::storage::lsm::read_manifest(storage.as_ref(), "auto") + .await + .unwrap(); + if !man.segments.is_empty() && man.uncompacted().count() < AUTO_COMPACT_FRAGMENT_THRESHOLD { + compacted = true; + break; + } + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + assert!( + compacted, + "auto-compaction should have folded the WAL into a segment" + ); + + // All data still present after auto-compaction (via materialize). + let (man, _) = crate::storage::lsm::read_manifest(storage.as_ref(), "auto") + .await + .unwrap(); + let mat = cloud::materialize(storage.as_ref(), "auto", &man) + .await + .unwrap(); + assert_eq!(mat.chunks.len(), batches, "no data lost in auto-compaction"); + + let _ = std::fs::remove_dir_all(&data_dir); +} + +// Negative: auto-compaction must NOT fire below the fragment threshold (a +// regression dropping the threshold to ~0 would compact on every ingest). +#[tokio::test] +async fn auto_compaction_does_not_fire_below_threshold() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let data_dir = unique_data_dir(); + std::fs::create_dir_all(&data_dir).unwrap(); + let storage: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m = CollectionManager::new_with_storage(&data_dir, storage.clone()) + .await + .unwrap(); + m.create_collection("below", None, Some(4), None) + .await + .unwrap(); + + // Well under the threshold: a handful of single-chunk ingests. + for i in 0..5u32 { + m.ingest("below", vec![ingest_chunk(i)], &embed) + .await + .unwrap(); + } + // Give any (wrongly) spawned compaction ample time to land a segment. + tokio::time::sleep(std::time::Duration::from_millis(500)).await; + + let (man, _) = crate::storage::lsm::read_manifest(storage.as_ref(), "below") + .await + .unwrap(); + assert!( + man.segments.is_empty(), + "auto-compaction must not fire below the threshold" + ); + assert_eq!(man.fragments.len(), 5, "all fragments still in the WAL"); + + let _ = std::fs::remove_dir_all(&data_dir); +} + +// F1 regression: compaction must physically GC old objects (deferred one +// cycle), not leak them forever. Ingest, compact twice, assert the first +// segment's object is deleted and the object count stays bounded. +#[tokio::test] +async fn compaction_gcs_old_objects() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let data_dir = unique_data_dir(); + std::fs::create_dir_all(&data_dir).unwrap(); + let storage: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m = CollectionManager::new_with_storage(&data_dir, storage.clone()) + .await + .unwrap(); + m.create_collection("gc", None, Some(4), None) + .await + .unwrap(); + m.ingest("gc", vec![ingest_chunk(0), ingest_chunk(1)], &embed) + .await + .unwrap(); + + // First compaction → segment S1, stages the 1 fragment for next-cycle GC. + m.compact_collection("gc").await.unwrap(); + let (man1, _) = crate::storage::lsm::read_manifest(storage.as_ref(), "gc") + .await + .unwrap(); + let seg1_id = man1.segments[0].id.clone(); + // The old WAL fragment object is staged (still present this cycle). + assert_eq!(man1.pending_deletes.len(), 1); + + // Drive enough tail-fold cycles to cross the merge threshold (8 + // segments) so a full merge runs; the merge (plus deferred GC) must + // physically delete S1 — the key point is it's GC'd, not leaked. + for i in 2..14u32 { + m.ingest("gc", vec![ingest_chunk(i)], &embed).await.unwrap(); + m.compact_collection("gc").await.unwrap(); + } + // Fold and merge are deliberately SEPARATE invocations (the merge + // never runs in the same call as a fold, preserving the one-cycle GC + // grace) — drive bare compactions so the merge and its deferred GC run. + m.compact_collection("gc").await.unwrap(); // merge (no tail) + m.ingest("gc", vec![ingest_chunk(99)], &embed) + .await + .unwrap(); + m.compact_collection("gc").await.unwrap(); // fold + m.compact_collection("gc").await.unwrap(); // merge + GC prior staged + m.ingest("gc", vec![ingest_chunk(100)], &embed) + .await + .unwrap(); + m.compact_collection("gc").await.unwrap(); // fold + m.compact_collection("gc").await.unwrap(); // merge + GC + let (man2, _) = crate::storage::lsm::read_manifest(storage.as_ref(), "gc") + .await + .unwrap(); + + // Segment S1 must be physically deleted (GC'd after being folded away). + let s1_key = format!("gc/segments/{seg1_id}"); + assert!( + !storage.exists(&s1_key).await.unwrap(), + "old segment must be GC'd, not leaked" + ); + // Object count stays BOUNDED across many compaction cycles — proving no + // unbounded leak (the F1 bug would grow this without limit). + let all = storage.list("gc/").await.unwrap(); + // Fixed per-namespace objects (manifest, collection.json, id-alloc) + // plus up to MERGE_SEGMENTS(8) tail segments and this-cycle staged + // objects — bounded, never growing with cycle count. + assert!( + all.len() <= 16, + "object count must stay bounded across cycles, got {}", + all.len() + ); + // Data intact. + let mat = cloud::materialize(storage.as_ref(), "gc", &man2) + .await + .unwrap(); + assert_eq!(mat.chunks.len(), 16); + + let _ = std::fs::remove_dir_all(&data_dir); +} + +// Relations must survive COMPACTION-then-restart (segment-relations path), +// not just the fragment-replay path. +#[tokio::test] +async fn relations_survive_compaction_then_restart() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + + let data_dir_a = unique_data_dir(); + std::fs::create_dir_all(&data_dir_a).unwrap(); + { + let storage: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m = CollectionManager::new_with_storage(&data_dir_a, storage) + .await + .unwrap(); + m.create_collection("rc", None, Some(4), None) + .await + .unwrap(); + m.ingest("rc", vec![ingest_chunk(0), ingest_chunk(1)], &embed) + .await + .unwrap(); + m.create_relations( + "rc", + vec![CreateRelation { + source_chunk_id: 0, + target_chunk_id: 1, + target_document_id: None, + relation_type: "cites".into(), + metadata: HashMap::new(), + }], + ) + .await + .unwrap(); + // Compact so the relation lives in the SEGMENT, not a WAL fragment. + m.compact_collection("rc").await.unwrap(); + } + std::fs::remove_dir_all(&data_dir_a).unwrap(); + + let data_dir_b = unique_data_dir(); + std::fs::create_dir_all(&data_dir_b).unwrap(); + let storage: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m2 = CollectionManager::new_with_storage(&data_dir_b, storage) + .await + .unwrap(); + let out = m2 + .get_chunk_relations("rc", 0, RelationDirection::Outgoing, None) + .await + .unwrap(); + assert_eq!(out.len(), 1, "relation must survive compaction+restart"); + assert_eq!(out[0].relation_type, "cites"); + + let _ = std::fs::remove_dir_all(&data_dir_b); +} + +// Concurrent ingests into the same collection: all chunks visible, all ids +// unique, no lost writes (stresses the lock drop/reacquire window). +#[tokio::test] +async fn concurrent_ingests_same_collection() { + let embed = std::sync::Arc::new(embed_state()); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let data_dir = unique_data_dir(); + std::fs::create_dir_all(&data_dir).unwrap(); + let storage: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m = CollectionManager::new_with_storage(&data_dir, storage.clone()) + .await + .unwrap(); + m.create_collection("conc", None, Some(4), None) + .await + .unwrap(); + + let n = 12usize; + let mut handles = Vec::new(); + for i in 0..n { + let m2 = m.clone(); + let e2 = embed.clone(); + handles.push(tokio::spawn(async move { + m2.ingest("conc", vec![ingest_chunk(i as u32)], &e2).await + })); + } + for h in handles { + h.await.unwrap().unwrap(); + } + + // All N chunks present, ids 0..N unique (no collision from the lock gap). + let (man, _) = crate::storage::lsm::read_manifest(storage.as_ref(), "conc") + .await + .unwrap(); + let mat = cloud::materialize(storage.as_ref(), "conc", &man) + .await + .unwrap(); + assert_eq!(mat.chunks.len(), n, "all concurrent ingests durable"); + let ids: std::collections::HashSet = mat.chunks.keys().copied().collect(); + assert_eq!(ids.len(), n, "no duplicate/lost ids"); + assert_eq!(ids, (0..n as u64).collect()); + + let _ = std::fs::remove_dir_all(&data_dir); +} + +// next_id must NEVER regress across compaction + cold restart. Compaction +// physically drops tombstoned chunks; without the segment's stored max_id +// high-water mark, a fresh-disk rebuild would recompute next_id from the +// live set only and REUSE the deleted ids for new chunks. +#[tokio::test] +async fn no_id_reuse_after_compaction_and_cold_restart() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let data_dir_a = unique_data_dir(); + std::fs::create_dir_all(&data_dir_a).unwrap(); + { + let storage: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m = CollectionManager::new_with_storage(&data_dir_a, storage) + .await + .unwrap(); + m.create_collection("idreuse", None, Some(4), None) + .await + .unwrap(); + // ids 0..3; delete the two HIGHEST, then compact them away. + m.ingest("idreuse", (0..4u32).map(ingest_chunk).collect(), &embed) + .await + .unwrap(); + m.delete_chunks("idreuse", &[2, 3]).await.unwrap(); + m.compact_collection("idreuse").await.unwrap(); + } + // Node loss: wipe local disk, cold-rebuild from S3 (max live id is 1). + std::fs::remove_dir_all(&data_dir_a).unwrap(); + let data_dir_b = unique_data_dir(); + std::fs::create_dir_all(&data_dir_b).unwrap(); + let storage: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m2 = CollectionManager::new_with_storage(&data_dir_b, storage.clone()) + .await + .unwrap(); + + // A new ingest must get a FRESH id (4), not reuse deleted id 2. + m2.ingest("idreuse", vec![ingest_chunk(9)], &embed) + .await + .unwrap(); + let (man, _) = crate::storage::lsm::read_manifest(storage.as_ref(), "idreuse") + .await + .unwrap(); + let mat = cloud::materialize(storage.as_ref(), "idreuse", &man) + .await + .unwrap(); + // Under block allocation the exact new id is an allocator detail (a + // fresh node claims a fresh block); the INVARIANT is that no previously + // assigned id — live or deleted — is ever reused. + let new_ids: Vec = mat.chunks.keys().copied().filter(|id| *id > 3).collect(); + assert_eq!( + new_ids.len(), + 1, + "exactly one new chunk with a never-before-assigned id, got {:?}", + mat.chunks.keys().collect::>() + ); + assert!( + !mat.chunks.contains_key(&2) && !mat.chunks.contains_key(&3), + "deleted ids must not be reused" + ); + + let _ = std::fs::remove_dir_all(&data_dir_b); +} + +// PERSISTENT-DISK restart path (the one the adversarial review flagged): +// in cloud mode, a node restarting with its local disk intact runs +// `load_collection` (rehydrate from redb) and SKIPS rebuild-from-S3 for +// already-loaded collections. A chunk tombstoned locally (redb) — which is +// exactly what delete AND the ingest-compensation path write — must stay +// masked after that restart, even though it's still physically in redb. +#[tokio::test] +async fn persistent_disk_restart_honors_local_tombstones() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let data_dir = unique_data_dir(); + std::fs::create_dir_all(&data_dir).unwrap(); + { + let storage: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m = CollectionManager::new_with_storage(&data_dir, storage) + .await + .unwrap(); + m.create_collection("pdisk", None, Some(4), None) + .await + .unwrap(); + m.ingest( + "pdisk", + vec![ingest_chunk(0), ingest_chunk(1), ingest_chunk(2)], + &embed, + ) + .await + .unwrap(); + // Writes the redb tombstone + RAM tombstone + S3 tombstone — the + // same three places the ingest-compensation path writes. + assert_eq!(m.delete_chunks("pdisk", &[1]).await.unwrap().0, 1); + } + + // Restart with the SAME data_dir (persistent disk — NOT wiped). This + // takes the load_collection-first, skip-cloud-rebuild path. + let storage: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m2 = CollectionManager::new_with_storage(&data_dir, storage) + .await + .unwrap(); + + let req = SearchRequest { + query: "chunk".to_string(), + mode: "semantic".to_string(), + vector_space: None, + top_k: 20, + query_vector: Some(vec![0.1, 0.2, 0.3, 0.4]), + filters: HashMap::new(), + score_weights: None, + recency: None, + recency_preset: None, + recency_field: None, + boosts: Vec::new(), + relationship_boost: None, + explain: false, + include_relations: false, + relation_types: None, + relation_direction: RelationDirection::Outgoing, + min_seq: None, + }; + let (hits, _, _, _) = m2.search("pdisk", &req, &embed).await.unwrap(); + let hit_ids: std::collections::HashSet = hits.iter().map(|(c, _, _, _, _)| c.id).collect(); + assert!( + !hit_ids.contains(&1), + "tombstoned chunk must stay masked after persistent-disk restart" + ); + assert!( + hit_ids.contains(&0) && hit_ids.contains(&2), + "live chunks must survive, got {:?}", + hit_ids + ); + + let _ = std::fs::remove_dir_all(&data_dir); +} + +// ── Warm-serverless: bucket config + id allocator + writer role ────── + +// The bucket collection.json is the source of truth on recovery: specs, +// created_at, and CollectionConfig must survive a cold rebuild instead of +// being re-inferred as model:"recovered" / defaults. +#[tokio::test] +async fn cold_rebuild_recovers_real_collection_config() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let data_dir_a = unique_data_dir(); + std::fs::create_dir_all(&data_dir_a).unwrap(); + let created; + { + let storage: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m = CollectionManager::new_with_storage(&data_dir_a, storage) + .await + .unwrap(); + let mut spaces = HashMap::new(); + spaces.insert( + "custom".to_string(), + VectorSpaceConfig { + dims: 4, + model: "my-real-model".to_string(), + status: "active".to_string(), + }, + ); + created = m + .create_collection("cfg", Some(spaces), None, None) + .await + .unwrap(); + m.ingest("cfg", vec![ingest_chunk(0)], &embed) + .await + .unwrap(); + } + std::fs::remove_dir_all(&data_dir_a).unwrap(); + + let data_dir_b = unique_data_dir(); + std::fs::create_dir_all(&data_dir_b).unwrap(); + let storage: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m2 = CollectionManager::new_with_storage(&data_dir_b, storage) + .await + .unwrap(); + let recovered = m2.get_collection("cfg").await.unwrap(); + let space = recovered.vector_spaces.get("custom").unwrap(); + assert_eq!( + space.model, "my-real-model", + "specs must not be re-inferred" + ); + assert_eq!(recovered.created_at, created.created_at); + let _ = std::fs::remove_dir_all(&data_dir_b); +} + +// A zero-ingest collection must be discoverable from a fresh disk (the +// create-only empty manifest + bucket config make the namespace exist). +#[tokio::test] +async fn empty_collection_survives_node_loss() { + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let data_dir_a = unique_data_dir(); + std::fs::create_dir_all(&data_dir_a).unwrap(); + { + let storage: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m = CollectionManager::new_with_storage(&data_dir_a, storage) + .await + .unwrap(); + m.create_collection("emptyns", None, Some(4), None) + .await + .unwrap(); + } + std::fs::remove_dir_all(&data_dir_a).unwrap(); + + let data_dir_b = unique_data_dir(); + std::fs::create_dir_all(&data_dir_b).unwrap(); + let storage: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m2 = CollectionManager::new_with_storage(&data_dir_b, storage) + .await + .unwrap(); + assert!( + m2.get_collection("emptyns").await.is_some(), + "zero-ingest collection must be rediscovered from the bucket" + ); + let _ = std::fs::remove_dir_all(&data_dir_b); +} + +// Writer role end-to-end: a node with NO local collection state ingests; +// a fresh serving node sees the data. Ids from writer and attached node +// never collide (both allocate from {ns}/id-alloc). +#[tokio::test] +async fn writer_role_ingest_is_stateless_and_ids_disjoint() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + + // Full node creates the collection and ingests two chunks. + let dir_full = unique_data_dir(); + std::fs::create_dir_all(&dir_full).unwrap(); + let storage_full: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m_full = CollectionManager::new_with_storage_role(&dir_full, storage_full, NodeRole::Full) + .await + .unwrap(); + m_full + .create_collection("wns", None, Some(4), None) + .await + .unwrap(); + m_full + .ingest("wns", vec![ingest_chunk(0), ingest_chunk(1)], &embed) + .await + .unwrap(); + + // Writer node: EMPTY data dir, writer role. Ingest must succeed with + // zero local collection state and never create local index files. + let dir_writer = unique_data_dir(); + std::fs::create_dir_all(&dir_writer).unwrap(); + let storage_writer: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m_writer = + CollectionManager::new_with_storage_role(&dir_writer, storage_writer, NodeRole::Writer) + .await + .unwrap(); + let (n, _, _) = m_writer + .ingest("wns", vec![ingest_chunk(2), ingest_chunk(3)], &embed) + .await + .unwrap(); + assert_eq!(n, 2); + assert!( + !dir_writer.join("wns").exists(), + "writer role must not create local collection state" + ); + // Reads are refused on the writer. + assert!(m_writer.get_facets("wns", "", &[]).await.is_err()); + + // A fresh serving node materializes ALL four chunks with unique ids. + let dir_read = unique_data_dir(); + std::fs::create_dir_all(&dir_read).unwrap(); + let storage_read: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m_read = CollectionManager::new_with_storage(&dir_read, storage_read.clone()) + .await + .unwrap(); + let (man, _) = crate::storage::lsm::read_manifest(storage_read.as_ref(), "wns") + .await + .unwrap(); + let mat = cloud::materialize(storage_read.as_ref(), "wns", &man) + .await + .unwrap(); + assert_eq!(mat.chunks.len(), 4, "all chunks durable"); + let ids: std::collections::HashSet = mat.chunks.keys().copied().collect(); + assert_eq!( + ids.len(), + 4, + "no id collisions between writer and full node" + ); + assert!(m_read.get_collection("wns").await.is_some()); + + let _ = std::fs::remove_dir_all(&dir_full); + let _ = std::fs::remove_dir_all(&dir_writer); + let _ = std::fs::remove_dir_all(&dir_read); +} + +// Pre-v0.4 migration: a namespace with data but NO id-alloc object seeds +// the allocator from the bucket-derived high-water mark — new ids never +// collide with existing ones. +#[tokio::test] +async fn id_alloc_migration_seeds_past_existing_ids() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let data_dir = unique_data_dir(); + std::fs::create_dir_all(&data_dir).unwrap(); + let storage: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m = CollectionManager::new_with_storage(&data_dir, storage.clone()) + .await + .unwrap(); + m.create_collection("mig", None, Some(4), None) + .await + .unwrap(); + m.ingest( + "mig", + vec![ingest_chunk(0), ingest_chunk(1), ingest_chunk(2)], + &embed, + ) + .await + .unwrap(); + // Simulate a pre-v0.4 namespace: remove the allocator object. + storage.delete("mig/id-alloc").await.unwrap(); + // Drain the local pool by restarting the manager (pool is in-RAM). + drop(m); + let m2 = CollectionManager::new_with_storage(&data_dir, storage.clone()) + .await + .unwrap(); + m2.ingest("mig", vec![ingest_chunk(9)], &embed) + .await + .unwrap(); + + let (man, _) = crate::storage::lsm::read_manifest(storage.as_ref(), "mig") + .await + .unwrap(); + let mat = cloud::materialize(storage.as_ref(), "mig", &man) + .await + .unwrap(); + assert_eq!(mat.chunks.len(), 4); + let ids: std::collections::HashSet = mat.chunks.keys().copied().collect(); + assert_eq!(ids.len(), 4, "migrated allocator must not reuse ids 0-2"); + assert!( + ids.contains(&3), + "first migrated id is one past the high-water" + ); + let _ = std::fs::remove_dir_all(&data_dir); +} + +// ── Warm-serverless: manifest refresh + read-your-writes ───────────── + +fn cloud_search_req(min_seq: Option) -> SearchRequest { + SearchRequest { + query: "chunk".to_string(), + mode: "semantic".to_string(), + vector_space: None, + top_k: 20, + query_vector: Some(vec![0.1, 0.2, 0.3, 0.4]), + filters: HashMap::new(), + score_weights: None, + recency: None, + recency_preset: None, + recency_field: None, + boosts: Vec::new(), + relationship_boost: None, + explain: false, + include_relations: false, + relation_types: None, + relation_direction: RelationDirection::Outgoing, + min_seq, + } +} + +// Two serving nodes on one bucket: writes on A become visible on B via +// refresh_collection — chunks, deletes, and relations all converge. +#[tokio::test] +async fn two_nodes_converge_via_refresh() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let dir_a = unique_data_dir(); + let dir_b = unique_data_dir(); + std::fs::create_dir_all(&dir_a).unwrap(); + std::fs::create_dir_all(&dir_b).unwrap(); + let sa: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let sb: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let a = CollectionManager::new_with_storage(&dir_a, sa) + .await + .unwrap(); + a.create_collection("conv", None, Some(4), None) + .await + .unwrap(); + a.ingest("conv", vec![ingest_chunk(0)], &embed) + .await + .unwrap(); + + // B boots AFTER the first write (rebuilds to seq frontier). + let b = CollectionManager::new_with_storage(&dir_b, sb) + .await + .unwrap(); + let (hits, _, _, _) = b + .search("conv", &cloud_search_req(None), &embed) + .await + .unwrap(); + assert_eq!(hits.len(), 1, "B rebuilt A's first write at boot"); + + // A writes more: a new chunk, a relation, and a delete of chunk id 0. + a.ingest("conv", vec![ingest_chunk(1), ingest_chunk(2)], &embed) + .await + .unwrap(); + let a_ids: Vec = { + let (hits, _, _, _) = a + .search("conv", &cloud_search_req(None), &embed) + .await + .unwrap(); + hits.iter().map(|(c, _, _, _, _)| c.id).collect() + }; + assert_eq!(a_ids.len(), 3); + let first_id = *a_ids.iter().min().unwrap(); + let others: Vec = a_ids.iter().copied().filter(|i| *i != first_id).collect(); + a.create_relations( + "conv", + vec![CreateRelation { + source_chunk_id: others[0], + target_chunk_id: others[1], + target_document_id: None, + relation_type: "cites".into(), + metadata: HashMap::new(), + }], + ) + .await + .unwrap(); + a.delete_chunks("conv", &[first_id]).await.unwrap(); + + // B converges via refresh (no restart, no rebuild). + b.refresh_collection("conv").await.unwrap(); + let (hits, _, _, _) = b + .search("conv", &cloud_search_req(None), &embed) + .await + .unwrap(); + let b_ids: std::collections::HashSet = hits.iter().map(|(c, _, _, _, _)| c.id).collect(); + assert!(!b_ids.contains(&first_id), "A's delete visible on B"); + assert_eq!(b_ids.len(), 2, "A's later chunks visible on B"); + let rels = b + .get_chunk_relations("conv", others[0], RelationDirection::Outgoing, None) + .await + .unwrap(); + assert_eq!(rels.len(), 1, "A's relation visible on B"); + assert_eq!(rels[0].relation_type, "cites"); + + let _ = std::fs::remove_dir_all(&dir_a); + let _ = std::fs::remove_dir_all(&dir_b); +} + +// The refresher must never double-apply a node's OWN fragments (the seq +// tracker covers them out-of-band). +#[tokio::test] +async fn refresh_never_double_applies_own_writes() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let dir = unique_data_dir(); + std::fs::create_dir_all(&dir).unwrap(); + let st: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m = CollectionManager::new_with_storage(&dir, st).await.unwrap(); + m.create_collection("own", None, Some(4), None) + .await + .unwrap(); + m.ingest("own", vec![ingest_chunk(0), ingest_chunk(1)], &embed) + .await + .unwrap(); + m.delete_chunks("own", &[0]).await.unwrap(); + + // Refresh repeatedly: state (incl. chunk_count) must not change. + let before = m.get_collection("own").await.unwrap().chunk_count; + for _ in 0..3 { + m.refresh_collection("own").await.unwrap(); + } + let after = m.get_collection("own").await.unwrap().chunk_count; + assert_eq!(before, after, "replay of own fragments must be a no-op"); + assert_eq!(after, 1); + let _ = std::fs::remove_dir_all(&dir); +} + +// Compaction two-branch rule: a node that saw everything skips segments; +// a node whose frontier is BEHIND the watermark re-attaches fully. +#[tokio::test] +async fn refresh_survives_remote_compaction() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let dir_a = unique_data_dir(); + let dir_b = unique_data_dir(); + std::fs::create_dir_all(&dir_a).unwrap(); + std::fs::create_dir_all(&dir_b).unwrap(); + let sa: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let sb: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let a = CollectionManager::new_with_storage(&dir_a, sa) + .await + .unwrap(); + a.create_collection("rc2", None, Some(4), None) + .await + .unwrap(); + a.ingest("rc2", vec![ingest_chunk(0)], &embed) + .await + .unwrap(); + + // B attaches at frontier 1 (one fragment applied). + let b = CollectionManager::new_with_storage(&dir_b, sb) + .await + .unwrap(); + + // Branch 1: A ingests + compacts; B's frontier is BEHIND the watermark + // (never saw seq 1) → refresh must full re-attach, not skip. + a.ingest("rc2", vec![ingest_chunk(1)], &embed) + .await + .unwrap(); + a.compact_collection("rc2").await.unwrap(); + b.refresh_collection("rc2").await.unwrap(); + let (hits, _, _, _) = b + .search("rc2", &cloud_search_req(None), &embed) + .await + .unwrap(); + assert_eq!(hits.len(), 2, "stale node re-attaches across compaction"); + + // Branch 2: B now has everything; another compaction (A side) must be + // a cheap no-op on refresh (no re-attach needed) and lose nothing. + a.ingest("rc2", vec![ingest_chunk(2)], &embed) + .await + .unwrap(); + b.refresh_collection("rc2").await.unwrap(); // B applies seq tail first + a.compact_collection("rc2").await.unwrap(); + b.refresh_collection("rc2").await.unwrap(); // wm <= frontier → skip + let (hits, _, _, _) = b + .search("rc2", &cloud_search_req(None), &embed) + .await + .unwrap(); + assert_eq!(hits.len(), 3); + + let _ = std::fs::remove_dir_all(&dir_a); + let _ = std::fs::remove_dir_all(&dir_b); +} + +// Read-your-writes across nodes: a write on A returns a seq; a search on B +// with min_seq=seq refreshes and serves the write. +#[tokio::test] +async fn min_seq_gives_read_your_writes_across_nodes() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let dir_a = unique_data_dir(); + let dir_b = unique_data_dir(); + std::fs::create_dir_all(&dir_a).unwrap(); + std::fs::create_dir_all(&dir_b).unwrap(); + let sa: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let sb: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let a = CollectionManager::new_with_storage(&dir_a, sa) + .await + .unwrap(); + a.create_collection("ryw", None, Some(4), None) + .await + .unwrap(); + a.ingest("ryw", vec![ingest_chunk(0)], &embed) + .await + .unwrap(); + let b = CollectionManager::new_with_storage(&dir_b, sb) + .await + .unwrap(); + + // A writes; B searches with min_seq — must see it without manual refresh. + let (_, _, seq) = a + .ingest("ryw", vec![ingest_chunk(1)], &embed) + .await + .unwrap(); + let seq = seq.expect("cloud ingest returns a seq"); + let (hits, _, _, _) = b + .search("ryw", &cloud_search_req(Some(seq)), &embed) + .await + .unwrap(); + assert_eq!(hits.len(), 2, "min_seq forces convergence before serving"); + + // A min_seq beyond the write history is rejected, not waited on. + assert!(b + .search("ryw", &cloud_search_req(Some(9_999)), &embed) + .await + .is_err()); + + let _ = std::fs::remove_dir_all(&dir_a); + let _ = std::fs::remove_dir_all(&dir_b); +} + +// ── Warm-serverless: lazy attach + LRU detach ───────────────────────── + +// Lazy boot registers namespaces without rebuilding; the first request +// attaches; a concurrent stampede attaches exactly once. +#[tokio::test] +async fn lazy_attach_on_first_request() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + // Seed the bucket with a collection via an eager node. + let dir_seed = unique_data_dir(); + std::fs::create_dir_all(&dir_seed).unwrap(); + { + let st: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m = CollectionManager::new_with_storage(&dir_seed, st) + .await + .unwrap(); + m.create_collection("lazy", None, Some(4), None) + .await + .unwrap(); + m.ingest("lazy", vec![ingest_chunk(0), ingest_chunk(1)], &embed) + .await + .unwrap(); + } + std::fs::remove_dir_all(&dir_seed).unwrap(); + + // Lazy node: boot must NOT rebuild (no local dir for the collection). + let dir = unique_data_dir(); + std::fs::create_dir_all(&dir).unwrap(); + let st: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m = CollectionManager::new_with_storage_opts(&dir, st, NodeRole::Full, true, 0, 0) + .await + .unwrap(); + assert!( + !dir.join("lazy").join("chunks.redb").exists(), + "lazy boot must not rebuild collections" + ); + + // Stampede: 8 concurrent first-requests; all succeed, attach happens once. + let mut handles = Vec::new(); + for _ in 0..8 { + let m2 = m.clone(); + let e2 = embed_state(); + handles.push(tokio::spawn(async move { + let (hits, _, _, _) = m2 + .search("lazy", &cloud_search_req(None), &e2) + .await + .unwrap(); + hits.len() + })); + } + for h in handles { + assert_eq!(h.await.unwrap(), 2); + } + assert!(dir.join("lazy").join("chunks.redb").exists()); + let _ = std::fs::remove_dir_all(&dir); +} + +// LRU detach: with a budget of 1, attaching a second collection evicts the +// least-recently-used one; the evicted collection re-attaches on demand +// with all its data (bucket is the source of truth). +#[tokio::test] +async fn lru_detach_and_reattach_roundtrip() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let dir_seed = unique_data_dir(); + std::fs::create_dir_all(&dir_seed).unwrap(); + { + let st: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m = CollectionManager::new_with_storage(&dir_seed, st) + .await + .unwrap(); + for name in ["one", "two"] { + m.create_collection(name, None, Some(4), None) + .await + .unwrap(); + m.ingest(name, vec![ingest_chunk(0)], &embed).await.unwrap(); + } + } + std::fs::remove_dir_all(&dir_seed).unwrap(); + + let dir = unique_data_dir(); + std::fs::create_dir_all(&dir).unwrap(); + let st: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m = CollectionManager::new_with_storage_opts(&dir, st, NodeRole::Full, true, 1, 0) + .await + .unwrap(); + + // Attach "one", then "two" — budget 1 evicts "one". + let (hits, _, _, _) = m + .search("one", &cloud_search_req(None), &embed) + .await + .unwrap(); + assert_eq!(hits.len(), 1); + let (hits, _, _, _) = m + .search("two", &cloud_search_req(None), &embed) + .await + .unwrap(); + assert_eq!(hits.len(), 1); + { + let attached = m.collections.read().await; + assert_eq!(attached.len(), 1, "LRU budget enforced"); + assert!(attached.contains_key("two")); + } + assert!(!dir.join("one").join("chunks.redb").exists()); + + // Evicted collection re-attaches on demand, data intact. + let (hits, _, _, _) = m + .search("one", &cloud_search_req(None), &embed) + .await + .unwrap(); + assert_eq!(hits.len(), 1, "re-attach after eviction serves all data"); + let _ = std::fs::remove_dir_all(&dir); +} + +// Lazy mode keeps metadata correct: list/get see registered collections; +// a collection created on ANOTHER node after boot attaches on demand. +#[tokio::test] +async fn lazy_attach_discovers_foreign_creates() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let dir_a = unique_data_dir(); + let dir_b = unique_data_dir(); + std::fs::create_dir_all(&dir_a).unwrap(); + std::fs::create_dir_all(&dir_b).unwrap(); + let sa: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let sb: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + // Lazy node boots FIRST (empty bucket). + let b = CollectionManager::new_with_storage_opts(&dir_b, sb, NodeRole::Full, true, 0, 0) + .await + .unwrap(); + // Another node creates + writes afterwards. + let a = CollectionManager::new_with_storage(&dir_a, sa) + .await + .unwrap(); + a.create_collection("late", None, Some(4), None) + .await + .unwrap(); + a.ingest("late", vec![ingest_chunk(0)], &embed) + .await + .unwrap(); + + // B never saw "late" at boot; first request attaches it anyway. + let (hits, _, _, _) = b + .search("late", &cloud_search_req(None), &embed) + .await + .unwrap(); + assert_eq!(hits.len(), 1, "foreign create attaches on demand"); + + let _ = std::fs::remove_dir_all(&dir_a); + let _ = std::fs::remove_dir_all(&dir_b); +} + +// ── Review-driven regression tests (adversarial round) ─────────────── + +#[test] +fn seq_tracker_semantics() { + let mut t = SeqTracker::default(); + assert!(!t.covers(0)); + t.mark(0); + assert_eq!(t.contiguous, 1); + // Out-of-band mark ahead of the frontier; contiguous holds. + t.mark(2); + assert!(t.covers(2) && !t.covers(1)); + assert_eq!(t.contiguous, 1); + // Filling the gap drains the whole out-of-band run. + t.mark(1); + assert_eq!(t.contiguous, 3); + assert!(t.out_of_band.is_empty()); + // Duplicate + below-frontier marks are no-ops (no unbounded growth). + t.mark(1); + t.mark(2); + assert_eq!(t.contiguous, 3); + assert!(t.out_of_band.is_empty()); + // starting_at seeds the frontier. + let t2 = SeqTracker::starting_at(7); + assert!(t2.covers(6) && !t2.covers(7)); +} + +// H4 regression: eviction must be least-recently-USED, not least-recently- +// attached. 3 collections, budget 2: attach a, attach b, USE a, attach c +// → b (not a) is evicted. +#[tokio::test] +async fn lru_evicts_least_recently_used() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let dir_seed = unique_data_dir(); + std::fs::create_dir_all(&dir_seed).unwrap(); + { + let st: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m = CollectionManager::new_with_storage(&dir_seed, st) + .await + .unwrap(); + for name in ["a", "b", "c"] { + m.create_collection(name, None, Some(4), None) + .await + .unwrap(); + m.ingest(name, vec![ingest_chunk(0)], &embed).await.unwrap(); + } + } + std::fs::remove_dir_all(&dir_seed).unwrap(); + let dir = unique_data_dir(); + std::fs::create_dir_all(&dir).unwrap(); + let st: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m = CollectionManager::new_with_storage_opts(&dir, st, NodeRole::Full, true, 2, 0) + .await + .unwrap(); + m.search("a", &cloud_search_req(None), &embed) + .await + .unwrap(); + m.search("b", &cloud_search_req(None), &embed) + .await + .unwrap(); + // USE a again — it is now hotter than b. + m.search("a", &cloud_search_req(None), &embed) + .await + .unwrap(); + m.search("c", &cloud_search_req(None), &embed) + .await + .unwrap(); + let attached = m.collections.read().await; + assert!(attached.contains_key("a"), "hot collection must survive"); + assert!(!attached.contains_key("b"), "cold collection is the victim"); + assert!(attached.contains_key("c")); +} + +// C1 regression: a vector space added on node A becomes visible on an +// already-attached node B via refresh (config is synced, not just +// fragments), so B never quarantines chunks carrying the new space. +#[tokio::test] +async fn vector_space_add_propagates_via_refresh() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let dir_a = unique_data_dir(); + let dir_b = unique_data_dir(); + std::fs::create_dir_all(&dir_a).unwrap(); + std::fs::create_dir_all(&dir_b).unwrap(); + let sa: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let sb: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let a = CollectionManager::new_with_storage(&dir_a, sa) + .await + .unwrap(); + a.create_collection("vsprop", None, Some(4), None) + .await + .unwrap(); + a.ingest("vsprop", vec![ingest_chunk(0)], &embed) + .await + .unwrap(); + let b = CollectionManager::new_with_storage(&dir_b, sb) + .await + .unwrap(); + + // A adds an 8-dim space, then ingests a chunk carrying it. + a.add_vector_space("vsprop", "wide", 8, "test-model") + .await + .unwrap(); + let mut ic = ingest_chunk(1); + ic.embeddings.insert("wide".to_string(), vec![0.1; 8]); + a.ingest("vsprop", vec![ic], &embed).await.unwrap(); + + // B refreshes: must learn the space AND apply the chunk (no quarantine). + b.refresh_collection("vsprop").await.unwrap(); + let bc = b.get_collection("vsprop").await.unwrap(); + assert!(bc.vector_spaces.contains_key("wide"), "config converged"); + let (hits, _, _, _) = b + .search("vsprop", &cloud_search_req(None), &embed) + .await + .unwrap(); + assert_eq!( + hits.len(), + 2, + "chunk with the new space applied, not quarantined" + ); + + let _ = std::fs::remove_dir_all(&dir_a); + let _ = std::fs::remove_dir_all(&dir_b); +} + +// Rank-1 regression: ingest racing a refresher loop never double-applies +// (chunk_count exact, no duplicate hits). +#[tokio::test] +async fn ingest_races_refresher_no_double_apply() { + let embed = std::sync::Arc::new(embed_state()); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let dir = unique_data_dir(); + std::fs::create_dir_all(&dir).unwrap(); + let st: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m = CollectionManager::new_with_storage(&dir, st).await.unwrap(); + m.create_collection("race", None, Some(4), None) + .await + .unwrap(); + + let n = 10usize; + let refresher = { + let m2 = m.clone(); + tokio::spawn(async move { + for _ in 0..200 { + let _ = m2.refresh_collection("race").await; + tokio::task::yield_now().await; + } + }) + }; + let mut handles = Vec::new(); + for i in 0..n { + let m2 = m.clone(); + let e2 = embed.clone(); + handles.push(tokio::spawn(async move { + m2.ingest("race", vec![ingest_chunk(i as u32)], &e2).await + })); + } + for h in handles { + h.await.unwrap().unwrap(); + } + refresher.await.unwrap(); + let _ = m.refresh_collection("race").await; + + let c = m.get_collection("race").await.unwrap(); + assert_eq!(c.chunk_count as usize, n, "no double-count under the race"); + let (hits, _, _, _) = m + .search("race", &cloud_search_req(None), &embed) + .await + .unwrap(); + assert_eq!(hits.len(), n, "no duplicate/lost chunks under the race"); + let _ = std::fs::remove_dir_all(&dir); +} + +// Rank-6: persistent-disk restart catches up the REMOTE delta via refresh +// instead of serving stale data (applied_seq persistence path). +#[tokio::test] +async fn persistent_restart_catches_up_remote_delta() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let dir_a = unique_data_dir(); + let dir_w = unique_data_dir(); + std::fs::create_dir_all(&dir_a).unwrap(); + std::fs::create_dir_all(&dir_w).unwrap(); + { + let sa: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let a = CollectionManager::new_with_storage(&dir_a, sa) + .await + .unwrap(); + a.create_collection("pd", None, Some(4), None) + .await + .unwrap(); + a.ingest("pd", vec![ingest_chunk(0)], &embed).await.unwrap(); + } // node A down; its disk PERSISTS. + { + let sw: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let w = CollectionManager::new_with_storage_role(&dir_w, sw, NodeRole::Writer) + .await + .unwrap(); + w.ingest("pd", vec![ingest_chunk(1)], &embed).await.unwrap(); + } + // A restarts on the SAME dir (load_collection path, not rebuild). + let sa: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let a = CollectionManager::new_with_storage(&dir_a, sa) + .await + .unwrap(); + a.refresh_collection("pd").await.unwrap(); + let (hits, _, _, _) = a + .search("pd", &cloud_search_req(None), &embed) + .await + .unwrap(); + assert_eq!( + hits.len(), + 2, + "restart + refresh catches up the writer's delta" + ); + let c = a.get_collection("pd").await.unwrap(); + assert_eq!(c.chunk_count, 2, "delta applied exactly once"); + let _ = std::fs::remove_dir_all(&dir_a); + let _ = std::fs::remove_dir_all(&dir_w); +} + +// Rank-8: a wrong-dims chunk inside a fragment is quarantined on replay +// without corrupting anything else. +#[tokio::test] +async fn refresh_quarantines_wrong_dims_without_corruption() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let dir = unique_data_dir(); + std::fs::create_dir_all(&dir).unwrap(); + let st: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m = CollectionManager::new_with_storage(&dir, st.clone()) + .await + .unwrap(); + m.create_collection("quar", None, Some(4), None) + .await + .unwrap(); + m.ingest("quar", vec![ingest_chunk(0)], &embed) + .await + .unwrap(); + + // Hand-craft a fragment with one bad (3-dim) and one good chunk, + // simulating a poisoned foreign writer. + let mut bad = DocumentChunk { + id: 500_000, + collection: "quar".into(), + file_id: "bad".into(), + chunk_index: 0, + page: None, + text: "bad chunk".into(), + metadata: HashMap::new(), + doc_type: "chunk".into(), + parent_id: None, + group_id: None, + embeddings: HashMap::new(), + embedding: None, + }; + bad.embeddings.insert("default".into(), vec![0.1, 0.2, 0.3]); + let mut good = bad.clone(); + good.id = 500_001; + good.file_id = "good".into(); + good.text = "good chunk".into(); + good.embeddings + .insert("default".into(), vec![0.1, 0.2, 0.3, 0.4]); + let payload = serde_json::to_vec(&vec![bad, good]).unwrap(); + crate::storage::lsm::append_fragment(st.as_ref(), "quar", bytes::Bytes::from(payload), 2) + .await + .unwrap(); + + m.refresh_collection("quar").await.unwrap(); + let (hits, _, _, _) = m + .search("quar", &cloud_search_req(None), &embed) + .await + .unwrap(); + let ids: std::collections::HashSet = hits.iter().map(|(c, _, _, _, _)| c.id).collect(); + assert!(ids.contains(&500_001), "good chunk applied"); + assert!(!ids.contains(&500_000), "bad chunk quarantined"); + // Post-quarantine ingest still works and searches correctly (mmap not shifted). + m.ingest("quar", vec![ingest_chunk(9)], &embed) + .await + .unwrap(); + let (hits, _, _, _) = m + .search("quar", &cloud_search_req(None), &embed) + .await + .unwrap(); + assert_eq!(hits.len(), 3); + let _ = std::fs::remove_dir_all(&dir); +} + +// Rank-9: min_seq is ignored in local mode; exact boundary at next_seq. +#[tokio::test] +async fn min_seq_local_mode_and_boundary() { + let embed = embed_state(); + // Local mode: min_seq must be ignored, not error. + let dir = unique_data_dir(); + std::fs::create_dir_all(&dir).unwrap(); + let m = CollectionManager::new(&dir).await.unwrap(); + m.create_collection("loc", None, Some(4), None) + .await + .unwrap(); + m.ingest("loc", vec![ingest_chunk(0)], &embed) + .await + .unwrap(); + let (hits, _, _, _) = m + .search("loc", &cloud_search_req(Some(999)), &embed) + .await + .unwrap(); + assert_eq!(hits.len(), 1, "local mode ignores min_seq"); + let _ = std::fs::remove_dir_all(&dir); + + // Cloud: last valid seq (next_seq-1) succeeds; next_seq is rejected. + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let dir2 = unique_data_dir(); + std::fs::create_dir_all(&dir2).unwrap(); + let st: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let m2 = CollectionManager::new_with_storage(&dir2, st) + .await + .unwrap(); + m2.create_collection("bnd", None, Some(4), None) + .await + .unwrap(); + let (_, _, seq) = m2 + .ingest("bnd", vec![ingest_chunk(0)], &embed) + .await + .unwrap(); + let seq = seq.unwrap(); + assert!(m2 + .search("bnd", &cloud_search_req(Some(seq)), &embed) + .await + .is_ok()); + assert!(m2 + .search("bnd", &cloud_search_req(Some(seq + 1)), &embed) + .await + .is_err()); + let _ = std::fs::remove_dir_all(&dir2); +} + +// Rank-4/H3: a writer delete against a bogus namespace must NOT create a +// phantom collection, and absurd ids are rejected by the allocator frontier. +#[tokio::test] +async fn writer_delete_validates_namespace_and_ids() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let dir = unique_data_dir(); + std::fs::create_dir_all(&dir).unwrap(); + let st: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let w = CollectionManager::new_with_storage_role(&dir, st.clone(), NodeRole::Writer) + .await + .unwrap(); + // Bogus namespace: error + nothing created in the bucket. + assert!(w.delete_chunks("ghost", &[1]).await.is_err()); + assert!( + !st.exists("ghost/manifest").await.unwrap(), + "no phantom namespace" + ); + + // Real collection: absurd id rejected (would poison max_id forever). + let dir_f = unique_data_dir(); + std::fs::create_dir_all(&dir_f).unwrap(); + let sf: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let f = CollectionManager::new_with_storage(&dir_f, sf) + .await + .unwrap(); + f.create_collection("real", None, Some(4), None) + .await + .unwrap(); + f.ingest("real", vec![ingest_chunk(0)], &embed) + .await + .unwrap(); + assert!(w.delete_chunks("real", &[u64::MAX]).await.is_err()); + // In-range delete works. + assert!(w.delete_chunks("real", &[0]).await.is_ok()); + let _ = std::fs::remove_dir_all(&dir); + let _ = std::fs::remove_dir_all(&dir_f); +} + +// H1-lite: delete+recreate on another node is detected via created_at and +// the stale node re-attaches to the NEW collection. +#[tokio::test] +async fn delete_recreate_detected_by_refresh() { + let embed = embed_state(); + let store = std::sync::Arc::new(object_store::memory::InMemory::new()); + let dir_a = unique_data_dir(); + let dir_b = unique_data_dir(); + std::fs::create_dir_all(&dir_a).unwrap(); + std::fs::create_dir_all(&dir_b).unwrap(); + let sa: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let sb: Arc = Arc::new(ObjectStoreBackend::from_store( + store.clone(), + "object-store:memory", + )); + let a = CollectionManager::new_with_storage(&dir_a, sa) + .await + .unwrap(); + a.create_collection("cycle", None, Some(4), None) + .await + .unwrap(); + a.ingest( + "cycle", + vec![ingest_chunk(0), ingest_chunk(1), ingest_chunk(2)], + &embed, + ) + .await + .unwrap(); + let b = CollectionManager::new_with_storage(&dir_b, sb) + .await + .unwrap(); + + // A deletes and recreates with different content. + a.delete_collection("cycle").await.unwrap(); + a.create_collection("cycle", None, Some(4), None) + .await + .unwrap(); + a.ingest("cycle", vec![ingest_chunk(9)], &embed) + .await + .unwrap(); + + // B refreshes: must serve the NEW collection (1 chunk), not the old 3. + b.refresh_collection("cycle").await.unwrap(); + let (hits, _, _, _) = b + .search("cycle", &cloud_search_req(None), &embed) + .await + .unwrap(); + assert_eq!( + hits.len(), + 1, + "stale node re-attached to the recreated collection" + ); + let _ = std::fs::remove_dir_all(&dir_a); + let _ = std::fs::remove_dir_all(&dir_b); +} + +// ── Scale harness (env-gated) ───────────────────────────────────────── +// COMPASS_SCALE_N= [COMPASS_SCALE_DIMS=] cargo test +// --features object-storage --release scale_envelope -- --nocapture +// Measures ingest throughput, attach (cold rebuild) time, and search +// latency against a local-disk Storage backend (same code paths as S3, +// disk-bound). Skips (passes) when COMPASS_SCALE_N is unset. +#[tokio::test] +async fn scale_envelope() { + let Ok(n) = std::env::var("COMPASS_SCALE_N") else { + eprintln!("skipped: COMPASS_SCALE_N not set"); + return; + }; + let n: usize = n.parse().unwrap(); + let dims: usize = std::env::var("COMPASS_SCALE_DIMS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(128); + let batch = 2_000usize; + let embed = embed_state(); + + let bucket_dir = unique_data_dir(); + std::fs::create_dir_all(&bucket_dir).unwrap(); + let storage: Arc = + Arc::new(crate::storage::local::LocalDiskStorage::new(&bucket_dir).unwrap()); + // local-disk backend reports "local-disk" => cloud_mode false. Wrap it + // to report as a cloud backend so the full S3-native path runs. + struct CloudyDisk(Arc); + #[async_trait::async_trait] + impl Storage for CloudyDisk { + async fn get(&self, k: &str) -> Result { + self.0.get(k).await + } + async fn get_range( + &self, + k: &str, + r: std::ops::Range, + ) -> Result { + self.0.get_range(k, r).await + } + async fn get_versioned( + &self, + k: &str, + ) -> Result<(bytes::Bytes, crate::storage::Version), crate::storage::StorageError> { + self.0.get_versioned(k).await + } + async fn put( + &self, + k: &str, + b: bytes::Bytes, + ) -> Result { + self.0.put(k, b).await + } + async fn put_if_match( + &self, + k: &str, + b: bytes::Bytes, + e: &crate::storage::Version, + ) -> Result { + self.0.put_if_match(k, b, e).await + } + async fn put_if_not_exists( + &self, + k: &str, + b: bytes::Bytes, + ) -> Result { + self.0.put_if_not_exists(k, b).await + } + async fn delete(&self, k: &str) -> Result<(), crate::storage::StorageError> { + self.0.delete(k).await + } + async fn put_large( + &self, + k: &str, + b: bytes::Bytes, + ) -> Result { + self.0.put_large(k, b).await + } + async fn list( + &self, + p: &str, + ) -> Result, crate::storage::StorageError> { + self.0.list(p).await + } + async fn list_dirs(&self, p: &str) -> Result, crate::storage::StorageError> { + self.0.list_dirs(p).await + } + fn backend_name(&self) -> &'static str { + "scale-disk" + } + } + let storage: Arc = Arc::new(CloudyDisk(storage)); + + let node_dir = unique_data_dir(); + std::fs::create_dir_all(&node_dir).unwrap(); + let m = CollectionManager::new_with_storage_opts( + &node_dir, + storage.clone(), + NodeRole::Full, + false, + 0, + 0, + ) + .await + .unwrap(); + let mut spaces = HashMap::new(); + spaces.insert( + "default".to_string(), + VectorSpaceConfig { + dims, + model: "scale".into(), + status: "active".into(), + }, + ); + m.create_collection("scale", Some(spaces), None, None) + .await + .unwrap(); + + // Deterministic pseudo-random embeddings (no Math.random / clock). + let mk_vec = |seed: usize| -> Vec { + let mut x = seed as u64 * 6364136223846793005 + 1442695040888963407; + (0..dims) + .map(|_| { + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + ((x % 2000) as f32 / 1000.0) - 1.0 + }) + .collect() + }; + let t0 = std::time::Instant::now(); + for b0 in (0..n).step_by(batch) { + let chunks: Vec = (b0..(b0 + batch).min(n)) + .map(|i| { + let mut embeddings = HashMap::new(); + embeddings.insert("default".to_string(), mk_vec(i)); + IngestChunk { + client_id: None, + file_id: format!("f{i}"), + chunk_index: 0, + page: None, + text: format!("scale test chunk number {i} lorem ipsum"), + metadata: HashMap::new(), + doc_type: "chunk".to_string(), + parent_id: None, + parent_ref: None, + group_id: None, + embeddings, + embedding: None, + } + }) + .collect(); + m.ingest("scale", chunks, &embed).await.unwrap(); + } + let ingest_s = t0.elapsed().as_secs_f64(); + + // Cold attach: fresh node dir, same bucket. + drop(m); + let node2 = unique_data_dir(); + std::fs::create_dir_all(&node2).unwrap(); + let t1 = std::time::Instant::now(); + let m2 = CollectionManager::new_with_storage_opts( + &node2, + storage.clone(), + NodeRole::Full, + false, + 0, + 0, + ) + .await + .unwrap(); + let attach_s = t1.elapsed().as_secs_f64(); + + // Search latency (semantic, 50 queries). + let mut req = cloud_search_req(None); + let t2 = std::time::Instant::now(); + let mut hits_total = 0usize; + for q in 0..50 { + req.query_vector = Some(mk_vec(q * 7919)); + let (hits, _, _, _) = m2.search("scale", &req, &embed).await.unwrap(); + hits_total += hits.len(); + } + let search_ms = t2.elapsed().as_secs_f64() * 1000.0 / 50.0; + assert!(hits_total > 0); + + eprintln!( + "SCALE n={n} dims={dims}: ingest {:.1}s ({:.0} chunks/s) | cold attach {:.1}s | search avg {:.1}ms", + ingest_s, n as f64 / ingest_s, attach_s, search_ms + ); + let _ = std::fs::remove_dir_all(&bucket_dir); + let _ = std::fs::remove_dir_all(&node_dir); + let _ = std::fs::remove_dir_all(&node2); +} diff --git a/crates/compass/src/collections/filter_aware_search_tests.rs b/crates/compass/src/collections/filter_aware_search_tests.rs new file mode 100644 index 0000000..5cb63c6 --- /dev/null +++ b/crates/compass/src/collections/filter_aware_search_tests.rs @@ -0,0 +1,557 @@ +// collections/filter_aware_search_tests.rs — extracted test module (was inline in mod.rs). +// Child module of `collections`: `super::*` sees the parent's private items. + +//! End-to-end test of filter-aware /search + /explain (a follow-up). +//! +//! Builds a real CollectionManager, ingests chunks with caller-provided +//! embeddings (skipping the in-process BGE model), runs filtered hybrid +//! search, and asserts: +//! 1. All hits respect the filter (filter-aware path, not post-filter). +//! 2. The /explain field is populated when requested and absent when not. +//! 3. Filter selectivity is reported correctly. + +use super::*; +use crate::embed::EmbedState; +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; + +fn unique_data_dir() -> std::path::PathBuf { + static N: AtomicU64 = AtomicU64::new(0); + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + std::env::temp_dir().join(format!( + "compass-filter-search-test-{}-{}-{}", + std::process::id(), + nanos, + N.fetch_add(1, Ordering::SeqCst) + )) +} + +fn embed_state() -> EmbedState { + EmbedState { + bge: None, + distilled: None, + } +} + +/// Deterministic 4-dim unit vector seeded from an integer. +fn pseudo_vec(seed: u64) -> Vec { + let mut state = seed + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + let mut v = Vec::with_capacity(4); + for _ in 0..4 { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + let f = (state >> 11) as f32 / (1u64 << 53) as f32 * 2.0 - 1.0; + v.push(f); + } + let norm: f32 = v.iter().map(|x| x * x).sum::().sqrt(); + if norm > 0.0 { + for x in &mut v { + *x /= norm; + } + } + v +} + +fn ingest_with(org: &str, idx: u32) -> IngestChunk { + let mut metadata = HashMap::new(); + metadata.insert("org_id".to_string(), MetadataValue::String(org.to_string())); + metadata.insert( + "created_at".to_string(), + MetadataValue::Int(1_700_000_000 + idx as i64), + ); + let mut embeddings = HashMap::new(); + embeddings.insert("default".to_string(), pseudo_vec(idx as u64 + 1)); + IngestChunk { + client_id: None, + file_id: format!("f{idx}"), + chunk_index: 0, + page: None, + text: format!("chunk-{idx}"), + metadata, + doc_type: "chunk".to_string(), + parent_id: None, + parent_ref: None, + group_id: None, + embeddings, + embedding: None, + } +} + +#[tokio::test] +async fn filter_aware_search_returns_only_matching_chunks() { + let data_dir = unique_data_dir(); + std::fs::create_dir_all(&data_dir).unwrap(); + let embed = embed_state(); + let manager = CollectionManager::new(&data_dir).await.unwrap(); + manager + .create_collection("filter-search", None, Some(4), None) + .await + .unwrap(); + + // 100 chunks: 20 from "acme", 80 from "widgets". + let mut chunks = Vec::new(); + for i in 0..100u32 { + let org = if i % 5 == 0 { "acme" } else { "widgets" }; + chunks.push(ingest_with(org, i)); + } + manager + .ingest("filter-search", chunks, &embed) + .await + .unwrap(); + + // Search with filter org_id=acme. ALL hits must come from acme. + let mut filters = HashMap::new(); + filters.insert( + "org_id".to_string(), + FilterValue::Exact(MetadataValue::String("acme".into())), + ); + let req = SearchRequest { + query: "chunk".to_string(), + mode: "hybrid".to_string(), + vector_space: None, + top_k: 10, + query_vector: Some(pseudo_vec(99_999)), + filters, + score_weights: None, + recency: None, + recency_preset: None, + recency_field: None, + boosts: Vec::new(), + relationship_boost: None, + explain: true, + include_relations: false, + relation_types: None, + relation_direction: RelationDirection::Outgoing, + min_seq: None, + }; + let (hits, total, _took_us, explain) = + manager.search("filter-search", &req, &embed).await.unwrap(); + + assert!(!hits.is_empty(), "search returned no hits"); + for (chunk, _, _, _, _) in &hits { + assert_eq!( + chunk.metadata.get("org_id"), + Some(&MetadataValue::String("acme".into())), + "all hits must satisfy the filter; got chunk {} with org_id {:?}", + chunk.id, + chunk.metadata.get("org_id") + ); + } + assert!( + total <= 20, + "no more than 20 hits possible at 20% selectivity" + ); + + // /explain should be populated. + let explain = explain.expect("explain plan requested but not returned"); + assert_eq!(explain.filter.eligible_count, 20); + assert_eq!(explain.filter.universe_count, 100); + assert!((explain.filter.selectivity - 0.20).abs() < 1e-9); + assert!( + matches!(explain.ann.engine.as_str(), "hnsw" | "brute_force"), + "ann engine reported as {}", + explain.ann.engine + ); + + let _ = std::fs::remove_dir_all(&data_dir); +} + +#[tokio::test] +async fn explain_absent_when_not_requested() { + let data_dir = unique_data_dir(); + std::fs::create_dir_all(&data_dir).unwrap(); + let embed = embed_state(); + let manager = CollectionManager::new(&data_dir).await.unwrap(); + manager + .create_collection("no-explain", None, Some(4), None) + .await + .unwrap(); + manager + .ingest("no-explain", vec![ingest_with("acme", 0)], &embed) + .await + .unwrap(); + + let req = SearchRequest { + query: "chunk".to_string(), + mode: "semantic".to_string(), + vector_space: None, + top_k: 1, + query_vector: Some(pseudo_vec(42)), + filters: HashMap::new(), + score_weights: None, + recency: None, + recency_preset: None, + recency_field: None, + boosts: Vec::new(), + relationship_boost: None, + explain: false, + include_relations: false, + relation_types: None, + relation_direction: RelationDirection::Outgoing, + min_seq: None, + }; + let (_hits, _total, _took, explain) = manager.search("no-explain", &req, &embed).await.unwrap(); + assert!(explain.is_none(), "explain must be None when not requested"); + + let _ = std::fs::remove_dir_all(&data_dir); +} + +#[tokio::test] +async fn relations_crud_and_search_enrichment() { + let data_dir = unique_data_dir(); + std::fs::create_dir_all(&data_dir).unwrap(); + let embed = embed_state(); + let manager = CollectionManager::new(&data_dir).await.unwrap(); + manager + .create_collection("rel-search", None, Some(4), None) + .await + .unwrap(); + + // Ingest 5 chunks -> ids 0..5 in order. + let chunks: Vec<_> = (0..5u32).map(|i| ingest_with("acme", i)).collect(); + manager.ingest("rel-search", chunks, &embed).await.unwrap(); + + // Create relations: 0 --cites--> 1, 0 --cites--> 2, 3 --supersedes--> 0. + let created = manager + .create_relations( + "rel-search", + vec![ + CreateRelation { + source_chunk_id: 0, + target_chunk_id: 1, + target_document_id: None, + relation_type: "cites".into(), + metadata: HashMap::new(), + }, + CreateRelation { + source_chunk_id: 0, + target_chunk_id: 2, + target_document_id: None, + relation_type: "cites".into(), + metadata: HashMap::new(), + }, + CreateRelation { + source_chunk_id: 3, + target_chunk_id: 0, + target_document_id: None, + relation_type: "supersedes".into(), + metadata: HashMap::new(), + }, + ], + ) + .await + .unwrap(); + assert_eq!(created.len(), 3); + assert!(created.iter().all(|r| !r.relation_id.is_empty())); + assert!(created.iter().all(|r| r.target_status == "found")); + + // Self-relation is rejected. + let bad = manager + .create_relations( + "rel-search", + vec![CreateRelation { + source_chunk_id: 1, + target_chunk_id: 1, + target_document_id: None, + relation_type: "cites".into(), + metadata: HashMap::new(), + }], + ) + .await; + assert!(bad.is_err()); + + // Direction filters. + let out = manager + .get_chunk_relations("rel-search", 0, RelationDirection::Outgoing, None) + .await + .unwrap(); + assert_eq!(out.len(), 2); + let inc = manager + .get_chunk_relations("rel-search", 0, RelationDirection::Incoming, None) + .await + .unwrap(); + assert_eq!(inc.len(), 1); + assert_eq!(inc[0].relation_type, "supersedes"); + + // Type filter. + let cites = manager + .get_chunk_relations( + "rel-search", + 0, + RelationDirection::Both, + Some(&["cites".to_string()]), + ) + .await + .unwrap(); + assert_eq!(cites.len(), 2); + + let base_req = |include: bool| SearchRequest { + query: "chunk".to_string(), + mode: "semantic".to_string(), + vector_space: None, + top_k: 10, + query_vector: Some(pseudo_vec(7)), + filters: HashMap::new(), + score_weights: None, + recency: None, + recency_preset: None, + recency_field: None, + boosts: Vec::new(), + relationship_boost: None, + explain: false, + include_relations: include, + relation_types: None, + relation_direction: RelationDirection::Outgoing, + min_seq: None, + }; + + // Without include_relations -> hits carry None. + let (hits_off, _, _, _) = manager + .search("rel-search", &base_req(false), &embed) + .await + .unwrap(); + assert!(hits_off.iter().all(|(_, _, _, _, rels)| rels.is_none())); + + // With include_relations -> chunk 0's hit carries its 2 outgoing cites. + let (hits_on, _, _, _) = manager + .search("rel-search", &base_req(true), &embed) + .await + .unwrap(); + let chunk0 = hits_on + .iter() + .find(|(c, _, _, _, _)| c.id == 0) + .expect("chunk 0 in results"); + let rels = chunk0.4.as_ref().expect("Some(relations) when requested"); + assert_eq!(rels.len(), 2, "chunk 0 has 2 outgoing cites"); + assert!(rels.iter().all(|r| r.relation_type == "cites")); + assert!(rels.iter().all(|r| r.target_status == "found")); + + // Delete one relation; outgoing from 0 drops to 1. + let rid = &created[0].relation_id; + assert!(manager.delete_relation("rel-search", rid).await.unwrap()); + let out2 = manager + .get_chunk_relations("rel-search", 0, RelationDirection::Outgoing, None) + .await + .unwrap(); + assert_eq!(out2.len(), 1); + + let _ = std::fs::remove_dir_all(&data_dir); +} + +#[tokio::test] +async fn delete_removes_from_search_and_survives_restart() { + let data_dir = unique_data_dir(); + std::fs::create_dir_all(&data_dir).unwrap(); + let embed = embed_state(); + + { + let manager = CollectionManager::new(&data_dir).await.unwrap(); + manager + .create_collection("del", None, Some(4), None) + .await + .unwrap(); + // Ingest 10 chunks (ids 0..10), org=acme. + let chunks: Vec<_> = (0..10u32).map(|i| ingest_with("acme", i)).collect(); + manager.ingest("del", chunks, &embed).await.unwrap(); + + // Delete chunk id 3 by id. + let (n, _) = manager.delete_chunks("del", &[3]).await.unwrap(); + assert_eq!(n, 1); + // Re-deleting is a no-op. + assert_eq!(manager.delete_chunks("del", &[3]).await.unwrap().0, 0); + + // A search must never return the deleted id. + let req = SearchRequest { + query: "chunk".to_string(), + mode: "semantic".to_string(), + vector_space: None, + top_k: 20, + query_vector: Some(pseudo_vec(4)), + filters: HashMap::new(), + score_weights: None, + recency: None, + recency_preset: None, + recency_field: None, + boosts: Vec::new(), + relationship_boost: None, + explain: false, + include_relations: false, + relation_types: None, + relation_direction: RelationDirection::Outgoing, + min_seq: None, + }; + let (hits, _, _, _) = manager.search("del", &req, &embed).await.unwrap(); + assert!( + hits.iter().all(|(c, _, _, _, _)| c.id != 3), + "deleted chunk must not appear in results" + ); + + // Delete-by-filter: delete everything with file_id f5 (chunk 5). + let mut filters = HashMap::new(); + filters.insert( + "file_id".to_string(), + FilterValue::Exact(MetadataValue::String("f5".into())), + ); + // ingest_with doesn't set file_id in metadata, so use a metadata field. + // org_id=acme matches all remaining -> delete the rest via a scan. + let mut org_filter = HashMap::new(); + org_filter.insert( + "org_id".to_string(), + FilterValue::Exact(MetadataValue::String("acme".into())), + ); + let deleted = manager.delete_by_filter("del", &org_filter).await.unwrap(); + // 10 ingested - 1 already deleted (id 3) = 9 remaining deleted now. + assert_eq!(deleted.0, 9); + let _ = filters; + } + + // Restart: tombstones must persist. Reopen the manager over the same dir. + { + let manager = CollectionManager::new(&data_dir).await.unwrap(); + let req = SearchRequest { + query: "chunk".to_string(), + mode: "semantic".to_string(), + vector_space: None, + top_k: 20, + query_vector: Some(pseudo_vec(4)), + filters: HashMap::new(), + score_weights: None, + recency: None, + recency_preset: None, + recency_field: None, + boosts: Vec::new(), + relationship_boost: None, + explain: false, + include_relations: false, + relation_types: None, + relation_direction: RelationDirection::Outgoing, + min_seq: None, + }; + let (hits, _, _, _) = manager.search("del", &req, &embed).await.unwrap(); + assert!( + hits.is_empty(), + "all chunks deleted; none should survive restart, got {}", + hits.len() + ); + } + + let _ = std::fs::remove_dir_all(&data_dir); +} + +// F6 coverage: deleting a chunk that participates in relations must prune +// those edges (both endpoints), not leave dangling references. +#[tokio::test] +async fn delete_chunk_prunes_its_relations() { + let data_dir = unique_data_dir(); + std::fs::create_dir_all(&data_dir).unwrap(); + let embed = embed_state(); + let manager = CollectionManager::new(&data_dir).await.unwrap(); + manager + .create_collection("delrel", None, Some(4), None) + .await + .unwrap(); + let chunks: Vec<_> = (0..3u32).map(|i| ingest_with("acme", i)).collect(); + manager.ingest("delrel", chunks, &embed).await.unwrap(); + + // 0 -> 1, 2 -> 0 (chunk 0 is both a source and a target). + manager + .create_relations( + "delrel", + vec![ + CreateRelation { + source_chunk_id: 0, + target_chunk_id: 1, + target_document_id: None, + relation_type: "cites".into(), + metadata: HashMap::new(), + }, + CreateRelation { + source_chunk_id: 2, + target_chunk_id: 0, + target_document_id: None, + relation_type: "cites".into(), + metadata: HashMap::new(), + }, + ], + ) + .await + .unwrap(); + + // Delete chunk 0 — both edges (as source and as target) must be pruned. + manager.delete_chunks("delrel", &[0]).await.unwrap(); + + assert!( + manager + .get_chunk_relations("delrel", 0, RelationDirection::Both, None) + .await + .unwrap() + .is_empty(), + "deleted chunk's own edges gone" + ); + // Chunk 2's outgoing edge (to deleted 0) must also be gone. + assert!( + manager + .get_chunk_relations("delrel", 2, RelationDirection::Outgoing, None) + .await + .unwrap() + .is_empty(), + "edge pointing AT the deleted chunk must be pruned" + ); + let _ = std::fs::remove_dir_all(&data_dir); +} + +// A stale tombstone must not suppress a NEWLY-ingested chunk. Since next_id +// is a monotonic high-water mark, re-ingest gets a fresh id that was never +// tombstoned, so it's fully searchable. +#[tokio::test] +async fn delete_then_reingest_new_chunk_is_searchable() { + let data_dir = unique_data_dir(); + std::fs::create_dir_all(&data_dir).unwrap(); + let embed = embed_state(); + let manager = CollectionManager::new(&data_dir).await.unwrap(); + manager + .create_collection("reing", None, Some(4), None) + .await + .unwrap(); + manager + .ingest("reing", vec![ingest_with("acme", 0)], &embed) + .await + .unwrap(); + manager.delete_chunks("reing", &[0]).await.unwrap(); + + // Re-ingest: gets id 1 (next_id advanced), NOT the tombstoned id 0. + manager + .ingest("reing", vec![ingest_with("acme", 9)], &embed) + .await + .unwrap(); + + let req = SearchRequest { + query: "chunk".to_string(), + mode: "semantic".to_string(), + vector_space: None, + top_k: 10, + query_vector: Some(pseudo_vec(10)), + filters: HashMap::new(), + score_weights: None, + recency: None, + recency_preset: None, + recency_field: None, + boosts: Vec::new(), + relationship_boost: None, + explain: false, + include_relations: false, + relation_types: None, + relation_direction: RelationDirection::Outgoing, + min_seq: None, + }; + let (hits, _, _, _) = manager.search("reing", &req, &embed).await.unwrap(); + assert_eq!(hits.len(), 1, "the re-ingested chunk must be searchable"); + assert_eq!(hits[0].0.id, 1, "re-ingest got a fresh (untombstoned) id"); + let _ = std::fs::remove_dir_all(&data_dir); +} diff --git a/crates/compass/src/collections/mod.rs b/crates/compass/src/collections/mod.rs index 5e50852..c84736d 100644 --- a/crates/compass/src/collections/mod.rs +++ b/crates/compass/src/collections/mod.rs @@ -18,6 +18,7 @@ pub mod store; use crate::embed::EmbedState; use crate::models::*; use crate::scoring::{self, ScoredCandidate}; +use crate::search::chunk_cache::ChunkCache; use crate::search::chunk_store::ChunkStore; use crate::search::filter_index::{selectivity, FilterIndex}; use crate::search::filter_pushdown::FilterExpr; @@ -55,21 +56,72 @@ pub(crate) fn validate_name_segment( } /// A loaded collection with all its search indices in memory. +/// Which manifest seqs this node has applied to its local indexes. +/// +/// `contiguous` is the count of contiguously-applied seqs (fragments +/// `0..contiguous` are reflected locally); `out_of_band` holds seqs this node +/// applied AHEAD of the contiguous frontier — its own appends land locally at +/// commit time while earlier REMOTE fragments may still be unapplied, so a +/// single watermark would silently skip those remote fragments forever. The +/// refresher advances `contiguous` in seq order, draining `out_of_band`. +#[derive(Debug, Default, Clone)] +struct SeqTracker { + contiguous: u64, + out_of_band: std::collections::BTreeSet, +} + +impl SeqTracker { + fn starting_at(contiguous: u64) -> Self { + Self { + contiguous, + ..Default::default() + } + } + + /// Has this seq been applied locally (either side of the frontier)? + fn covers(&self, seq: u64) -> bool { + seq < self.contiguous || self.out_of_band.contains(&seq) + } + + /// Record a locally-applied seq and advance the contiguous frontier. + fn mark(&mut self, seq: u64) { + if seq < self.contiguous { + return; + } + self.out_of_band.insert(seq); + while self.out_of_band.remove(&self.contiguous) { + self.contiguous += 1; + } + } +} + struct LoadedCollection { metadata: Collection, + /// Per-space count of ingest batches since the HNSW index was last saved + /// (saving rewrites the whole index file — O(index) per batch was a scale + /// wall). A stale on-disk index is detected at load (size < keymap) and + /// rebuilt from the mmap file. u32::MAX means "no mutable in-RAM index". + hnsw_unsaved: HashMap, + /// LRU stamp for lazy-attach eviction (process-monotonic tick). + last_used: std::sync::atomic::AtomicU64, + /// Manifest seqs applied to this node's local indexes (see [`SeqTracker`]). + applied: SeqTracker, + /// Cloud-mode id pool: ranges CAS-leased from `{ns}/id-alloc`. In cloud + /// mode ids are ONLY taken from here (never from `next_id`, which becomes + /// a diagnostic high-water mark) so attached nodes and stateless writers + /// can never mint colliding ids. + id_pool: std::collections::VecDeque>, fts: FtsState, /// Named vector spaces, each with its own USearch HNSW index. /// Arc-wrapped so search can clone cheaply and run in spawn_blocking. vector_spaces: HashMap>, /// Document relationships (parent-child + sibling groups) relationships: RelationshipStore, - /// All chunks in memory, keyed by chunk ID for O(1) retrieval. This is a - /// hot cache; the disk source of truth is `chunk_store`. Populated on - /// startup from `chunk_store.for_each` and kept in sync on every ingest. - chunks: HashMap, - /// Disk-backed chunk metadata. Every ingest writes through to this redb - /// database so chunks survive process restarts and crashes. - chunk_store: ChunkStore, + /// Bounded read-through cache over the disk-backed chunk store. Chunks + /// are NOT held wholesale in RAM anymore — serving memory is O(cache + /// budget), not O(collection). Existence checks go through the filter + /// index universe (live ids as a treemap). + chunk_store: ChunkCache, /// Disk-backed typed many-to-many chunk relations. Source of truth on disk; /// read on demand at search time (never rehydrated into RAM). relation_store: RelationStore, @@ -84,7 +136,8 @@ struct LoadedCollection { /// and HNSW indexes on every ingest batch, and on load from the rehydrated /// chunks. Powers filter-aware ANN: queries with `filters={...}` compile /// to a `FilterExpr`, resolve to an eligible bitmap, and route through - /// USearch's `filtered_search`. Planned follow-up. + /// USearch's `filtered_search`. Also the live-id universe for facets and + /// delete-by-filter. filter_index: FilterIndex, } @@ -104,11 +157,49 @@ pub struct CollectionManager { /// background compaction so concurrent triggers don't each write (and, on /// CAS loss, leak) a full segment. compacting: Arc>>, + /// Node role (COMPASS_ROLE). Writer = durable-append-only ingest with no + /// local indexes; Full = today's behavior. Cloud mode only. + role: NodeRole, + /// Stateless-writer id pools, keyed by namespace (attached collections + /// pool on `LoadedCollection.id_pool` instead). + writer_pools: + tokio::sync::Mutex>>>, + /// Cache of bucket collection configs for stateless-writer validation. + bucket_configs: tokio::sync::RwLock>, + /// Lazy attach (COMPASS_LAZY_ATTACH): namespaces discovered in the bucket + /// but not yet attached. Attach happens on first request. + registered: tokio::sync::RwLock>, + /// Per-namespace attach mutexes: a request stampede on a cold namespace + /// rebuilds ONCE, without holding the global collections lock. + attach_locks: tokio::sync::Mutex>>>, + /// Lazy attach enabled (cloud mode + COMPASS_LAZY_ATTACH=true). + lazy_attach: bool, + /// LRU budget for attached collections (COMPASS_MAX_ATTACHED; 0 = unbounded). + max_attached: usize, +} + +/// What this node does. Parsed from `COMPASS_ROLE` (default `full`). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum NodeRole { + /// Serve reads and writes with full local indexes (default). + Full, + /// Durable-append-only writes; no local indexes, no read serving. + Writer, +} + +impl NodeRole { + fn from_env() -> Self { + match std::env::var("COMPASS_ROLE").as_deref() { + Ok("writer") => NodeRole::Writer, + _ => NodeRole::Full, + } + } } impl CollectionManager { /// Create a manager with local-disk storage (the default embedded mode). - /// Convenience wrapper used by tests and local-only callers. + /// Test-only convenience; `main.rs` goes through `new_with_storage_opts`. + #[cfg(test)] pub async fn new( data_dir: &Path, ) -> Result, Box> { @@ -122,6 +213,50 @@ impl CollectionManager { pub async fn new_with_storage( data_dir: &Path, storage: Arc, + ) -> Result, Box> { + let role = NodeRole::from_env(); + Self::new_with_storage_role(data_dir, storage, role).await + } + + /// Like [`new_with_storage`] with an explicit node role (used by tests; + /// `new_with_storage` parses `COMPASS_ROLE`). + pub async fn new_with_storage_role( + data_dir: &Path, + storage: Arc, + role: NodeRole, + ) -> Result, Box> { + let lazy = std::env::var("COMPASS_LAZY_ATTACH") + .map(|v| v == "true" || v == "1") + .unwrap_or(false); + let max_attached = std::env::var("COMPASS_MAX_ATTACHED") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(0); + let refresh_interval_secs = std::env::var("COMPASS_REFRESH_INTERVAL") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(5); + Self::new_with_storage_opts( + data_dir, + storage, + role, + lazy, + max_attached, + refresh_interval_secs, + ) + .await + } + + /// Fully-explicit constructor (role + lazy-attach + LRU budget), used by + /// tests to avoid process-global env races and by callers embedding + /// Compass as a library. + pub async fn new_with_storage_opts( + data_dir: &Path, + storage: Arc, + role: NodeRole, + lazy_attach: bool, + max_attached: usize, + refresh_interval_secs: u64, ) -> Result, Box> { std::fs::create_dir_all(data_dir)?; @@ -129,6 +264,12 @@ impl CollectionManager { rebuild::cleanup_stale_rebuilds(data_dir); let cloud_mode = storage.backend_name() != "local-disk"; + // The writer role is meaningless without a shared bucket; force Full + // in local mode so a stray COMPASS_ROLE can't disable local serving. + let role = if cloud_mode { role } else { NodeRole::Full }; + if role == NodeRole::Writer { + tracing::info!("Node role: writer (durable-append-only; no read serving)"); + } let manager = Arc::new(Self { data_dir: data_dir.to_path_buf(), collections: RwLock::new(HashMap::new()), @@ -136,8 +277,22 @@ impl CollectionManager { storage, cloud_mode, compacting: Arc::new(std::sync::Mutex::new(std::collections::HashSet::new())), + role, + writer_pools: tokio::sync::Mutex::new(HashMap::new()), + bucket_configs: tokio::sync::RwLock::new(HashMap::new()), + registered: tokio::sync::RwLock::new(std::collections::HashSet::new()), + attach_locks: tokio::sync::Mutex::new(HashMap::new()), + lazy_attach: cloud_mode && lazy_attach, + max_attached, }); + // Writer role: no local collections, no recovery — the node serves + // durable appends only, validated against bucket configs. Boot is + // instant regardless of how much data lives in the bucket. + if manager.role == NodeRole::Writer { + return Ok(manager); + } + // Load existing collections from local disk. let names = store::list_collection_names(data_dir)?; for name in &names { @@ -166,6 +321,12 @@ impl CollectionManager { if already { continue; } + if manager.lazy_attach { + // Lazy mode: register only — attach on first + // request. Boot cost is O(namespaces), not O(data). + manager.registered.write().await.insert(ns.clone()); + continue; + } match manager.rebuild_collection_from_storage(ns).await { Ok(n) => { recovered += 1; @@ -183,11 +344,35 @@ impl CollectionManager { if recovered > 0 { tracing::info!("Recovered {} collection(s) from object storage", recovered); } + if manager.lazy_attach { + let n = manager.registered.read().await.len(); + if n > 0 { + tracing::info!("Registered {} collection(s) for lazy attach", n); + } + } } Err(e) => tracing::error!("Could not list collections from object storage: {}", e), } } + // Background manifest refresher: keeps this node's local indexes + // converged with fragments written by OTHER nodes (stateless writers, + // other serving nodes). COMPASS_REFRESH_INTERVAL seconds, default 5, + // 0 disables. Holds only a Weak — the task dies with the manager. + if cloud_mode { + let interval_secs = refresh_interval_secs; + if interval_secs > 0 { + let weak = Arc::downgrade(&manager); + tokio::spawn(async move { + loop { + tokio::time::sleep(std::time::Duration::from_secs(interval_secs)).await; + let Some(m) = weak.upgrade() else { break }; + m.refresh_all().await; + } + }); + } + } + Ok(manager) } @@ -201,10 +386,10 @@ impl CollectionManager { let vectors_dir = store::vectors_dir(&self.data_dir, name); // Open the Tantivy FTS index - let fts = if tantivy_dir.join("meta.json").exists() { + let mut fts = if tantivy_dir.join("meta.json").exists() { tantivy_fts::open_index(&tantivy_dir)? } else { - tantivy_fts::build_index(&tantivy_dir, &[], 0)? + tantivy_fts::build_index(&tantivy_dir, &[])? }; // Load each named vector space from disk @@ -235,7 +420,6 @@ impl CollectionManager { key_to_chunk_id: Vec::new(), mmap_vectors: None, vectors: Vec::new(), - dims: space_config.dims, }), ); } @@ -254,16 +438,26 @@ impl CollectionManager { if let Some(parent) = chunks_db.parent() { std::fs::create_dir_all(parent)?; } - let chunk_store = ChunkStore::open(&chunks_db)?; - let mut chunks: HashMap = HashMap::new(); + let chunk_store = ChunkCache::new(ChunkStore::open(&chunks_db)?); let mut max_seen_id: u64 = 0; + let mut rehydrated_count: usize = 0; + let mut filter_index = FilterIndex::new(); + let tombstones_vec = chunk_store.load_tombstones()?; + let tombstones: std::collections::HashSet = tombstones_vec.into_iter().collect(); + let mut facet_rebuild = tantivy_fts::FacetBitsets::default(); chunk_store.for_each(|id, chunk| { if id >= max_seen_id { max_seen_id = id; } - chunks.insert(id, chunk); + rehydrated_count += 1; + if !tombstones.contains(&id) { + filter_index.insert(id, &filter_meta(&chunk)); + // Facets were EMPTY after every restart (open_index returns + // none and nothing rebuilt them) — rebuild here, same pass. + facet_rebuild.insert_chunk(&chunk); + } })?; - let rehydrated_count = chunks.len(); + fts.facet_bitsets = facet_rebuild; // next_id is a MONOTONIC high-water mark that must never regress or reuse // an id. Take the max of: the persisted metadata.next_id (survives even // when the local chunk store is empty on a cold restart), and one past @@ -278,18 +472,21 @@ impl CollectionManager { let next_id = metadata.next_id.max(from_disk).max(metadata.chunk_count); let chunk_count = metadata.chunk_count; - let tombstones: std::collections::HashSet = - chunk_store.load_tombstones()?.into_iter().collect(); - let filter_index = build_filter_index_from_chunks(&chunks, &tombstones); let relations_db = store::relations_db_path(&self.data_dir, name); let relation_store = RelationStore::open(&relations_db)?; let loaded = LoadedCollection { + id_pool: Default::default(), + hnsw_unsaved: HashMap::new(), + last_used: std::sync::atomic::AtomicU64::new(next_lru_tick()), + // Persistent-disk restart: local indexes reflect fragments + // 0..applied_seq (persisted on every apply); the refresher applies + // the delta instead of a full rebuild. + applied: SeqTracker::starting_at(metadata.applied_seq), next_id, metadata, fts, vector_spaces, relationships, - chunks, chunk_store, relation_store, tombstones, @@ -320,97 +517,207 @@ impl CollectionManager { embedding_dims: Option, config: Option, ) -> Result> { + if self.role == NodeRole::Writer { + return Err( + "this node runs in writer role; create collections via a serving node".into(), + ); + } validate_name_segment(name, "Collection")?; - let mut collections = self.collections.write().await; - if collections.contains_key(name) { - return Err(format!("Collection '{}' already exists", name).into()); - } + let collection = { + let mut collections = self.collections.write().await; + if collections.contains_key(name) { + return Err(format!("Collection '{}' already exists", name).into()); + } - // Build vector spaces config: use explicit spaces, or create a "default" space - let spaces = vector_spaces.unwrap_or_else(|| { - let dims = embedding_dims.unwrap_or(384); - let mut m = HashMap::new(); - m.insert( - "default".to_string(), - VectorSpaceConfig { - dims, - model: "bge-small-en-v1.5".to_string(), - status: "active".to_string(), - }, - ); - m - }); + // Build vector spaces config: use explicit spaces, or create a "default" space + let spaces = vector_spaces.unwrap_or_else(|| { + let dims = embedding_dims.unwrap_or(384); + let mut m = HashMap::new(); + m.insert( + "default".to_string(), + VectorSpaceConfig { + dims, + model: "bge-small-en-v1.5".to_string(), + status: "active".to_string(), + }, + ); + m + }); + + let default_space = spaces.keys().next().cloned(); + let dims = spaces.values().next().map(|s| s.dims).unwrap_or(384); + + let collection = Collection { + name: name.to_string(), + created_at: Utc::now(), + vector_spaces: spaces, + default_vector_space: default_space, + embedding_dims: dims, + chunk_count: 0, + next_id: 0, + config: config.unwrap_or_default(), + applied_seq: 0, + }; - let default_space = spaces.keys().next().cloned(); - let dims = spaces.values().next().map(|s| s.dims).unwrap_or(384); + store::save_metadata(&self.data_dir, &collection)?; - let collection = Collection { - name: name.to_string(), - created_at: Utc::now(), - vector_spaces: spaces, - default_vector_space: default_space, - embedding_dims: dims, - chunk_count: 0, - next_id: 0, - config: config.unwrap_or_default(), - }; + // Build empty FTS index + let tantivy_dir = store::tantivy_dir(&self.data_dir, name); + let fts = tantivy_fts::build_index(&tantivy_dir, &[])?; - store::save_metadata(&self.data_dir, &collection)?; + // Create empty vector spaces + let mut vs_map = HashMap::new(); + for sname in collection.vector_spaces.keys() { + vs_map.insert( + sname.clone(), + Arc::new(VectorState { + index: None, + key_to_chunk_id: Vec::new(), + mmap_vectors: None, + vectors: Vec::new(), + }), + ); + } - // Build empty FTS index - let tantivy_dir = store::tantivy_dir(&self.data_dir, name); - let fts = tantivy_fts::build_index(&tantivy_dir, &[], 0)?; + // Open the disk-backed chunk store for the new collection. Empty + // database file is created at //chunks.redb. + let chunks_db = store::chunks_db_path(&self.data_dir, name); + if let Some(parent) = chunks_db.parent() { + std::fs::create_dir_all(parent)?; + } + let chunk_store = ChunkCache::new(ChunkStore::open(&chunks_db)?); + let relations_db = store::relations_db_path(&self.data_dir, name); + let relation_store = RelationStore::open(&relations_db)?; + + let loaded = LoadedCollection { + id_pool: Default::default(), + hnsw_unsaved: HashMap::new(), + last_used: std::sync::atomic::AtomicU64::new(next_lru_tick()), + applied: SeqTracker::default(), + metadata: collection.clone(), + fts, + vector_spaces: vs_map, + relationships: RelationshipStore::new(), + chunk_store, + relation_store, + tombstones: std::collections::HashSet::new(), + next_id: 0, + filter_index: FilterIndex::new(), + }; - // Create empty vector spaces - let mut vs_map = HashMap::new(); - for (sname, sconfig) in &collection.vector_spaces { - vs_map.insert( - sname.clone(), - Arc::new(VectorState { - index: None, - key_to_chunk_id: Vec::new(), - mmap_vectors: None, - vectors: Vec::new(), - dims: sconfig.dims, - }), - ); - } + collections.insert(name.to_string(), loaded); + collection + }; // write lock released — never hold it across S3 round-trips. - // Open the disk-backed chunk store for the new collection. Empty - // database file is created at //chunks.redb. - let chunks_db = store::chunks_db_path(&self.data_dir, name); - if let Some(parent) = chunks_db.parent() { - std::fs::create_dir_all(parent)?; + // Cloud mode: make the collection exist DURABLY in the bucket — + // create-only config object + empty manifest — so a zero-ingest + // collection is discoverable from a fresh disk and stateless writers + // can validate against its config. On any bucket failure, roll the + // local creation back so local and bucket state agree (= absent). + if self.cloud_mode { + let rollback_local = || async { + self.collections.write().await.remove(name); + let _ = store::delete_collection_data(&self.data_dir, name); + }; + let bucket_cfg = cloud::BucketConfig::from_collection(&collection); + match cloud::write_bucket_config_if_absent(self.storage.as_ref(), name, &bucket_cfg) + .await + { + Ok(()) => {} + Err(crate::storage::StorageError::AlreadyExists(_)) => { + rollback_local().await; + return Err( + format!("Collection '{}' already exists in object storage", name).into(), + ); + } + Err(e) => { + rollback_local().await; + return Err(format!("bucket config write failed: {e}").into()); + } + } + match crate::storage::lsm::init_namespace(self.storage.as_ref(), name).await { + Ok(()) => { + // Fresh namespace: seed the id allocator at 0 so every + // ingest path (attached or stateless) can claim blocks. + if let Err(e) = + crate::storage::id_alloc::seed(self.storage.as_ref(), name, 0).await + { + let _ = crate::storage::lsm::delete_namespace(self.storage.as_ref(), name) + .await; + rollback_local().await; + return Err(format!("id allocator seed failed: {e}").into()); + } + } + Err(crate::storage::StorageError::AlreadyExists(_)) => { + // Data exists in the bucket without a config (pre-v0.4 + // namespace): this create collides with real data. Remove + // the config we just wrote and refuse. + let _ = self.storage.delete(&cloud::config_key(name)).await; + rollback_local().await; + return Err( + format!("namespace '{}' already has data in object storage", name).into(), + ); + } + Err(e) => { + let _ = self.storage.delete(&cloud::config_key(name)).await; + rollback_local().await; + return Err(format!("bucket manifest init failed: {e}").into()); + } + } } - let chunk_store = ChunkStore::open(&chunks_db)?; - let relations_db = store::relations_db_path(&self.data_dir, name); - let relation_store = RelationStore::open(&relations_db)?; - - let loaded = LoadedCollection { - metadata: collection.clone(), - fts, - vector_spaces: vs_map, - relationships: RelationshipStore::new(), - chunks: HashMap::new(), - chunk_store, - relation_store, - tombstones: std::collections::HashSet::new(), - next_id: 0, - filter_index: FilterIndex::new(), - }; - collections.insert(name.to_string(), loaded); tracing::info!("Created collection '{}'", name); Ok(collection) } - pub async fn list_collections(&self) -> Vec { + /// Metadata of ATTACHED collections only — no bucket round-trips (used + /// by /metrics; lazy-registered namespaces are intentionally excluded). + pub async fn attached_collections(&self) -> Vec { let collections = self.collections.read().await; collections.values().map(|c| c.metadata.clone()).collect() } + pub async fn list_collections(&self) -> Vec { + let mut out: Vec = { + let collections = self.collections.read().await; + collections.values().map(|c| c.metadata.clone()).collect() + }; + if self.lazy_attach { + let attached: std::collections::HashSet = + out.iter().map(|c| c.name.clone()).collect(); + let names: Vec = { + let reg = self.registered.read().await; + reg.iter() + .filter(|n| !attached.contains(*n)) + .cloned() + .collect() + }; + for name in names { + if let Ok(Some(cfg)) = cloud::read_bucket_config(self.storage.as_ref(), &name).await + { + out.push(Collection { + name: cfg.name, + created_at: cfg.created_at, + vector_spaces: cfg.vector_spaces, + default_vector_space: cfg.default_vector_space, + embedding_dims: cfg.embedding_dims, + // Live counts are known only once attached. + chunk_count: 0, + next_id: 0, + config: cfg.config, + applied_seq: 0, + }); + } + } + } + out + } + pub async fn get_collection(&self, name: &str) -> Option { + // Lazy mode: a registered-but-unattached collection attaches on its + // first request — including a metadata read. + let _ = self.ensure_attached(name).await; let collections = self.collections.read().await; collections.get(name).map(|c| c.metadata.clone()) } @@ -419,11 +726,34 @@ impl CollectionManager { &self, name: &str, ) -> Result<(), Box> { - let mut collections = self.collections.write().await; - if collections.remove(name).is_none() { - return Err(format!("Collection '{}' not found", name).into()); + // Lazy mode: the collection may be registered-but-unattached (or LRU + // evicted) — deleting it must still purge the bucket. + let attached = { + let mut collections = self.collections.write().await; + collections.remove(name).is_some() + }; // write lock released BEFORE any filesystem/S3 work. + let registered = self.registered.write().await.remove(name); + if !attached && !registered { + // Not known locally; in cloud mode it may still exist in the bucket + // (created by another node). + let in_bucket = self.cloud_mode + && cloud::read_bucket_config(self.storage.as_ref(), name) + .await + .ok() + .flatten() + .is_some(); + if !in_bucket { + return Err(not_found(format_args!("Collection \'{}\' not found", name))); + } } - store::delete_collection_data(&self.data_dir, name)?; + if attached { + store::delete_collection_data(&self.data_dir, name)?; + } + // Purge every node-local cache tied to the namespace so a later + // recreate can't consume stale pooled ids or stale configs. + self.bucket_configs.write().await.remove(name); + self.writer_pools.lock().await.remove(name); + self.attach_locks.lock().await.remove(name); // Cloud mode: also purge the collection's objects from storage, so it // can't be resurrected from S3 on a later cold start (and so a racing // ingest's orphan fragment doesn't bring a "deleted" collection back). @@ -453,36 +783,69 @@ impl CollectionManager { // character set as collection names. validate_name_segment(space_name, "Vector space")?; - let mut collections = self.collections.write().await; - let loaded = collections - .get_mut(collection_name) - .ok_or_else(|| format!("Collection '{}' not found", collection_name))?; - - if loaded.metadata.vector_spaces.contains_key(space_name) { - return Err(format!("Vector space '{}' already exists", space_name).into()); + if self.role == NodeRole::Writer { + return Err( + "this node runs in writer role; manage vector spaces via a serving node".into(), + ); + } + self.ensure_attached(collection_name).await?; + // Phase 1 (short read lock): preconditions only. + { + let collections = self.collections.read().await; + let loaded = collections.get(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; + if loaded.metadata.vector_spaces.contains_key(space_name) { + return Err(format!("Vector space '{}' already exists", space_name).into()); + } } - loaded.metadata.vector_spaces.insert( - space_name.to_string(), - VectorSpaceConfig { - dims, - model: model.to_string(), - status: "building".to_string(), - }, - ); - - loaded.vector_spaces.insert( - space_name.to_string(), - Arc::new(VectorState { - index: None, - key_to_chunk_id: Vec::new(), - mmap_vectors: None, - vectors: Vec::new(), - dims, - }), - ); + // Phase 2 (NO lock): bucket-first CAS — the bucket config is the source + // of truth in cloud mode; the mutate closure revalidates against the + // LATEST doc so a racing add loses cleanly. + if self.cloud_mode { + cloud::cas_update_bucket_config(self.storage.as_ref(), collection_name, |cfg| { + if cfg.vector_spaces.contains_key(space_name) { + return Err(format!("Vector space '{}' already exists", space_name).into()); + } + cfg.vector_spaces.insert( + space_name.to_string(), + VectorSpaceConfig { + dims, + model: model.to_string(), + status: "building".to_string(), + }, + ); + Ok(()) + }) + .await?; + } - store::save_metadata(&self.data_dir, &loaded.metadata)?; + // Phase 3 (write lock): apply locally. + let mut collections = self.collections.write().await; + let loaded = collections.get_mut(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; + if !loaded.metadata.vector_spaces.contains_key(space_name) { + loaded.metadata.vector_spaces.insert( + space_name.to_string(), + VectorSpaceConfig { + dims, + model: model.to_string(), + status: "building".to_string(), + }, + ); + loaded.vector_spaces.insert( + space_name.to_string(), + Arc::new(VectorState { + index: None, + key_to_chunk_id: Vec::new(), + mmap_vectors: None, + vectors: Vec::new(), + }), + ); + store::save_metadata(&self.data_dir, &loaded.metadata)?; + } Ok(()) } @@ -496,16 +859,42 @@ impl CollectionManager { // `remove_file` calls below. validate_name_segment(space_name, "Vector space")?; - let mut collections = self.collections.write().await; - let loaded = collections - .get_mut(collection_name) - .ok_or_else(|| format!("Collection '{}' not found", collection_name))?; + if self.role == NodeRole::Writer { + return Err( + "this node runs in writer role; manage vector spaces via a serving node".into(), + ); + } + self.ensure_attached(collection_name).await?; + // Phase 1 (short read lock): preconditions. + { + let collections = self.collections.read().await; + let loaded = collections.get(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; + if loaded.metadata.default_vector_space.as_deref() == Some(space_name) { + return Err("Cannot delete the default vector space. Switch default first.".into()); + } + } - // Don't delete the default vector space - if loaded.metadata.default_vector_space.as_deref() == Some(space_name) { - return Err("Cannot delete the default vector space. Switch default first.".into()); + // Phase 2 (NO lock): bucket-first CAS, revalidating against the latest doc. + if self.cloud_mode { + cloud::cas_update_bucket_config(self.storage.as_ref(), collection_name, |cfg| { + if cfg.default_vector_space.as_deref() == Some(space_name) { + return Err( + "Cannot delete the default vector space. Switch default first.".into(), + ); + } + cfg.vector_spaces.remove(space_name); + Ok(()) + }) + .await?; } + // Phase 3 (write lock): apply locally. + let mut collections = self.collections.write().await; + let loaded = collections.get_mut(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; loaded.metadata.vector_spaces.remove(space_name); loaded.vector_spaces.remove(space_name); @@ -525,31 +914,80 @@ impl CollectionManager { collection_name: &str, space_name: &str, ) -> Result<(), Box> { - let mut collections = self.collections.write().await; - let loaded = collections - .get_mut(collection_name) - .ok_or_else(|| format!("Collection '{}' not found", collection_name))?; + if self.role == NodeRole::Writer { + return Err( + "this node runs in writer role; manage vector spaces via a serving node".into(), + ); + } + self.ensure_attached(collection_name).await?; + // Phase 1 (short read lock): preconditions. + { + let collections = self.collections.read().await; + let loaded = collections.get(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; + if !loaded.metadata.vector_spaces.contains_key(space_name) { + return Err(not_found(format_args!( + "Vector space \'{}\' not found", + space_name + ))); + } + } - if !loaded.metadata.vector_spaces.contains_key(space_name) { - return Err(format!("Vector space '{}' not found", space_name).into()); + // Phase 2 (NO lock): bucket-first CAS. + if self.cloud_mode { + cloud::cas_update_bucket_config(self.storage.as_ref(), collection_name, |cfg| { + if !cfg.vector_spaces.contains_key(space_name) { + return Err(not_found(format_args!( + "Vector space \'{}\' not found", + space_name + ))); + } + cfg.default_vector_space = Some(space_name.to_string()); + Ok(()) + }) + .await?; } + // Phase 3 (write lock): apply locally. + let mut collections = self.collections.write().await; + let loaded = collections.get_mut(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; loaded.metadata.default_vector_space = Some(space_name.to_string()); store::save_metadata(&self.data_dir, &loaded.metadata)?; Ok(()) } - /// Mark a vector space as active (called when rebuild completes). - #[allow(dead_code)] + /// Mark a vector space as active: flip the persisted status (bucket-first + /// CAS in cloud mode) and hot-load the rebuilt index into the serving + /// collection. Called by the rebuild job on completion. pub async fn mark_vector_space_active( &self, collection_name: &str, space_name: &str, ) -> Result<(), Box> { + if self.role == NodeRole::Writer { + return Err( + "this node runs in writer role; manage vector spaces via a serving node".into(), + ); + } + self.ensure_attached(collection_name).await?; + // Bucket-first status flip (NO lock during the CAS). + if self.cloud_mode { + cloud::cas_update_bucket_config(self.storage.as_ref(), collection_name, |cfg| { + if let Some(space) = cfg.vector_spaces.get_mut(space_name) { + space.status = "active".to_string(); + } + Ok(()) + }) + .await?; + } + let mut collections = self.collections.write().await; - let loaded = collections - .get_mut(collection_name) - .ok_or_else(|| format!("Collection '{}' not found", collection_name))?; + let loaded = collections.get_mut(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; if let Some(config) = loaded.metadata.vector_spaces.get_mut(space_name) { config.status = "active".to_string(); @@ -585,33 +1023,352 @@ impl CollectionManager { // ── Ingest ─────────────────────────────────────────────────────────── /// Ingest chunks with batch parent resolution, named embeddings, and relationships. - pub async fn ingest( + /// Claim an id block, migrating a pre-v0.4 namespace on first use: if the + /// allocator object is absent, seed it from the bucket-derived high-water + /// mark (create-only, race-safe — no new ids can be minted while the + /// allocator is absent because every cloud ingest path requires it). + async fn claim_ids_or_migrate( + &self, + ns: &str, + count: u64, + ) -> Result, Box> { + use crate::storage::id_alloc; + match id_alloc::claim(self.storage.as_ref(), ns, count).await { + Ok(r) => Ok(r), + Err(crate::storage::StorageError::NotFound(_)) => { + let (manifest, _) = + crate::storage::lsm::read_manifest(self.storage.as_ref(), ns).await?; + let mat = cloud::materialize(self.storage.as_ref(), ns, &manifest).await?; + let start = if mat.max_id > 0 || !mat.chunks.is_empty() { + mat.max_id + 1 + } else { + 0 + }; + id_alloc::seed(self.storage.as_ref(), ns, start).await?; + Ok(id_alloc::claim(self.storage.as_ref(), ns, count).await?) + } + Err(e) => Err(e.into()), + } + } + + /// Take `count` ids for an ATTACHED collection from its pooled blocks, + /// refilling via CAS with the collections lock RELEASED (never hold the + /// global lock across an S3 round-trip). Racing refills both push their + /// ranges — nothing leaks, no extra mutex. + async fn take_ids_cloud( + &self, + collection_name: &str, + count: usize, + ) -> Result, Box> { + loop { + { + let mut collections = self.collections.write().await; + let loaded = collections.get_mut(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; + let available: u64 = loaded.id_pool.iter().map(|r| r.end - r.start).sum(); + if available >= count as u64 { + let mut ids = Vec::with_capacity(count); + while ids.len() < count { + let front = loaded + .id_pool + .front_mut() + .expect("available >= count guarantees a range"); + ids.push(front.start); + front.start += 1; + if front.start == front.end { + loaded.id_pool.pop_front(); + } + } + return Ok(ids); + } + } // lock released before the S3 round-trip below. + let range = self + .claim_ids_or_migrate(collection_name, count as u64) + .await?; + let mut collections = self.collections.write().await; + match collections.get_mut(collection_name) { + Some(loaded) => loaded.id_pool.push_back(range), + // Collection deleted mid-claim: the block leaks (gaps are fine). + None => { + return Err(not_found(format_args!( + "Collection \'{}\' not found", + collection_name + ))) + } + } + } + } + + /// Bucket collection config, cached. `refresh` forces a re-fetch (used + /// once on validation failure, so a just-added vector space is seen + /// without restarting the writer). + async fn bucket_config( + &self, + ns: &str, + refresh: bool, + ) -> Result> { + if !refresh { + if let Some(cfg) = self.bucket_configs.read().await.get(ns) { + return Ok(cfg.clone()); + } + } + let cfg = cloud::read_bucket_config(self.storage.as_ref(), ns) + .await? + .ok_or_else(|| { + not_found(format_args!( + "Collection \'{}\' not found in object storage", + ns + )) + })?; + self.bucket_configs + .write() + .await + .insert(ns.to_string(), cfg.clone()); + Ok(cfg) + } + + /// Writer-role ingest: validate against the bucket config, claim ids from + /// the shared allocator, append ONE durable WAL fragment, return. No + /// collections lock, no local indexes — the batch becomes searchable on + /// serving nodes after their manifest refresh (or attach). + async fn ingest_stateless( &self, collection_name: &str, ingest_chunks: Vec, embed_state: &EmbedState, - ) -> Result<(usize, HashMap), Box> { - let mut collections = self.collections.write().await; - let loaded = collections - .get_mut(collection_name) - .ok_or_else(|| format!("Collection '{}' not found", collection_name))?; - + ) -> Result<(usize, HashMap, Option), Box> + { + validate_name_segment(collection_name, "Collection")?; let count = ingest_chunks.len(); + if count == 0 { + return Ok((0, HashMap::new(), None)); + } + let cfg = self.bucket_config(collection_name, false).await?; - // Phase 1: Assign IDs and build client_id -> chunk_id map - let mut client_id_map: HashMap = HashMap::new(); - let mut assigned_ids: Vec = Vec::with_capacity(count); + // Ids from the writer-side pool (same allocator as attached nodes). + // The pool mutex is NEVER held across the S3 claim: drain what's + // available, release, claim, push, repeat. Ids already drained are + // kept across iterations (a failed later claim leaks them — fine). + let mut ids: Vec = Vec::with_capacity(count); + loop { + { + let mut pools = self.writer_pools.lock().await; + let pool = pools.entry(collection_name.to_string()).or_default(); + while ids.len() < count { + let Some(front) = pool.front_mut() else { break }; + if front.start < front.end { + ids.push(front.start); + front.start += 1; + } + if front.start >= front.end { + pool.pop_front(); + } + } + if ids.len() == count { + break; + } + } // pool mutex released before the S3 round-trip. + let need = (count - ids.len()) as u64; + let range = self.claim_ids_or_migrate(collection_name, need).await?; + let mut pools = self.writer_pools.lock().await; + pools + .entry(collection_name.to_string()) + .or_default() + .push_back(range); + } + + // Build chunks with the same embedding rules as the attached path, + // validating dims against the bucket config. On a validation failure, + // refresh the config once (a space may have just been added) before + // rejecting — a stale cache must never poison a durable fragment. + let build = |cfg: &cloud::BucketConfig| -> Result< + (Vec, HashMap), + Box, + > { + let default_space = cfg + .default_vector_space + .clone() + .unwrap_or_else(|| "default".into()); + let mut client_id_map: HashMap = HashMap::new(); + for (ic, &id) in ingest_chunks.iter().zip(ids.iter()) { + if let Some(ref cid) = ic.client_id { + client_id_map.insert(cid.clone(), id); + } + } + let parent_ids: Vec> = + ingest_chunks.iter().map(|ic| ic.parent_id).collect(); + let parent_refs: Vec> = ingest_chunks + .iter() + .map(|ic| ic.parent_ref.clone()) + .collect(); + let group_ids: Vec> = + ingest_chunks.iter().map(|ic| ic.group_id.clone()).collect(); + let resolved = RelationshipStore::resolve_batch_refs( + &client_id_map, + &parent_ids, + &parent_refs, + &group_ids, + ); - for ic in &ingest_chunks { - let id = loaded.next_id; - loaded.next_id += 1; - assigned_ids.push(id); - if let Some(ref cid) = ic.client_id { - client_id_map.insert(cid.clone(), id); + let mut chunks: Vec = Vec::with_capacity(count); + for (i, ic) in ingest_chunks.iter().enumerate() { + let id = ids[i]; + let (parent_id, group_id) = resolved[i].clone(); + let mut embeddings = ic.embeddings.clone(); + if let Some(emb) = ic.embedding.clone() { + embeddings.entry(default_space.clone()).or_insert(emb); + } + if embeddings.is_empty() { + if let Ok(emb) = embed_state.embed_query(&ic.text) { + let expected = cfg + .vector_spaces + .get(&default_space) + .map(|c| c.dims) + .unwrap_or(cfg.embedding_dims); + if emb.len() == expected { + embeddings.insert(default_space.clone(), emb); + } else { + tracing::warn!( + "writer ingest: built-in embedder produces {} dims but space \ + '{}' expects {} — chunk {} will be FTS-only", + emb.len(), + default_space, + expected, + i + ); + } + } + } + for (space_name, vec) in &embeddings { + let expected = cfg + .vector_spaces + .get(space_name) + .map(|c| c.dims) + .unwrap_or(cfg.embedding_dims); + if vec.len() != expected { + return Err(format!( + "chunk {i}: embedding for vector space '{space_name}' has {} dims, \ + expected {expected}", + vec.len() + ) + .into()); + } + } + chunks.push(DocumentChunk { + id, + collection: collection_name.to_string(), + file_id: ic.file_id.clone(), + chunk_index: ic.chunk_index, + page: ic.page, + text: ic.text.clone(), + metadata: ic.metadata.clone(), + doc_type: ic.doc_type.clone(), + parent_id, + group_id, + embeddings, + embedding: None, + }); } - } + Ok((chunks, client_id_map)) + }; + let (chunks, client_id_map) = match build(&cfg) { + Ok(out) => out, + Err(first_err) => { + let fresh = self.bucket_config(collection_name, true).await?; + build(&fresh).map_err(|_| first_err)? + } + }; - // Phase 2: Resolve batch parent references + // ONE durable append; searchable on serving nodes after refresh. + let payload = serde_json::to_vec(&chunks)?; + let records = chunks.len() as u64; + let seq = crate::storage::lsm::append_fragment( + self.storage.as_ref(), + collection_name, + bytes::Bytes::from(payload), + records, + ) + .await + .map_err(|e| format!("cloud WAL append failed: {e}"))?; + maybe_auto_compact( + self.storage.clone(), + collection_name.to_string(), + self.compacting.clone(), + ); + tracing::info!( + "Writer ingest: WAL fragment seq={} ({} chunks) durable for '{}'", + seq, + records, + collection_name + ); + Ok((count, client_id_map, Some(seq))) + } + + pub async fn ingest( + &self, + collection_name: &str, + ingest_chunks: Vec, + embed_state: &EmbedState, + ) -> Result<(usize, HashMap, Option), Box> + { + crate::metrics::inc(&crate::metrics::INGEST_REQUESTS_TOTAL); + crate::metrics::add( + &crate::metrics::INGEST_CHUNKS_TOTAL, + ingest_chunks.len() as u64, + ); + // Writer role: durable-append-only ingest, no local state required. + if self.role == NodeRole::Writer { + return self + .ingest_stateless(collection_name, ingest_chunks, embed_state) + .await; + } + + let count = ingest_chunks.len(); + self.ensure_attached(collection_name).await?; + + // Cloud mode: ids come from CAS-leased blocks (storage/id_alloc.rs) so + // they can NEVER collide with a stateless writer's ids. This happens + // BEFORE taking the write lock (its refill path does S3 round-trips). + // A failed ingest after this point leaks the taken ids — gaps are fine; + // the invariant is no-reuse, not density. + let cloud_ids: Option> = if self.cloud_mode && count > 0 { + Some(self.take_ids_cloud(collection_name, count).await?) + } else { + None + }; + + let mut collections = self.collections.write().await; + let loaded = collections.get_mut(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; + + // Phase 1: Assign IDs and build client_id -> chunk_id map + let mut client_id_map: HashMap = HashMap::new(); + let assigned_ids: Vec = match cloud_ids { + Some(ids) => { + // Keep the local counter as a diagnostic high-water mark only. + if let Some(&max) = ids.iter().max() { + loaded.next_id = loaded.next_id.max(max + 1); + } + ids + } + None => { + let mut ids = Vec::with_capacity(count); + for _ in 0..count { + ids.push(loaded.next_id); + loaded.next_id += 1; + } + ids + } + }; + for (ic, &id) in ingest_chunks.iter().zip(assigned_ids.iter()) { + if let Some(ref cid) = ic.client_id { + client_id_map.insert(cid.clone(), id); + } + } + + // Phase 2: Resolve batch parent references let parent_ids: Vec> = ingest_chunks.iter().map(|ic| ic.parent_id).collect(); let parent_refs: Vec> = ingest_chunks .iter() @@ -727,6 +1484,7 @@ impl CollectionManager { // Phase 3a (cloud): DURABLE S3 WAL append FIRST, before any local commit // (fixes F14 split-brain — a failed append leaves nothing local, clean retry). + let mut appended_seq: Option = None; if self.cloud_mode { let payload = serde_json::to_vec(&chunks)?; let records = chunks.len() as u64; @@ -738,6 +1496,7 @@ impl CollectionManager { ) .await .map_err(|e| format!("cloud WAL append failed, ingest not applied: {e}"))?; + appended_seq = Some(seq); tracing::info!( "Cloud ingest: WAL fragment seq={} ({} chunks) durable for '{}'", seq, @@ -756,11 +1515,28 @@ impl CollectionManager { // so no concurrent ingest can collide; applying by id is order-independent. let mut collections = self.collections.write().await; let loaded = match collections.get_mut(collection_name) { - Some(l) => l, + Some(l) => { + l.last_used + .store(next_lru_tick(), std::sync::atomic::Ordering::Relaxed); + l + } None => { - // Collection was deleted in the lock gap. `delete_collection` - // purges S3, but our fragment may have landed after that purge — - // append a tombstone so a re-materialize (which would recreate a + // Missing from the map: either DELETED or merely LRU-EVICTED in + // the lock gap. If the bucket still has the collection, the + // append is healthy and durable — do NOT erase it; the next + // attach/refresh applies it. + if self.cloud_mode + && cloud::read_bucket_config(self.storage.as_ref(), collection_name) + .await + .ok() + .flatten() + .is_some() + { + return Ok((count, client_id_map, appended_seq)); + } + // Genuinely deleted: `delete_collection` purges S3, but our + // fragment may have landed after that purge — append a + // tombstone so a re-materialize (which would recreate a // manifest referencing only our orphan fragment) yields nothing. if self.cloud_mode { if let Err(te) = crate::storage::lsm::append_tombstone( @@ -777,9 +1553,27 @@ impl CollectionManager { ); } } - return Err(format!("Collection '{}' not found", collection_name).into()); + return Err(not_found(format_args!( + "Collection \'{}\' not found", + collection_name + ))); } }; + // Double-apply guard: between our S3 append and this reacquire, the + // manifest refresher may have polled and applied OUR fragment. The + // tracker is the single source of truth for "already reflected + // locally" — skip the local apply if it covers our seq. + if let Some(seq) = appended_seq { + if loaded.applied.covers(seq) { + tracing::debug!( + "ingest seq={} for '{}' already applied by refresher; skipping local apply", + seq, + collection_name + ); + return Ok((count, client_id_map, appended_seq)); + } + } + // Apply all local state (chunks map, redb, FTS, HNSW, metadata, filter // index) in one fallible step. On ANY failure in cloud mode we've already // written a durable S3 fragment for these ids, so we compensate with a @@ -794,6 +1588,12 @@ impl CollectionManager { space_vectors, count, ); + if commit_result.is_ok() { + if let Some(seq) = appended_seq { + loaded.applied.mark(seq); + loaded.metadata.applied_seq = loaded.applied.contiguous; + } + } if let Err(e) = commit_result { if self.cloud_mode { // Compensate for the durable S3 fragment whose local commit @@ -817,10 +1617,10 @@ impl CollectionManager { } for id in &assigned_ids { loaded.tombstones.insert(*id); - loaded.chunks.remove(id); + if let Ok(Some(c)) = loaded.chunk_store.get(*id) { + loaded.filter_index.remove(*id, &filter_meta(&c)); + } } - loaded.filter_index = - build_filter_index_from_chunks(&loaded.chunks, &loaded.tombstones); drop(collections); if let Err(te) = crate::storage::lsm::append_tombstone( self.storage.as_ref(), @@ -847,7 +1647,7 @@ impl CollectionManager { tracing::info!("Ingested {} chunks into '{}'", count, collection_name); - Ok((count, client_id_map)) + Ok((count, client_id_map, appended_seq)) } /// Apply an ingest batch's local state (chunk map, redb, FTS, HNSW, metadata, @@ -867,9 +1667,6 @@ impl CollectionManager { for (id, parent_id, group_id) in rel_adds { loaded.relationships.add(id, parent_id, group_id); } - for chunk in chunks { - loaded.chunks.insert(chunk.id, chunk.clone()); - } // Phase 3b: Persist chunks to the disk-backed store BEFORE updating // FTS/HNSW. If this write fails we error out before any index commits, @@ -879,9 +1676,13 @@ impl CollectionManager { chunks.iter().map(|c| (c.id, c.clone())).collect(); loaded.chunk_store.insert_batch(&to_persist)?; - // Phase 4: Update Tantivy FTS index + // Phase 4: Update Tantivy FTS index. build_index returns facet state + // for THIS batch only — absorb the prior batches' facets (replacing + // them wholesale was the latent since-v0.2 facet bug). let tantivy_dir = store::tantivy_dir(data_dir, collection_name); - loaded.fts = tantivy_fts::build_index(&tantivy_dir, chunks, loaded.metadata.chunk_count)?; + let mut new_fts = tantivy_fts::build_index(&tantivy_dir, chunks)?; + new_fts.facet_bitsets.absorb(&loaded.fts.facet_bitsets); + loaded.fts = new_fts; // Phase 5: Update each vector space's HNSW index let vectors_dir = store::vectors_dir(data_dir, collection_name); @@ -930,6 +1731,17 @@ impl CollectionManager { continue; }; + // Save-batching state lives on the collection (the closure + // owns only the unwrapped VectorState). + let prev = loaded + .hnsw_unsaved + .get(&space_name) + .copied() + .unwrap_or(u32::MAX); + let mut mutable_flag = prev != u32::MAX; + let mut unsaved_ctr = if mutable_flag { prev } else { 0 }; + let unsaved = &mut unsaved_ctr; + let mutable_now = &mut mutable_flag; // Run the fallible updates in a closure so the space is ALWAYS // re-inserted into `vector_spaces` afterward — an early `?` here // used to drop the unwrapped space entirely, silently disabling @@ -951,39 +1763,76 @@ impl CollectionManager { let map_path = index_path.with_extension("keymap"); vector::save_key_map(&map_path, &vs.key_to_chunk_id)?; - // Add to HNSW index (use load() for mutability, not view()) + // Add to HNSW index. The in-RAM index stays mutable across + // batches (first mutation loads from disk once); the FILE is + // rewritten only every HNSW_SAVE_EVERY batches — per-batch + // saves were O(index size), a scale wall. A crash between + // saves leaves a stale file, detected and rebuilt from the + // mmap at next load (vectors are already durable there). + const HNSW_SAVE_EVERY: u32 = 16; let total = vs.key_to_chunk_id.len(); - if total >= 1000 && (vs.index.is_none() || index_path.exists()) { + if total >= 1000 { let index_path_str = index_path .to_str() .ok_or("USearch index path is not valid UTF-8")?; - let index = vector::create_index(dims, total)?; - if index_path.exists() { - index - .load(index_path_str) - .map_err(|e| format!("Failed to load USearch index: {}", e))?; - } - // Reserve for new vectors - let threads = 128.max(rayon::current_num_threads()); + let (index, was_fresh) = match (*mutable_now, vs.index.take()) { + (true, Some(idx)) => (idx, false), + _ => { + let idx = vector::create_index(dims, total)?; + if index_path.exists() { + idx.load(index_path_str).map_err(|e| { + format!("Failed to load USearch index: {}", e) + })?; + } + // A prior batch may have errored after its + // in-RAM adds but before a save: the on-disk + // file is STALE (missing committed batches + // whose vectors live in the mmap). Adding only + // the new batch and saving would bake that + // hole in permanently — heal from the mmap + // first (rows idx.size()..base_key). + if (idx.size() as usize) < base_key { + if let Some(m) = &vs.mmap_vectors { + let threads = vector::index_threads(); + idx.reserve_capacity_and_threads(total, threads) + .map_err(|e| format!("Reserve failed: {}", e))?; + for i in (idx.size() as usize)..base_key.min(m.len()) { + idx.add(i as u64, m.get(i)).map_err(|e| { + format!("Failed to heal index: {}", e) + })?; + } + } + } + (idx, true) + } + }; + let threads = vector::index_threads(); index .reserve_capacity_and_threads(total, threads) .map_err(|e| format!("Reserve failed: {}", e))?; - // Add new vectors incrementally for (i, (_, vec)) in new_vecs.iter().enumerate() { index .add((base_key + i) as u64, vec) .map_err(|e| format!("Failed to add vector: {}", e))?; } - index - .save(index_path_str) - .map_err(|e| format!("Failed to save index: {}", e))?; + *unsaved += 1; + if was_fresh || *unsaved >= HNSW_SAVE_EVERY { + index + .save(index_path_str) + .map_err(|e| format!("Failed to save index: {}", e))?; + *unsaved = 0; + } vs.index = Some(index); + *mutable_now = true; } Ok(()) })(); // Space goes back in whatever happened; a partial update is // recoverable (caller compensates the batch), a vanished space // is a silent outage. + if mutable_flag { + loaded.hnsw_unsaved.insert(space_name.clone(), unsaved_ctr); + } loaded.vector_spaces.insert(space_name, Arc::new(vs)); result?; } else { @@ -1013,7 +1862,11 @@ impl CollectionManager { store::save_metadata(data_dir, &loaded.metadata)?; let rel_path = store::collection_dir(data_dir, collection_name).join("relationships.bin"); loaded.relationships.save(&rel_path)?; - loaded.filter_index = build_filter_index_from_chunks(&loaded.chunks, &loaded.tombstones); + // Incremental: O(batch), not O(collection) — a full index rebuild here + // made every ingest/replay cost scale with the whole collection. + for c in chunks { + loaded.filter_index.insert(c.id, &filter_meta(c)); + } Ok(()) } @@ -1044,11 +1897,58 @@ impl CollectionManager { ), Box, > { + if self.role == NodeRole::Writer { + return Err("this node runs in writer role and does not serve queries".into()); + } + crate::metrics::inc(&crate::metrics::SEARCH_REQUESTS_TOTAL); + self.ensure_attached(collection_name).await?; + + // Read-your-writes: wait (bounded) until fragments up to `min_seq` are + // applied locally, refreshing on demand. A `min_seq` beyond the + // manifest is rejected rather than waited on forever. + if let Some(min_seq) = req.min_seq { + if self.cloud_mode { + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(2); + loop { + let covered = { + let collections = self.collections.read().await; + let loaded = collections.get(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; + loaded + .last_used + .store(next_lru_tick(), std::sync::atomic::Ordering::Relaxed); + loaded.applied.covers(min_seq) + }; + if covered { + break; + } + let next_seq = self.refresh_collection(collection_name).await?; + if min_seq >= next_seq { + return Err(format!( + "min_seq {} is beyond the collection's write history ({})", + min_seq, next_seq + ) + .into()); + } + if std::time::Instant::now() >= deadline { + return Err(format!( + "timed out waiting for min_seq {min_seq} to be applied" + ) + .into()); + } + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + } + } let start = std::time::Instant::now(); let collections = self.collections.read().await; - let loaded = collections - .get(collection_name) - .ok_or_else(|| format!("Collection '{}' not found", collection_name))?; + let loaded = collections.get(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; + loaded + .last_used + .store(next_lru_tick(), std::sync::atomic::Ordering::Relaxed); let mode = SearchMode::from_str_param(&req.mode); let rerank_k = req.top_k * 3; // fetch extra candidates for scoring @@ -1076,8 +1976,7 @@ impl CollectionManager { // ── Step 1: Retrieve candidates (filter-aware) ─────────────────── let fts_results = if matches!(mode, SearchMode::Fts | SearchMode::Hybrid) { - let (raw, _, _) = - tantivy_fts::search(&loaded.fts, &req.query, &HashMap::new(), rerank_k)?; + let (raw, _, _) = tantivy_fts::search(&loaded.fts, &req.query, rerank_k)?; // FTS doesn't yet have predicate pushdown; post-filter results // against the same eligible bitmap so the merged top-k respects // the filter exactly the same way the semantic path does. @@ -1228,14 +2127,13 @@ impl CollectionManager { recency_config.is_some() || !req.boosts.is_empty() || req.relationship_boost.is_some(); if has_scoring && !candidates.is_empty() { - let chunk_metadata: HashMap> = candidates - .iter() - .filter_map(|c| { - loaded - .chunks - .get(&c.chunk_id) - .map(|chunk| (c.chunk_id, chunk.metadata.clone())) - }) + let candidate_ids: Vec = candidates.iter().map(|c| c.chunk_id).collect(); + let chunk_metadata: HashMap> = loaded + .chunk_store + .get_batch(&candidate_ids) + .unwrap_or_default() + .into_iter() + .map(|chunk| (chunk.id, chunk.metadata.clone())) .collect(); let candidate_ids: Vec = candidates.iter().map(|c| c.chunk_id).collect(); @@ -1263,7 +2161,8 @@ impl CollectionManager { // segments sharing the same parent_id pay for one HashMap lookup, // not N. No additional I/O; the chunk map is already in memory. let candidate_chunk_ids: Vec = candidates.iter().map(|c| c.chunk_id).collect(); - let parent_meta_cache = build_parent_metadata_cache(&candidate_chunk_ids, &loaded.chunks); + let parent_meta_cache = + build_parent_metadata_cache(&candidate_chunk_ids, &loaded.chunk_store); // ── Step 5b: Relation enrichment (opt-in) ─────────────────────── // When include_relations is set, fetch each hit's edges in ONE batched, @@ -1280,9 +2179,7 @@ impl CollectionManager { )?; for edges in relations_by_chunk.values_mut() { for edge in edges.iter_mut() { - edge.target_status = if loaded.chunks.contains_key(&edge.target_chunk_id) - && !loaded.tombstones.contains(&edge.target_chunk_id) - { + edge.target_status = if loaded.filter_index.contains(edge.target_chunk_id) { "found".to_string() } else { "missing".to_string() @@ -1300,28 +2197,33 @@ impl CollectionManager { )> = candidates .iter() .filter_map(|c| { - loaded.chunks.get(&c.chunk_id).map(|chunk| { - let parent_metadata = parent_metadata_for(chunk, &parent_meta_cache); - // Some(vec) when requested (possibly empty), None when not — - // mirrors the parent_metadata Option discipline. - let relations = if req.include_relations { - Some( - relations_by_chunk - .get(&c.chunk_id) - .cloned() - .unwrap_or_default(), + loaded + .chunk_store + .get(c.chunk_id) + .ok() + .flatten() + .map(|chunk| { + let parent_metadata = parent_metadata_for(&chunk, &parent_meta_cache); + // Some(vec) when requested (possibly empty), None when not — + // mirrors the parent_metadata Option discipline. + let relations = if req.include_relations { + Some( + relations_by_chunk + .get(&c.chunk_id) + .cloned() + .unwrap_or_default(), + ) + } else { + None + }; + ( + chunk.clone(), + c.final_score, + c.source.clone(), + parent_metadata, + relations, ) - } else { - None - }; - ( - chunk.clone(), - c.final_score, - c.source.clone(), - parent_metadata, - relations, - ) - }) + }) }) .collect(); @@ -1359,15 +2261,51 @@ impl CollectionManager { collection_name: &str, new: Vec, ) -> Result, Box> { + // Writer role: build the edges without local state — target_status is + // stored as "missing" and re-resolved against the live chunk set at + // every read on serving nodes — and append ONE durable fragment. + if self.role == NodeRole::Writer { + let now = Utc::now(); + let mut built: Vec = Vec::with_capacity(new.len()); + for r in new { + if r.source_chunk_id == r.target_chunk_id { + return Err("A relation's source and target chunk must differ".into()); + } + built.push(ChunkRelation { + relation_id: uuid::Uuid::new_v4().to_string(), + source_chunk_id: r.source_chunk_id, + target_chunk_id: r.target_chunk_id, + target_document_id: r.target_document_id, + relation_type: r.relation_type, + target_status: "missing".to_string(), + metadata: r.metadata, + created_at: now, + }); + } + if !built.is_empty() { + let payload = serde_json::to_vec(&built)?; + let records = built.len() as u64; + crate::storage::lsm::append_relation_upsert( + self.storage.as_ref(), + collection_name, + bytes::Bytes::from(payload), + records, + ) + .await + .map_err(|e| format!("cloud relation-upsert append failed: {e}"))?; + } + return Ok(built); + } + self.ensure_attached(collection_name).await?; // Phase 1 (read lock): build the edges, resolving target_status against // the chunk map. Then release the lock BEFORE the S3 round-trip (#2/#4). let now = Utc::now(); let mut built: Vec = Vec::with_capacity(new.len()); { let collections = self.collections.read().await; - let loaded = collections - .get(collection_name) - .ok_or_else(|| format!("Collection '{}' not found", collection_name))?; + let loaded = collections.get(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; for r in new { if r.source_chunk_id == r.target_chunk_id { return Err("A relation's source and target chunk must differ".into()); @@ -1378,9 +2316,7 @@ impl CollectionManager { target_chunk_id: r.target_chunk_id, target_document_id: r.target_document_id, relation_type: r.relation_type, - target_status: if loaded.chunks.contains_key(&r.target_chunk_id) - && !loaded.tombstones.contains(&r.target_chunk_id) - { + target_status: if loaded.filter_index.contains(r.target_chunk_id) { "found".to_string() } else { "missing".to_string() @@ -1392,10 +2328,11 @@ impl CollectionManager { } // read lock released before S3 I/O. // Phase 2 (NO lock): durable S3 relation-upsert FIRST (S3-first ordering). + let mut appended_seq: Option = None; if self.cloud_mode && !built.is_empty() { let payload = serde_json::to_vec(&built)?; let records = built.len() as u64; - crate::storage::lsm::append_relation_upsert( + let seq = crate::storage::lsm::append_relation_upsert( self.storage.as_ref(), collection_name, bytes::Bytes::from(payload), @@ -1403,6 +2340,7 @@ impl CollectionManager { ) .await .map_err(|e| format!("cloud relation-upsert append failed: {e}"))?; + appended_seq = Some(seq); } // Phase 3 (read lock): apply locally (durable S3 record already written). @@ -1411,10 +2349,30 @@ impl CollectionManager { // ids, so the durable upsert fragment can't resurrect orphan edges on a // later materialize (the same discipline ingest applies to chunks). let apply_result: Result<(), Box> = { - let collections = self.collections.read().await; - match collections.get(collection_name) { - Some(loaded) => loaded.relation_store.insert_batch(&built), - None => Err(format!("Collection '{}' not found", collection_name).into()), + let mut collections = self.collections.write().await; + match collections.get_mut(collection_name) { + Some(loaded) => { + // Double-apply guard vs the manifest refresher; replay of a + // relation upsert is idempotent anyway (same relation_ids). + if appended_seq.map(|s| loaded.applied.covers(s)) == Some(true) { + Ok(()) + } else { + let r = loaded.relation_store.insert_batch(&built); + if r.is_ok() { + // Mark only after a successful apply (see C2). + if let Some(seq) = appended_seq { + loaded.applied.mark(seq); + loaded.metadata.applied_seq = loaded.applied.contiguous; + let _ = store::save_metadata(&self.data_dir, &loaded.metadata); + } + } + r + } + } + None => Err(not_found(format_args!( + "Collection \'{}\' not found", + collection_name + ))), } }; if let Err(e) = apply_result { @@ -1446,33 +2404,62 @@ impl CollectionManager { collection_name: &str, relation_id: &str, ) -> Result> { + // Writer role: durable relation-delete only (idempotent on replay). + if self.role == NodeRole::Writer { + crate::storage::lsm::append_relation_delete( + self.storage.as_ref(), + collection_name, + std::slice::from_ref(&relation_id.to_string()), + ) + .await + .map_err(|e| format!("cloud relation-delete append failed: {e}"))?; + return Ok(true); + } + self.ensure_attached(collection_name).await?; // Existence check under a short read lock, then release before S3 I/O. { let collections = self.collections.read().await; if !collections.contains_key(collection_name) { - return Err(format!("Collection '{}' not found", collection_name).into()); + return Err(not_found(format_args!( + "Collection \'{}\' not found", + collection_name + ))); } } // Cloud mode: durable S3 relation-delete FIRST — NO lock held across the // S3 round-trip (#4). Replay drops the id; deleting an absent id is an // idempotent no-op on materialize. + let mut appended_seq: Option = None; if self.cloud_mode { - crate::storage::lsm::append_relation_delete( + let seq = crate::storage::lsm::append_relation_delete( self.storage.as_ref(), collection_name, std::slice::from_ref(&relation_id.to_string()), ) .await .map_err(|e| format!("cloud relation-delete append failed: {e}"))?; + appended_seq = Some(seq); } - // Apply locally. - let collections = self.collections.read().await; - let loaded = collections - .get(collection_name) - .ok_or_else(|| format!("Collection '{}' not found", collection_name))?; - loaded.relation_store.delete(relation_id) + // Apply locally (write lock: the seq tracker needs &mut). + let mut collections = self.collections.write().await; + let loaded = collections.get_mut(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; + if appended_seq.map(|s| loaded.applied.covers(s)) == Some(true) { + return Ok(true); // refresher already applied our delete + } + let r = loaded.relation_store.delete(relation_id); + if r.is_ok() { + // Mark only after a successful apply (see C2). + if let Some(seq) = appended_seq { + loaded.applied.mark(seq); + loaded.metadata.applied_seq = loaded.applied.contiguous; + let _ = store::save_metadata(&self.data_dir, &loaded.metadata); + } + } + r } /// List a single chunk's relations, with `target_status` resolved against @@ -1484,17 +2471,22 @@ impl CollectionManager { direction: RelationDirection, types: Option<&[String]>, ) -> Result, Box> { + if self.role == NodeRole::Writer { + return Err("this node runs in writer role and does not serve queries".into()); + } + self.ensure_attached(collection_name).await?; let collections = self.collections.read().await; - let loaded = collections - .get(collection_name) - .ok_or_else(|| format!("Collection '{}' not found", collection_name))?; + let loaded = collections.get(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; + loaded + .last_used + .store(next_lru_tick(), std::sync::atomic::Ordering::Relaxed); let mut edges = loaded .relation_store .for_chunk(chunk_id, direction, types)?; for edge in edges.iter_mut() { - edge.target_status = if loaded.chunks.contains_key(&edge.target_chunk_id) - && !loaded.tombstones.contains(&edge.target_chunk_id) - { + edge.target_status = if loaded.filter_index.contains(edge.target_chunk_id) { "found".to_string() } else { "missing".to_string() @@ -1511,68 +2503,16 @@ impl CollectionManager { /// fragment so the deletion is durably S3-native. The vectors physically /// remain in the HNSW/FTS indexes until the next rebuild/compaction; search /// filters them out in the meantime. Returns the number newly deleted. - pub async fn delete_chunks( - &self, - collection_name: &str, - ids: &[u64], - ) -> Result> { - // Phase 1 (read lock): determine which ids are actually deletable. - // DEDUP the input — `{"ids":[5,5,5]}` must count (and decrement - // chunk_count by) ONE delete, not three. - let newly: Vec = { - let collections = self.collections.read().await; - let loaded = collections - .get(collection_name) - .ok_or_else(|| format!("Collection '{}' not found", collection_name))?; - let mut seen = std::collections::HashSet::new(); - ids.iter() - .copied() - .filter(|id| { - seen.insert(*id) - && loaded.chunks.contains_key(id) - && !loaded.tombstones.contains(id) - }) - .collect() - }; // read lock released here. - if newly.is_empty() { - return Ok(0); - } - - // Phase 2 (NO lock held): DURABLE S3 tombstone FIRST. This is the S3 - // network round-trip; doing it without the collections lock means a slow - // S3 call no longer stalls every other collection's reads/writes (#2). - // S3-first also fixes the F5 split-brain: on failure nothing local is - // committed, so the caller retries cleanly. - if self.cloud_mode { - crate::storage::lsm::append_tombstone(self.storage.as_ref(), collection_name, &newly) - .await - .map_err(|e| format!("LSM tombstone append failed (delete not applied): {e}"))?; - maybe_auto_compact( - self.storage.clone(), - collection_name.to_string(), - self.compacting.clone(), - ); - } - - // Phase 3 (write lock): apply local state. Re-check membership under the - // lock (a concurrent delete could have tombstoned some ids meanwhile); - // a redundant S3 tombstone for an already-deleted id is a harmless - // idempotent no-op on replay. - let mut collections = self.collections.write().await; - let loaded = collections - .get_mut(collection_name) - .ok_or_else(|| format!("Collection '{}' not found", collection_name))?; - let apply: Vec = newly - .iter() - .copied() - .filter(|id| loaded.chunks.contains_key(id) && !loaded.tombstones.contains(id)) - .collect(); - if apply.is_empty() { - return Ok(0); - } - - loaded.chunk_store.tombstone_batch(&apply)?; - for id in &apply { + /// Local tombstone apply — shared by the delete path and fragment replay. + /// Idempotent: already-deleted / absent ids are filtered by the caller (or + /// harmlessly re-tombstoned in redb). + fn apply_tombstones_locally( + data_dir: &Path, + loaded: &mut LoadedCollection, + apply: &[u64], + ) -> Result<(), Box> { + loaded.chunk_store.tombstone_batch(apply)?; + for id in apply { loaded.tombstones.insert(*id); } // Persist the corrected live count IMMEDIATELY after the tombstones — @@ -1580,15 +2520,22 @@ impl CollectionManager { // chunk_count permanently overstated. let removed = apply.len() as u64; loaded.metadata.chunk_count = loaded.metadata.chunk_count.saturating_sub(removed); - store::save_metadata(&self.data_dir, &loaded.metadata)?; + store::save_metadata(data_dir, &loaded.metadata)?; // Keep the filter index in step with the tombstones so `eligible` / - // selectivity don't count deleted chunks (which would underfill top-k - // on deleted-heavy collections). - loaded.filter_index = build_filter_index_from_chunks(&loaded.chunks, &loaded.tombstones); + // selectivity don't count deleted chunks — incrementally (O(batch)). + for id in apply { + match loaded.chunk_store.get(*id) { + Ok(Some(c)) => loaded.filter_index.remove(*id, &filter_meta(&c)), + other => tracing::error!( + "filter-index removal skipped for chunk {id}: {other:?} — universe may \ + overcount until re-attach (results stay correct via tombstone masking)" + ), + } + } // Prune relations incident on the deleted chunks (F6: propagate errors; // on failure the edges are orphaned but target_status reports their // endpoints as missing, and cloud replay prunes them independently). - for &id in &apply { + for &id in apply { let edges = loaded .relation_store .for_chunk(id, RelationDirection::Both, None)?; @@ -1596,72 +2543,635 @@ impl CollectionManager { loaded.relation_store.delete(&e.relation_id)?; } } - - tracing::info!( - "Deleted {} chunk(s) from '{}' (tombstoned{})", - apply.len(), - collection_name, - if self.cloud_mode { - " + WAL tombstone" - } else { - "" - } - ); - Ok(apply.len()) + Ok(()) } - /// Soft-delete every chunk matching a metadata filter (e.g. all chunks of a - /// file_id, or a metadata predicate). Resolves matching ids, then delegates - /// to `delete_chunks`. - pub async fn delete_by_filter( - &self, + /// Replay one WAL fragment into the local indexes (the refresher's apply + /// path — mirrors `cloud::materialize`'s kind dispatch exactly). MUST be + /// idempotent: fragments may race the originating node's own local apply. + fn apply_fragment_locally( + data_dir: &Path, collection_name: &str, - filters: &HashMap, - ) -> Result> { - // Collect matching, not-yet-deleted ids under a read lock first. - let ids: Vec = { - let collections = self.collections.read().await; - let loaded = collections - .get(collection_name) - .ok_or_else(|| format!("Collection '{}' not found", collection_name))?; - loaded - .chunks - .values() - .filter(|c| !loaded.tombstones.contains(&c.id)) - .filter(|c| crate::filter::matches_filters(c, filters)) - .map(|c| c.id) - .collect() - }; - if ids.is_empty() { - return Ok(0); + loaded: &mut LoadedCollection, + kind: crate::storage::lsm::FragmentKind, + payload: &[u8], + ) -> Result<(), Box> { + use crate::storage::lsm::FragmentKind; + match kind { + FragmentKind::Data => { + let chunks: Vec = serde_json::from_slice(payload)?; + // Defensive dims validation: a foreign writer's stale config + // could have let a wrong-length vector into a durable fragment; + // appending it would corrupt the mmap file for every vector + // after it. Quarantine (skip + loud error), never apply. + let mut fresh: Vec = Vec::with_capacity(chunks.len()); + 'chunk: for c in chunks { + // Idempotent replay: skip ids already present so + // chunk_count can't double-count. + if loaded.filter_index.contains(c.id) || loaded.tombstones.contains(&c.id) { + continue; + } + for (space, emb) in &c.embeddings { + let expected = loaded + .metadata + .vector_spaces + .get(space) + .map(|v| v.dims) + .unwrap_or(loaded.metadata.embedding_dims); + if emb.len() != expected { + tracing::error!( + "replay: chunk {} in '{}' has {}-dim embedding for space '{}' \ + (expected {}); quarantined", + c.id, + collection_name, + emb.len(), + space, + expected + ); + crate::metrics::inc(&crate::metrics::QUARANTINED_CHUNKS_TOTAL); + continue 'chunk; + } + } + fresh.push(c); + } + if fresh.is_empty() { + return Ok(()); + } + let rel_adds: Vec<(u64, Option, Option)> = fresh + .iter() + .map(|c| (c.id, c.parent_id, c.group_id.clone())) + .collect(); + let mut space_vectors: HashMap)>> = HashMap::new(); + for c in &fresh { + for (space, emb) in &c.embeddings { + space_vectors + .entry(space.clone()) + .or_default() + .push((c.id, emb.clone())); + } + } + let count = fresh.len(); + Self::apply_ingest_commit( + data_dir, + collection_name, + loaded, + rel_adds, + &fresh, + space_vectors, + count, + ) + } + FragmentKind::Tombstone => { + let ids: Vec = serde_json::from_slice(payload)?; + let apply: Vec = ids + .into_iter() + .filter(|id| loaded.filter_index.contains(*id)) + .collect(); + if apply.is_empty() { + return Ok(()); + } + Self::apply_tombstones_locally(data_dir, loaded, &apply) + } + FragmentKind::RelationUpsert => { + let rels: Vec = serde_json::from_slice(payload)?; + loaded.relation_store.insert_batch(&rels) + } + FragmentKind::RelationDelete => { + let ids: Vec = serde_json::from_slice(payload)?; + for id in &ids { + loaded.relation_store.delete(id)?; + } + Ok(()) + } } - self.delete_chunks(collection_name, &ids).await } - // ── Cloud compaction (object-storage mode) ─────────────────────────── + /// The per-namespace attach mutex (created on demand). Serializes every + /// destructive local-state transition for a namespace: attach, refresh- + /// triggered full re-attach, and LRU detach — so two of them can never run + /// concurrently on the same live index directory. + async fn attach_lock(&self, ns: &str) -> Arc> { + let mut locks = self.attach_locks.lock().await; + locks.entry(ns.to_string()).or_default().clone() + } - /// Compact a collection's S3 LSM: fold all segments + WAL fragments into a - /// single new segment (applying deletes), then rewrite the manifest to - /// reference only it. Reclaims space for tombstoned data. No-op in local - /// mode. Returns the number of live records in the resulting segment. - /// - /// This CAS-retries against concurrent appends: if the manifest changed - /// under us, we re-materialize the fresh state and try again. - pub async fn compact_collection( + /// Lazy attach: make sure `ns` is attached (rebuilt from the bucket) before + /// serving a request against it. No-op when already attached or when lazy + /// attach is off. A request stampede on a cold namespace rebuilds ONCE via + /// the per-namespace mutex; the global collections lock is never held + /// across the rebuild. + async fn ensure_attached( &self, - collection_name: &str, - ) -> Result> { - if !self.cloud_mode { - return Ok(0); - } - // Verify the collection exists (under a short read lock). - { - let collections = self.collections.read().await; - if !collections.contains_key(collection_name) { - return Err(format!("Collection '{}' not found", collection_name).into()); + ns: &str, + ) -> Result<(), Box> { + if !self.lazy_attach { + return Ok(()); + } + if self.collections.read().await.contains_key(ns) { + return Ok(()); + } + let lock = self.attach_lock(ns).await; + let _guard = lock.lock().await; + // Double-check under the attach mutex: a racer may have attached. + if self.collections.read().await.contains_key(ns) { + return Ok(()); + } + // Confirm the namespace exists in the bucket. Check the registry first + // (boot-time discovery), then the bucket itself — a collection created + // by ANOTHER node after our boot is attachable too. + let known = self.registered.read().await.contains(ns); + if !known { + validate_name_segment(ns, "Collection")?; + let exists = cloud::read_bucket_config(self.storage.as_ref(), ns) + .await? + .is_some() + || self + .storage + .exists(&format!("{ns}/manifest")) + .await + .unwrap_or(false); + if !exists { + // Don't leak an attach-lock entry per garbage name probed. + self.attach_locks.lock().await.remove(ns); + return Err(not_found(format_args!("Collection \'{}\' not found", ns))); } + self.registered.write().await.insert(ns.to_string()); } - + let start = std::time::Instant::now(); + let n = self.rebuild_collection_from_storage(ns).await?; + crate::metrics::inc(&crate::metrics::ATTACH_TOTAL); + crate::metrics::add( + &crate::metrics::ATTACH_SECONDS_SUM_MILLIS, + start.elapsed().as_millis() as u64, + ); + tracing::info!( + "Attached '{}' on demand ({} chunks in {:.2}s)", + ns, + n, + start.elapsed().as_secs_f64() + ); + self.maybe_evict_lru(ns).await; + Ok(()) + } + + /// Enforce the attached-collection budget: detach the least-recently-used + /// collection (never `just_attached`). Detach is safe — the bucket is the + /// source of truth — and local files are deleted only AFTER the global + /// lock is released (never filesystem I/O under the lock). The evicted + /// namespace stays registered for future re-attach. + async fn maybe_evict_lru(&self, just_attached: &str) { + if self.max_attached == 0 { + return; + } + // Pick the victim under a short read lock. + let victim: Option = { + let collections = self.collections.read().await; + if collections.len() <= self.max_attached { + None + } else { + collections + .iter() + .filter(|(name, _)| name.as_str() != just_attached) + .min_by_key(|(_, l)| l.last_used.load(std::sync::atomic::Ordering::Relaxed)) + .map(|(name, _)| name.clone()) + } + }; + let Some(name) = victim else { return }; + // Serialize with attach/re-attach on the same namespace: file deletion + // must never race a rebuild into the same directory. + let lock = self.attach_lock(&name).await; + let _guard = lock.lock().await; + { + let mut collections = self.collections.write().await; + // Re-check under the attach mutex (a racer may have evicted or the + // budget may have been satisfied meanwhile). + if collections.len() <= self.max_attached || !collections.contains_key(&name) { + return; + } + collections.remove(&name); + } // global lock released before any filesystem work. + self.registered.write().await.insert(name.clone()); + if let Err(e) = store::delete_collection_data(&self.data_dir, &name) { + tracing::warn!("detach '{}': local cleanup failed: {}", name, e); + } + tracing::info!("Detached '{}' (LRU, budget {})", name, self.max_attached); + } + + /// Converge this node's local indexes with the bucket manifest: apply + /// fragments this node hasn't seen (a remote writer's, or another serving + /// node's), in seq order, idempotently. Returns the manifest's `next_seq`. + /// + /// Two-branch compaction rule: if the compaction watermark has passed our + /// contiguous frontier, fragments we NEVER applied were folded into the + /// segment — the only correct recovery is a full re-attach. Otherwise the + /// folded fragments are ones we already applied, and only the live tail + /// needs replay. + pub async fn refresh_collection( + &self, + collection_name: &str, + ) -> Result> { + if !self.cloud_mode { + return Ok(0); + } + // Deleted-collection detection (no namespace generations yet): a + // manifest that has VANISHED means the collection was deleted on + // another node — detach instead of warning forever while serving dead + // data. + let (manifest, _) = match crate::storage::lsm::read_manifest( + self.storage.as_ref(), + collection_name, + ) + .await + { + Ok(m) => m, + Err(e) => { + let manifest_gone = !self + .storage + .exists(&format!("{collection_name}/manifest")) + .await + .unwrap_or(true); + if manifest_gone { + self.detach_deleted(collection_name).await; + return Err(format!( + "collection '{collection_name}' was deleted in object storage" + ) + .into()); + } + return Err(e.into()); + } + }; + let next_seq = manifest.next_seq; + + // Config convergence: vector-space adds/removes/default switches are + // CAS'd into {ns}/collection.json, NOT written as fragments — sync them + // here so already-attached nodes learn about them. A recreate (config + // created_at differs from ours) forces a full re-attach. + let bucket_cfg = cloud::read_bucket_config(self.storage.as_ref(), collection_name).await?; + let mut force_reattach = false; + if let Some(cfg) = &bucket_cfg { + let mut collections = self.collections.write().await; + if let Some(loaded) = collections.get_mut(collection_name) { + if loaded.metadata.created_at != cfg.created_at { + // Same name, different collection: it was deleted and + // recreated while we were attached. + force_reattach = true; + } else if loaded.metadata.vector_spaces != cfg.vector_spaces + || loaded.metadata.default_vector_space != cfg.default_vector_space + { + for name in cfg.vector_spaces.keys() { + if !loaded.vector_spaces.contains_key(name) { + loaded.vector_spaces.insert( + name.clone(), + Arc::new(VectorState { + index: None, + key_to_chunk_id: Vec::new(), + mmap_vectors: None, + vectors: Vec::new(), + }), + ); + } + } + loaded + .vector_spaces + .retain(|name, _| cfg.vector_spaces.contains_key(name)); + loaded.metadata.vector_spaces = cfg.vector_spaces.clone(); + loaded.metadata.default_vector_space = cfg.default_vector_space.clone(); + loaded.metadata.config = cfg.config.clone(); + store::save_metadata(&self.data_dir, &loaded.metadata)?; + tracing::info!( + "refresh '{}': synced vector-space config from bucket", + collection_name + ); + } + } + } + + let contiguous = { + let collections = self.collections.read().await; + let loaded = collections.get(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; + loaded.applied.contiguous + }; + + let needs_reattach = force_reattach + || manifest + .compaction_watermark + .map(|wm| wm + 1 > contiguous) + .unwrap_or(false); + if needs_reattach { + // Full re-attach, SERIALIZED on the per-namespace attach mutex so + // concurrent refresh ticks / min_seq waiters can't run destructive + // rebuilds into the same live directory. + let lock = self.attach_lock(collection_name).await; + let _guard = lock.lock().await; + // Re-check under the mutex: a racer may have already re-attached. + let still_needed = { + let collections = self.collections.read().await; + match collections.get(collection_name) { + Some(loaded) => { + force_reattach + && loaded.metadata.created_at + != bucket_cfg + .as_ref() + .map(|c| c.created_at) + .unwrap_or(loaded.metadata.created_at) + || manifest + .compaction_watermark + .map(|wm| wm + 1 > loaded.applied.contiguous) + .unwrap_or(false) + } + None => true, + } + }; + if still_needed { + tracing::info!( + "refresh '{}': full re-attach (compaction passed local frontier or recreate)", + collection_name + ); + crate::metrics::inc(&crate::metrics::REFRESH_REATTACHES_TOTAL); + self.rebuild_collection_from_storage(collection_name) + .await?; + } + return Ok(next_seq); + } + + // Filter fragment REFS first, fetch only what we need (a caught-up + // node fetches nothing), then apply per-fragment with the lock + // RELEASED between fragments so a large backlog can't cause a + // node-wide read outage. + let pending_refs: Vec = manifest + .uncompacted() + .filter(|r| r.seq >= contiguous) + .cloned() + .collect(); + if pending_refs.is_empty() { + return Ok(next_seq); + } + let mut applied_any = false; + for fref in pending_refs { + let bytes = crate::storage::lsm::read_fragment( + self.storage.as_ref(), + collection_name, + &fref.id, + ) + .await?; + let mut collections = self.collections.write().await; + let loaded = collections.get_mut(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; + if loaded.applied.covers(fref.seq) { + continue; + } + Self::apply_fragment_locally( + &self.data_dir, + collection_name, + loaded, + fref.kind, + &bytes, + )?; + loaded.applied.mark(fref.seq); + crate::metrics::inc(&crate::metrics::REFRESH_FRAGMENTS_APPLIED_TOTAL); + applied_any = true; + } + if applied_any { + let mut collections = self.collections.write().await; + if let Some(loaded) = collections.get_mut(collection_name) { + loaded.metadata.applied_seq = loaded.applied.contiguous; + store::save_metadata(&self.data_dir, &loaded.metadata)?; + } + } + Ok(next_seq) + } + + /// Detach a collection whose bucket namespace disappeared (deleted by + /// another node): drop it from the map + registry + caches and remove + /// local files, serialized on the attach mutex. + async fn detach_deleted(&self, ns: &str) { + let lock = self.attach_lock(ns).await; + let _guard = lock.lock().await; + let removed = { + let mut collections = self.collections.write().await; + collections.remove(ns).is_some() + }; + self.registered.write().await.remove(ns); + self.bucket_configs.write().await.remove(ns); + self.writer_pools.lock().await.remove(ns); + if removed { + let _ = store::delete_collection_data(&self.data_dir, ns); + tracing::info!("Detached '{}': deleted in object storage", ns); + } + } + + /// Refresh every attached collection /// Refresh every attached collection (the background refresher's tick). + pub async fn refresh_all(&self) { + let names: Vec = { + let collections = self.collections.read().await; + collections.keys().cloned().collect() + }; + for name in names { + if let Err(e) = self.refresh_collection(&name).await { + tracing::warn!("refresh of '{}' failed: {}", name, e); + } + } + } + + pub async fn delete_chunks( + &self, + collection_name: &str, + ids: &[u64], + ) -> Result<(usize, Option), Box> { + crate::metrics::inc(&crate::metrics::DELETE_REQUESTS_TOTAL); + // Writer role: durable tombstone only. Without local indexes we can't + // filter to ids-that-exist; a tombstone for an absent id is an + // idempotent no-op on replay, so append the deduped set as-is. + if self.role == NodeRole::Writer { + validate_name_segment(collection_name, "Collection")?; + // Existence check: without it, a tombstone for a bogus namespace + // would CREATE that namespace in the bucket (phantom collection). + self.bucket_config(collection_name, false).await?; + let mut seen = std::collections::HashSet::new(); + let newly: Vec = ids.iter().copied().filter(|id| seen.insert(*id)).collect(); + if newly.is_empty() { + return Ok((0, None)); + } + // Ids can never legitimately reach the allocator frontier; a bogus + // huge id would otherwise poison max_id forever (rebuilds compute + // next_id = max_id + 1 → overflow / id reuse). + let frontier = + crate::storage::id_alloc::frontier(self.storage.as_ref(), collection_name).await?; + if let Some(bad) = newly.iter().find(|id| **id >= frontier) { + return Err(format!( + "chunk id {bad} was never allocated in '{collection_name}' \ + (allocator frontier {frontier})" + ) + .into()); + } + let seq = crate::storage::lsm::append_tombstone( + self.storage.as_ref(), + collection_name, + &newly, + ) + .await + .map_err(|e| format!("LSM tombstone append failed: {e}"))?; + maybe_auto_compact( + self.storage.clone(), + collection_name.to_string(), + self.compacting.clone(), + ); + return Ok((newly.len(), Some(seq))); + } + self.ensure_attached(collection_name).await?; + + // Phase 1 (read lock): determine which ids are actually deletable. + // DEDUP the input — `{"ids":[5,5,5]}` must count (and decrement + // chunk_count by) ONE delete, not three. + let newly: Vec = { + let collections = self.collections.read().await; + let loaded = collections.get(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; + let mut seen = std::collections::HashSet::new(); + ids.iter() + .copied() + .filter(|id| seen.insert(*id) && loaded.filter_index.contains(*id)) + .collect() + }; // read lock released here. + if newly.is_empty() { + return Ok((0, None)); + } + + // Phase 2 (NO lock held): DURABLE S3 tombstone FIRST. This is the S3 + // network round-trip; doing it without the collections lock means a slow + // S3 call no longer stalls every other collection's reads/writes (#2). + // S3-first also fixes the F5 split-brain: on failure nothing local is + // committed, so the caller retries cleanly. + let mut appended_seq: Option = None; + if self.cloud_mode { + let seq = crate::storage::lsm::append_tombstone( + self.storage.as_ref(), + collection_name, + &newly, + ) + .await + .map_err(|e| format!("LSM tombstone append failed (delete not applied): {e}"))?; + appended_seq = Some(seq); + maybe_auto_compact( + self.storage.clone(), + collection_name.to_string(), + self.compacting.clone(), + ); + } + + // Phase 3 (write lock): apply local state. Re-check membership under the + // lock (a concurrent delete could have tombstoned some ids meanwhile); + // a redundant S3 tombstone for an already-deleted id is a harmless + // idempotent no-op on replay. + let mut collections = self.collections.write().await; + let loaded = collections.get_mut(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; + // Double-apply guard: the refresher may have applied OUR tombstone + // fragment between the append and this reacquire. + if let Some(seq) = appended_seq { + if loaded.applied.covers(seq) { + return Ok((newly.len(), appended_seq)); + } + } + let apply: Vec = newly + .iter() + .copied() + .filter(|id| loaded.filter_index.contains(*id)) + .collect(); + if apply.is_empty() { + return Ok((0, appended_seq)); + } + + Self::apply_tombstones_locally(&self.data_dir, loaded, &apply)?; + // Mark ONLY after the apply succeeded: marking first would make a + // failed apply invisible to the refresher forever (the node would keep + // serving deleted data). + if let Some(seq) = appended_seq { + loaded.applied.mark(seq); + loaded.metadata.applied_seq = loaded.applied.contiguous; + store::save_metadata(&self.data_dir, &loaded.metadata)?; + } + + tracing::info!( + "Deleted {} chunk(s) from '{}' (tombstoned{})", + apply.len(), + collection_name, + if self.cloud_mode { + " + WAL tombstone" + } else { + "" + } + ); + Ok((apply.len(), appended_seq)) + } + + /// Soft-delete every chunk matching a metadata filter (e.g. all chunks of a + /// file_id, or a metadata predicate). Resolves matching ids, then delegates + /// to `delete_chunks`. + pub async fn delete_by_filter( + &self, + collection_name: &str, + filters: &HashMap, + ) -> Result<(usize, Option), Box> { + if self.role == NodeRole::Writer { + return Err( + "delete-by-filter needs a serving node's indexes; this node runs in writer role \ + (delete by explicit ids instead)" + .into(), + ); + } + self.ensure_attached(collection_name).await?; + // Resolve matching live ids from the roaring filter index — the same + // pushdown search uses, so delete-by-filter and search can never + // disagree about what a filter matches. (This replaced a second, + // chunk-scanning filter implementation.) + let ids: Vec = { + let collections = self.collections.read().await; + let loaded = collections.get(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; + let expr = crate::search::filter_pushdown::FilterExpr::compile(filters); + loaded.filter_index.eligible(&expr).iter().collect() + }; + if ids.is_empty() { + return Ok((0, None)); + } + self.delete_chunks(collection_name, &ids).await + } + + // ── Cloud compaction (object-storage mode) ─────────────────────────── + + /// Compact a collection's S3 LSM: fold all segments + WAL fragments into a + /// single new segment (applying deletes), then rewrite the manifest to + /// reference only it. Reclaims space for tombstoned data. No-op in local + /// mode. Returns the number of live records in the resulting segment. + /// + /// This CAS-retries against concurrent appends: if the manifest changed + /// under us, we re-materialize the fresh state and try again. + pub async fn compact_collection( + &self, + collection_name: &str, + ) -> Result> { + if !self.cloud_mode { + return Ok(0); + } + self.ensure_attached(collection_name).await?; + // Verify the collection exists (under a short read lock). + { + let collections = self.collections.read().await; + if !collections.contains_key(collection_name) { + return Err(not_found(format_args!( + "Collection \'{}\' not found", + collection_name + ))); + } + } + Ok(compact_storage(self.storage.as_ref(), collection_name).await?) } @@ -1682,31 +3192,55 @@ impl CollectionManager { let chunks: Vec = materialized.chunks.values().cloned().collect(); let live_count = chunks.len(); - // Infer vector-space config from the recovered chunks' embeddings. - let mut vector_spaces: HashMap = HashMap::new(); - for c in &chunks { - for (space, emb) in &c.embeddings { - vector_spaces - .entry(space.clone()) - .or_insert(VectorSpaceConfig { - dims: emb.len(), - model: "recovered".to_string(), - status: "active".to_string(), - }); + // Collection config: the bucket `{ns}/collection.json` is authoritative + // (carries the user's real vector-space specs, created_at, and + // CollectionConfig — the old inference fabricated all three and lost + // `embed_model`). Fall back to inference only for pre-v0.4 namespaces, + // and back-fill the bucket config so the fallback runs at most once. + let bucket_cfg = cloud::read_bucket_config(self.storage.as_ref(), collection_name).await?; + let (vector_spaces, default_space, dims, created_at, coll_config) = match &bucket_cfg { + Some(cfg) => ( + cfg.vector_spaces.clone(), + cfg.default_vector_space.clone(), + cfg.embedding_dims, + cfg.created_at, + cfg.config.clone(), + ), + None => { + // Legacy inference from recovered embeddings. + let mut vector_spaces: HashMap = HashMap::new(); + for c in &chunks { + for (space, emb) in &c.embeddings { + vector_spaces + .entry(space.clone()) + .or_insert(VectorSpaceConfig { + dims: emb.len(), + model: "recovered".to_string(), + status: "active".to_string(), + }); + } + } + if vector_spaces.is_empty() { + vector_spaces.insert( + "default".to_string(), + VectorSpaceConfig { + dims: 384, + model: "recovered".to_string(), + status: "active".to_string(), + }, + ); + } + let default_space = vector_spaces.keys().next().cloned(); + let dims = vector_spaces.values().next().map(|s| s.dims).unwrap_or(384); + ( + vector_spaces, + default_space, + dims, + Utc::now(), + CollectionConfig::default(), + ) } - } - if vector_spaces.is_empty() { - vector_spaces.insert( - "default".to_string(), - VectorSpaceConfig { - dims: 384, - model: "recovered".to_string(), - status: "active".to_string(), - }, - ); - } - let default_space = vector_spaces.keys().next().cloned(); - let dims = vector_spaces.values().next().map(|s| s.dims).unwrap_or(384); + }; // next_id must never regress or reuse an id: one past the high-water // mark whenever ANY id was ever assigned (max_id covers tombstoned ids // via the WAL and the segment's stored max_id). The old @@ -1720,15 +3254,36 @@ impl CollectionManager { let metadata = Collection { name: collection_name.to_string(), - created_at: Utc::now(), + created_at, vector_spaces: vector_spaces.clone(), default_vector_space: default_space.clone(), embedding_dims: dims, chunk_count: live_count as u64, next_id, - config: CollectionConfig::default(), + config: coll_config, + applied_seq: manifest.next_seq, }; store::save_metadata(&self.data_dir, &metadata)?; + // Organic migration: back-fill the bucket config for pre-v0.4 + // namespaces (create-only, race-safe; best-effort). + if bucket_cfg.is_none() { + let backfill = cloud::BucketConfig::from_collection(&metadata); + if let Err(e) = cloud::write_bucket_config_if_absent( + self.storage.as_ref(), + collection_name, + &backfill, + ) + .await + { + if !matches!(e, crate::storage::StorageError::AlreadyExists(_)) { + tracing::warn!( + "bucket config back-fill for '{}' failed: {}", + collection_name, + e + ); + } + } + } // Build local stores from the materialized chunks. Start from a CLEAN // chunk store: S3 is the source of truth on recovery, so any pre-existing @@ -1739,14 +3294,14 @@ impl CollectionManager { std::fs::create_dir_all(parent)?; } let _ = std::fs::remove_file(&chunks_db); - let chunk_store = ChunkStore::open(&chunks_db)?; + let chunk_store = ChunkCache::new(ChunkStore::open(&chunks_db)?); let to_persist: Vec<(u64, DocumentChunk)> = chunks.iter().map(|c| (c.id, c.clone())).collect(); chunk_store.insert_batch(&to_persist)?; // FTS index. let tantivy_dir = store::tantivy_dir(&self.data_dir, collection_name); - let fts = tantivy_fts::build_index(&tantivy_dir, &chunks, 0)?; + let fts = tantivy_fts::build_index(&tantivy_dir, &chunks)?; // Vector spaces (HNSW) from embeddings. let vectors_dir = store::vectors_dir(&self.data_dir, collection_name); @@ -1771,11 +3326,12 @@ impl CollectionManager { for c in &chunks { relationships.add(c.id, c.parent_id, c.group_id.clone()); } - let chunk_map: HashMap = - chunks.iter().map(|c| (c.id, c.clone())).collect(); - // Materialized state is already live-only (tombstones applied on replay). - let filter_index = - build_filter_index_from_chunks(&chunk_map, &std::collections::HashSet::new()); + // Materialized state is already live-only (tombstones applied on + // replay); build the index streaming, no full map in RAM. + let mut filter_index = FilterIndex::new(); + for c in &chunks { + filter_index.insert(c.id, &filter_meta(c)); + } // Reconstruct the typed-relation store from the materialized relations // (recovered from the S3 WAL/segments) — so relations survive a cold @@ -1790,11 +3346,15 @@ impl CollectionManager { } let loaded = LoadedCollection { + id_pool: Default::default(), + hnsw_unsaved: HashMap::new(), + last_used: std::sync::atomic::AtomicU64::new(next_lru_tick()), + // A rebuild materialized EVERYTHING in the manifest it read. + applied: SeqTracker::starting_at(manifest.next_seq), metadata, fts, vector_spaces: vs_map, relationships, - chunks: chunk_map, chunk_store, relation_store, tombstones: std::collections::HashSet::new(), @@ -1816,11 +3376,18 @@ impl CollectionManager { (HashMap>, u64), Box, > { + if self.role == NodeRole::Writer { + return Err("this node runs in writer role and does not serve queries".into()); + } + self.ensure_attached(collection_name).await?; let collections = self.collections.read().await; - let loaded = collections - .get(collection_name) - .ok_or_else(|| format!("Collection '{}' not found", collection_name))?; - tantivy_fts::get_facets(&loaded.fts, query, fields) + let loaded = collections.get(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; + loaded + .last_used + .store(next_lru_tick(), std::sync::atomic::Ordering::Relaxed); + tantivy_fts::get_facets(&loaded.fts, query, fields, loaded.filter_index.universe()) } /// Get all chunk texts and IDs for rebuild jobs. @@ -1829,16 +3396,18 @@ impl CollectionManager { collection_name: &str, ) -> Result<(Vec, Vec), Box> { let collections = self.collections.read().await; - let loaded = collections - .get(collection_name) - .ok_or_else(|| format!("Collection '{}' not found", collection_name))?; + let loaded = collections.get(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; let mut texts = Vec::new(); let mut ids = Vec::new(); - for (&id, chunk) in &loaded.chunks { - ids.push(id); - texts.push(chunk.text.clone()); - } + loaded.chunk_store.for_each(|id, chunk| { + if !loaded.tombstones.contains(&id) { + ids.push(id); + texts.push(chunk.text); + } + })?; Ok((texts, ids)) } @@ -1860,19 +3429,21 @@ impl CollectionManager { time_end_ms: Option, ) -> Result, Box> { let collections = self.collections.read().await; - let loaded = collections - .get(collection_name) - .ok_or_else(|| format!("Collection '{}' not found", collection_name))?; - - let mut results: Vec = loaded - .chunks - .values() - .filter(|c| !loaded.tombstones.contains(&c.id)) - .filter(|c| c.doc_type == "segment") - .filter(|c| c.group_id.as_deref() == Some(asset)) - .filter(|c| segment_in_time_window(c, time_ms, time_start_ms, time_end_ms)) - .cloned() - .collect(); + let loaded = collections.get(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; + + let mut collected: Vec = Vec::new(); + loaded.chunk_store.for_each(|id, c| { + if !loaded.tombstones.contains(&id) + && c.doc_type == "segment" + && c.group_id.as_deref() == Some(asset) + && segment_in_time_window(&c, time_ms, time_start_ms, time_end_ms) + { + collected.push(c); + } + })?; + let mut results: Vec = collected.into_iter().collect(); // Sort ascending by timerange_start_ms. Segments missing the metadata // sort to the end (f64::INFINITY) instead of position 0, so callers @@ -1942,186 +3513,30 @@ pub(crate) fn segment_in_time_window( } #[cfg(test)] -mod segments_at_tests { - use super::*; - - fn make_segment(group_id: &str, ts_ms: f64, te_ms: f64) -> DocumentChunk { - let mut metadata = HashMap::new(); - metadata.insert( - "timerange_start_ms".to_string(), - MetadataValue::Float(ts_ms), - ); - metadata.insert("timerange_end_ms".to_string(), MetadataValue::Float(te_ms)); - DocumentChunk { - id: 1, - collection: "test".to_string(), - file_id: "f1".to_string(), - chunk_index: 0, - page: None, - text: String::new(), - metadata, - doc_type: "segment".to_string(), - parent_id: None, - group_id: Some(group_id.to_string()), - embeddings: HashMap::new(), - embedding: None, - } - } +mod segments_at_tests; - /// Make a zero-duration "instant" segment, the convention for sidecar - /// events that have a single timestamp (e.g. standout_timestamps). - fn make_instant(group_id: &str, t_ms: f64) -> DocumentChunk { - make_segment(group_id, t_ms, t_ms) - } +/// Typed "does not exist" error. The API layer downcasts to map these to +/// HTTP 404; every other engine error keeps the handler's default status. +/// (Previously a missing collection surfaced as 500 from /search and 400 +/// from /ingest — stringly errors carried no classification.) +#[derive(Debug)] +pub struct NotFound(pub String); - #[test] - fn point_inside_window() { - let c = make_segment("a", 100.0, 200.0); - assert!(segment_in_time_window(&c, Some(150.0), None, None)); +impl std::fmt::Display for NotFound { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(&self.0) } +} - #[test] - fn point_outside_window() { - let c = make_segment("a", 100.0, 200.0); - assert!(!segment_in_time_window(&c, Some(250.0), None, None)); - } +impl std::error::Error for NotFound {} - #[test] - fn point_boundaries_inclusive() { - let c = make_segment("a", 100.0, 200.0); - assert!(segment_in_time_window(&c, Some(100.0), None, None)); - assert!(segment_in_time_window(&c, Some(200.0), None, None)); - } +fn not_found(what: impl std::fmt::Display) -> Box { + Box::new(NotFound(what.to_string())) +} - #[test] - fn range_overlap_matches() { - let c = make_segment("a", 100.0, 200.0); - assert!(segment_in_time_window(&c, None, Some(180.0), Some(300.0))); - } - - #[test] - fn range_no_overlap() { - let c = make_segment("a", 100.0, 200.0); - assert!(!segment_in_time_window(&c, None, Some(250.0), Some(400.0))); - } - - #[test] - fn range_open_lower_bound() { - let c = make_segment("a", 100.0, 200.0); - assert!(segment_in_time_window(&c, None, None, Some(150.0))); - assert!(!segment_in_time_window(&c, None, None, Some(50.0))); - } - - #[test] - fn range_open_upper_bound() { - let c = make_segment("a", 100.0, 200.0); - assert!(segment_in_time_window(&c, None, Some(150.0), None)); - assert!(!segment_in_time_window(&c, None, Some(300.0), None)); - } - - #[test] - fn missing_metadata_with_filter_excludes() { - let c = DocumentChunk { - id: 2, - collection: "test".to_string(), - file_id: "f2".to_string(), - chunk_index: 0, - page: None, - text: String::new(), - metadata: HashMap::new(), - doc_type: "segment".to_string(), - parent_id: None, - group_id: Some("a".to_string()), - embeddings: HashMap::new(), - embedding: None, - }; - assert!(!segment_in_time_window(&c, Some(100.0), None, None)); - assert!(!segment_in_time_window(&c, None, Some(0.0), Some(1000.0))); - } - - #[test] - fn no_filter_matches_all() { - let with_meta = make_segment("a", 100.0, 200.0); - assert!(segment_in_time_window(&with_meta, None, None, None)); - - let without_meta = DocumentChunk { - id: 3, - collection: "test".to_string(), - file_id: "f3".to_string(), - chunk_index: 0, - page: None, - text: String::new(), - metadata: HashMap::new(), - doc_type: "segment".to_string(), - parent_id: None, - group_id: Some("a".to_string()), - embeddings: HashMap::new(), - embedding: None, - }; - assert!(segment_in_time_window(&without_meta, None, None, None)); - } - - // When both `time_ms` and `time_start_ms`/`time_end_ms` are provided, - // `time_ms` wins. Documented in the segments.rs handler comment; this - // test asserts it. - #[test] - fn point_lookup_takes_precedence_over_range() { - let c = make_segment("a", 100.0, 200.0); - // Point=150 is inside [100, 200], but the range [300, 400] is outside. - // If `time_ms` correctly takes precedence, this must return true. - assert!(segment_in_time_window( - &c, - Some(150.0), - Some(300.0), - Some(400.0) - )); - // Point=250 is outside, but the range [100, 300] would match. - // If `time_ms` correctly takes precedence, this must return false. - assert!(!segment_in_time_window( - &c, - Some(250.0), - Some(100.0), - Some(300.0) - )); - } - - // Instants (zero-duration events like a standout_timestamp) match a - // point query at their exact timestamp and any range that overlaps it. - // Critical for ingesting sidecar fields like - // `gemini.response.standout_timestamps[]` which only carry a single ms. - #[test] - fn instant_matches_exact_point_query() { - let c = make_instant("a", 5200.0); - assert!(segment_in_time_window(&c, Some(5200.0), None, None)); - assert!(!segment_in_time_window(&c, Some(5199.0), None, None)); - assert!(!segment_in_time_window(&c, Some(5201.0), None, None)); - } - - #[test] - fn instant_matches_overlapping_range_query() { - let c = make_instant("a", 5200.0); - assert!(segment_in_time_window(&c, None, Some(5000.0), Some(6000.0))); - assert!(segment_in_time_window(&c, None, Some(5200.0), Some(5200.0))); - assert!(!segment_in_time_window( - &c, - None, - Some(5201.0), - Some(6000.0) - )); - } -} - -/// Build a roaring-bitmap FilterIndex over a chunk map. Synthesizes a -/// `doc_type` metadata entry from the struct field so filter expressions -/// can target it without requiring callers to duplicate `doc_type` into -/// `chunk.metadata`. Matches the semantics of `filter::matches_filters`. -/// -/// Called on collection load (over rehydrated chunks) and after every -/// ingest batch (alongside FTS/HNSW rebuild). The index lives in-memory -/// only for now; persistence lands when chunk metadata migrates off redb. -/// Uncompacted-fragment count above which a cloud collection is auto-compacted. -/// Keeps the WAL bounded and reclaims tombstoned data without operator action. -pub(crate) const AUTO_COMPACT_FRAGMENT_THRESHOLD: usize = 32; +/// Uncompacted-fragment count above which a cloud collection is auto-compacted. +/// Keeps the WAL bounded and reclaims tombstoned data without operator action. +pub(crate) const AUTO_COMPACT_FRAGMENT_THRESHOLD: usize = 32; /// Storage-only compaction (no local manager state touched): materialize the /// full live set from S3, write it as one new segment, and CAS-rewrite the @@ -2132,18 +3547,72 @@ pub(crate) async fn compact_storage( storage: &dyn Storage, ns: &str, ) -> Result { + /// Segments tolerated before a full merge. Tail folds are O(batch); only + /// the merge is O(live set), and it runs 1/K as often. + const MERGE_SEGMENTS: usize = 8; const MAX_RETRIES: u32 = 10; + + // Phase 1: fold the WAL tail into an APPENDED segment (bounded work). + let mut folded_this_run = false; for _ in 0..MAX_RETRIES { let (manifest, version) = crate::storage::lsm::read_manifest(storage, ns).await?; - if manifest.segments.is_empty() && manifest.fragments.is_empty() { - return Ok(0); + let tail: Vec<_> = manifest.uncompacted().cloned().collect(); + if tail.is_empty() { + break; + } + let folded_through = tail.iter().map(|f| f.seq).max().unwrap(); + // STRICT reads: folding advances the watermark past these fragments; + // a tolerated NotFound here would be silent data loss. + let frags = + crate::storage::lsm::read_uncompacted_fragments_strict(storage, ns, &manifest).await?; + let segment = cloud::fold_tail(&frags)?; + let records = segment.chunks.len() as u64; + let bytes = cloud::encode_segment_v2(&segment)?; + match crate::storage::lsm::append_segment( + storage, + ns, + &version, + &manifest, + bytes::Bytes::from(bytes), + records, + folded_through, + ) + .await + { + Ok(()) => { + crate::metrics::inc(&crate::metrics::COMPACTIONS_TOTAL); + tracing::info!( + "Compacted '{}': folded WAL tail through seq {} ({} live records)", + ns, + folded_through, + records + ); + folded_this_run = true; + break; + } + Err(crate::storage::StorageError::VersionConflict { .. }) => continue, + Err(e) => return Err(e), + } + } + + // Phase 2: merge segments when they pile up (the only O(live-set) step). + // NEVER in the same invocation as a fold: phase 1 staged the folded + // fragments for next-cycle GC, and an immediate merge would GC them out + // from under readers still holding the pre-fold manifest. + if folded_this_run { + let (manifest, _) = crate::storage::lsm::read_manifest(storage, ns).await?; + return Ok(manifest.segments.iter().map(|s| s.records).sum()); + } + for _ in 0..MAX_RETRIES { + let (manifest, version) = crate::storage::lsm::read_manifest(storage, ns).await?; + if manifest.segments.len() <= MERGE_SEGMENTS { + return Ok(manifest.segments.iter().map(|s| s.records).sum()); } let materialized = cloud::materialize(storage, ns, &manifest).await?; let chunks: Vec = materialized.chunks.values().cloned().collect(); let relations: Vec = materialized.relations.values().cloned().collect(); let records = chunks.len() as u64; let segment_bytes = cloud::encode_segment(&chunks, &relations, materialized.max_id)?; - match crate::storage::lsm::replace_with_single_segment( storage, ns, @@ -2155,11 +3624,7 @@ pub(crate) async fn compact_storage( .await { Ok(()) => { - tracing::info!( - "Compacted '{}': {} live records in one segment", - ns, - records - ); + tracing::info!("Merged '{}' segments: {} live records", ns, records); return Ok(records); } Err(crate::storage::StorageError::VersionConflict { .. }) => continue, @@ -2214,28 +3679,21 @@ fn maybe_auto_compact( /// `eligible`/selectivity agree with what search may actually return, on every /// load path (local load, ingest rebuild, cloud recovery). Freshly-deleted ids /// are additionally masked post-retrieval until the next rebuild. -pub(crate) fn build_filter_index_from_chunks( - chunks: &HashMap, - tombstones: &std::collections::HashSet, -) -> FilterIndex { - let mut idx = FilterIndex::new(); - for (&chunk_id, chunk) in chunks { - if tombstones.contains(&chunk_id) { - continue; - } - let mut effective = chunk.metadata.clone(); - // doc_type is a struct field, not a metadata key, but the filter - // language treats it as one. Mirror it here so the bitmap covers it. - effective.insert( - "doc_type".to_string(), - MetadataValue::String(chunk.doc_type.clone()), - ); - // chunk_id is the full u64; the treemap-backed FilterIndex indexes the - // whole id space, so no chunk is dropped regardless of id magnitude. - idx.insert(chunk_id, &effective); - } - idx.finalize(); - idx +/// Process-monotonic LRU tick (no wall clock — avoids Date-based flakiness). +fn next_lru_tick() -> u64 { + static TICK: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1); + TICK.fetch_add(1, std::sync::atomic::Ordering::Relaxed) +} + +/// The metadata view the filter index sees: chunk metadata plus the mirrored +/// doc_type field (the filter language treats it as metadata). +fn filter_meta(chunk: &DocumentChunk) -> HashMap { + let mut m = chunk.metadata.clone(); + m.insert( + "doc_type".to_string(), + MetadataValue::String(chunk.doc_type.clone()), + ); + m } /// Build a deduplicated cache of parent chunk metadata for a set of candidate @@ -2253,11 +3711,11 @@ pub(crate) fn build_filter_index_from_chunks( /// from "parent exists with empty metadata." pub(crate) fn build_parent_metadata_cache( candidate_chunk_ids: &[u64], - chunks: &HashMap, + chunks: &ChunkCache, ) -> HashMap> { let mut cache: HashMap> = HashMap::new(); for cid in candidate_chunk_ids { - let Some(chunk) = chunks.get(cid) else { + let Ok(Some(chunk)) = chunks.get(*cid) else { continue; }; if chunk.doc_type != "segment" { @@ -2271,7 +3729,7 @@ pub(crate) fn build_parent_metadata_cache( } // Only cache parents that actually exist. Missing parents stay out // of the cache so `parent_metadata_for` returns None for them. - if let Some(parent) = chunks.get(&pid) { + if let Ok(Some(parent)) = chunks.get(pid) { cache.insert(pid, parent.metadata.clone()); } } @@ -2294,1745 +3752,16 @@ pub(crate) fn parent_metadata_for( } #[cfg(test)] -mod parent_metadata_tests { - use super::*; - - fn segment(id: u64, parent_id: Option) -> DocumentChunk { - DocumentChunk { - id, - collection: "test".to_string(), - file_id: format!("f{}", id), - chunk_index: 0, - page: None, - text: String::new(), - metadata: HashMap::new(), - doc_type: "segment".to_string(), - parent_id, - group_id: None, - embeddings: HashMap::new(), - embedding: None, - } - } - - fn source_with_meta(id: u64, key: &str, val: &str) -> DocumentChunk { - let mut metadata = HashMap::new(); - metadata.insert(key.to_string(), MetadataValue::String(val.to_string())); - DocumentChunk { - id, - collection: "test".to_string(), - file_id: format!("f{}", id), - chunk_index: 0, - page: None, - text: String::new(), - metadata, - doc_type: "source".to_string(), - parent_id: None, - group_id: None, - embeddings: HashMap::new(), - embedding: None, - } - } - - fn into_map(chunks: Vec) -> HashMap { - chunks.into_iter().map(|c| (c.id, c)).collect() - } - - #[test] - fn segment_with_parent_gets_metadata() { - let chunks = into_map(vec![ - source_with_meta(1, "title", "Keynote"), - segment(2, Some(1)), - ]); - let cache = build_parent_metadata_cache(&[2], &chunks); - let meta = parent_metadata_for(chunks.get(&2).unwrap(), &cache); - assert_eq!( - meta.unwrap().get("title"), - Some(&MetadataValue::String("Keynote".to_string())) - ); - } - - #[test] - fn source_hit_gets_none() { - let chunks = into_map(vec![source_with_meta(1, "title", "Keynote")]); - let cache = build_parent_metadata_cache(&[1], &chunks); - let meta = parent_metadata_for(chunks.get(&1).unwrap(), &cache); - assert!(meta.is_none()); - } - - #[test] - fn segment_without_parent_gets_none() { - let chunks = into_map(vec![segment(2, None)]); - let cache = build_parent_metadata_cache(&[2], &chunks); - let meta = parent_metadata_for(chunks.get(&2).unwrap(), &cache); - assert!(meta.is_none()); - } - - #[test] - fn dedup_one_lookup_per_unique_parent() { - // Three segments, all pointing at parent_id=10. The cache should - // contain exactly one entry (for pid=10), proving the dedup. - let chunks = into_map(vec![ - source_with_meta(10, "source_id", "src-001"), - segment(11, Some(10)), - segment(12, Some(10)), - segment(13, Some(10)), - ]); - let cache = build_parent_metadata_cache(&[11, 12, 13], &chunks); - assert_eq!( - cache.len(), - 1, - "expected one cache entry for the shared parent" - ); - assert!(cache.contains_key(&10)); - // All three segments resolve to the same parent metadata. - for cid in [11, 12, 13] { - let meta = parent_metadata_for(chunks.get(&cid).unwrap(), &cache); - assert_eq!( - meta.unwrap().get("source_id"), - Some(&MetadataValue::String("src-001".to_string())) - ); - } - } - - #[test] - fn orphan_segment_yields_none() { - // parent_id=99 not in chunks. The cache must NOT contain pid=99, - // and parent_metadata_for must return None. This distinguishes - // "parent exists with empty metadata" (Some({})) from "parent - // doesn't exist" (None). - let chunks = into_map(vec![segment(5, Some(99))]); - let cache = build_parent_metadata_cache(&[5], &chunks); - assert!(!cache.contains_key(&99), "orphan parent must not be cached"); - let meta = parent_metadata_for(chunks.get(&5).unwrap(), &cache); - assert!(meta.is_none(), "orphan segment must yield None"); - } - - #[test] - fn parent_exists_with_empty_metadata_yields_some_empty() { - // Parent chunk exists but has no metadata fields. Must return Some({}) - // so callers can distinguish from the orphan case (None). - let parent_no_meta = DocumentChunk { - id: 20, - collection: "test".to_string(), - file_id: "f20".to_string(), - chunk_index: 0, - page: None, - text: String::new(), - metadata: HashMap::new(), - doc_type: "source".to_string(), - parent_id: None, - group_id: None, - embeddings: HashMap::new(), - embedding: None, - }; - let chunks = into_map(vec![parent_no_meta, segment(21, Some(20))]); - let cache = build_parent_metadata_cache(&[21], &chunks); - let meta = parent_metadata_for(chunks.get(&21).unwrap(), &cache); - assert!(meta.is_some()); - assert!(meta.unwrap().is_empty()); - } - - #[test] - fn parent_metadata_for_cache_miss_returns_none() { - // Defensive: if the cache was built with a different set of IDs than - // the one we're looking up, the function must return None (not panic, - // not return stale data). Catches regressions where someone "optimizes" - // parent_metadata_for to assume the cache is always complete. - let parent = source_with_meta(1, "title", "Keynote"); - let seg = segment(2, Some(1)); - let chunks = into_map(vec![parent, seg]); - // Build cache against an empty candidate list, then look up segment 2. - let cache = build_parent_metadata_cache(&[], &chunks); - assert!(cache.is_empty()); - let meta = parent_metadata_for(chunks.get(&2).unwrap(), &cache); - assert!(meta.is_none()); - } -} +mod parent_metadata_tests; #[cfg(test)] -mod persistence_tests { - //! End-to-end durability test. Builds a CollectionManager in a temp dir, - //! ingests chunks, drops the manager (closing the chunk store), creates a - //! new manager pointing at the same dir, and asserts the chunks come back. - //! - //! This is the test that proves Compass survives process restarts. Without - //! the disk-backed ChunkStore wiring, this test would fail because - //! `loaded.chunks` would be empty after the manager restart. - - use super::*; - use crate::embed::EmbedState; - use std::collections::HashMap; - use std::sync::atomic::{AtomicU64, Ordering}; - - fn unique_data_dir() -> std::path::PathBuf { - static N: AtomicU64 = AtomicU64::new(0); - let nanos = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_nanos(); - std::env::temp_dir().join(format!( - "compass-persist-test-{}-{}-{}", - std::process::id(), - nanos, - N.fetch_add(1, Ordering::SeqCst) - )) - } - - fn empty_embed_state() -> EmbedState { - // No embedding models loaded. Safe for the persistence test because - // we provide chunks without text-only embedding requirements. Any - // call to embed_query returns Err and the ingest path tolerates that. - EmbedState { - bge: None, - distilled: None, - } - } - - fn make_ingest_chunk(file_id: &str, text: &str) -> IngestChunk { - IngestChunk { - client_id: None, - file_id: file_id.to_string(), - chunk_index: 0, - page: None, - text: text.to_string(), - metadata: HashMap::new(), - doc_type: "chunk".to_string(), - parent_id: None, - parent_ref: None, - group_id: None, - embeddings: HashMap::new(), - embedding: None, - } - } - - #[tokio::test] - async fn chunks_persist_across_manager_restart() { - let data_dir = unique_data_dir(); - std::fs::create_dir_all(&data_dir).unwrap(); - let embed = empty_embed_state(); - - // First manager lifetime: create collection, ingest three chunks, - // then drop the manager to close all file handles (including redb). - { - let manager = CollectionManager::new(&data_dir).await.unwrap(); - manager - .create_collection("persist-test", None, None, None) - .await - .unwrap(); - let to_ingest = vec![ - make_ingest_chunk("f1", "first chunk"), - make_ingest_chunk("f2", "second chunk"), - make_ingest_chunk("f3", "third chunk"), - ]; - let (ingested, _) = manager - .ingest("persist-test", to_ingest, &embed) - .await - .unwrap(); - assert_eq!(ingested, 3, "ingest call reports 3 chunks written"); - // manager dropped here - } - - // Second manager: same data dir, must rehydrate chunks from disk. - let manager2 = CollectionManager::new(&data_dir).await.unwrap(); - let (texts, ids) = manager2.get_all_chunk_data("persist-test").await.unwrap(); - - assert_eq!( - ids.len(), - 3, - "expected 3 chunks rehydrated from disk after manager restart, got {}", - ids.len() - ); - let mut sorted_texts = texts.clone(); - sorted_texts.sort(); - assert_eq!( - sorted_texts, - vec![ - "first chunk".to_string(), - "second chunk".to_string(), - "third chunk".to_string(), - ], - "chunk texts should match what was ingested before the restart" - ); - - // Cleanup - let _ = std::fs::remove_dir_all(&data_dir); - } - - #[tokio::test] - async fn next_id_advances_correctly_after_rehydration() { - // After rehydration, next_id should be max(seen) + 1 so new ingests - // don't collide with persisted IDs. Verify by ingesting again after - // restart and checking the new chunk got a fresh ID. - let data_dir = unique_data_dir(); - std::fs::create_dir_all(&data_dir).unwrap(); - let embed = empty_embed_state(); - - // Round 1: ingest two chunks (IDs 0, 1) - { - let manager = CollectionManager::new(&data_dir).await.unwrap(); - manager - .create_collection("next-id-test", None, None, None) - .await - .unwrap(); - manager - .ingest( - "next-id-test", - vec![ - make_ingest_chunk("f0", "round-one-a"), - make_ingest_chunk("f1", "round-one-b"), - ], - &embed, - ) - .await - .unwrap(); - } - - // Round 2: restart and ingest one more chunk. The new chunk's ID - // should be 2, not 0. - let manager2 = CollectionManager::new(&data_dir).await.unwrap(); - manager2 - .ingest( - "next-id-test", - vec![make_ingest_chunk("f2", "round-two")], - &embed, - ) - .await - .unwrap(); - let (_, ids) = manager2.get_all_chunk_data("next-id-test").await.unwrap(); - let mut sorted_ids = ids.clone(); - sorted_ids.sort(); - assert_eq!( - sorted_ids, - vec![0, 1, 2], - "next_id must advance past max persisted id, got ids: {:?}", - sorted_ids - ); - - let _ = std::fs::remove_dir_all(&data_dir); - } -} +mod persistence_tests; #[cfg(test)] -mod validate_name_segment_tests { - use super::validate_name_segment; - - #[test] - fn accepts_simple_names() { - assert!(validate_name_segment("my-collection", "Collection").is_ok()); - assert!(validate_name_segment("harrier", "Vector space").is_ok()); - assert!(validate_name_segment("qwen3-vl", "Vector space").is_ok()); - assert!(validate_name_segment("a", "Collection").is_ok()); - } - - #[test] - fn rejects_empty() { - let err = validate_name_segment("", "Vector space").expect_err("empty name should error"); - assert!(err.to_string().contains("Vector space")); - } - - #[test] - fn rejects_path_traversal() { - // The whole reason this validator exists: a vector space name flows - // into on-disk paths like `/.bin`. A `../` segment - // must never be accepted. - for bad in [ - "../etc/passwd", - "..", - "foo/bar", - "foo\\bar", - "/abs", - "name with space", - "name.with.dot", - "name_with_underscore", // hyphens only, no underscores - "tab\there", - "name\nwith\nnewline", - ] { - assert!( - validate_name_segment(bad, "Vector space").is_err(), - "validator must reject {bad:?}" - ); - } - } - - #[test] - fn rejects_unicode_lookalikes() { - // Cyrillic 'а' (U+0430) looks like 'a' but is not ASCII. - assert!(validate_name_segment("\u{0430}bc", "Collection").is_err()); - assert!(validate_name_segment("emoji-🚀", "Collection").is_err()); - } -} +mod validate_name_segment_tests; #[cfg(test)] -mod filter_aware_search_tests { - //! End-to-end test of filter-aware /search + /explain (a follow-up). - //! - //! Builds a real CollectionManager, ingests chunks with caller-provided - //! embeddings (skipping the in-process BGE model), runs filtered hybrid - //! search, and asserts: - //! 1. All hits respect the filter (filter-aware path, not post-filter). - //! 2. The /explain field is populated when requested and absent when not. - //! 3. Filter selectivity is reported correctly. - - use super::*; - use crate::embed::EmbedState; - use std::collections::HashMap; - use std::sync::atomic::{AtomicU64, Ordering}; - - fn unique_data_dir() -> std::path::PathBuf { - static N: AtomicU64 = AtomicU64::new(0); - let nanos = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_nanos(); - std::env::temp_dir().join(format!( - "compass-filter-search-test-{}-{}-{}", - std::process::id(), - nanos, - N.fetch_add(1, Ordering::SeqCst) - )) - } - - fn embed_state() -> EmbedState { - EmbedState { - bge: None, - distilled: None, - } - } - - /// Deterministic 4-dim unit vector seeded from an integer. - fn pseudo_vec(seed: u64) -> Vec { - let mut state = seed - .wrapping_mul(6364136223846793005) - .wrapping_add(1442695040888963407); - let mut v = Vec::with_capacity(4); - for _ in 0..4 { - state = state - .wrapping_mul(6364136223846793005) - .wrapping_add(1442695040888963407); - let f = (state >> 11) as f32 / (1u64 << 53) as f32 * 2.0 - 1.0; - v.push(f); - } - let norm: f32 = v.iter().map(|x| x * x).sum::().sqrt(); - if norm > 0.0 { - for x in &mut v { - *x /= norm; - } - } - v - } - - fn ingest_with(org: &str, idx: u32) -> IngestChunk { - let mut metadata = HashMap::new(); - metadata.insert("org_id".to_string(), MetadataValue::String(org.to_string())); - metadata.insert( - "created_at".to_string(), - MetadataValue::Int(1_700_000_000 + idx as i64), - ); - let mut embeddings = HashMap::new(); - embeddings.insert("default".to_string(), pseudo_vec(idx as u64 + 1)); - IngestChunk { - client_id: None, - file_id: format!("f{idx}"), - chunk_index: 0, - page: None, - text: format!("chunk-{idx}"), - metadata, - doc_type: "chunk".to_string(), - parent_id: None, - parent_ref: None, - group_id: None, - embeddings, - embedding: None, - } - } - - #[tokio::test] - async fn filter_aware_search_returns_only_matching_chunks() { - let data_dir = unique_data_dir(); - std::fs::create_dir_all(&data_dir).unwrap(); - let embed = embed_state(); - let manager = CollectionManager::new(&data_dir).await.unwrap(); - manager - .create_collection("filter-search", None, Some(4), None) - .await - .unwrap(); - - // 100 chunks: 20 from "acme", 80 from "widgets". - let mut chunks = Vec::new(); - for i in 0..100u32 { - let org = if i % 5 == 0 { "acme" } else { "widgets" }; - chunks.push(ingest_with(org, i)); - } - manager - .ingest("filter-search", chunks, &embed) - .await - .unwrap(); - - // Search with filter org_id=acme. ALL hits must come from acme. - let mut filters = HashMap::new(); - filters.insert( - "org_id".to_string(), - FilterValue::Exact(MetadataValue::String("acme".into())), - ); - let req = SearchRequest { - query: "chunk".to_string(), - mode: "hybrid".to_string(), - vector_space: None, - top_k: 10, - query_vector: Some(pseudo_vec(99_999)), - filters, - score_weights: None, - recency: None, - recency_preset: None, - recency_field: None, - boosts: Vec::new(), - relationship_boost: None, - explain: true, - include_relations: false, - relation_types: None, - relation_direction: RelationDirection::Outgoing, - }; - let (hits, total, _took_us, explain) = - manager.search("filter-search", &req, &embed).await.unwrap(); - - assert!(!hits.is_empty(), "search returned no hits"); - for (chunk, _, _, _, _) in &hits { - assert_eq!( - chunk.metadata.get("org_id"), - Some(&MetadataValue::String("acme".into())), - "all hits must satisfy the filter; got chunk {} with org_id {:?}", - chunk.id, - chunk.metadata.get("org_id") - ); - } - assert!( - total <= 20, - "no more than 20 hits possible at 20% selectivity" - ); - - // /explain should be populated. - let explain = explain.expect("explain plan requested but not returned"); - assert_eq!(explain.filter.eligible_count, 20); - assert_eq!(explain.filter.universe_count, 100); - assert!((explain.filter.selectivity - 0.20).abs() < 1e-9); - assert!( - matches!(explain.ann.engine.as_str(), "hnsw" | "brute_force"), - "ann engine reported as {}", - explain.ann.engine - ); - - let _ = std::fs::remove_dir_all(&data_dir); - } - - #[tokio::test] - async fn explain_absent_when_not_requested() { - let data_dir = unique_data_dir(); - std::fs::create_dir_all(&data_dir).unwrap(); - let embed = embed_state(); - let manager = CollectionManager::new(&data_dir).await.unwrap(); - manager - .create_collection("no-explain", None, Some(4), None) - .await - .unwrap(); - manager - .ingest("no-explain", vec![ingest_with("acme", 0)], &embed) - .await - .unwrap(); - - let req = SearchRequest { - query: "chunk".to_string(), - mode: "semantic".to_string(), - vector_space: None, - top_k: 1, - query_vector: Some(pseudo_vec(42)), - filters: HashMap::new(), - score_weights: None, - recency: None, - recency_preset: None, - recency_field: None, - boosts: Vec::new(), - relationship_boost: None, - explain: false, - include_relations: false, - relation_types: None, - relation_direction: RelationDirection::Outgoing, - }; - let (_hits, _total, _took, explain) = - manager.search("no-explain", &req, &embed).await.unwrap(); - assert!(explain.is_none(), "explain must be None when not requested"); - - let _ = std::fs::remove_dir_all(&data_dir); - } - - #[tokio::test] - async fn relations_crud_and_search_enrichment() { - let data_dir = unique_data_dir(); - std::fs::create_dir_all(&data_dir).unwrap(); - let embed = embed_state(); - let manager = CollectionManager::new(&data_dir).await.unwrap(); - manager - .create_collection("rel-search", None, Some(4), None) - .await - .unwrap(); - - // Ingest 5 chunks -> ids 0..5 in order. - let chunks: Vec<_> = (0..5u32).map(|i| ingest_with("acme", i)).collect(); - manager.ingest("rel-search", chunks, &embed).await.unwrap(); - - // Create relations: 0 --cites--> 1, 0 --cites--> 2, 3 --supersedes--> 0. - let created = manager - .create_relations( - "rel-search", - vec![ - CreateRelation { - source_chunk_id: 0, - target_chunk_id: 1, - target_document_id: None, - relation_type: "cites".into(), - metadata: HashMap::new(), - }, - CreateRelation { - source_chunk_id: 0, - target_chunk_id: 2, - target_document_id: None, - relation_type: "cites".into(), - metadata: HashMap::new(), - }, - CreateRelation { - source_chunk_id: 3, - target_chunk_id: 0, - target_document_id: None, - relation_type: "supersedes".into(), - metadata: HashMap::new(), - }, - ], - ) - .await - .unwrap(); - assert_eq!(created.len(), 3); - assert!(created.iter().all(|r| !r.relation_id.is_empty())); - assert!(created.iter().all(|r| r.target_status == "found")); - - // Self-relation is rejected. - let bad = manager - .create_relations( - "rel-search", - vec![CreateRelation { - source_chunk_id: 1, - target_chunk_id: 1, - target_document_id: None, - relation_type: "cites".into(), - metadata: HashMap::new(), - }], - ) - .await; - assert!(bad.is_err()); - - // Direction filters. - let out = manager - .get_chunk_relations("rel-search", 0, RelationDirection::Outgoing, None) - .await - .unwrap(); - assert_eq!(out.len(), 2); - let inc = manager - .get_chunk_relations("rel-search", 0, RelationDirection::Incoming, None) - .await - .unwrap(); - assert_eq!(inc.len(), 1); - assert_eq!(inc[0].relation_type, "supersedes"); - - // Type filter. - let cites = manager - .get_chunk_relations( - "rel-search", - 0, - RelationDirection::Both, - Some(&["cites".to_string()]), - ) - .await - .unwrap(); - assert_eq!(cites.len(), 2); - - let base_req = |include: bool| SearchRequest { - query: "chunk".to_string(), - mode: "semantic".to_string(), - vector_space: None, - top_k: 10, - query_vector: Some(pseudo_vec(7)), - filters: HashMap::new(), - score_weights: None, - recency: None, - recency_preset: None, - recency_field: None, - boosts: Vec::new(), - relationship_boost: None, - explain: false, - include_relations: include, - relation_types: None, - relation_direction: RelationDirection::Outgoing, - }; - - // Without include_relations -> hits carry None. - let (hits_off, _, _, _) = manager - .search("rel-search", &base_req(false), &embed) - .await - .unwrap(); - assert!(hits_off.iter().all(|(_, _, _, _, rels)| rels.is_none())); - - // With include_relations -> chunk 0's hit carries its 2 outgoing cites. - let (hits_on, _, _, _) = manager - .search("rel-search", &base_req(true), &embed) - .await - .unwrap(); - let chunk0 = hits_on - .iter() - .find(|(c, _, _, _, _)| c.id == 0) - .expect("chunk 0 in results"); - let rels = chunk0.4.as_ref().expect("Some(relations) when requested"); - assert_eq!(rels.len(), 2, "chunk 0 has 2 outgoing cites"); - assert!(rels.iter().all(|r| r.relation_type == "cites")); - assert!(rels.iter().all(|r| r.target_status == "found")); - - // Delete one relation; outgoing from 0 drops to 1. - let rid = &created[0].relation_id; - assert!(manager.delete_relation("rel-search", rid).await.unwrap()); - let out2 = manager - .get_chunk_relations("rel-search", 0, RelationDirection::Outgoing, None) - .await - .unwrap(); - assert_eq!(out2.len(), 1); - - let _ = std::fs::remove_dir_all(&data_dir); - } - - #[tokio::test] - async fn delete_removes_from_search_and_survives_restart() { - let data_dir = unique_data_dir(); - std::fs::create_dir_all(&data_dir).unwrap(); - let embed = embed_state(); - - { - let manager = CollectionManager::new(&data_dir).await.unwrap(); - manager - .create_collection("del", None, Some(4), None) - .await - .unwrap(); - // Ingest 10 chunks (ids 0..10), org=acme. - let chunks: Vec<_> = (0..10u32).map(|i| ingest_with("acme", i)).collect(); - manager.ingest("del", chunks, &embed).await.unwrap(); - - // Delete chunk id 3 by id. - let n = manager.delete_chunks("del", &[3]).await.unwrap(); - assert_eq!(n, 1); - // Re-deleting is a no-op. - assert_eq!(manager.delete_chunks("del", &[3]).await.unwrap(), 0); - - // A search must never return the deleted id. - let req = SearchRequest { - query: "chunk".to_string(), - mode: "semantic".to_string(), - vector_space: None, - top_k: 20, - query_vector: Some(pseudo_vec(4)), - filters: HashMap::new(), - score_weights: None, - recency: None, - recency_preset: None, - recency_field: None, - boosts: Vec::new(), - relationship_boost: None, - explain: false, - include_relations: false, - relation_types: None, - relation_direction: RelationDirection::Outgoing, - }; - let (hits, _, _, _) = manager.search("del", &req, &embed).await.unwrap(); - assert!( - hits.iter().all(|(c, _, _, _, _)| c.id != 3), - "deleted chunk must not appear in results" - ); - - // Delete-by-filter: delete everything with file_id f5 (chunk 5). - let mut filters = HashMap::new(); - filters.insert( - "file_id".to_string(), - FilterValue::Exact(MetadataValue::String("f5".into())), - ); - // ingest_with doesn't set file_id in metadata, so use a metadata field. - // org_id=acme matches all remaining -> delete the rest via a scan. - let mut org_filter = HashMap::new(); - org_filter.insert( - "org_id".to_string(), - FilterValue::Exact(MetadataValue::String("acme".into())), - ); - let deleted = manager.delete_by_filter("del", &org_filter).await.unwrap(); - // 10 ingested - 1 already deleted (id 3) = 9 remaining deleted now. - assert_eq!(deleted, 9); - let _ = filters; - } - - // Restart: tombstones must persist. Reopen the manager over the same dir. - { - let manager = CollectionManager::new(&data_dir).await.unwrap(); - let req = SearchRequest { - query: "chunk".to_string(), - mode: "semantic".to_string(), - vector_space: None, - top_k: 20, - query_vector: Some(pseudo_vec(4)), - filters: HashMap::new(), - score_weights: None, - recency: None, - recency_preset: None, - recency_field: None, - boosts: Vec::new(), - relationship_boost: None, - explain: false, - include_relations: false, - relation_types: None, - relation_direction: RelationDirection::Outgoing, - }; - let (hits, _, _, _) = manager.search("del", &req, &embed).await.unwrap(); - assert!( - hits.is_empty(), - "all chunks deleted; none should survive restart, got {}", - hits.len() - ); - } - - let _ = std::fs::remove_dir_all(&data_dir); - } - - // F6 coverage: deleting a chunk that participates in relations must prune - // those edges (both endpoints), not leave dangling references. - #[tokio::test] - async fn delete_chunk_prunes_its_relations() { - let data_dir = unique_data_dir(); - std::fs::create_dir_all(&data_dir).unwrap(); - let embed = embed_state(); - let manager = CollectionManager::new(&data_dir).await.unwrap(); - manager - .create_collection("delrel", None, Some(4), None) - .await - .unwrap(); - let chunks: Vec<_> = (0..3u32).map(|i| ingest_with("acme", i)).collect(); - manager.ingest("delrel", chunks, &embed).await.unwrap(); - - // 0 -> 1, 2 -> 0 (chunk 0 is both a source and a target). - manager - .create_relations( - "delrel", - vec![ - CreateRelation { - source_chunk_id: 0, - target_chunk_id: 1, - target_document_id: None, - relation_type: "cites".into(), - metadata: HashMap::new(), - }, - CreateRelation { - source_chunk_id: 2, - target_chunk_id: 0, - target_document_id: None, - relation_type: "cites".into(), - metadata: HashMap::new(), - }, - ], - ) - .await - .unwrap(); - - // Delete chunk 0 — both edges (as source and as target) must be pruned. - manager.delete_chunks("delrel", &[0]).await.unwrap(); - - assert!( - manager - .get_chunk_relations("delrel", 0, RelationDirection::Both, None) - .await - .unwrap() - .is_empty(), - "deleted chunk's own edges gone" - ); - // Chunk 2's outgoing edge (to deleted 0) must also be gone. - assert!( - manager - .get_chunk_relations("delrel", 2, RelationDirection::Outgoing, None) - .await - .unwrap() - .is_empty(), - "edge pointing AT the deleted chunk must be pruned" - ); - let _ = std::fs::remove_dir_all(&data_dir); - } - - // A stale tombstone must not suppress a NEWLY-ingested chunk. Since next_id - // is a monotonic high-water mark, re-ingest gets a fresh id that was never - // tombstoned, so it's fully searchable. - #[tokio::test] - async fn delete_then_reingest_new_chunk_is_searchable() { - let data_dir = unique_data_dir(); - std::fs::create_dir_all(&data_dir).unwrap(); - let embed = embed_state(); - let manager = CollectionManager::new(&data_dir).await.unwrap(); - manager - .create_collection("reing", None, Some(4), None) - .await - .unwrap(); - manager - .ingest("reing", vec![ingest_with("acme", 0)], &embed) - .await - .unwrap(); - manager.delete_chunks("reing", &[0]).await.unwrap(); - - // Re-ingest: gets id 1 (next_id advanced), NOT the tombstoned id 0. - manager - .ingest("reing", vec![ingest_with("acme", 9)], &embed) - .await - .unwrap(); - - let req = SearchRequest { - query: "chunk".to_string(), - mode: "semantic".to_string(), - vector_space: None, - top_k: 10, - query_vector: Some(pseudo_vec(10)), - filters: HashMap::new(), - score_weights: None, - recency: None, - recency_preset: None, - recency_field: None, - boosts: Vec::new(), - relationship_boost: None, - explain: false, - include_relations: false, - relation_types: None, - relation_direction: RelationDirection::Outgoing, - }; - let (hits, _, _, _) = manager.search("reing", &req, &embed).await.unwrap(); - assert_eq!(hits.len(), 1, "the re-ingested chunk must be searchable"); - assert_eq!(hits[0].0.id, 1, "re-ingest got a fresh (untombstoned) id"); - let _ = std::fs::remove_dir_all(&data_dir); - } -} +mod filter_aware_search_tests; #[cfg(all(test, feature = "object-storage"))] -mod cloud_ingest_tests { - //! Verifies that in object-storage (cloud) mode, ingest mirrors the batch - //! into the LSM as a WAL fragment + CAS-committed manifest — the S3-native - //! path. Uses the in-memory object_store backend, which - //! exercises the identical `Storage`/`ObjectStoreBackend` code an S3 bucket - //! would, without needing real credentials. - - use super::*; - use crate::embed::EmbedState; - use crate::storage::object_store_backend::ObjectStoreBackend; - use std::collections::HashMap; - use std::sync::atomic::{AtomicU64, Ordering}; - - fn unique_data_dir() -> std::path::PathBuf { - static N: AtomicU64 = AtomicU64::new(0); - std::env::temp_dir().join(format!( - "compass-cloud-ingest-{}-{}", - std::process::id(), - N.fetch_add(1, Ordering::Relaxed) - )) - } - - fn embed_state() -> EmbedState { - // No models needed: chunks carry precomputed embeddings. - EmbedState { - bge: None, - distilled: None, - } - } - - fn ingest_chunk(idx: u32) -> IngestChunk { - let mut metadata = HashMap::new(); - metadata.insert( - "org_id".to_string(), - MetadataValue::String("acme".to_string()), - ); - let mut embeddings = HashMap::new(); - embeddings.insert("default".to_string(), vec![0.1, 0.2, 0.3, 0.4]); - IngestChunk { - client_id: None, - file_id: format!("f{idx}"), - chunk_index: 0, - page: None, - text: format!("chunk-{idx}"), - metadata, - doc_type: "chunk".to_string(), - parent_id: None, - parent_ref: None, - group_id: None, - embeddings, - embedding: None, - } - } - - #[tokio::test] - async fn ingest_writes_wal_fragment_and_manifest_to_object_storage() { - let data_dir = unique_data_dir(); - std::fs::create_dir_all(&data_dir).unwrap(); - let embed = embed_state(); - - // In-memory object storage backend (same code path as s3://). - let storage: Arc = Arc::new(ObjectStoreBackend::from_store( - std::sync::Arc::new(object_store::memory::InMemory::new()), - "object-store:memory", - )); - let manager = CollectionManager::new_with_storage(&data_dir, storage.clone()) - .await - .unwrap(); - manager - .create_collection("cloudcoll", None, Some(4), None) - .await - .unwrap(); - - // Ingest two batches. - manager - .ingest("cloudcoll", vec![ingest_chunk(0), ingest_chunk(1)], &embed) - .await - .unwrap(); - manager - .ingest("cloudcoll", vec![ingest_chunk(2)], &embed) - .await - .unwrap(); - - // The manifest exists and records two WAL fragments. - let (manifest, version) = crate::storage::lsm::read_manifest(storage.as_ref(), "cloudcoll") - .await - .unwrap(); - assert!(version.is_some(), "manifest must exist in object storage"); - assert_eq!(manifest.fragments.len(), 2, "one fragment per ingest batch"); - assert_eq!(manifest.next_seq, 2); - - // The WAL fragment objects exist and decode back to the ingested chunks. - let frags = crate::storage::lsm::read_uncompacted_fragments( - storage.as_ref(), - "cloudcoll", - &manifest, - ) - .await - .unwrap(); - assert_eq!(frags.len(), 2); - - let batch0: Vec = serde_json::from_slice(&frags[0].1).unwrap(); - assert_eq!(batch0.len(), 2); - assert_eq!(batch0[0].text, "chunk-0"); - let batch1: Vec = serde_json::from_slice(&frags[1].1).unwrap(); - assert_eq!(batch1.len(), 1); - assert_eq!(batch1[0].text, "chunk-2"); - - // Total records across fragments == total chunks ingested. - let total: u64 = manifest.fragments.iter().map(|f| f.records).sum(); - assert_eq!(total, 3); - - let _ = std::fs::remove_dir_all(&data_dir); - } - - #[tokio::test] - async fn local_mode_writes_no_wal() { - // Sanity: a local-disk manager must NOT create any WAL/manifest objects. - let data_dir = unique_data_dir(); - std::fs::create_dir_all(&data_dir).unwrap(); - let embed = embed_state(); - let manager = CollectionManager::new(&data_dir).await.unwrap(); - manager - .create_collection("localcoll", None, Some(4), None) - .await - .unwrap(); - manager - .ingest("localcoll", vec![ingest_chunk(0)], &embed) - .await - .unwrap(); - - // No manifest object should exist under the collection prefix. - let manifest_path = data_dir.join("localcoll").join("manifest"); - assert!( - !manifest_path.exists(), - "local mode must not write an LSM manifest" - ); - let _ = std::fs::remove_dir_all(&data_dir); - } - - #[tokio::test] - async fn delete_writes_tombstone_wal_fragment() { - use crate::storage::lsm::{read_manifest, read_uncompacted_fragments, FragmentKind}; - - let data_dir = unique_data_dir(); - std::fs::create_dir_all(&data_dir).unwrap(); - let embed = embed_state(); - let storage: Arc = Arc::new(ObjectStoreBackend::from_store( - std::sync::Arc::new(object_store::memory::InMemory::new()), - "object-store:memory", - )); - let manager = CollectionManager::new_with_storage(&data_dir, storage.clone()) - .await - .unwrap(); - manager - .create_collection("delcloud", None, Some(4), None) - .await - .unwrap(); - manager - .ingest( - "delcloud", - vec![ingest_chunk(0), ingest_chunk(1), ingest_chunk(2)], - &embed, - ) - .await - .unwrap(); - - // Delete chunk 1 -> a tombstone WAL fragment lands in object storage. - let n = manager.delete_chunks("delcloud", &[1]).await.unwrap(); - assert_eq!(n, 1); - - let (manifest, _) = read_manifest(storage.as_ref(), "delcloud").await.unwrap(); - // seq 0 = data fragment (the ingest), seq 1 = tombstone fragment. - assert_eq!(manifest.fragments.len(), 2); - assert_eq!(manifest.fragments[0].kind, FragmentKind::Data); - assert_eq!(manifest.fragments[1].kind, FragmentKind::Tombstone); - - // The tombstone fragment decodes to the deleted id [1]. - let frags = read_uncompacted_fragments(storage.as_ref(), "delcloud", &manifest) - .await - .unwrap(); - let deleted_ids: Vec = serde_json::from_slice(&frags[1].1).unwrap(); - assert_eq!(deleted_ids, vec![1]); - - let _ = std::fs::remove_dir_all(&data_dir); - } - - // THE structural fix: a cloud collection must survive a restart on a FRESH - // local disk by rebuilding from S3. Ingest, delete one, then drop the manager - // AND wipe the local data dir, then reload from the SAME object store — the - // data (minus the deleted chunk) must come back. - #[tokio::test] - async fn cloud_restart_rehydrates_from_object_storage() { - let embed = embed_state(); - // Shared object store persists across the "restart". - let store = std::sync::Arc::new(object_store::memory::InMemory::new()); - - let data_dir_a = unique_data_dir(); - std::fs::create_dir_all(&data_dir_a).unwrap(); - { - let storage: Arc = Arc::new(ObjectStoreBackend::from_store( - store.clone(), - "object-store:memory", - )); - let m = CollectionManager::new_with_storage(&data_dir_a, storage) - .await - .unwrap(); - m.create_collection("survive", None, Some(4), None) - .await - .unwrap(); - m.ingest( - "survive", - vec![ingest_chunk(0), ingest_chunk(1), ingest_chunk(2)], - &embed, - ) - .await - .unwrap(); - m.delete_chunks("survive", &[1]).await.unwrap(); - } - // Simulate node loss: wipe the local disk entirely. - std::fs::remove_dir_all(&data_dir_a).unwrap(); - - // Restart on a BRAND-NEW empty local dir, same object store. - let data_dir_b = unique_data_dir(); - std::fs::create_dir_all(&data_dir_b).unwrap(); - let storage: Arc = Arc::new(ObjectStoreBackend::from_store( - store.clone(), - "object-store:memory", - )); - let m2 = CollectionManager::new_with_storage(&data_dir_b, storage) - .await - .unwrap(); - - // The collection is back, recovered from S3. - let info = m2.get_collection("survive").await; - assert!(info.is_some(), "collection must be recovered from S3"); - - // Search finds the surviving chunks (0 and 2), not the deleted one (1). - let req = SearchRequest { - query: "chunk".to_string(), - mode: "semantic".to_string(), - vector_space: None, - top_k: 10, - query_vector: Some(vec![0.1, 0.2, 0.3, 0.4]), - filters: HashMap::new(), - score_weights: None, - recency: None, - recency_preset: None, - recency_field: None, - boosts: Vec::new(), - relationship_boost: None, - explain: false, - include_relations: false, - relation_types: None, - relation_direction: RelationDirection::Outgoing, - }; - let (hits, _, _, _) = m2.search("survive", &req, &embed).await.unwrap(); - let ids: std::collections::HashSet = hits.iter().map(|(c, _, _, _, _)| c.id).collect(); - assert!(ids.contains(&0), "chunk 0 recovered"); - assert!(ids.contains(&2), "chunk 2 recovered"); - assert!(!ids.contains(&1), "deleted chunk 1 must NOT reappear"); - - let _ = std::fs::remove_dir_all(&data_dir_b); - } - - // Compaction folds segments+fragments into one segment, dropping tombstoned - // records so they can never resurrect. - #[tokio::test] - async fn compaction_reclaims_tombstoned_data() { - use crate::storage::lsm::read_manifest; - let embed = embed_state(); - let store = std::sync::Arc::new(object_store::memory::InMemory::new()); - let data_dir = unique_data_dir(); - std::fs::create_dir_all(&data_dir).unwrap(); - let storage: Arc = Arc::new(ObjectStoreBackend::from_store( - store.clone(), - "object-store:memory", - )); - let m = CollectionManager::new_with_storage(&data_dir, storage.clone()) - .await - .unwrap(); - m.create_collection("comp", None, Some(4), None) - .await - .unwrap(); - m.ingest( - "comp", - vec![ingest_chunk(0), ingest_chunk(1), ingest_chunk(2)], - &embed, - ) - .await - .unwrap(); - m.delete_chunks("comp", &[1]).await.unwrap(); - - // Before: manifest has data + tombstone fragments, no segment. - let (before, _) = read_manifest(storage.as_ref(), "comp").await.unwrap(); - assert!(before.segments.is_empty()); - assert_eq!(before.fragments.len(), 2); - - // Compact. - let live = m.compact_collection("comp").await.unwrap(); - assert_eq!(live, 2, "2 live records (0 and 2) after dropping deleted 1"); - - // After: one segment, no fragments. - let (after, _) = read_manifest(storage.as_ref(), "comp").await.unwrap(); - assert_eq!(after.segments.len(), 1); - assert!(after.fragments.is_empty()); - - // The compacted segment contains only live chunks (0, 2) — deleted 1 gone. - let seg = - crate::storage::lsm::read_segment(storage.as_ref(), "comp", &after.segments[0].id) - .await - .unwrap(); - let segment: cloud::Segment = serde_json::from_slice(&seg).unwrap(); - let ids: std::collections::HashSet = segment.chunks.iter().map(|c| c.id).collect(); - assert!(ids.contains(&0) && ids.contains(&2)); - assert!( - !ids.contains(&1), - "compaction must drop the tombstoned chunk" - ); - - let _ = std::fs::remove_dir_all(&data_dir); - } - - // #1 regression: typed RELATIONS must survive a cold restart from S3 (the bug - // where relation_store was local-redb-only and vanished on rebuild). Create - // relations, wipe the local disk, restart on a fresh dir, relations return. - #[tokio::test] - async fn cloud_restart_recovers_relations() { - let embed = embed_state(); - let store = std::sync::Arc::new(object_store::memory::InMemory::new()); - - let data_dir_a = unique_data_dir(); - std::fs::create_dir_all(&data_dir_a).unwrap(); - { - let storage: Arc = Arc::new(ObjectStoreBackend::from_store( - store.clone(), - "object-store:memory", - )); - let m = CollectionManager::new_with_storage(&data_dir_a, storage) - .await - .unwrap(); - m.create_collection("relsurv", None, Some(4), None) - .await - .unwrap(); - m.ingest( - "relsurv", - vec![ingest_chunk(0), ingest_chunk(1), ingest_chunk(2)], - &embed, - ) - .await - .unwrap(); - // Create two relations, then delete one — only the survivor should - // come back. - let created = m - .create_relations( - "relsurv", - vec![ - CreateRelation { - source_chunk_id: 0, - target_chunk_id: 1, - target_document_id: None, - relation_type: "cites".into(), - metadata: HashMap::new(), - }, - CreateRelation { - source_chunk_id: 0, - target_chunk_id: 2, - target_document_id: None, - relation_type: "supersedes".into(), - metadata: HashMap::new(), - }, - ], - ) - .await - .unwrap(); - m.delete_relation("relsurv", &created[1].relation_id) - .await - .unwrap(); - } - // Node loss: wipe local disk. - std::fs::remove_dir_all(&data_dir_a).unwrap(); - - // Restart on a fresh local dir, same object store. - let data_dir_b = unique_data_dir(); - std::fs::create_dir_all(&data_dir_b).unwrap(); - let storage: Arc = Arc::new(ObjectStoreBackend::from_store( - store.clone(), - "object-store:memory", - )); - let m2 = CollectionManager::new_with_storage(&data_dir_b, storage) - .await - .unwrap(); - - // The surviving relation (0 --cites--> 1) must be recovered from S3; - // the deleted one (0 --supersedes--> 2) must NOT reappear. - let out = m2 - .get_chunk_relations("relsurv", 0, RelationDirection::Outgoing, None) - .await - .unwrap(); - assert_eq!(out.len(), 1, "exactly one relation should survive restart"); - assert_eq!(out[0].relation_type, "cites"); - assert_eq!(out[0].target_chunk_id, 1); - - let _ = std::fs::remove_dir_all(&data_dir_b); - } - - // #3: auto-compaction. Ingest enough batches to cross the fragment threshold; - // the background trigger should fold them into a segment. We poll briefly for - // the detached task to run, then assert the WAL is bounded. - #[tokio::test] - async fn auto_compaction_bounds_the_wal() { - let embed = embed_state(); - let store = std::sync::Arc::new(object_store::memory::InMemory::new()); - let data_dir = unique_data_dir(); - std::fs::create_dir_all(&data_dir).unwrap(); - let storage: Arc = Arc::new(ObjectStoreBackend::from_store( - store.clone(), - "object-store:memory", - )); - let m = CollectionManager::new_with_storage(&data_dir, storage.clone()) - .await - .unwrap(); - m.create_collection("auto", None, Some(4), None) - .await - .unwrap(); - - // One chunk per ingest = one fragment per ingest. Cross the threshold. - let batches = AUTO_COMPACT_FRAGMENT_THRESHOLD + 2; - for i in 0..batches { - m.ingest("auto", vec![ingest_chunk(i as u32)], &embed) - .await - .unwrap(); - } - - // Poll up to ~3s for the detached auto-compaction to land a segment and - // shrink the uncompacted fragment set. - let mut compacted = false; - for _ in 0..30 { - let (man, _) = crate::storage::lsm::read_manifest(storage.as_ref(), "auto") - .await - .unwrap(); - if !man.segments.is_empty() - && man.uncompacted().count() < AUTO_COMPACT_FRAGMENT_THRESHOLD - { - compacted = true; - break; - } - tokio::time::sleep(std::time::Duration::from_millis(100)).await; - } - assert!( - compacted, - "auto-compaction should have folded the WAL into a segment" - ); - - // All data still present after auto-compaction (via materialize). - let (man, _) = crate::storage::lsm::read_manifest(storage.as_ref(), "auto") - .await - .unwrap(); - let mat = cloud::materialize(storage.as_ref(), "auto", &man) - .await - .unwrap(); - assert_eq!(mat.chunks.len(), batches, "no data lost in auto-compaction"); - - let _ = std::fs::remove_dir_all(&data_dir); - } - - // Negative: auto-compaction must NOT fire below the fragment threshold (a - // regression dropping the threshold to ~0 would compact on every ingest). - #[tokio::test] - async fn auto_compaction_does_not_fire_below_threshold() { - let embed = embed_state(); - let store = std::sync::Arc::new(object_store::memory::InMemory::new()); - let data_dir = unique_data_dir(); - std::fs::create_dir_all(&data_dir).unwrap(); - let storage: Arc = Arc::new(ObjectStoreBackend::from_store( - store.clone(), - "object-store:memory", - )); - let m = CollectionManager::new_with_storage(&data_dir, storage.clone()) - .await - .unwrap(); - m.create_collection("below", None, Some(4), None) - .await - .unwrap(); - - // Well under the threshold: a handful of single-chunk ingests. - for i in 0..5u32 { - m.ingest("below", vec![ingest_chunk(i)], &embed) - .await - .unwrap(); - } - // Give any (wrongly) spawned compaction ample time to land a segment. - tokio::time::sleep(std::time::Duration::from_millis(500)).await; - - let (man, _) = crate::storage::lsm::read_manifest(storage.as_ref(), "below") - .await - .unwrap(); - assert!( - man.segments.is_empty(), - "auto-compaction must not fire below the threshold" - ); - assert_eq!(man.fragments.len(), 5, "all fragments still in the WAL"); - - let _ = std::fs::remove_dir_all(&data_dir); - } - - // F1 regression: compaction must physically GC old objects (deferred one - // cycle), not leak them forever. Ingest, compact twice, assert the first - // segment's object is deleted and the object count stays bounded. - #[tokio::test] - async fn compaction_gcs_old_objects() { - let embed = embed_state(); - let store = std::sync::Arc::new(object_store::memory::InMemory::new()); - let data_dir = unique_data_dir(); - std::fs::create_dir_all(&data_dir).unwrap(); - let storage: Arc = Arc::new(ObjectStoreBackend::from_store( - store.clone(), - "object-store:memory", - )); - let m = CollectionManager::new_with_storage(&data_dir, storage.clone()) - .await - .unwrap(); - m.create_collection("gc", None, Some(4), None) - .await - .unwrap(); - m.ingest("gc", vec![ingest_chunk(0), ingest_chunk(1)], &embed) - .await - .unwrap(); - - // First compaction → segment S1, stages the 1 fragment for next-cycle GC. - m.compact_collection("gc").await.unwrap(); - let (man1, _) = crate::storage::lsm::read_manifest(storage.as_ref(), "gc") - .await - .unwrap(); - let seg1_id = man1.segments[0].id.clone(); - // The old WAL fragment object is staged (still present this cycle). - assert_eq!(man1.pending_deletes.len(), 1); - - // Compact twice more (each cycle GCs the PRIOR cycle's staged objects, - // deferred one cycle for in-flight readers). After enough cycles, S1 is - // physically gone — the key point is it's GC'd, not leaked forever. - for i in 2..5u32 { - m.ingest("gc", vec![ingest_chunk(i)], &embed).await.unwrap(); - m.compact_collection("gc").await.unwrap(); - } - let (man2, _) = crate::storage::lsm::read_manifest(storage.as_ref(), "gc") - .await - .unwrap(); - - // Segment S1 must be physically deleted (GC'd after being folded away). - let s1_key = format!("gc/segments/{seg1_id}"); - assert!( - !storage.exists(&s1_key).await.unwrap(), - "old segment must be GC'd, not leaked" - ); - // Object count stays BOUNDED across many compaction cycles — proving no - // unbounded leak (the F1 bug would grow this without limit). - let all = storage.list("gc/").await.unwrap(); - assert!( - all.len() <= 5, - "object count must stay bounded across cycles, got {}", - all.len() - ); - // Data intact. - let mat = cloud::materialize(storage.as_ref(), "gc", &man2) - .await - .unwrap(); - assert_eq!(mat.chunks.len(), 5); - - let _ = std::fs::remove_dir_all(&data_dir); - } - - // Relations must survive COMPACTION-then-restart (segment-relations path), - // not just the fragment-replay path. - #[tokio::test] - async fn relations_survive_compaction_then_restart() { - let embed = embed_state(); - let store = std::sync::Arc::new(object_store::memory::InMemory::new()); - - let data_dir_a = unique_data_dir(); - std::fs::create_dir_all(&data_dir_a).unwrap(); - { - let storage: Arc = Arc::new(ObjectStoreBackend::from_store( - store.clone(), - "object-store:memory", - )); - let m = CollectionManager::new_with_storage(&data_dir_a, storage) - .await - .unwrap(); - m.create_collection("rc", None, Some(4), None) - .await - .unwrap(); - m.ingest("rc", vec![ingest_chunk(0), ingest_chunk(1)], &embed) - .await - .unwrap(); - m.create_relations( - "rc", - vec![CreateRelation { - source_chunk_id: 0, - target_chunk_id: 1, - target_document_id: None, - relation_type: "cites".into(), - metadata: HashMap::new(), - }], - ) - .await - .unwrap(); - // Compact so the relation lives in the SEGMENT, not a WAL fragment. - m.compact_collection("rc").await.unwrap(); - } - std::fs::remove_dir_all(&data_dir_a).unwrap(); - - let data_dir_b = unique_data_dir(); - std::fs::create_dir_all(&data_dir_b).unwrap(); - let storage: Arc = Arc::new(ObjectStoreBackend::from_store( - store.clone(), - "object-store:memory", - )); - let m2 = CollectionManager::new_with_storage(&data_dir_b, storage) - .await - .unwrap(); - let out = m2 - .get_chunk_relations("rc", 0, RelationDirection::Outgoing, None) - .await - .unwrap(); - assert_eq!(out.len(), 1, "relation must survive compaction+restart"); - assert_eq!(out[0].relation_type, "cites"); - - let _ = std::fs::remove_dir_all(&data_dir_b); - } - - // Concurrent ingests into the same collection: all chunks visible, all ids - // unique, no lost writes (stresses the lock drop/reacquire window). - #[tokio::test] - async fn concurrent_ingests_same_collection() { - let embed = std::sync::Arc::new(embed_state()); - let store = std::sync::Arc::new(object_store::memory::InMemory::new()); - let data_dir = unique_data_dir(); - std::fs::create_dir_all(&data_dir).unwrap(); - let storage: Arc = Arc::new(ObjectStoreBackend::from_store( - store.clone(), - "object-store:memory", - )); - let m = CollectionManager::new_with_storage(&data_dir, storage.clone()) - .await - .unwrap(); - m.create_collection("conc", None, Some(4), None) - .await - .unwrap(); - - let n = 12usize; - let mut handles = Vec::new(); - for i in 0..n { - let m2 = m.clone(); - let e2 = embed.clone(); - handles.push(tokio::spawn(async move { - m2.ingest("conc", vec![ingest_chunk(i as u32)], &e2).await - })); - } - for h in handles { - h.await.unwrap().unwrap(); - } - - // All N chunks present, ids 0..N unique (no collision from the lock gap). - let (man, _) = crate::storage::lsm::read_manifest(storage.as_ref(), "conc") - .await - .unwrap(); - let mat = cloud::materialize(storage.as_ref(), "conc", &man) - .await - .unwrap(); - assert_eq!(mat.chunks.len(), n, "all concurrent ingests durable"); - let ids: std::collections::HashSet = mat.chunks.keys().copied().collect(); - assert_eq!(ids.len(), n, "no duplicate/lost ids"); - assert_eq!(ids, (0..n as u64).collect()); - - let _ = std::fs::remove_dir_all(&data_dir); - } - - // next_id must NEVER regress across compaction + cold restart. Compaction - // physically drops tombstoned chunks; without the segment's stored max_id - // high-water mark, a fresh-disk rebuild would recompute next_id from the - // live set only and REUSE the deleted ids for new chunks. - #[tokio::test] - async fn no_id_reuse_after_compaction_and_cold_restart() { - let embed = embed_state(); - let store = std::sync::Arc::new(object_store::memory::InMemory::new()); - let data_dir_a = unique_data_dir(); - std::fs::create_dir_all(&data_dir_a).unwrap(); - { - let storage: Arc = Arc::new(ObjectStoreBackend::from_store( - store.clone(), - "object-store:memory", - )); - let m = CollectionManager::new_with_storage(&data_dir_a, storage) - .await - .unwrap(); - m.create_collection("idreuse", None, Some(4), None) - .await - .unwrap(); - // ids 0..3; delete the two HIGHEST, then compact them away. - m.ingest("idreuse", (0..4u32).map(ingest_chunk).collect(), &embed) - .await - .unwrap(); - m.delete_chunks("idreuse", &[2, 3]).await.unwrap(); - m.compact_collection("idreuse").await.unwrap(); - } - // Node loss: wipe local disk, cold-rebuild from S3 (max live id is 1). - std::fs::remove_dir_all(&data_dir_a).unwrap(); - let data_dir_b = unique_data_dir(); - std::fs::create_dir_all(&data_dir_b).unwrap(); - let storage: Arc = Arc::new(ObjectStoreBackend::from_store( - store.clone(), - "object-store:memory", - )); - let m2 = CollectionManager::new_with_storage(&data_dir_b, storage.clone()) - .await - .unwrap(); - - // A new ingest must get a FRESH id (4), not reuse deleted id 2. - m2.ingest("idreuse", vec![ingest_chunk(9)], &embed) - .await - .unwrap(); - let (man, _) = crate::storage::lsm::read_manifest(storage.as_ref(), "idreuse") - .await - .unwrap(); - let mat = cloud::materialize(storage.as_ref(), "idreuse", &man) - .await - .unwrap(); - assert!( - mat.chunks.contains_key(&4), - "new chunk must take id 4 (one past the pre-compaction high-water), got ids {:?}", - mat.chunks.keys().collect::>() - ); - assert!( - !mat.chunks.contains_key(&2) && !mat.chunks.contains_key(&3), - "deleted ids must not be reused" - ); - - let _ = std::fs::remove_dir_all(&data_dir_b); - } - - // PERSISTENT-DISK restart path (the one the adversarial review flagged): - // in cloud mode, a node restarting with its local disk intact runs - // `load_collection` (rehydrate from redb) and SKIPS rebuild-from-S3 for - // already-loaded collections. A chunk tombstoned locally (redb) — which is - // exactly what delete AND the ingest-compensation path write — must stay - // masked after that restart, even though it's still physically in redb. - #[tokio::test] - async fn persistent_disk_restart_honors_local_tombstones() { - let embed = embed_state(); - let store = std::sync::Arc::new(object_store::memory::InMemory::new()); - let data_dir = unique_data_dir(); - std::fs::create_dir_all(&data_dir).unwrap(); - { - let storage: Arc = Arc::new(ObjectStoreBackend::from_store( - store.clone(), - "object-store:memory", - )); - let m = CollectionManager::new_with_storage(&data_dir, storage) - .await - .unwrap(); - m.create_collection("pdisk", None, Some(4), None) - .await - .unwrap(); - m.ingest( - "pdisk", - vec![ingest_chunk(0), ingest_chunk(1), ingest_chunk(2)], - &embed, - ) - .await - .unwrap(); - // Writes the redb tombstone + RAM tombstone + S3 tombstone — the - // same three places the ingest-compensation path writes. - assert_eq!(m.delete_chunks("pdisk", &[1]).await.unwrap(), 1); - } - - // Restart with the SAME data_dir (persistent disk — NOT wiped). This - // takes the load_collection-first, skip-cloud-rebuild path. - let storage: Arc = Arc::new(ObjectStoreBackend::from_store( - store.clone(), - "object-store:memory", - )); - let m2 = CollectionManager::new_with_storage(&data_dir, storage) - .await - .unwrap(); - - let req = SearchRequest { - query: "chunk".to_string(), - mode: "semantic".to_string(), - vector_space: None, - top_k: 20, - query_vector: Some(vec![0.1, 0.2, 0.3, 0.4]), - filters: HashMap::new(), - score_weights: None, - recency: None, - recency_preset: None, - recency_field: None, - boosts: Vec::new(), - relationship_boost: None, - explain: false, - include_relations: false, - relation_types: None, - relation_direction: RelationDirection::Outgoing, - }; - let (hits, _, _, _) = m2.search("pdisk", &req, &embed).await.unwrap(); - let hit_ids: std::collections::HashSet = - hits.iter().map(|(c, _, _, _, _)| c.id).collect(); - assert!( - !hit_ids.contains(&1), - "tombstoned chunk must stay masked after persistent-disk restart" - ); - assert!( - hit_ids.contains(&0) && hit_ids.contains(&2), - "live chunks must survive, got {:?}", - hit_ids - ); - - let _ = std::fs::remove_dir_all(&data_dir); - } -} +mod cloud_ingest_tests; diff --git a/crates/compass/src/collections/parent_metadata_tests.rs b/crates/compass/src/collections/parent_metadata_tests.rs new file mode 100644 index 0000000..5a98da1 --- /dev/null +++ b/crates/compass/src/collections/parent_metadata_tests.rs @@ -0,0 +1,167 @@ +// collections/parent_metadata_tests.rs — extracted test module (was inline in mod.rs). +// Child module of `collections`: `super::*` sees the parent's private items. + +use super::*; + +fn segment(id: u64, parent_id: Option) -> DocumentChunk { + DocumentChunk { + id, + collection: "test".to_string(), + file_id: format!("f{}", id), + chunk_index: 0, + page: None, + text: String::new(), + metadata: HashMap::new(), + doc_type: "segment".to_string(), + parent_id, + group_id: None, + embeddings: HashMap::new(), + embedding: None, + } +} + +fn source_with_meta(id: u64, key: &str, val: &str) -> DocumentChunk { + let mut metadata = HashMap::new(); + metadata.insert(key.to_string(), MetadataValue::String(val.to_string())); + DocumentChunk { + id, + collection: "test".to_string(), + file_id: format!("f{}", id), + chunk_index: 0, + page: None, + text: String::new(), + metadata, + doc_type: "source".to_string(), + parent_id: None, + group_id: None, + embeddings: HashMap::new(), + embedding: None, + } +} + +fn into_map(chunks: Vec) -> ChunkCache { + use std::sync::atomic::{AtomicU64, Ordering}; + static N: AtomicU64 = AtomicU64::new(0); + let mut p = std::env::temp_dir(); + p.push(format!( + "compass_pmc_{}_{}.redb", + std::process::id(), + N.fetch_add(1, Ordering::Relaxed) + )); + let _ = std::fs::remove_file(&p); + let cache = ChunkCache::new(ChunkStore::open(&p).unwrap()); + let batch: Vec<(u64, DocumentChunk)> = chunks.into_iter().map(|c| (c.id, c)).collect(); + cache.insert_batch(&batch).unwrap(); + cache +} + +#[test] +fn segment_with_parent_gets_metadata() { + let chunks = into_map(vec![ + source_with_meta(1, "title", "Keynote"), + segment(2, Some(1)), + ]); + let cache = build_parent_metadata_cache(&[2], &chunks); + let meta = parent_metadata_for(&chunks.get(2).unwrap().unwrap(), &cache); + assert_eq!( + meta.unwrap().get("title"), + Some(&MetadataValue::String("Keynote".to_string())) + ); +} + +#[test] +fn source_hit_gets_none() { + let chunks = into_map(vec![source_with_meta(1, "title", "Keynote")]); + let cache = build_parent_metadata_cache(&[1], &chunks); + let meta = parent_metadata_for(&chunks.get(1).unwrap().unwrap(), &cache); + assert!(meta.is_none()); +} + +#[test] +fn segment_without_parent_gets_none() { + let chunks = into_map(vec![segment(2, None)]); + let cache = build_parent_metadata_cache(&[2], &chunks); + let meta = parent_metadata_for(&chunks.get(2).unwrap().unwrap(), &cache); + assert!(meta.is_none()); +} + +#[test] +fn dedup_one_lookup_per_unique_parent() { + // Three segments, all pointing at parent_id=10. The cache should + // contain exactly one entry (for pid=10), proving the dedup. + let chunks = into_map(vec![ + source_with_meta(10, "source_id", "src-001"), + segment(11, Some(10)), + segment(12, Some(10)), + segment(13, Some(10)), + ]); + let cache = build_parent_metadata_cache(&[11, 12, 13], &chunks); + assert_eq!( + cache.len(), + 1, + "expected one cache entry for the shared parent" + ); + assert!(cache.contains_key(&10)); + // All three segments resolve to the same parent metadata. + for cid in [11, 12, 13] { + let meta = parent_metadata_for(&chunks.get(cid).unwrap().unwrap(), &cache); + assert_eq!( + meta.unwrap().get("source_id"), + Some(&MetadataValue::String("src-001".to_string())) + ); + } +} + +#[test] +fn orphan_segment_yields_none() { + // parent_id=99 not in chunks. The cache must NOT contain pid=99, + // and parent_metadata_for must return None. This distinguishes + // "parent exists with empty metadata" (Some({})) from "parent + // doesn't exist" (None). + let chunks = into_map(vec![segment(5, Some(99))]); + let cache = build_parent_metadata_cache(&[5], &chunks); + assert!(!cache.contains_key(&99), "orphan parent must not be cached"); + let meta = parent_metadata_for(&chunks.get(5).unwrap().unwrap(), &cache); + assert!(meta.is_none(), "orphan segment must yield None"); +} + +#[test] +fn parent_exists_with_empty_metadata_yields_some_empty() { + // Parent chunk exists but has no metadata fields. Must return Some({}) + // so callers can distinguish from the orphan case (None). + let parent_no_meta = DocumentChunk { + id: 20, + collection: "test".to_string(), + file_id: "f20".to_string(), + chunk_index: 0, + page: None, + text: String::new(), + metadata: HashMap::new(), + doc_type: "source".to_string(), + parent_id: None, + group_id: None, + embeddings: HashMap::new(), + embedding: None, + }; + let chunks = into_map(vec![parent_no_meta, segment(21, Some(20))]); + let cache = build_parent_metadata_cache(&[21], &chunks); + let meta = parent_metadata_for(&chunks.get(21).unwrap().unwrap(), &cache); + assert!(meta.is_some()); + assert!(meta.unwrap().is_empty()); +} + +#[test] +fn parent_metadata_for_cache_miss_returns_none() { + // Defensive: if the cache was built with a different set of IDs than + // the one we're looking up, the function must return None (not panic, + // not return stale data). Catches regressions where someone "optimizes" + // parent_metadata_for to assume the cache is always complete. + let parent = source_with_meta(1, "title", "Keynote"); + let seg = segment(2, Some(1)); + let chunks = into_map(vec![parent, seg]); + // Build cache against an empty candidate list, then look up segment 2. + let cache = build_parent_metadata_cache(&[], &chunks); + assert!(cache.is_empty()); + let meta = parent_metadata_for(&chunks.get(2).unwrap().unwrap(), &cache); + assert!(meta.is_none()); +} diff --git a/crates/compass/src/collections/persistence_tests.rs b/crates/compass/src/collections/persistence_tests.rs new file mode 100644 index 0000000..a604bed --- /dev/null +++ b/crates/compass/src/collections/persistence_tests.rs @@ -0,0 +1,246 @@ +// collections/persistence_tests.rs — extracted test module (was inline in mod.rs). +// Child module of `collections`: `super::*` sees the parent's private items. + +//! End-to-end durability test. Builds a CollectionManager in a temp dir, +//! ingests chunks, drops the manager (closing the chunk store), creates a +//! new manager pointing at the same dir, and asserts the chunks come back. +//! +//! This is the test that proves Compass survives process restarts. Without +//! the disk-backed ChunkStore wiring, this test would fail because +//! `loaded.chunks` would be empty after the manager restart. + +use super::*; +use crate::embed::EmbedState; +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; + +fn unique_data_dir() -> std::path::PathBuf { + static N: AtomicU64 = AtomicU64::new(0); + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + std::env::temp_dir().join(format!( + "compass-persist-test-{}-{}-{}", + std::process::id(), + nanos, + N.fetch_add(1, Ordering::SeqCst) + )) +} + +fn empty_embed_state() -> EmbedState { + // No embedding models loaded. Safe for the persistence test because + // we provide chunks without text-only embedding requirements. Any + // call to embed_query returns Err and the ingest path tolerates that. + EmbedState { + bge: None, + distilled: None, + } +} + +fn make_ingest_chunk(file_id: &str, text: &str) -> IngestChunk { + IngestChunk { + client_id: None, + file_id: file_id.to_string(), + chunk_index: 0, + page: None, + text: text.to_string(), + metadata: HashMap::new(), + doc_type: "chunk".to_string(), + parent_id: None, + parent_ref: None, + group_id: None, + embeddings: HashMap::new(), + embedding: None, + } +} + +// Regression for the three facet bugs the live E2E harness caught: +// (1) a second ingest batch replaced facet state instead of accumulating +// (latent since v0.2 — build_index returned new-batch-only bitsets); +// (2) facets came back empty after a restart (open_index returns empty +// state and nothing rebuilt it); +// (3) deleted chunks kept inflating counts (facets never saw tombstones). +#[tokio::test] +async fn facets_accumulate_survive_restart_and_exclude_deleted() { + let data_dir = unique_data_dir(); + std::fs::create_dir_all(&data_dir).unwrap(); + let embed = empty_embed_state(); + let tagged = |file: &str, text: &str, kind: &str| { + let mut c = make_ingest_chunk(file, text); + c.metadata.insert( + "kind".to_string(), + crate::models::MetadataValue::String(kind.to_string()), + ); + c + }; + let field = ["kind".to_string()]; + + { + let manager = CollectionManager::new(&data_dir).await.unwrap(); + manager + .create_collection("facet-test", None, None, None) + .await + .unwrap(); + manager + .ingest( + "facet-test", + vec![ + tagged("a", "alpha doc", "report"), + tagged("b", "beta doc", "memo"), + ], + &embed, + ) + .await + .unwrap(); + // Bug 1: this second batch must ADD to the first, not replace it. + manager + .ingest( + "facet-test", + vec![tagged("c", "gamma doc", "report")], + &embed, + ) + .await + .unwrap(); + let (facets, _) = manager.get_facets("facet-test", "", &field).await.unwrap(); + let kind = facets.get("kind").expect("facets survive a second batch"); + assert_eq!(kind.get("report"), Some(&2)); + assert_eq!(kind.get("memo"), Some(&1)); + } + + // Bug 2: facets must be rebuilt from the chunk store on restart. + let manager2 = CollectionManager::new(&data_dir).await.unwrap(); + let (facets, _) = manager2.get_facets("facet-test", "", &field).await.unwrap(); + let kind = facets.get("kind").expect("facets survive a restart"); + assert_eq!(kind.get("report"), Some(&2)); + assert_eq!(kind.get("memo"), Some(&1)); + + // Bug 3: deleting a chunk must drop it from counts immediately. + let (_, ids) = manager2.get_all_chunk_data("facet-test").await.unwrap(); + let (texts, _) = manager2.get_all_chunk_data("facet-test").await.unwrap(); + let memo_id = ids + .iter() + .zip(texts.iter()) + .find(|(_, t)| t.contains("beta")) + .map(|(id, _)| *id) + .unwrap(); + manager2 + .delete_chunks("facet-test", &[memo_id]) + .await + .unwrap(); + let (facets, _) = manager2.get_facets("facet-test", "", &field).await.unwrap(); + let kind = facets.get("kind").unwrap(); + assert_eq!(kind.get("report"), Some(&2)); + assert!( + kind.get("memo").is_none() || kind.get("memo") == Some(&0), + "deleted chunk still counted in facets: {kind:?}" + ); + + let _ = std::fs::remove_dir_all(&data_dir); +} + +#[tokio::test] +async fn chunks_persist_across_manager_restart() { + let data_dir = unique_data_dir(); + std::fs::create_dir_all(&data_dir).unwrap(); + let embed = empty_embed_state(); + + // First manager lifetime: create collection, ingest three chunks, + // then drop the manager to close all file handles (including redb). + { + let manager = CollectionManager::new(&data_dir).await.unwrap(); + manager + .create_collection("persist-test", None, None, None) + .await + .unwrap(); + let to_ingest = vec![ + make_ingest_chunk("f1", "first chunk"), + make_ingest_chunk("f2", "second chunk"), + make_ingest_chunk("f3", "third chunk"), + ]; + let (ingested, _, _) = manager + .ingest("persist-test", to_ingest, &embed) + .await + .unwrap(); + assert_eq!(ingested, 3, "ingest call reports 3 chunks written"); + // manager dropped here + } + + // Second manager: same data dir, must rehydrate chunks from disk. + let manager2 = CollectionManager::new(&data_dir).await.unwrap(); + let (texts, ids) = manager2.get_all_chunk_data("persist-test").await.unwrap(); + + assert_eq!( + ids.len(), + 3, + "expected 3 chunks rehydrated from disk after manager restart, got {}", + ids.len() + ); + let mut sorted_texts = texts.clone(); + sorted_texts.sort(); + assert_eq!( + sorted_texts, + vec![ + "first chunk".to_string(), + "second chunk".to_string(), + "third chunk".to_string(), + ], + "chunk texts should match what was ingested before the restart" + ); + + // Cleanup + let _ = std::fs::remove_dir_all(&data_dir); +} + +#[tokio::test] +async fn next_id_advances_correctly_after_rehydration() { + // After rehydration, next_id should be max(seen) + 1 so new ingests + // don't collide with persisted IDs. Verify by ingesting again after + // restart and checking the new chunk got a fresh ID. + let data_dir = unique_data_dir(); + std::fs::create_dir_all(&data_dir).unwrap(); + let embed = empty_embed_state(); + + // Round 1: ingest two chunks (IDs 0, 1) + { + let manager = CollectionManager::new(&data_dir).await.unwrap(); + manager + .create_collection("next-id-test", None, None, None) + .await + .unwrap(); + manager + .ingest( + "next-id-test", + vec![ + make_ingest_chunk("f0", "round-one-a"), + make_ingest_chunk("f1", "round-one-b"), + ], + &embed, + ) + .await + .unwrap(); + } + + // Round 2: restart and ingest one more chunk. The new chunk's ID + // should be 2, not 0. + let manager2 = CollectionManager::new(&data_dir).await.unwrap(); + manager2 + .ingest( + "next-id-test", + vec![make_ingest_chunk("f2", "round-two")], + &embed, + ) + .await + .unwrap(); + let (_, ids) = manager2.get_all_chunk_data("next-id-test").await.unwrap(); + let mut sorted_ids = ids.clone(); + sorted_ids.sort(); + assert_eq!( + sorted_ids, + vec![0, 1, 2], + "next_id must advance past max persisted id, got ids: {:?}", + sorted_ids + ); + + let _ = std::fs::remove_dir_all(&data_dir); +} diff --git a/crates/compass/src/collections/rebuild.rs b/crates/compass/src/collections/rebuild.rs index 1fa3c02..2c0bd90 100644 --- a/crates/compass/src/collections/rebuild.rs +++ b/crates/compass/src/collections/rebuild.rs @@ -78,6 +78,7 @@ pub async fn start_rebuild( _batch_size: usize, tracker: RebuildTracker, collection_name: String, + manager: Arc, ) -> Result<(), String> { let key = format!("{}/{}", collection_name, space_name); @@ -113,15 +114,14 @@ pub async fn start_rebuild( let rt = tokio::runtime::Handle::current(); let mut all_vectors: Vec> = Vec::with_capacity(texts.len()); + // External embedding endpoints are accepted in the request but not yet + // dispatched to — every rebuild embeds with the built-in models. Kept + // as a field (not a branch) so clients sending it keep working. + let _ = embed_endpoint; + // Embed each chunk's text for (i, text) in texts.iter().enumerate() { - let vec = if let Some(ref _endpoint) = embed_endpoint { - // TODO: HTTP POST to external endpoint for GPU embedding - // For now, fall back to built-in embedder - embed_state - .embed_query(text) - .unwrap_or_else(|_| vec![0.0; dims]) - } else { + let vec = { // Use built-in Candle embedder embed_state .embed_query(text) @@ -145,16 +145,31 @@ pub async fn start_rebuild( let result = vector::build_vector_index(&index_path, &vectors_path, &chunk_ids, &all_vectors, dims); - // Update final status + // Update final status. On success, ALSO flip the space's persisted + // status and hot-load the new index into the serving collection — + // without this the space stayed "building" (and the rebuilt index + // unused) until the next restart, even though the progress endpoint + // reported active. let progress = progress.clone(); let key = key.clone(); rt.block_on(async { let mut p = progress.write().await; match result { Ok(_) => { - p.status = "active".to_string(); - p.embedded = p.total; - tracing::info!("Rebuild complete for {}", key); + match manager + .mark_vector_space_active(&collection_name, &space_name) + .await + { + Ok(()) => { + p.status = "active".to_string(); + p.embedded = p.total; + tracing::info!("Rebuild complete for {}", key); + } + Err(e) => { + p.status = format!("failed: activation: {}", e); + tracing::error!("Rebuild activation failed for {}: {}", key, e); + } + } } Err(e) => { p.status = format!("failed: {}", e); diff --git a/crates/compass/src/collections/relation_store.rs b/crates/compass/src/collections/relation_store.rs index a2d6a7b..fc87a31 100644 --- a/crates/compass/src/collections/relation_store.rs +++ b/crates/compass/src/collections/relation_store.rs @@ -235,6 +235,7 @@ impl RelationStore { } /// Total number of stored edges. Used by tests + diagnostics. + #[cfg(test)] pub fn count(&self) -> Result { use redb::ReadableTableMetadata; let txn = self.db.begin_read()?; diff --git a/crates/compass/src/collections/relationships.rs b/crates/compass/src/collections/relationships.rs index f8c14de..9db9c1c 100644 --- a/crates/compass/src/collections/relationships.rs +++ b/crates/compass/src/collections/relationships.rs @@ -129,11 +129,6 @@ impl RelationshipStore { .collect() } - /// Total number of tracked relationships. - pub fn len(&self) -> usize { - self.forward.len() - } - // ── Disk persistence ───────────────────────────────────────────────── // Simple binary format: // [u32 count] diff --git a/crates/compass/src/collections/segments_at_tests.rs b/crates/compass/src/collections/segments_at_tests.rs new file mode 100644 index 0000000..9e7f3fd --- /dev/null +++ b/crates/compass/src/collections/segments_at_tests.rs @@ -0,0 +1,169 @@ +// collections/segments_at_tests.rs — extracted test module (was inline in mod.rs). +// Child module of `collections`: `super::*` sees the parent's private items. + +use super::*; + +fn make_segment(group_id: &str, ts_ms: f64, te_ms: f64) -> DocumentChunk { + let mut metadata = HashMap::new(); + metadata.insert( + "timerange_start_ms".to_string(), + MetadataValue::Float(ts_ms), + ); + metadata.insert("timerange_end_ms".to_string(), MetadataValue::Float(te_ms)); + DocumentChunk { + id: 1, + collection: "test".to_string(), + file_id: "f1".to_string(), + chunk_index: 0, + page: None, + text: String::new(), + metadata, + doc_type: "segment".to_string(), + parent_id: None, + group_id: Some(group_id.to_string()), + embeddings: HashMap::new(), + embedding: None, + } +} + +/// Make a zero-duration "instant" segment, the convention for sidecar +/// events that have a single timestamp (e.g. standout_timestamps). +fn make_instant(group_id: &str, t_ms: f64) -> DocumentChunk { + make_segment(group_id, t_ms, t_ms) +} + +#[test] +fn point_inside_window() { + let c = make_segment("a", 100.0, 200.0); + assert!(segment_in_time_window(&c, Some(150.0), None, None)); +} + +#[test] +fn point_outside_window() { + let c = make_segment("a", 100.0, 200.0); + assert!(!segment_in_time_window(&c, Some(250.0), None, None)); +} + +#[test] +fn point_boundaries_inclusive() { + let c = make_segment("a", 100.0, 200.0); + assert!(segment_in_time_window(&c, Some(100.0), None, None)); + assert!(segment_in_time_window(&c, Some(200.0), None, None)); +} + +#[test] +fn range_overlap_matches() { + let c = make_segment("a", 100.0, 200.0); + assert!(segment_in_time_window(&c, None, Some(180.0), Some(300.0))); +} + +#[test] +fn range_no_overlap() { + let c = make_segment("a", 100.0, 200.0); + assert!(!segment_in_time_window(&c, None, Some(250.0), Some(400.0))); +} + +#[test] +fn range_open_lower_bound() { + let c = make_segment("a", 100.0, 200.0); + assert!(segment_in_time_window(&c, None, None, Some(150.0))); + assert!(!segment_in_time_window(&c, None, None, Some(50.0))); +} + +#[test] +fn range_open_upper_bound() { + let c = make_segment("a", 100.0, 200.0); + assert!(segment_in_time_window(&c, None, Some(150.0), None)); + assert!(!segment_in_time_window(&c, None, Some(300.0), None)); +} + +#[test] +fn missing_metadata_with_filter_excludes() { + let c = DocumentChunk { + id: 2, + collection: "test".to_string(), + file_id: "f2".to_string(), + chunk_index: 0, + page: None, + text: String::new(), + metadata: HashMap::new(), + doc_type: "segment".to_string(), + parent_id: None, + group_id: Some("a".to_string()), + embeddings: HashMap::new(), + embedding: None, + }; + assert!(!segment_in_time_window(&c, Some(100.0), None, None)); + assert!(!segment_in_time_window(&c, None, Some(0.0), Some(1000.0))); +} + +#[test] +fn no_filter_matches_all() { + let with_meta = make_segment("a", 100.0, 200.0); + assert!(segment_in_time_window(&with_meta, None, None, None)); + + let without_meta = DocumentChunk { + id: 3, + collection: "test".to_string(), + file_id: "f3".to_string(), + chunk_index: 0, + page: None, + text: String::new(), + metadata: HashMap::new(), + doc_type: "segment".to_string(), + parent_id: None, + group_id: Some("a".to_string()), + embeddings: HashMap::new(), + embedding: None, + }; + assert!(segment_in_time_window(&without_meta, None, None, None)); +} + +// When both `time_ms` and `time_start_ms`/`time_end_ms` are provided, +// `time_ms` wins. Documented in the segments.rs handler comment; this +// test asserts it. +#[test] +fn point_lookup_takes_precedence_over_range() { + let c = make_segment("a", 100.0, 200.0); + // Point=150 is inside [100, 200], but the range [300, 400] is outside. + // If `time_ms` correctly takes precedence, this must return true. + assert!(segment_in_time_window( + &c, + Some(150.0), + Some(300.0), + Some(400.0) + )); + // Point=250 is outside, but the range [100, 300] would match. + // If `time_ms` correctly takes precedence, this must return false. + assert!(!segment_in_time_window( + &c, + Some(250.0), + Some(100.0), + Some(300.0) + )); +} + +// Instants (zero-duration events like a standout_timestamp) match a +// point query at their exact timestamp and any range that overlaps it. +// Critical for ingesting sidecar fields like +// `gemini.response.standout_timestamps[]` which only carry a single ms. +#[test] +fn instant_matches_exact_point_query() { + let c = make_instant("a", 5200.0); + assert!(segment_in_time_window(&c, Some(5200.0), None, None)); + assert!(!segment_in_time_window(&c, Some(5199.0), None, None)); + assert!(!segment_in_time_window(&c, Some(5201.0), None, None)); +} + +#[test] +fn instant_matches_overlapping_range_query() { + let c = make_instant("a", 5200.0); + assert!(segment_in_time_window(&c, None, Some(5000.0), Some(6000.0))); + assert!(segment_in_time_window(&c, None, Some(5200.0), Some(5200.0))); + assert!(!segment_in_time_window( + &c, + None, + Some(5201.0), + Some(6000.0) + )); +} diff --git a/crates/compass/src/collections/validate_name_segment_tests.rs b/crates/compass/src/collections/validate_name_segment_tests.rs new file mode 100644 index 0000000..4ec5d67 --- /dev/null +++ b/crates/compass/src/collections/validate_name_segment_tests.rs @@ -0,0 +1,49 @@ +// collections/validate_name_segment_tests.rs — extracted test module (was inline in mod.rs). +// Child module of `collections`: `super::*` sees the parent's private items. + +use super::validate_name_segment; + +#[test] +fn accepts_simple_names() { + assert!(validate_name_segment("my-collection", "Collection").is_ok()); + assert!(validate_name_segment("harrier", "Vector space").is_ok()); + assert!(validate_name_segment("qwen3-vl", "Vector space").is_ok()); + assert!(validate_name_segment("a", "Collection").is_ok()); +} + +#[test] +fn rejects_empty() { + let err = validate_name_segment("", "Vector space").expect_err("empty name should error"); + assert!(err.to_string().contains("Vector space")); +} + +#[test] +fn rejects_path_traversal() { + // The whole reason this validator exists: a vector space name flows + // into on-disk paths like `/.bin`. A `../` segment + // must never be accepted. + for bad in [ + "../etc/passwd", + "..", + "foo/bar", + "foo\\bar", + "/abs", + "name with space", + "name.with.dot", + "name_with_underscore", // hyphens only, no underscores + "tab\there", + "name\nwith\nnewline", + ] { + assert!( + validate_name_segment(bad, "Vector space").is_err(), + "validator must reject {bad:?}" + ); + } +} + +#[test] +fn rejects_unicode_lookalikes() { + // Cyrillic 'а' (U+0430) looks like 'a' but is not ASCII. + assert!(validate_name_segment("\u{0430}bc", "Collection").is_err()); + assert!(validate_name_segment("emoji-🚀", "Collection").is_err()); +} diff --git a/crates/compass/src/filter.rs b/crates/compass/src/filter.rs deleted file mode 100644 index 7c5da02..0000000 --- a/crates/compass/src/filter.rs +++ /dev/null @@ -1,269 +0,0 @@ -// filter.rs — Post-retrieval metadata filtering with operator support. -// -// Operators: -// exact match — "department": "Legal" -// range — "timerange_start": {"gte": 2040.0} -// contains — "tags": {"contains": "sports"} -// set member — "doc_type": {"in": ["segment", "flow"]} -// -// "doc_type" is special-cased: it reads from chunk.doc_type (struct field) -// instead of chunk.metadata, so existing data on disk works without migration. - -use crate::models::{DocumentChunk, FilterCondition, FilterValue, MetadataValue}; -use std::collections::HashMap; - -pub fn matches_filters(chunk: &DocumentChunk, filters: &HashMap) -> bool { - filters.iter().all(|(key, filter_val)| { - let meta_val = if key == "doc_type" { - Some(MetadataValue::String(chunk.doc_type.clone())) - } else { - chunk.metadata.get(key).cloned() - }; - - match filter_val { - FilterValue::Exact(expected) => meta_val.as_ref().map_or(false, |v| v == expected), - FilterValue::Condition(cond) => eval_condition(meta_val.as_ref(), cond), - } - }) -} - -fn eval_condition(val: Option<&MetadataValue>, cond: &FilterCondition) -> bool { - if cond.gte.is_some() || cond.lte.is_some() { - match val.and_then(|v| v.as_f64()) { - None => return false, - Some(n) => { - if let Some(g) = cond.gte { - if n < g { - return false; - } - } - if let Some(l) = cond.lte { - if n > l { - return false; - } - } - } - } - } - - if let Some(ref target) = cond.contains { - match val { - Some(MetadataValue::StringList(list)) => { - if !list.iter().any(|s| s == target) { - return false; - } - } - Some(MetadataValue::String(s)) => { - if s != target { - return false; - } - } - _ => return false, - } - } - - if let Some(ref allowed) = cond.in_values { - match val { - Some(MetadataValue::String(s)) => { - if !allowed.contains(s) { - return false; - } - } - _ => return false, - } - } - - true -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::models::DocumentChunk; - use std::collections::HashMap; - - fn make_chunk(doc_type: &str, metadata: HashMap) -> DocumentChunk { - DocumentChunk { - id: 1, - collection: "test".to_string(), - file_id: "f1".to_string(), - chunk_index: 0, - page: None, - text: "test".to_string(), - metadata, - doc_type: doc_type.to_string(), - parent_id: None, - group_id: None, - embeddings: HashMap::new(), - embedding: None, - } - } - - #[test] - fn exact_match_backward_compat() { - let mut meta = HashMap::new(); - meta.insert( - "department".to_string(), - MetadataValue::String("Legal".to_string()), - ); - let chunk = make_chunk("chunk", meta); - - let mut filters = HashMap::new(); - filters.insert( - "department".to_string(), - FilterValue::Exact(MetadataValue::String("Legal".to_string())), - ); - assert!(matches_filters(&chunk, &filters)); - - filters.insert( - "department".to_string(), - FilterValue::Exact(MetadataValue::String("HR".to_string())), - ); - assert!(!matches_filters(&chunk, &filters)); - } - - #[test] - fn numeric_range_gte() { - let mut meta = HashMap::new(); - meta.insert("timerange_start".to_string(), MetadataValue::Float(2040.0)); - let chunk = make_chunk("segment", meta); - - let mut filters = HashMap::new(); - filters.insert( - "timerange_start".to_string(), - FilterValue::Condition(FilterCondition { - gte: Some(2000.0), - lte: None, - contains: None, - in_values: None, - }), - ); - assert!(matches_filters(&chunk, &filters)); - - filters.insert( - "timerange_start".to_string(), - FilterValue::Condition(FilterCondition { - gte: Some(2100.0), - lte: None, - contains: None, - in_values: None, - }), - ); - assert!(!matches_filters(&chunk, &filters)); - } - - #[test] - fn numeric_range_combined() { - let mut meta = HashMap::new(); - meta.insert("priority".to_string(), MetadataValue::Int(5)); - let chunk = make_chunk("chunk", meta); - - let mut filters = HashMap::new(); - filters.insert( - "priority".to_string(), - FilterValue::Condition(FilterCondition { - gte: Some(3.0), - lte: Some(10.0), - contains: None, - in_values: None, - }), - ); - assert!(matches_filters(&chunk, &filters)); - - filters.insert( - "priority".to_string(), - FilterValue::Condition(FilterCondition { - gte: Some(6.0), - lte: Some(10.0), - contains: None, - in_values: None, - }), - ); - assert!(!matches_filters(&chunk, &filters)); - } - - #[test] - fn contains_string_list() { - let mut meta = HashMap::new(); - meta.insert( - "tags".to_string(), - MetadataValue::StringList(vec!["sports".to_string(), "goals".to_string()]), - ); - let chunk = make_chunk("chunk", meta); - - let mut filters = HashMap::new(); - filters.insert( - "tags".to_string(), - FilterValue::Condition(FilterCondition { - gte: None, - lte: None, - contains: Some("sports".to_string()), - in_values: None, - }), - ); - assert!(matches_filters(&chunk, &filters)); - - filters.insert( - "tags".to_string(), - FilterValue::Condition(FilterCondition { - gte: None, - lte: None, - contains: Some("music".to_string()), - in_values: None, - }), - ); - assert!(!matches_filters(&chunk, &filters)); - } - - #[test] - fn in_set_membership() { - let chunk = make_chunk("segment", HashMap::new()); - - let mut filters = HashMap::new(); - filters.insert( - "doc_type".to_string(), - FilterValue::Condition(FilterCondition { - gte: None, - lte: None, - contains: None, - in_values: Some(vec!["segment".to_string(), "flow".to_string()]), - }), - ); - assert!(matches_filters(&chunk, &filters)); - - filters.insert( - "doc_type".to_string(), - FilterValue::Condition(FilterCondition { - gte: None, - lte: None, - contains: None, - in_values: Some(vec!["source".to_string()]), - }), - ); - assert!(!matches_filters(&chunk, &filters)); - } - - #[test] - fn doc_type_special_case() { - let chunk = make_chunk("segment", HashMap::new()); - - let mut filters = HashMap::new(); - filters.insert( - "doc_type".to_string(), - FilterValue::Exact(MetadataValue::String("segment".to_string())), - ); - assert!(matches_filters(&chunk, &filters)); - } - - #[test] - fn missing_field_returns_false() { - let chunk = make_chunk("chunk", HashMap::new()); - - let mut filters = HashMap::new(); - filters.insert( - "nonexistent".to_string(), - FilterValue::Exact(MetadataValue::String("value".to_string())), - ); - assert!(!matches_filters(&chunk, &filters)); - } -} diff --git a/crates/compass/src/main.rs b/crates/compass/src/main.rs index ad7ff9f..2c87211 100644 --- a/crates/compass/src/main.rs +++ b/crates/compass/src/main.rs @@ -1,6 +1,5 @@ // Pre-existing clippy lints from newer toolchain — will be cleaned up separately. #![allow( - dead_code, clippy::too_many_arguments, clippy::type_complexity, clippy::collapsible_if, @@ -32,15 +31,13 @@ mod api; mod collections; mod embed; -mod filter; +mod metrics; mod models; mod scoring; mod search; -// Storage abstraction (Storage trait + LocalDiskStorage + object-storage backend -// + LSM). main() selects and verifies the backend at startup; full engine -// persistence through it is the follow-on. `allow(dead_code)` covers the parts -// (LSM, chunk cache, filter-index serde) not yet on the hot path. -#[allow(dead_code)] +// Storage abstraction: Storage trait + LocalDiskStorage + object-storage +// backend + the LSM (WAL fragments, manifest, segments) — the cloud-mode +// persistence layer. mod storage; mod telemetry; diff --git a/crates/compass/src/metrics.rs b/crates/compass/src/metrics.rs new file mode 100644 index 0000000..fe89c98 --- /dev/null +++ b/crates/compass/src/metrics.rs @@ -0,0 +1,55 @@ +//! Minimal Prometheus-text metrics — no external deps, atomic counters only. +//! +//! Operating a multi-node deployment blind was a production blocker: you +//! could not see ingest/search rates, refresh convergence, or attach costs. +//! This is deliberately tiny; a full metrics facade can replace it later +//! without touching call sites (they go through these free functions). + +use std::sync::atomic::{AtomicU64, Ordering}; + +macro_rules! counters { + ($($name:ident),* $(,)?) => { + $(pub static $name: AtomicU64 = AtomicU64::new(0);)* + fn render_counters(out: &mut String) { + $( + out.push_str(&format!( + "compass_{} {}\n", + stringify!($name).to_lowercase(), + $name.load(Ordering::Relaxed) + )); + )* + } + }; +} + +counters!( + INGEST_REQUESTS_TOTAL, + INGEST_CHUNKS_TOTAL, + SEARCH_REQUESTS_TOTAL, + DELETE_REQUESTS_TOTAL, + REFRESH_FRAGMENTS_APPLIED_TOTAL, + REFRESH_REATTACHES_TOTAL, + ATTACH_TOTAL, + ATTACH_SECONDS_SUM_MILLIS, + COMPACTIONS_TOTAL, + QUARANTINED_CHUNKS_TOTAL, +); + +#[inline] +pub fn inc(counter: &AtomicU64) { + counter.fetch_add(1, Ordering::Relaxed); +} + +#[inline] +pub fn add(counter: &AtomicU64, n: u64) { + counter.fetch_add(n, Ordering::Relaxed); +} + +/// Render every counter plus caller-supplied gauge lines (e.g. per-collection +/// state the manager owns). +pub fn render(extra_gauges: &str) -> String { + let mut out = String::with_capacity(1024); + render_counters(&mut out); + out.push_str(extra_gauges); + out +} diff --git a/crates/compass/src/models.rs b/crates/compass/src/models.rs index 37e8ea0..85fbf87 100644 --- a/crates/compass/src/models.rs +++ b/crates/compass/src/models.rs @@ -52,7 +52,7 @@ impl MetadataValue { // same collection (e.g. BGE-small for text, CLIP for images) and swap models // without re-indexing everything at once. -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct VectorSpaceConfig { /// Dimensionality of vectors in this space (e.g. 384 for BGE-small) pub dims: usize, @@ -137,6 +137,14 @@ pub struct Collection { pub next_id: u64, #[serde(default)] pub config: CollectionConfig, + /// Cloud mode: count of contiguously-applied manifest seqs (the local + /// indexes reflect fragments 0..applied_seq). Persisted so a + /// persistent-disk restart knows how fresh its local state is and the + /// refresher can catch up the delta instead of a full rebuild. May LAG + /// the true applied state (relation applies don't force a save); replay + /// of already-applied fragments is idempotent. + #[serde(default)] + pub applied_seq: u64, } fn default_dims() -> usize { @@ -262,6 +270,11 @@ pub struct SearchRequest { /// Which edges to include per hit when `include_relations` is set. #[serde(default)] pub relation_direction: RelationDirection, + /// Cloud mode read-your-writes: only serve once fragments up to this seq + /// (returned by a write) are applied locally, refreshing if needed + /// (bounded wait). Ignored in local mode. + #[serde(default)] + pub min_seq: Option, } // ── Chunk Relations ─────────────────────────────────────────────────────── @@ -351,6 +364,10 @@ pub struct DeleteRequest { pub struct DeleteResponse { /// Number of chunks newly soft-deleted (excludes already-deleted/missing). pub deleted: usize, + /// Cloud mode: manifest seq of the durable tombstone fragment (for + /// read-your-writes via `min_seq`). + #[serde(skip_serializing_if = "Option::is_none")] + pub seq: Option, } fn default_search_mode() -> String { @@ -614,6 +631,10 @@ pub struct IngestResponse { #[serde(skip_serializing_if = "HashMap::is_empty")] pub id_map: HashMap, pub took_ms: u64, + /// Cloud mode: manifest seq of the durable WAL fragment for this batch. + /// Pass as `min_seq` on a later search for read-your-writes. + #[serde(skip_serializing_if = "Option::is_none")] + pub seq: Option, } #[derive(Debug, Serialize)] diff --git a/crates/compass/src/search/backend.rs b/crates/compass/src/search/backend.rs deleted file mode 100644 index 26a7f81..0000000 --- a/crates/compass/src/search/backend.rs +++ /dev/null @@ -1,200 +0,0 @@ -//! Vector index backend abstraction. -//! -//! Wraps the existing USearch HNSW path in a [`VectorIndex`] implementation so -//! that callers can swap to the GPU-accelerated [`compass_vector_gpu::CuvsHnswIndex`] -//! transparently. The trait itself lives in [`compass_index_api`]. -//! -//! # Why a trait -//! -//! Compass has historically used USearch directly. As we add a GPU backend -//! (cuVS / CAGRA→HNSW), and as we anticipate IVF-PQ for very large corpora, -//! the call sites benefit from binding to a stable trait instead of the -//! USearch types. New backends slot in without touching `collections/`, -//! `api/`, or the rebuild path. -//! -//! # Backend selection -//! -//! Construction goes through [`build_backend`], which inspects environment -//! variables and feature flags to pick: -//! -//! - `COMPASS_BACKEND=cpu` (default): [`UsearchHnswIndex`]. -//! - `COMPASS_BACKEND=gpu`: requires the `gpu` feature; returns -//! `CuvsHnswIndex` from `compass-vector-gpu`. Falls back to CPU with a -//! `tracing::warn!` if CUDA is unavailable at runtime. -//! - `COMPASS_BACKEND=auto`: probe GPU first, fall back to CPU. - -use std::path::Path; - -pub use compass_index_api::{IndexError, IndexParams, LoadableIndex, VectorIndex, VectorMatch}; - -use super::vector; - -/// CPU-backed HNSW via USearch. Wraps the existing `vector::VectorState` so -/// the in-tree code keeps working while new code can bind to the trait. -pub struct UsearchHnswIndex { - state: vector::VectorState, - /// Where on disk the persisted index lives. Set by `build` or `load`. - persisted_at: Option, - vectors_path: Option, -} - -impl UsearchHnswIndex { - /// Empty index ready to receive a build. - pub fn new(params: IndexParams) -> Self { - Self { - state: vector::VectorState { - index: None, - key_to_chunk_id: Vec::new(), - mmap_vectors: None, - vectors: Vec::new(), - dims: params.dims, - }, - persisted_at: None, - vectors_path: None, - } - } - - /// Mount an existing index that's already on disk. The companion - /// `vectors_path` holds the raw float buffer for brute-force fallback. - pub fn from_paths( - index_path: &Path, - vectors_path: &Path, - dims: usize, - ) -> Result { - let state = vector::load_vector_index(index_path, vectors_path, dims) - .map_err(|e| IndexError::Io(e.to_string()))?; - Ok(Self { - state, - persisted_at: Some(index_path.to_path_buf()), - vectors_path: Some(vectors_path.to_path_buf()), - }) - } - - /// Direct accessor for code that still uses the legacy `VectorState` shape. - /// New code should go through the [`VectorIndex`] methods. - pub fn state(&self) -> &vector::VectorState { - &self.state - } -} - -impl VectorIndex for UsearchHnswIndex { - fn build(&mut self, vectors: &[Vec], chunk_ids: &[u64]) -> Result<(), IndexError> { - let index_path = self - .persisted_at - .clone() - .unwrap_or_else(|| std::path::PathBuf::from("./data/.compass-tmp.usearch")); - let vectors_path = self - .vectors_path - .clone() - .unwrap_or_else(|| std::path::PathBuf::from("./data/.compass-tmp.vectors")); - let state = vector::build_vector_index( - &index_path, - &vectors_path, - chunk_ids, - vectors, - self.state.dims, - ) - .map_err(|e| IndexError::Backend(e.to_string()))?; - self.state = state; - self.persisted_at = Some(index_path); - self.vectors_path = Some(vectors_path); - Ok(()) - } - - fn add(&mut self, _chunk_id: u64, _vector: &[f32]) -> Result<(), IndexError> { - // USearch does support incremental insert; wiring it here means - // re-saving the index after each add or batching at the rebuild layer. - // Today, ingestion goes through `build_vector_index` via the rebuild - // path. Surface this when the streaming-ingest API lands. - Err(IndexError::Unsupported( - "incremental add via VectorIndex trait not wired yet; use rebuild()".into(), - )) - } - - fn search(&self, query: &[f32], top_k: usize) -> Result, IndexError> { - if query.len() != self.state.dims { - return Err(IndexError::DimMismatch { - expected: self.state.dims, - actual: query.len(), - }); - } - let results = vector::search_vectors(query, &self.state, top_k); - Ok(results - .into_iter() - .map(|r| VectorMatch { - chunk_id: r.chunk_id, - score: r.score, - }) - .collect()) - } - - fn len(&self) -> usize { - self.state.vectors.len() - } - - fn dims(&self) -> usize { - self.state.dims - } - - fn save(&self, _path: &Path) -> Result<(), IndexError> { - // USearch saves at build time via `build_vector_index`. Re-saving an - // already-mmap'd index requires `index.save()` which the `Index` type - // exposes; we can wire it when downstream callers need atomic snapshot. - Ok(()) - } - - fn backend_name(&self) -> &'static str { - "usearch" - } -} - -impl LoadableIndex for UsearchHnswIndex { - fn load(path: &Path, params: IndexParams) -> Result { - let vectors_path = path.with_extension("vectors"); - Self::from_paths(path, &vectors_path, params.dims) - } -} - -/// Backend selection at startup. Reads `COMPASS_BACKEND` and feature flags. -/// -/// Returns a `Box` so the call site stays backend-agnostic. -/// Callers can downcast via [`std::any::Any`] if they need the concrete type -/// for backend-specific tuning. -pub fn build_backend(params: IndexParams) -> Box { - let preference = std::env::var("COMPASS_BACKEND").unwrap_or_else(|_| "cpu".into()); - match preference.as_str() { - "gpu" => build_gpu_or_warn(params), - "auto" => { - #[cfg(feature = "gpu")] - { - if compass_vector_gpu::cuda_available() { - return build_gpu_or_warn(params); - } - } - Box::new(UsearchHnswIndex::new(params)) - } - _ => Box::new(UsearchHnswIndex::new(params)), - } -} - -#[cfg(feature = "gpu")] -fn build_gpu_or_warn(params: IndexParams) -> Box { - match compass_vector_gpu::CuvsHnswIndex::new(params) { - Ok(idx) => { - tracing::info!("vector backend = cuvs-hnsw (GPU)"); - Box::new(idx) - } - Err(e) => { - tracing::warn!("GPU backend requested but unavailable ({e}); falling back to USearch"); - Box::new(UsearchHnswIndex::new(params)) - } - } -} - -#[cfg(not(feature = "gpu"))] -fn build_gpu_or_warn(params: IndexParams) -> Box { - tracing::warn!( - "COMPASS_BACKEND=gpu but binary built without --features gpu; falling back to USearch" - ); - Box::new(UsearchHnswIndex::new(params)) -} diff --git a/crates/compass/src/search/chunk_cache.rs b/crates/compass/src/search/chunk_cache.rs index e8478d1..d23dda8 100644 --- a/crates/compass/src/search/chunk_cache.rs +++ b/crates/compass/src/search/chunk_cache.rs @@ -113,7 +113,23 @@ impl ChunkCache { Ok(()) } + /// Tombstone passthrough (evicts tombstoned ids from the cache too). + pub fn tombstone_batch(&self, ids: &[u64]) -> Result<(), BoxErr> { + self.store.tombstone_batch(ids)?; + let mut cache = self.cache.lock().unwrap_or_else(|e| e.into_inner()); + for id in ids { + cache.pop(id); + } + Ok(()) + } + + /// Load persisted tombstones (passthrough). + pub fn load_tombstones(&self) -> Result, BoxErr> { + self.store.load_tombstones() + } + /// Number of chunks durably stored (not the cache size). + #[cfg(test)] pub fn count(&self) -> Result { self.store.count() } @@ -125,11 +141,13 @@ impl ChunkCache { } /// Current number of resident (cached) chunks — for tests/metrics. + #[cfg(test)] pub fn resident(&self) -> usize { self.cache.lock().unwrap_or_else(|e| e.into_inner()).len() } /// Access the underlying store (for code paths that must bypass the cache). + #[cfg(test)] pub fn store(&self) -> &ChunkStore { &self.store } diff --git a/crates/compass/src/search/chunk_store.rs b/crates/compass/src/search/chunk_store.rs index 03dd6a1..a852fb2 100644 --- a/crates/compass/src/search/chunk_store.rs +++ b/crates/compass/src/search/chunk_store.rs @@ -4,7 +4,7 @@ //! Point lookups by u64 ID, batch inserts, full scans for rebuild. use crate::models::DocumentChunk; -use redb::{Database, DatabaseError, ReadableTable, ReadableTableMetadata, TableDefinition}; +use redb::{Database, DatabaseError, ReadableTable, TableDefinition}; use std::path::Path; use std::time::Duration; @@ -131,6 +131,7 @@ impl ChunkStore { Ok(results) } + #[cfg(test)] pub fn insert( &self, id: u64, @@ -162,7 +163,9 @@ impl ChunkStore { Ok(()) } + #[cfg(test)] pub fn count(&self) -> Result> { + use redb::ReadableTableMetadata; let txn = self.db.begin_read()?; let table = txn.open_table(CHUNKS_TABLE)?; Ok(table.len()?) diff --git a/crates/compass/src/search/filter_bench.rs b/crates/compass/src/search/filter_bench.rs index 71515b3..b60968c 100644 --- a/crates/compass/src/search/filter_bench.rs +++ b/crates/compass/src/search/filter_bench.rs @@ -76,13 +76,11 @@ fn build_corpus(n: u32) -> (VectorState, FilterIndex) { ); filter_index.insert(i as u64, &metadata); } - filter_index.finalize(); let state = VectorState { index: Some(index), key_to_chunk_id: chunk_ids, mmap_vectors: None, vectors, - dims: DIMS, }; (state, filter_index) } diff --git a/crates/compass/src/search/filter_index.rs b/crates/compass/src/search/filter_index.rs index ca0d2d9..2d4f4f3 100644 --- a/crates/compass/src/search/filter_index.rs +++ b/crates/compass/src/search/filter_index.rs @@ -12,7 +12,7 @@ // // Storage shape (v0): // - equality: (field, canonical_string) -> RoaringTreemap of chunk_ids -// - numeric: field -> Vec<(value: f64, chunk_id)> sorted by value +// - numeric: field -> BTreeMap (O(log N) ops) // - string_list:(field, element) -> RoaringTreemap (for `contains`) // - present: field -> RoaringTreemap of chunk_ids that have any value // @@ -55,8 +55,10 @@ pub struct FilterIndex { equality: HashMap>, /// field -> string value -> chunk_ids (for `in` semantics on strings). equality_strings: HashMap>, - /// field -> sorted (value, chunk_id) for range predicates. - numeric: HashMap>, + /// field -> total-order-encoded f64 -> ids, for range predicates. + /// BTreeMap keys let inserts/removes stay O(log N) (a sorted Vec made + /// every incremental update O(N) — disqualifying at scale). + numeric: HashMap>, /// field -> element -> chunk_ids whose StringList contains the element. string_list_contains: HashMap>, /// field -> chunk_ids that have any value for this field. @@ -66,6 +68,17 @@ pub struct FilterIndex { universe: RoaringTreemap, } +/// Map f64 to a u64 preserving total order (IEEE-754 bit trick; NaNs are +/// filtered before insertion by `as_f64`). +fn f64_ord_key(x: f64) -> u64 { + let b = x.to_bits(); + if b >> 63 == 1 { + !b + } else { + b | (1 << 63) + } +} + impl FilterIndex { pub fn new() -> Self { Self::default() @@ -75,8 +88,17 @@ impl FilterIndex { self.universe.len() } - pub fn is_empty(&self) -> bool { - self.universe.is_empty() + /// The live-id universe (treemap) — shared with facet counting so deleted + /// chunks never inflate counts. + pub fn universe(&self) -> &RoaringTreemap { + &self.universe + } + + /// Is this id live (inserted and not removed)? The universe excludes + /// tombstoned ids on every maintenance path, so this doubles as the + /// existence check now that chunks are not held in RAM. + pub fn contains(&self, id: u64) -> bool { + self.universe.contains(id) } /// Insert a single chunk with its metadata. `chunk_id` is the full u64 @@ -107,7 +129,9 @@ impl FilterIndex { self.numeric .entry(field.clone()) .or_default() - .push((n, chunk_id)); + .entry(f64_ord_key(n)) + .or_default() + .insert(chunk_id); } if let MetadataValue::StringList(xs) = value { for x in xs { @@ -122,10 +146,56 @@ impl FilterIndex { } } - /// Call after all inserts so range scans are O(log N) per bound. - pub fn finalize(&mut self) { - for v in self.numeric.values_mut() { - v.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal)); + /// Remove one chunk (reverse of `insert`). O(log N) per field value — + /// deletes no longer trigger an O(collection) index rebuild. + pub fn remove(&mut self, chunk_id: u64, metadata: &HashMap) { + self.universe.remove(chunk_id); + for (field, value) in metadata { + if let Some(tm) = self.present.get_mut(field) { + tm.remove(chunk_id); + } + if let Some(vals) = self.equality.get_mut(field) { + let key = MetadataKey::from_metadata(value); + if let Some(tm) = vals.get_mut(&key) { + tm.remove(chunk_id); + if tm.is_empty() { + vals.remove(&key); + } + } + } + if let MetadataValue::String(sv) = value { + if let Some(vals) = self.equality_strings.get_mut(field) { + if let Some(tm) = vals.get_mut(sv) { + tm.remove(chunk_id); + if tm.is_empty() { + vals.remove(sv); + } + } + } + } + if let Some(n) = value.as_f64() { + if let Some(vals) = self.numeric.get_mut(field) { + let key = f64_ord_key(n); + if let Some(tm) = vals.get_mut(&key) { + tm.remove(chunk_id); + if tm.is_empty() { + vals.remove(&key); + } + } + } + } + if let MetadataValue::StringList(xs) = value { + if let Some(vals) = self.string_list_contains.get_mut(field) { + for x in xs { + if let Some(tm) = vals.get_mut(x) { + tm.remove(chunk_id); + if tm.is_empty() { + vals.remove(x); + } + } + } + } + } } } @@ -181,17 +251,14 @@ impl FilterIndex { } fn range(&self, field: &str, gte: Option, lte: Option) -> RoaringTreemap { - let Some(sorted) = self.numeric.get(field) else { + let Some(vals) = self.numeric.get(field) else { return RoaringTreemap::new(); }; - let lo = gte.unwrap_or(f64::NEG_INFINITY); - let hi = lte.unwrap_or(f64::INFINITY); - // sorted is by value; binary-search the bounds. - let start = sorted.partition_point(|(v, _)| *v < lo); - let end = sorted.partition_point(|(v, _)| *v <= hi); + let lo = f64_ord_key(gte.unwrap_or(f64::NEG_INFINITY)); + let hi = f64_ord_key(lte.unwrap_or(f64::INFINITY)); let mut out = RoaringTreemap::new(); - for (_, id) in &sorted[start..end] { - out.insert(*id); + for (_, tm) in vals.range(lo..=hi) { + out |= tm; } out } @@ -206,247 +273,6 @@ pub fn selectivity(eligible: &RoaringTreemap, universe_len: u64) -> f64 { eligible.len() as f64 / universe_len as f64 } -// ── Persistence ──────────────────────────────────────────────────────────── -// Serialize the whole index to bytes. NOT YET WIRED: every load path currently -// rebuilds the index from the chunk map; persisting/reloading it through the -// Storage trait (skipping the O(N) rebuild on startup) is future work. -// RoaringTreemaps use their native portable format; the container framing is a -// small length-prefixed encoding. Note the length prefixes are u32 — per-field -// entry counts are bounded by that (fine in practice; the u64 work in this -// module is about CHUNK IDS, not per-field entry counts). - -impl MetadataKey { - // Type-tagged encoding. Tag byte + payload. - fn encode(&self, buf: &mut Vec) { - match self { - MetadataKey::Bool(b) => { - buf.push(0); - buf.push(*b as u8); - } - MetadataKey::Int(i) => { - buf.push(1); - buf.extend_from_slice(&i.to_le_bytes()); - } - MetadataKey::Float(bits) => { - buf.push(2); - buf.extend_from_slice(&bits.to_le_bytes()); - } - MetadataKey::String(s) => { - buf.push(3); - write_str(buf, s); - } - MetadataKey::StringList(xs) => { - buf.push(4); - buf.extend_from_slice(&(xs.len() as u32).to_le_bytes()); - for x in xs { - write_str(buf, x); - } - } - } - } - - fn decode(buf: &[u8], pos: &mut usize) -> Option { - let tag = *buf.get(*pos)?; - *pos += 1; - Some(match tag { - 0 => { - let b = *buf.get(*pos)? != 0; - *pos += 1; - MetadataKey::Bool(b) - } - 1 => MetadataKey::Int(read_i64(buf, pos)?), - 2 => MetadataKey::Float(read_u64(buf, pos)?), - 3 => MetadataKey::String(read_str(buf, pos)?), - 4 => { - let n = read_u32(buf, pos)? as usize; - let mut xs = Vec::with_capacity(n); - for _ in 0..n { - xs.push(read_str(buf, pos)?); - } - MetadataKey::StringList(xs) - } - _ => return None, - }) - } -} - -fn write_str(buf: &mut Vec, s: &str) { - buf.extend_from_slice(&(s.len() as u32).to_le_bytes()); - buf.extend_from_slice(s.as_bytes()); -} - -fn read_u32(buf: &[u8], pos: &mut usize) -> Option { - let end = *pos + 4; - let v = u32::from_le_bytes(buf.get(*pos..end)?.try_into().ok()?); - *pos = end; - Some(v) -} - -fn read_u64(buf: &[u8], pos: &mut usize) -> Option { - let end = *pos + 8; - let v = u64::from_le_bytes(buf.get(*pos..end)?.try_into().ok()?); - *pos = end; - Some(v) -} - -fn read_i64(buf: &[u8], pos: &mut usize) -> Option { - Some(read_u64(buf, pos)? as i64) -} - -fn read_str(buf: &[u8], pos: &mut usize) -> Option { - let len = read_u32(buf, pos)? as usize; - let end = *pos + len; - let s = String::from_utf8(buf.get(*pos..end)?.to_vec()).ok()?; - *pos = end; - Some(s) -} - -fn write_treemap(buf: &mut Vec, t: &RoaringTreemap) { - let mut tmp = Vec::new(); - // RoaringTreemap::serialize_into writes the portable format. - t.serialize_into(&mut tmp).expect("treemap serialize"); - buf.extend_from_slice(&(tmp.len() as u32).to_le_bytes()); - buf.extend_from_slice(&tmp); -} - -fn read_treemap(buf: &[u8], pos: &mut usize) -> Option { - let len = read_u32(buf, pos)? as usize; - let end = *pos + len; - let slice = buf.get(*pos..end)?; - let t = RoaringTreemap::deserialize_from(slice).ok()?; - *pos = end; - Some(t) -} - -// field -> RoaringTreemap -fn write_map_tm(buf: &mut Vec, m: &HashMap) { - buf.extend_from_slice(&(m.len() as u32).to_le_bytes()); - for (k, v) in m { - write_str(buf, k); - write_treemap(buf, v); - } -} - -fn read_map_tm(buf: &[u8], pos: &mut usize) -> Option> { - let n = read_u32(buf, pos)? as usize; - let mut m = HashMap::with_capacity(n); - for _ in 0..n { - let k = read_str(buf, pos)?; - let v = read_treemap(buf, pos)?; - m.insert(k, v); - } - Some(m) -} - -// field -> (string -> RoaringTreemap) -fn write_map_str_tm(buf: &mut Vec, m: &HashMap>) { - buf.extend_from_slice(&(m.len() as u32).to_le_bytes()); - for (k, inner) in m { - write_str(buf, k); - write_map_tm(buf, inner); - } -} - -fn read_map_str_tm( - buf: &[u8], - pos: &mut usize, -) -> Option>> { - let n = read_u32(buf, pos)? as usize; - let mut m = HashMap::with_capacity(n); - for _ in 0..n { - let k = read_str(buf, pos)?; - let inner = read_map_tm(buf, pos)?; - m.insert(k, inner); - } - Some(m) -} - -impl FilterIndex { - /// Format version for the serialized index (bump on any framing change). - const FORMAT_VERSION: u8 = 1; - - /// Serialize the whole index to a byte buffer. - pub fn serialize(&self) -> Vec { - let mut buf = Vec::new(); - buf.push(Self::FORMAT_VERSION); - - // equality: field -> (MetadataKey -> treemap) - buf.extend_from_slice(&(self.equality.len() as u32).to_le_bytes()); - for (field, inner) in &self.equality { - write_str(&mut buf, field); - buf.extend_from_slice(&(inner.len() as u32).to_le_bytes()); - for (key, tm) in inner { - key.encode(&mut buf); - write_treemap(&mut buf, tm); - } - } - - write_map_str_tm(&mut buf, &self.equality_strings); - - // numeric: field -> Vec<(f64 bits, u64)> - buf.extend_from_slice(&(self.numeric.len() as u32).to_le_bytes()); - for (field, vals) in &self.numeric { - write_str(&mut buf, field); - buf.extend_from_slice(&(vals.len() as u32).to_le_bytes()); - for (v, id) in vals { - buf.extend_from_slice(&v.to_bits().to_le_bytes()); - buf.extend_from_slice(&id.to_le_bytes()); - } - } - - write_map_str_tm(&mut buf, &self.string_list_contains); - write_map_tm(&mut buf, &self.present); - write_treemap(&mut buf, &self.universe); - buf - } - - /// Reconstruct an index from bytes produced by [`FilterIndex::serialize`]. - pub fn deserialize(buf: &[u8]) -> Option { - let mut pos = 0usize; - let version = *buf.get(pos)?; - pos += 1; - if version != Self::FORMAT_VERSION { - return None; - } - - let mut idx = FilterIndex::new(); - - let n_eq = read_u32(buf, &mut pos)? as usize; - for _ in 0..n_eq { - let field = read_str(buf, &mut pos)?; - let n_inner = read_u32(buf, &mut pos)? as usize; - let mut inner = HashMap::with_capacity(n_inner); - for _ in 0..n_inner { - let key = MetadataKey::decode(buf, &mut pos)?; - let tm = read_treemap(buf, &mut pos)?; - inner.insert(key, tm); - } - idx.equality.insert(field, inner); - } - - idx.equality_strings = read_map_str_tm(buf, &mut pos)?; - - let n_num = read_u32(buf, &mut pos)? as usize; - for _ in 0..n_num { - let field = read_str(buf, &mut pos)?; - let n_vals = read_u32(buf, &mut pos)? as usize; - let mut vals = Vec::with_capacity(n_vals); - for _ in 0..n_vals { - let v = f64::from_bits(read_u64(buf, &mut pos)?); - let id = read_u64(buf, &mut pos)?; - vals.push((v, id)); - } - idx.numeric.insert(field, vals); - } - - idx.string_list_contains = read_map_str_tm(buf, &mut pos)?; - idx.present = read_map_tm(buf, &mut pos)?; - idx.universe = read_treemap(buf, &mut pos)?; - - Some(idx) - } -} - #[cfg(test)] mod tests { use super::*; @@ -481,7 +307,6 @@ mod tests { ]), ); } - idx.finalize(); idx } @@ -583,7 +408,6 @@ mod tests { ("tags", MetadataValue::StringList(vec!["even".into()])), ]), ); - idx.finalize(); assert_eq!(idx.len(), 1); @@ -623,57 +447,4 @@ mod tests { ); assert!(idx.eligible(&FilterExpr::compile(&con)).contains(big)); } - - #[test] - fn serialize_roundtrip_preserves_queries() { - let idx = build_index(); - let bytes = idx.serialize(); - let restored = FilterIndex::deserialize(&bytes).expect("deserialize ok"); - - assert_eq!(restored.len(), idx.len()); - - // Equality query matches identically. - let mut eq = HashMap::new(); - eq.insert( - "org_id".into(), - FilterValue::Exact(MetadataValue::String("acme".into())), - ); - let expr = FilterExpr::compile(&eq); - assert_eq!(idx.eligible(&expr).len(), restored.eligible(&expr).len()); - assert_eq!(restored.eligible(&expr).len(), 10); - - // Range query matches identically after restore. - let mut rng = HashMap::new(); - rng.insert( - "created_at".into(), - FilterValue::Condition(FilterCondition { - gte: Some(100.0), - lte: Some(200.0), - contains: None, - in_values: None, - }), - ); - let rexpr = FilterExpr::compile(&rng); - assert_eq!(idx.eligible(&rexpr).len(), restored.eligible(&rexpr).len()); - - // Contains query matches identically. - let mut con = HashMap::new(); - con.insert( - "tags".into(), - FilterValue::Condition(FilterCondition { - gte: None, - lte: None, - contains: Some("even".into()), - in_values: None, - }), - ); - let cexpr = FilterExpr::compile(&con); - assert_eq!(restored.eligible(&cexpr).len(), 500); - } - - #[test] - fn deserialize_rejects_bad_version() { - assert!(FilterIndex::deserialize(&[]).is_none()); - assert!(FilterIndex::deserialize(&[99]).is_none()); // bad version byte - } } diff --git a/crates/compass/src/search/filter_pushdown.rs b/crates/compass/src/search/filter_pushdown.rs index 849ee9c..2bf2db6 100644 --- a/crates/compass/src/search/filter_pushdown.rs +++ b/crates/compass/src/search/filter_pushdown.rs @@ -55,11 +55,6 @@ impl FilterExpr { } FilterExpr { predicates } } - - /// Evaluate the expression against a chunk's metadata. AND across predicates. - pub fn eval(&self, metadata: &HashMap) -> bool { - self.predicates.iter().all(|p| eval_predicate(p, metadata)) - } } fn push_condition(out: &mut Vec, field: &str, cond: &FilterCondition) { @@ -84,77 +79,20 @@ fn push_condition(out: &mut Vec, field: &str, cond: &FilterCondition) } } -fn eval_predicate(p: &Predicate, metadata: &HashMap) -> bool { - match p { - Predicate::Eq { field, value } => match metadata.get(field) { - Some(mv) => mv == value, - None => false, - }, - Predicate::Range { field, gte, lte } => { - match metadata.get(field).and_then(|m| m.as_f64()) { - Some(n) => { - gte.map(|g| n >= g).unwrap_or(true) && lte.map(|l| n <= l).unwrap_or(true) - } - None => false, - } - } - Predicate::Contains { field, value } => match metadata.get(field) { - Some(MetadataValue::StringList(xs)) => xs.iter().any(|x| x == value), - Some(MetadataValue::String(s)) => s == value, - _ => false, - }, - Predicate::In { field, values } => match metadata.get(field) { - Some(MetadataValue::String(s)) => values.contains(s), - None => false, - _ => false, - }, - } -} - -/// Canonical string form for an equality / set-membership key. Booleans and -/// numbers normalize to a stable string so that the filter index can key on -/// `(field, string)` without juggling typed variants. -pub fn stringify_metadata(mv: &MetadataValue) -> String { - match mv { - MetadataValue::Bool(b) => b.to_string(), - MetadataValue::Int(i) => i.to_string(), - MetadataValue::Float(f) => f.to_string(), - MetadataValue::String(s) => s.clone(), - MetadataValue::StringList(xs) => xs.join(","), - } -} - #[cfg(test)] mod tests { use super::*; - fn meta(pairs: &[(&str, MetadataValue)]) -> HashMap { - pairs - .iter() - .cloned() - .map(|(k, v)| (k.to_string(), v)) - .collect() - } - + // Semantics (eq / range / contains / in, AND across fields) are covered + // end-to-end in filter_index.rs tests via FilterIndex::eligible — the one + // live evaluator. These only pin the compile() shape. #[test] - fn eq_matches_string() { + fn compile_shapes() { let mut f = HashMap::new(); f.insert( "org_id".into(), FilterValue::Exact(MetadataValue::String("acme".into())), ); - let expr = FilterExpr::compile(&f); - assert!(expr.eval(&meta(&[("org_id", MetadataValue::String("acme".into()))]))); - assert!(!expr.eval(&meta(&[( - "org_id", - MetadataValue::String("widgets".into()) - )]))); - assert!(!expr.eval(&meta(&[]))); - } - - #[test] - fn range_inclusive_bounds() { - let mut f = HashMap::new(); f.insert( "created_at".into(), FilterValue::Condition(FilterCondition { @@ -165,67 +103,8 @@ mod tests { }), ); let expr = FilterExpr::compile(&f); - assert!(expr.eval(&meta(&[("created_at", MetadataValue::Int(100))]))); - assert!(expr.eval(&meta(&[("created_at", MetadataValue::Float(150.5))]))); - assert!(expr.eval(&meta(&[("created_at", MetadataValue::Int(200))]))); - assert!(!expr.eval(&meta(&[("created_at", MetadataValue::Int(99))]))); - assert!(!expr.eval(&meta(&[("created_at", MetadataValue::Int(201))]))); - } - - #[test] - fn and_of_eq_and_range() { - let mut f = HashMap::new(); - f.insert( - "org_id".into(), - FilterValue::Exact(MetadataValue::String("acme".into())), - ); - f.insert( - "created_at".into(), - FilterValue::Condition(FilterCondition { - gte: Some(100.0), - lte: None, - contains: None, - in_values: None, - }), - ); - let expr = FilterExpr::compile(&f); - let ok = meta(&[ - ("org_id", MetadataValue::String("acme".into())), - ("created_at", MetadataValue::Int(150)), - ]); - let wrong_org = meta(&[ - ("org_id", MetadataValue::String("widgets".into())), - ("created_at", MetadataValue::Int(150)), - ]); - let too_old = meta(&[ - ("org_id", MetadataValue::String("acme".into())), - ("created_at", MetadataValue::Int(50)), - ]); - assert!(expr.eval(&ok)); - assert!(!expr.eval(&wrong_org)); - assert!(!expr.eval(&too_old)); - } - - #[test] - fn contains_on_string_list() { - let mut f = HashMap::new(); - f.insert( - "tags".into(), - FilterValue::Condition(FilterCondition { - gte: None, - lte: None, - contains: Some("sports".into()), - in_values: None, - }), - ); - let expr = FilterExpr::compile(&f); - assert!(expr.eval(&meta(&[( - "tags", - MetadataValue::StringList(vec!["sports".into(), "goals".into()]), - )]))); - assert!(!expr.eval(&meta(&[( - "tags", - MetadataValue::StringList(vec!["news".into()]), - )]))); + assert_eq!(expr.predicates.len(), 2); + assert!(!expr.is_empty()); + assert!(FilterExpr::compile(&HashMap::new()).is_empty()); } } diff --git a/crates/compass/src/search/mod.rs b/crates/compass/src/search/mod.rs index 0a302d4..f233d6b 100644 --- a/crates/compass/src/search/mod.rs +++ b/crates/compass/src/search/mod.rs @@ -5,33 +5,17 @@ // Semantic: USearch HNSW approximate nearest neighbor search // Hybrid: Both combined via Reciprocal Rank Fusion (RRF, k=60) -#[allow(dead_code)] -pub mod backend; -#[allow(dead_code)] pub mod chunk_cache; -#[allow(dead_code)] pub mod chunk_store; -// Filter-aware ANN modules . Not yet wired into the API -// surface; `search_vectors_filtered` below is the prototype call site. #[cfg(test)] mod filter_bench; -#[allow(dead_code)] pub mod filter_index; -#[allow(dead_code)] pub mod filter_pushdown; pub mod hybrid; -#[allow(dead_code)] pub mod mmap_vectors; pub mod tantivy_fts; pub mod vector; -// Re-export the stable trait surface for external consumers and future use. -#[allow(unused_imports)] -pub use backend::{ - build_backend, IndexError, IndexParams, LoadableIndex, UsearchHnswIndex, VectorIndex, - VectorMatch, -}; - /// Search mode — determines which search engines are used for a query. #[derive(Debug, Clone, Copy)] pub enum SearchMode { diff --git a/crates/compass/src/search/tantivy_fts.rs b/crates/compass/src/search/tantivy_fts.rs index 8c74f3a..b9771d9 100644 --- a/crates/compass/src/search/tantivy_fts.rs +++ b/crates/compass/src/search/tantivy_fts.rs @@ -1,19 +1,12 @@ -// search/tantivy_fts.rs — Full-text search + precomputed bitset faceting via Tantivy. +// search/tantivy_fts.rs — Full-text search + precomputed facet treemaps via Tantivy. // // Performance architecture: // - Full-text search: Tantivy's inverted index (BM25 scoring, sub-ms for any dataset size) -// - Facet counting: PRECOMPUTED BITSETS. At index time, we build one bitset per unique -// metadata value (e.g. one bitset for department="Legal"). At query time, we AND the -// query's result bitset with each precomputed bitset and popcount. -// This gives microsecond faceting even at millions of documents. -// -// How bitset faceting works: -// At index time: -// facetBitsets = { "department": { "Legal": BitSet([0,1,4,7,...]), "Eng": BitSet([2,3,...]) } } -// At query time: -// queryBits = search(query) // bitset of matching doc IDs -// legalCount = (queryBits AND facetBitsets["department"]["Legal"]).popcount() -// // ^ This is ~20 microseconds for 250K documents +// - Facet counting: precomputed roaring treemaps keyed by CHUNK ID, one per unique +// metadata value. At query time each value's treemap is intersected with the +// live-id universe (and the query's hit set, if any) and popcounted — +// microsecond faceting independent of collection size, correct under +// sparse/block-allocated ids and deletes. use crate::models::{DocumentChunk, MetadataValue}; use std::collections::HashMap; @@ -26,93 +19,51 @@ use tantivy::tokenizer::{ }; use tantivy::{Index, IndexWriter, ReloadPolicy}; -// ── Bitset implementation ──────────────────────────────────────────────────── -// A compact bitset stored as a Vec. Each u64 holds 64 bits. -// This is the core data structure that makes faceting fast. +// ── Precomputed facet bitsets ──────────────────────────────────────────────── +// Built once at index time, reused for every facet query. +// Structure: { "department" => { "Legal" => RoaringTreemap(chunk ids), ... } } -#[derive(Clone, Debug)] -pub struct BitSet { - /// Each u64 stores 64 bits. words[0] covers bits 0-63, words[1] covers 64-127, etc. - words: Vec, - /// Total number of bits (= total number of documents in the collection) - len: usize, +#[derive(Clone, Debug, Default)] +pub struct FacetBitsets { + /// Nested map: field_name -> { value -> treemap of matching CHUNK IDS }. + /// Keyed by chunk id (not doc position): ids are u64 and non-dense once + /// block-allocated, and the query side has always intersected on the + /// stored id — position-keyed dense bitsets silently misaligned. + pub groups: HashMap>, } -impl BitSet { - /// Create a new bitset with all bits set to 0 (nothing matches). - fn new(num_bits: usize) -> Self { - // Ceiling division: how many u64 words we need to cover all bits - let num_words = (num_bits + 63) / 64; - Self { - words: vec![0u64; num_words], - len: num_bits, - } - } - - /// Create a bitset with ALL bits set to 1 (everything matches). - /// Used for unfiltered facet queries where every document counts. - #[allow(dead_code)] - fn all(num_bits: usize) -> Self { - let num_words = (num_bits + 63) / 64; - let mut words = vec![u64::MAX; num_words]; - // Clear the extra trailing bits in the last word so popcount stays accurate - let trailing = num_bits % 64; - if trailing > 0 && !words.is_empty() { - let last = words.len() - 1; - words[last] = (1u64 << trailing) - 1; - } - Self { - words, - len: num_bits, - } - } - - /// Set a single bit to 1 (mark document at this position as matching). - #[inline] - fn set(&mut self, bit: usize) { - if bit < self.len { - // bit >> 6 = which u64 word (dividing by 64) - // bit & 63 = which bit within that word (modulo 64) - self.words[bit >> 6] |= 1u64 << (bit & 63); - } - } - - /// AND two bitsets together, producing a new bitset. - /// This is the hot path — called once per facet value per query. - /// Each iteration processes 64 documents in a single CPU instruction. - #[inline] - fn and(&self, other: &BitSet) -> BitSet { - let min_len = self.words.len().min(other.words.len()); - let mut result = Vec::with_capacity(min_len); - for i in 0..min_len { - // Compiles down to a single AND instruction per 64 documents - result.push(self.words[i] & other.words[i]); - } - BitSet { - words: result, - len: self.len.min(other.len), +impl FacetBitsets { + /// Union another (older) facet map into this one. Appending a batch used + /// to REPLACE the facet state with new-batch-only bitsets — facets went + /// wrong after the second ingest batch, latent since v0.2. + pub fn absorb(&mut self, older: &FacetBitsets) { + for (field, vals) in &older.groups { + let dst = self.groups.entry(field.clone()).or_default(); + for (value, tm) in vals { + *dst.entry(value.clone()).or_default() |= tm; + } } } - /// Count the number of set bits (1s) in the entire bitset. - /// Uses the CPU's native POPCNT instruction for maximum speed. - #[inline] - fn popcount(&self) -> u64 { - // count_ones() compiles to hardware POPCNT — processes 64 bits per clock cycle - self.words.iter().map(|w| w.count_ones() as u64).sum() + /// Record one chunk's facet values (used by streaming rebuilds at load). + pub fn insert_chunk(&mut self, chunk: &DocumentChunk) { + insert_facets_for(&mut self.groups, chunk); } } -// ── Precomputed facet bitsets ──────────────────────────────────────────────── -// Built once at index time, reused for every facet query. -// Structure: { "department" => { "Legal" => BitSet, "Eng" => BitSet }, ... } - -#[derive(Clone, Debug)] -pub struct FacetBitsets { - /// Nested map: field_name -> { value -> bitset of matching doc positions } - pub groups: HashMap>, - /// Total number of documents (needed to create "all" bitsets for unfiltered queries) - pub total_docs: usize, +fn insert_facets_for( + groups: &mut HashMap>, + chunk: &DocumentChunk, +) { + for (field, value) in &chunk.metadata { + let repr = metadata_value_repr(value); + groups + .entry(field.clone()) + .or_default() + .entry(repr) + .or_default() + .insert(chunk.id); + } } // ── FtsState ───────────────────────────────────────────────────────────────── @@ -123,15 +74,11 @@ pub struct FtsState { pub index: Index, /// Cached reader — created once, reused for all queries (avoids ~1ms overhead per query) pub reader: tantivy::IndexReader, - // Field handles for the Tantivy schema + // Field handles the query paths read. (The schema defines more columns — + // collection/file_id/chunk_index/page/metadata — written at index time via + // FtsFields; only these two are read back.) pub id_field: Field, - pub collection_field: Field, - pub file_id_field: Field, - pub chunk_index_field: Field, - pub page_field: Field, pub text_field: Field, - /// We store arbitrary metadata as a JSON string field (indexed per-key via facet bitsets) - pub metadata_field: Field, /// Precomputed bitsets for microsecond faceting pub facet_bitsets: FacetBitsets, } @@ -202,7 +149,6 @@ fn register_tokenizers(index: &Index) { pub fn build_index( dir: &Path, chunks: &[DocumentChunk], - existing_count: u64, ) -> Result> { let (schema, fields) = build_schema(); @@ -245,12 +191,9 @@ pub fn build_index( writer.commit()?; - // ── Precompute facet bitsets ────────────────────────────────────────────── - // We need ALL chunks in the collection (existing + new) to build accurate bitsets. - // For now, we rebuild bitsets from the chunks we have. On reload from disk, - // the collection manager will call rebuild_facets() with all chunks. - let total_docs = (existing_count as usize) + chunks.len(); - let facet_bitsets = build_facet_bitsets(chunks, existing_count as usize, total_docs); + // Facet treemaps for THIS batch only. Callers accumulate: ingest absorbs + // the prior state; the load/rebuild scans reconstruct from all live chunks. + let facet_bitsets = build_facet_bitsets(chunks); // Create a reader once, reuse for all queries let reader = index @@ -262,12 +205,7 @@ pub fn build_index( index, reader, id_field: fields.id, - collection_field: fields.collection, - file_id_field: fields.file_id, - chunk_index_field: fields.chunk_index, - page_field: fields.page, text_field: fields.text, - metadata_field: fields.metadata, facet_bitsets, }) } @@ -279,63 +217,38 @@ pub fn open_index(dir: &Path) -> Result FacetBitsets { - let mut groups: HashMap> = HashMap::new(); - - // Scan all chunks and set bits for each metadata key-value pair. - // MetadataValue is converted to a string for facet grouping (e.g. Float(9.5) -> "9.5"). - for (i, chunk) in chunks.iter().enumerate() { - let bit_pos = offset + i; - for (key, value) in &chunk.metadata { - let value_str = metadata_to_facet_string(value); - groups - .entry(key.clone()) - .or_default() - .entry(value_str) - .or_insert_with(|| BitSet::new(total_docs)) - .set(bit_pos); - } +fn build_facet_bitsets(chunks: &[DocumentChunk]) -> FacetBitsets { + let mut fb = FacetBitsets::default(); + for chunk in chunks { + fb.insert_chunk(chunk); } - - FacetBitsets { groups, total_docs } + fb } /// Convert a MetadataValue to a string for facet grouping. -fn metadata_to_facet_string(val: &MetadataValue) -> String { +fn metadata_value_repr(val: &MetadataValue) -> String { match val { MetadataValue::String(s) => s.clone(), MetadataValue::Int(i) => i.to_string(), @@ -347,12 +260,11 @@ fn metadata_to_facet_string(val: &MetadataValue) -> String { /// Run a full-text search query. Returns (matching doc IDs + scores, total count, microseconds). /// -/// Metadata filtering is handled post-search by the collection manager (using the scoring -/// pipeline), so this function only does text-based search. +/// Metadata filtering is handled by the collection manager (roaring filter-index +/// pushdown + scoring pipeline), so this function only does text-based search. pub fn search( state: &FtsState, query_str: &str, - filters: &HashMap, limit: usize, ) -> Result<(Vec<(u64, f32)>, usize, u64), Box> { let start = std::time::Instant::now(); @@ -378,22 +290,7 @@ pub fn search( } }; - // If there are metadata filters, combine them with the text query using BooleanQuery - let query: Box = if filters.is_empty() { - text_query - } else { - // Each filter becomes a MUST clause — all must match - let mut clauses: Vec<(tantivy::query::Occur, Box)> = Vec::new(); - clauses.push((tantivy::query::Occur::Must, text_query)); - - // Metadata filters are matched against stored text fields. - // Since metadata is stored as JSON, we can't filter directly in Tantivy. - // Instead, we apply metadata filtering post-search using the bitsets. - // For now, we include the text query only and let the caller handle filtering. - // TODO: implement metadata filtering via bitset post-filtering - - Box::new(tantivy::query::BooleanQuery::new(clauses)) - }; + let query: Box = text_query; // Execute search: get top results + total count in a single pass let (top_docs, total_count) = searcher.search(&query, &(TopDocs::with_limit(limit), Count))?; @@ -423,60 +320,54 @@ pub fn get_facets( state: &FtsState, query_str: &str, requested_fields: &[String], + live: &roaring::RoaringTreemap, ) -> Result<(HashMap>, u64), Box> { let start = std::time::Instant::now(); let bs = &state.facet_bitsets; - // For unfiltered queries, every document matches — use "all ones" bitset - let query_bitset = if query_str.is_empty() || query_str == "*" { - None // fast path: skip query execution entirely + // Text-filtered queries build a treemap of matching CHUNK IDS; unfiltered + // queries skip query execution entirely. Counts always intersect with the + // LIVE universe, so soft-deleted chunks never inflate facets. + let query_ids: Option = if query_str.is_empty() || query_str == "*" { + None } else { - // Execute the text query and build a bitset from matching doc IDs let searcher = state.reader.searcher(); let query_parser = QueryParser::for_index(&state.index, vec![state.text_field]); - let query: Box = match query_parser.parse_query(query_str) { Ok(q) => q, Err(_) => Box::new(tantivy::query::AllQuery), }; - - let top_docs = searcher.search(&query, &TopDocs::with_limit(bs.total_docs))?; - - let mut result_bits = BitSet::new(bs.total_docs); + let top_docs = searcher.search(&query, &TopDocs::with_limit(usize::MAX >> 32))?; + let mut ids = roaring::RoaringTreemap::new(); for (_score, doc_address) in &top_docs { let doc: tantivy::TantivyDocument = searcher.doc(*doc_address)?; if let Some(tantivy::schema::OwnedValue::U64(id)) = doc.get_first(state.id_field) { - result_bits.set(*id as usize); + ids.insert(*id); } } - Some(result_bits) + Some(ids) }; - // THE HOT PATH: bitset AND + popcount for each facet value - let mut facets: HashMap> = HashMap::new(); - - for (group_name, value_bitsets) in &bs.groups { - // If specific fields were requested, skip fields not in the list - if !requested_fields.is_empty() && !requested_fields.contains(group_name) { + let mut out: HashMap> = HashMap::new(); + for (field, values) in &bs.groups { + if !requested_fields.is_empty() && !requested_fields.contains(field) { continue; } - let mut counts: HashMap = HashMap::new(); - for (value, value_bits) in value_bitsets { - let count = match &query_bitset { - // Unfiltered: just popcount the precomputed bitset directly - None => value_bits.popcount(), - // Filtered: AND with query results, then popcount the intersection - Some(qb) => qb.and(value_bits).popcount(), - }; - if count > 0 { - counts.insert(value.clone(), count); + for (value, tm) in values { + let mut hit = tm & live; + if let Some(q) = &query_ids { + hit &= q; } + let n = hit.len(); + if n > 0 { + counts.insert(value.clone(), n); + } + } + if !counts.is_empty() { + out.insert(field.clone(), counts); } - facets.insert(group_name.clone(), counts); } - - let took_us = start.elapsed().as_micros() as u64; - Ok((facets, took_us)) + Ok((out, start.elapsed().as_micros() as u64)) } diff --git a/crates/compass/src/search/vector.rs b/crates/compass/src/search/vector.rs index 62ff4a1..e30490b 100644 --- a/crates/compass/src/search/vector.rs +++ b/crates/compass/src/search/vector.rs @@ -46,8 +46,6 @@ pub struct VectorState { pub mmap_vectors: Option, /// Legacy in-memory vectors for datasets without an mmap file (e.g. first build). pub vectors: Vec>, - /// Embedding dimensionality (e.g. 384 for BGE-small) - pub dims: usize, } unsafe impl Send for VectorState {} @@ -78,10 +76,10 @@ pub fn create_index( let index = Index::new(&opts).map_err(|e| format!("Failed to create USearch index: {}", e))?; if capacity > 0 { // Reserve enough concurrent search slots for the spawn_blocking pool. - // Default rayon threads (=CPU count) is too low when search runs on - // tokio's blocking pool. 128 slots costs ~256KB and avoids the - // "No available threads to lock" fallback to brute-force. - let threads = 128.max(rayon::current_num_threads()); + // CPU count is too low when search runs on tokio's blocking pool. + // 128 slots costs ~256KB and avoids the "No available threads to + // lock" fallback to brute-force. + let threads = index_threads(); index .reserve_capacity_and_threads(capacity, threads) .map_err(|e| format!("Failed to reserve USearch capacity: {}", e))?; @@ -105,7 +103,6 @@ pub fn build_vector_index( key_to_chunk_id: Vec::new(), mmap_vectors: None, vectors: Vec::new(), - dims, }); } @@ -115,21 +112,25 @@ pub fn build_vector_index( } let mmap = super::mmap_vectors::MmapVectors::create(vectors_path, dims, vectors)?; - // For small datasets, skip HNSW and use brute-force search + // For small datasets, skip HNSW and use brute-force search. The keymap + // must STILL be persisted: without it a restart loses key->chunk-id + // mapping and falls back to identity, which silently returns wrong ids + // once ids are non-dense (block-allocated ids exposed this). if vectors.len() < HNSW_THRESHOLD { + let map_path = index_path.with_extension("keymap"); + save_key_map(&map_path, chunk_ids)?; return Ok(VectorState { index: None, key_to_chunk_id: chunk_ids.to_vec(), mmap_vectors: Some(mmap), vectors: Vec::new(), - dims, }); } // Build the HNSW index let index = create_index(dims, vectors.len())?; - // Insert vectors using parallel threads via rayon + // Insert vectors using usearch's internal thread slots for (key, vec) in vectors.iter().enumerate() { index .add(key as u64, vec) @@ -156,7 +157,6 @@ pub fn build_vector_index( key_to_chunk_id: chunk_ids.to_vec(), mmap_vectors: Some(mmap), vectors: Vec::new(), - dims, }) } @@ -198,7 +198,18 @@ pub fn load_vector_index( // Load the key-to-chunk-id mapping let map_path = index_path.with_extension("keymap"); - let key_to_chunk_id = load_key_map(&map_path)?; + let mut key_to_chunk_id = load_key_map(&map_path)?; + // Pre-fix local dirs never persisted the keymap for small datasets and + // relied implicitly on identity mapping (dense ids from 0). Make that + // explicit so a later incremental append can't push new ids onto an empty + // keymap and misalign every existing vector. + if key_to_chunk_id.is_empty() { + if let Some(m) = &mmap { + if !m.is_empty() { + key_to_chunk_id = (0..m.len() as u64).collect(); + } + } + } // For small datasets, skip HNSW if count < HNSW_THRESHOLD { @@ -207,7 +218,6 @@ pub fn load_vector_index( key_to_chunk_id, mmap_vectors: mmap, vectors: Vec::new(), - dims, }); } @@ -221,12 +231,54 @@ pub fn load_vector_index( .view(index_path_str) .map_err(|e| format!("Failed to mmap USearch index: {}", e))?; + // Crash/shutdown recovery for batched HNSW saves: vectors are durable + // in the mmap file per batch, but the index file is rewritten only + // every N batches — the file can be missing up to N-1 batches' rows. + // Heal by APPENDING the missing tail from the mmap (same as the + // runtime heal in apply_ingest_commit); a full rebuild here made every + // warm restart O(collection) instead of O(unsaved tail). The keymap is + // saved per batch, so the index is only ever behind it, never ahead. + let index = if index.size() < key_to_chunk_id.len() { + tracing::warn!( + "HNSW index at {} is stale ({} < {}); appending missing rows from mmap", + index_path.display(), + index.size(), + key_to_chunk_id.len() + ); + let healed = create_index(dims, key_to_chunk_id.len())?; + healed + .load(index_path_str) + .map_err(|e| format!("Failed to load USearch index for heal: {}", e))?; + let threads = index_threads(); + healed + .reserve_capacity_and_threads(key_to_chunk_id.len(), threads) + .map_err(|e| format!("Reserve failed: {}", e))?; + if let Some(m) = &mmap { + for i in (healed.size() as usize)..key_to_chunk_id.len().min(m.len()) { + healed + .add(i as u64, m.get(i)) + .map_err(|e| format!("Failed to add vector: {}", e))?; + } + } + healed + .save(index_path_str) + .map_err(|e| format!("Failed to save healed index: {}", e))?; + // Serve the healed file mmap-backed like the clean path, instead + // of keeping the whole graph resident. + let viewed = create_index(dims, 0)?; + viewed + .view(index_path_str) + .map_err(|e| format!("Failed to mmap healed USearch index: {}", e))?; + viewed + } else { + index + }; + Ok(VectorState { index: Some(index), key_to_chunk_id, mmap_vectors: mmap, vectors: Vec::new(), - dims, }) } else { Ok(VectorState { @@ -234,7 +286,6 @@ pub fn load_vector_index( key_to_chunk_id, mmap_vectors: mmap, vectors: Vec::new(), - dims, }) } } @@ -244,12 +295,6 @@ pub fn load_vector_index( /// + selectivity story end-to-end. #[derive(Debug, Clone, Default)] pub struct FilteredSearchExplain { - /// |eligible| at query time. - pub eligible_count: u64, - /// |universe| at query time. - pub universe_count: u64, - /// eligible / universe. - pub selectivity: f64, /// Whether the HNSW filtered walk was used (vs. brute force fallback). pub used_hnsw: bool, /// Number of HNSW candidates inspected. Counted via the filter closure @@ -272,20 +317,11 @@ pub fn search_vectors_filtered( top_k: usize, eligible: &RoaringTreemap, ) -> (Vec, FilteredSearchExplain) { - let universe = state.key_to_chunk_id.len() as u64; - let eligible_count = eligible.len(); let mut explain = FilteredSearchExplain { - eligible_count, - universe_count: universe, - selectivity: if universe == 0 { - 1.0 - } else { - eligible_count as f64 / universe as f64 - }, used_hnsw: false, candidates_inspected: 0, }; - if eligible_count == 0 { + if eligible.is_empty() { return (Vec::new(), explain); } @@ -440,30 +476,18 @@ pub fn search_vectors(query_vec: &[f32], state: &VectorState, top_k: usize) -> V } // ── Persistence helpers ────────────────────────────────────────────────────── -// Simple binary formats for saving/loading vectors and key maps to disk. - -/// Save vectors to a binary file (legacy format, kept for migration). -/// Format: [u32 count] [u32 dims] [count * dims * f32 values] -#[allow(dead_code)] -fn save_vectors( - path: &Path, - vectors: &[Vec], - dims: usize, -) -> Result<(), Box> { - if let Some(parent) = path.parent() { - std::fs::create_dir_all(parent)?; - } - let count = vectors.len(); - let mut buf: Vec = Vec::with_capacity(8 + count * dims * 4); - buf.extend_from_slice(&(count as u32).to_le_bytes()); - buf.extend_from_slice(&(dims as u32).to_le_bytes()); - for vec in vectors { - for &val in vec { - buf.extend_from_slice(&val.to_le_bytes()); - } - } - std::fs::write(path, buf)?; - Ok(()) +// Binary formats for loading vectors and key maps from disk. (The legacy +// vector WRITER is gone — only the mmap format is written; the legacy reader +// below survives for migration.) + +/// Thread-slot count for usearch reserve calls: at least 128 (tokio's +/// blocking pool can run more concurrent searches than there are cores). +pub(crate) fn index_threads() -> usize { + 128.max( + std::thread::available_parallelism() + .map(|n| n.get()) + .unwrap_or(1), + ) } /// Load vectors from a binary file. diff --git a/crates/compass/src/storage/id_alloc.rs b/crates/compass/src/storage/id_alloc.rs new file mode 100644 index 0000000..8d8a9bc --- /dev/null +++ b/crates/compass/src/storage/id_alloc.rs @@ -0,0 +1,179 @@ +//! CAS-leased chunk-id block allocator — `{ns}/id-alloc` in object storage. +//! +//! In cloud mode, EVERY ingest path allocates chunk ids from blocks claimed +//! here (attached serving nodes pool a block; stateless writers claim per +//! batch). A local `next_id` counter cannot be the allocation source in cloud +//! mode: recovery computes `max_id + 1`, which can land inside another +//! writer's active, partially-used block — colliding with ids that writer +//! will mint next. The invariant this module defends is *no id is ever handed +//! out twice*; global ordering is irrelevant (replay order comes from the +//! manifest `seq`, not from ids), and gaps from crashed writers are fine. +//! +//! Local mode never touches this module (`next_id` remains the allocator). + +use super::{Storage, StorageError}; +use serde::{Deserialize, Serialize}; +use std::ops::Range; + +/// Ids claimed per CAS round-trip. Large enough that an attached node's pool +/// refill is rare; small enough that a crashed writer leaks little. +pub const BLOCK: u64 = 10_000; + +const MAX_CAS_RETRIES: u32 = 10; + +fn alloc_key(ns: &str) -> String { + format!("{ns}/id-alloc") +} + +#[derive(Debug, Serialize, Deserialize)] +struct AllocState { + next_block_start: u64, +} + +/// Create-only seed of the allocator. `start` must be one past the highest id +/// ever assigned in the namespace (0 for a fresh collection). Losing the +/// create race is fine — the winner's value is equally valid because no new +/// ids can be minted while the allocator is absent (all cloud ingest paths +/// require it), so concurrent seeders compute the same high-water mark. +pub async fn seed(storage: &dyn Storage, ns: &str, start: u64) -> Result<(), StorageError> { + let state = AllocState { + next_block_start: start, + }; + let bytes = serde_json::to_vec(&state) + .map_err(|e| StorageError::Io(format!("id-alloc encode: {e}")))?; + match storage + .put_if_not_exists(&alloc_key(ns), bytes::Bytes::from(bytes)) + .await + { + Ok(_) => Ok(()), + Err(StorageError::AlreadyExists(_)) => Ok(()), // racer seeded it — fine + Err(e) => Err(e), + } +} + +/// The current allocation frontier: every legitimately-minted id is < this. +/// `NotFound` when the allocator was never seeded (pre-v0.4 namespace). +pub async fn frontier(storage: &dyn Storage, ns: &str) -> Result { + let bytes = storage.get(&alloc_key(ns)).await?; + let state: AllocState = serde_json::from_slice(&bytes) + .map_err(|e| StorageError::Io(format!("id-alloc decode for '{ns}': {e}")))?; + Ok(state.next_block_start) +} + +/// Claim a block of at least `count` ids (min [`BLOCK`]) via CAS. Returns the +/// claimed half-open range. `NotFound` means the allocator was never seeded +/// (pre-v0.4 namespace) — the caller migrates via [`seed`] and retries. +pub async fn claim( + storage: &dyn Storage, + ns: &str, + count: u64, +) -> Result, StorageError> { + let want = count.max(BLOCK); + let key = alloc_key(ns); + for _ in 0..MAX_CAS_RETRIES { + let (bytes, version) = storage.get_versioned(&key).await?; + let state: AllocState = serde_json::from_slice(&bytes) + .map_err(|e| StorageError::Io(format!("id-alloc decode for '{ns}': {e}")))?; + let start = state.next_block_start; + let end = start.checked_add(want).ok_or_else(|| { + StorageError::Io(format!("id space exhausted for '{ns}' (u64 overflow)")) + })?; + let next = AllocState { + next_block_start: end, + }; + let encoded = serde_json::to_vec(&next) + .map_err(|e| StorageError::Io(format!("id-alloc encode: {e}")))?; + match storage + .put_if_match(&key, bytes::Bytes::from(encoded), &version) + .await + { + Ok(_) => return Ok(start..end), + Err(StorageError::VersionConflict { .. }) => continue, + Err(e) => return Err(e), + } + } + Err(StorageError::Io(format!( + "id-alloc CAS failed after {MAX_CAS_RETRIES} retries for '{ns}' (persistent contention)" + ))) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::storage::local::LocalDiskStorage; + use std::sync::Arc; + + fn storage(name: &str) -> Arc { + use std::sync::atomic::{AtomicU64, Ordering}; + static N: AtomicU64 = AtomicU64::new(0); + let mut root = std::env::temp_dir(); + root.push(format!( + "compass_idalloc_test_{}_{}_{}", + name, + std::process::id(), + N.fetch_add(1, Ordering::Relaxed) + )); + let _ = std::fs::remove_dir_all(&root); + Arc::new(LocalDiskStorage::new(root).unwrap()) + } + + #[tokio::test] + async fn seed_then_claim_advances() { + let s = storage("basic"); + seed(s.as_ref(), "ns", 0).await.unwrap(); + let a = claim(s.as_ref(), "ns", 5).await.unwrap(); + assert_eq!(a, 0..BLOCK); // min block size applies + let b = claim(s.as_ref(), "ns", 25_000).await.unwrap(); + assert_eq!(b, BLOCK..BLOCK + 25_000); // large batches claim exactly enough + } + + #[tokio::test] + async fn seed_race_is_idempotent() { + let s = storage("seedrace"); + seed(s.as_ref(), "ns", 42).await.unwrap(); + // A racing seeder (same computed high-water) loses silently. + seed(s.as_ref(), "ns", 42).await.unwrap(); + let a = claim(s.as_ref(), "ns", 1).await.unwrap(); + assert_eq!(a.start, 42); + } + + #[tokio::test] + async fn claim_before_seed_is_not_found() { + let s = storage("unseeded"); + assert!(matches!( + claim(s.as_ref(), "ns", 1).await, + Err(StorageError::NotFound(_)) + )); + } + + // The allocator's whole job under concurrency: N racing claimants must + // receive disjoint ranges (modeled on the LSM's concurrent-appends test). + #[tokio::test] + async fn concurrent_claims_are_disjoint() { + let s = storage("concurrent"); + seed(s.as_ref(), "ns", 0).await.unwrap(); + let n = 16; + let mut handles = Vec::new(); + for _ in 0..n { + let s2 = s.clone(); + handles.push(tokio::spawn( + async move { claim(s2.as_ref(), "ns", 1).await }, + )); + } + let mut ranges: Vec> = Vec::new(); + for h in handles { + ranges.push(h.await.unwrap().unwrap()); + } + ranges.sort_by_key(|r| r.start); + for w in ranges.windows(2) { + assert!( + w[0].end <= w[1].start, + "overlapping claims: {:?} vs {:?}", + w[0], + w[1] + ); + } + // No holes either: 16 min-size blocks tile exactly. + assert_eq!(ranges.last().unwrap().end, n as u64 * BLOCK); + } +} diff --git a/crates/compass/src/storage/lsm.rs b/crates/compass/src/storage/lsm.rs index c920c3a..deddd6d 100644 --- a/crates/compass/src/storage/lsm.rs +++ b/crates/compass/src/storage/lsm.rs @@ -21,7 +21,6 @@ use super::{Storage, StorageError, Version}; use bytes::Bytes; use serde::{Deserialize, Serialize}; -use std::sync::Arc; /// Max CAS attempts when committing the manifest before giving up. const MAX_CAS_RETRIES: u32 = 10; @@ -150,6 +149,25 @@ async fn commit_manifest( } } +/// Fetch one WAL fragment's payload by id. +pub async fn read_fragment( + storage: &dyn Storage, + ns: &str, + id: &str, +) -> Result { + storage.get(&fragment_key(ns, id)).await +} + +/// Create-only commit of an EMPTY manifest for a new namespace, making a +/// zero-ingest collection discoverable (`list_namespaces` keys off +/// `{ns}/manifest`). `AlreadyExists` bubbles up — it means the namespace +/// already has data in the bucket (e.g. a pre-existing collection). +pub async fn init_namespace(storage: &dyn Storage, ns: &str) -> Result<(), StorageError> { + commit_manifest(storage, ns, &Manifest::default(), &None) + .await + .map(|_| ()) +} + /// Append a data WAL fragment. Returns the assigned sequence number. pub async fn append_fragment( storage: &dyn Storage, @@ -274,7 +292,7 @@ pub async fn read_uncompacted_fragments( /// MUST be readable. A missing fragment here is corruption — folding a subset /// and then advancing the watermark past the missing one would silently drop /// that batch. So we fail loudly instead of skipping. -async fn read_uncompacted_fragments_strict( +pub async fn read_uncompacted_fragments_strict( storage: &dyn Storage, ns: &str, manifest: &Manifest, @@ -296,97 +314,8 @@ async fn read_uncompacted_fragments_strict( Ok(out) } -/// Compact: fold all uncompacted WAL fragments into ONE new segment object via -/// `merge`, advance the watermark, and CAS-commit. `merge` receives the ordered -/// `(FragmentRef, payload)` list — so it can see each fragment's `kind` and -/// apply latest-wins + drop tombstoned records — and returns the segment bytes + -/// record count. -/// -/// Fixes vs. the earlier version: -/// - **F3**: the segment gets a UNIQUE id (UUID), so two concurrent compactions -/// never overwrite each other's segment object. -/// - **F2/F4**: deferred-delete GC runs only AFTER the manifest CAS succeeds, so -/// a losing retry never deletes objects the committed manifest still needs. -/// - Strict read: a missing fragment aborts (no silent data loss). -pub async fn compact(storage: &dyn Storage, ns: &str, merge: F) -> Result -where - // `Fn` (not `FnOnce`) because the CAS retry loop may call it more than once. - F: Fn(&[(FragmentRef, Bytes)]) -> Result<(Bytes, u64), StorageError>, -{ - let mut attempt = 0u32; - loop { - attempt += 1; - let (mut manifest, version) = read_manifest(storage, ns).await?; - - // Snapshot the previous cycle's deferred deletes; we only physically - // delete these AFTER our CAS commit succeeds (below). - let carried_deletes = manifest.pending_deletes.clone(); - - // Strict read: every listed uncompacted fragment must be present, so we - // only ever fold a complete set (never a subset with a hole). - let frags = read_uncompacted_fragments_strict(storage, ns, &manifest).await?; - if frags.is_empty() { - // Nothing to compact. Still commit if we have deletes to drain. - if carried_deletes.is_empty() { - return Ok(false); - } - manifest.pending_deletes.clear(); - match commit_manifest(storage, ns, &manifest, &version).await { - Ok(_) => { - gc_keys(storage, &carried_deletes).await; - return Ok(false); - } - Err(StorageError::VersionConflict { .. }) if attempt < MAX_CAS_RETRIES => continue, - Err(e) => return Err(e), - } - } - - let (segment_bytes, records) = merge(&frags)?; - // Unique segment id (F3): concurrent compactions can't clobber. - let segment_id = uuid::Uuid::new_v4().to_string(); - storage - .put(&segment_key(ns, &segment_id), segment_bytes) - .await?; - - // Safe: `frags` is the COMPLETE uncompacted set, so max seq covers - // exactly what we merged. - let new_watermark = frags.iter().map(|(fref, _)| fref.seq).max().unwrap_or(0); - - // The fragment objects we just compacted away — stage them for deletion - // NEXT cycle (keyed by unique id), so any in-flight reader still on the - // old manifest can read them for one more cycle. - let newly_staged: Vec = manifest - .fragments - .iter() - .filter(|f| f.seq <= new_watermark) - .map(|f| fragment_key(ns, &f.id)) - .collect(); - - manifest.compaction_watermark = Some(new_watermark); - manifest.fragments.retain(|f| f.seq > new_watermark); - manifest.segments.push(SegmentRef { - id: segment_id, - records, - }); - manifest.pending_deletes = newly_staged; - - match commit_manifest(storage, ns, &manifest, &version).await { - Ok(_) => { - // Commit succeeded: NOW physically delete the carried (previous - // cycle's) objects. The just-compacted fragments stay one cycle - // in pending_deletes so any in-flight reader on the old manifest - // can still read them. - gc_keys(storage, &carried_deletes).await; - return Ok(true); - } - Err(StorageError::VersionConflict { .. }) if attempt < MAX_CAS_RETRIES => continue, - Err(e) => return Err(e), - } - } -} - -/// Physically delete a set of object keys, best-effort (a transient failure is -/// logged; the key stays referenced only if it was still in pending_deletes). +/// Best-effort deferred GC: delete the prior cycle's staged objects. A failed +/// delete only leaks an orphan object (retried next cycle via pending_deletes). async fn gc_keys(storage: &dyn Storage, keys: &[String]) { for key in keys { if let Err(e) = storage.delete(key).await { @@ -437,6 +366,64 @@ pub async fn list_namespaces(storage: &dyn Storage) -> Result, Stora Ok(names) } +/// Partitioned compaction commit: APPEND a segment folding only the WAL tail +/// (fragments with seq <= `folded_through`) and advance the watermark. Old +/// segments stay; the folded fragments are staged for next-cycle GC and the +/// PRIOR cycle's staged keys are deleted now. On CAS conflict the just-written +/// orphan segment is removed. Bounded work: O(tail), never O(collection). +pub async fn append_segment( + storage: &dyn Storage, + ns: &str, + expected: &Option, + prior: &Manifest, + segment_bytes: Bytes, + records: u64, + folded_through: u64, +) -> Result<(), StorageError> { + let segment_id = uuid::Uuid::new_v4().to_string(); + let new_segment_key = segment_key(ns, &segment_id); + storage.put_large(&new_segment_key, segment_bytes).await?; + + let folded: Vec = prior + .fragments + .iter() + .filter(|f| f.seq <= folded_through) + .map(|f| fragment_key(ns, &f.id)) + .collect(); + let mut segments = prior.segments.clone(); + segments.push(SegmentRef { + id: segment_id, + records, + }); + let new_manifest = Manifest { + fragments: prior + .fragments + .iter() + .filter(|f| f.seq > folded_through) + .cloned() + .collect(), + segments, + next_seq: prior.next_seq, + compaction_watermark: Some( + prior + .compaction_watermark + .map(|w| w.max(folded_through)) + .unwrap_or(folded_through), + ), + pending_deletes: folded, + }; + match commit_manifest(storage, ns, &new_manifest, expected).await { + Ok(_) => { + gc_keys(storage, &prior.pending_deletes).await; + Ok(()) + } + Err(e) => { + let _ = storage.delete(&new_segment_key).await; + Err(e) + } + } +} + /// Full compaction: replace the ENTIRE manifest state (all segments + all /// uncompacted fragments) with a single new segment containing `segment_bytes` /// (the fully-materialized live set, deletes already applied). This is the @@ -467,7 +454,7 @@ pub async fn replace_with_single_segment( ) -> Result<(), StorageError> { let segment_id = uuid::Uuid::new_v4().to_string(); let new_segment_key = segment_key(ns, &segment_id); - storage.put(&new_segment_key, segment_bytes).await?; + storage.put_large(&new_segment_key, segment_bytes).await?; // Objects we're folding away THIS cycle (old segments + all fragments) — stage // for deletion NEXT cycle. @@ -504,13 +491,11 @@ pub async fn replace_with_single_segment( } } -/// Convenience to share a storage handle into the async helpers. -pub type SharedStorage = Arc; - #[cfg(test)] mod tests { use super::*; use crate::storage::local::LocalDiskStorage; + use std::sync::Arc; fn storage(name: &str) -> Arc { use std::sync::atomic::{AtomicU64, Ordering}; @@ -651,18 +636,28 @@ mod tests { .await .unwrap(); } - // Merge concatenates fragment payloads. - let did = compact(s.as_ref(), "ns", |frags| { - let mut out = Vec::new(); - for (_, b) in frags { - out.extend_from_slice(b); - } - let records = frags.len() as u64; - Ok((Bytes::from(out), records)) - }) + // Fold via the live primitives: strict tail read -> append_segment + // (what compact_storage does), concatenating fragment payloads. + let (m0, v0) = read_manifest(s.as_ref(), "ns").await.unwrap(); + let frags = read_uncompacted_fragments_strict(s.as_ref(), "ns", &m0) + .await + .unwrap(); + let mut out = Vec::new(); + for (_, b) in &frags { + out.extend_from_slice(b); + } + let folded_through = m0.fragments.iter().map(|f| f.seq).max().unwrap(); + append_segment( + s.as_ref(), + "ns", + &v0, + &m0, + Bytes::from(out), + frags.len() as u64, + folded_through, + ) .await .unwrap(); - assert!(did); let (m, _) = read_manifest(s.as_ref(), "ns").await.unwrap(); assert_eq!(m.segments.len(), 1); @@ -675,14 +670,12 @@ mod tests { let seg = read_segment(s.as_ref(), "ns", &seg_id).await.unwrap(); assert_eq!(&seg[..], b"012"); - // A subsequent compaction with nothing new is a no-op. - let did2 = compact(s.as_ref(), "ns", |frags| { - assert!(frags.is_empty()); - Ok((Bytes::new(), 0)) - }) - .await - .unwrap(); - assert!(!did2); + // Nothing new to fold: the strict tail read comes back empty. + let (m1, _) = read_manifest(s.as_ref(), "ns").await.unwrap(); + let tail = read_uncompacted_fragments_strict(s.as_ref(), "ns", &m1) + .await + .unwrap(); + assert!(tail.is_empty()); } #[tokio::test] @@ -691,11 +684,10 @@ mod tests { append_fragment(s.as_ref(), "ns", Bytes::from_static(b"old"), 1) .await .unwrap(); - compact(s.as_ref(), "ns", |frags| { - Ok((Bytes::from_static(b"seg"), frags.len() as u64)) - }) - .await - .unwrap(); + let (m0, v0) = read_manifest(s.as_ref(), "ns").await.unwrap(); + append_segment(s.as_ref(), "ns", &v0, &m0, Bytes::from_static(b"seg"), 1, 0) + .await + .unwrap(); let seq = append_fragment(s.as_ref(), "ns", Bytes::from_static(b"new"), 1) .await @@ -749,13 +741,13 @@ mod tests { let victim = format!("ns/wal/{}.frag", m0.fragments[0].id); s.delete(&victim).await.unwrap(); - let result = compact(s.as_ref(), "ns", |frags| { - Ok((Bytes::from_static(b"seg"), frags.len() as u64)) - }) - .await; + // The live compaction path reads the tail STRICTLY before folding; a + // missing fragment must error out (never silently skip lost data). + let (m1, _) = read_manifest(s.as_ref(), "ns").await.unwrap(); + let result = read_uncompacted_fragments_strict(s.as_ref(), "ns", &m1).await; assert!( result.is_err(), - "compaction must abort on a missing fragment" + "strict tail read must abort on a missing fragment" ); // Manifest is untouched: watermark did NOT advance, fragment 1 still live. diff --git a/crates/compass/src/storage/mod.rs b/crates/compass/src/storage/mod.rs index 2f0a666..0e5608b 100644 --- a/crates/compass/src/storage/mod.rs +++ b/crates/compass/src/storage/mod.rs @@ -19,6 +19,7 @@ //! Nothing routes through this trait yet — it is introduced standalone and //! wired into the engine incrementally in later steps. +pub mod id_alloc; pub mod local; pub mod lsm; #[cfg(feature = "object-storage")] @@ -52,6 +53,9 @@ impl Version { } /// True when the token carries no usable precondition (CAS must refuse it). + // Used by the object-store backend at runtime and by s3_integration + // tests — all behind the feature; default builds never reference it. + #[cfg(feature = "object-storage")] pub fn is_empty(&self) -> bool { self.e_tag.is_empty() && self.version.as_deref().is_none_or(str::is_empty) } @@ -59,6 +63,9 @@ impl Version { /// Metadata about a stored object, returned by `list`. #[derive(Debug, Clone)] +// size/version are part of the listing contract; current callers key off +// `key` only. Kept — deleting them would change every backend's list(). +#[allow(dead_code)] pub struct ObjectMeta { pub key: String, pub size: u64, @@ -103,6 +110,9 @@ pub trait Storage: Send + Sync { /// Range read — fetch only `range` bytes of the object. The primitive that /// makes large segments servable without loading the whole object. + // Range reads are the sectioned-segment read primitive (v2 TOC points at + // byte ranges); both backends implement it, callers land with Phase 5/6. + #[allow(dead_code)] async fn get_range(&self, key: &str, range: Range) -> Result; /// Read the object together with its current version, for a CAS cycle. @@ -162,6 +172,13 @@ pub trait Storage: Send + Sync { Ok(dirs) } + /// Large-object write. Default delegates to `put`; the object-store + /// backend overrides with multipart upload (S3 caps single PUTs at 5GB — + /// compacted segments can exceed that). + async fn put_large(&self, key: &str, bytes: Bytes) -> Result { + self.put(key, bytes).await + } + /// Whether an object exists. async fn exists(&self, key: &str) -> Result { match self.get_versioned(key).await { diff --git a/crates/compass/src/storage/object_store_backend.rs b/crates/compass/src/storage/object_store_backend.rs index 70c16ec..05c9a82 100644 --- a/crates/compass/src/storage/object_store_backend.rs +++ b/crates/compass/src/storage/object_store_backend.rs @@ -132,6 +132,7 @@ impl ObjectStoreBackend { } /// Construct directly from an existing object store (used by tests). + #[cfg(test)] pub fn from_store(inner: Arc, label: &'static str) -> Self { Self { inner, label } } @@ -261,6 +262,29 @@ impl Storage for ObjectStoreBackend { }) } + async fn put_large(&self, key: &str, bytes: Bytes) -> Result { + // Multipart for anything past a conservative threshold; small objects + // take the single-PUT fast path. + const PART: usize = 16 * 1024 * 1024; + if bytes.len() <= PART { + return self.put(key, bytes).await; + } + let path = OsPath::from(key); + let upload = self + .inner + .put_multipart(&path) + .await + .map_err(|e| map_os_err(key, e))?; + let mut w = object_store::WriteMultipart::new(upload); + for part in bytes.chunks(PART) { + w.write(part); + } + w.finish().await.map_err(|e| map_os_err(key, e))?; + // Multipart results don't return an ETag through this helper; segments + // are immutable + UUID-keyed, so no CAS token is needed on them. + Ok(Version::etag(String::new())) + } + async fn delete(&self, key: &str) -> Result<(), StorageError> { let path = OsPath::from(key); match self.inner.delete(&path).await { diff --git a/docs/scale-envelope.md b/docs/scale-envelope.md new file mode 100644 index 0000000..74e541e --- /dev/null +++ b/docs/scale-envelope.md @@ -0,0 +1,49 @@ +# Scale Envelope (measured, not claimed) + +Every number here comes from the env-gated harness in +`collections/mod.rs::scale_envelope`: + +```bash +COMPASS_SCALE_N=250000 cargo test -p compass --features object-storage \ + --release scale_envelope -- --nocapture +``` + +Runs use a local-disk Storage backend (identical code paths to S3, disk-bound) +inside a linux/amd64 container **under ARM emulation** — native x86 hardware +runs meaningfully faster; treat these as conservative floors. + +| chunks | dims | ingest | cold attach | search (semantic, avg) | +|---|---|---|---|---| +| 250,000 | 128 | 156s (1,603 chunks/s) | 134.5s | 5.5ms | +| 500,000 | 128 | 359s (1,392 chunks/s) | 369.3s | 11.2ms | + +## What the envelope means + +- **RAM is O(cache budget)** since chunks moved out-of-core (bounded LRU over + redb); segment format v2 + multipart removed the 5GB object ceiling; + partitioned compaction is O(batch) per cycle; per-write index costs are + O(batch). None of the previous hard walls bind below ~100M chunks. +- **Merge and attach still need O(live set) RAM on the node doing them** + (the periodic full merge clones the live set to encode the merged segment; + attach materializes it). Steady-state serving RAM is bounded; the + compacting/attaching moment is not — budget worker memory for your largest + collection, or shard. +- **The binding constraint is cold-attach time** (HNSW rebuild from the mmap + file — roughly linear in collection size). Lazy attach + LRU keep this a + first-request cost per namespace, not a boot cost, but a 100M-chunk + collection still takes tens of minutes to attach on first use. +- **Billion-vector serving therefore remains out of envelope** until + serve-from-storage indexes land (roadmap Phase 6: centroid routing over + range-readable segments — attach becomes "fetch centroids", milliseconds). + Do not deploy a single collection past ~10–50M chunks and expect + sub-minute cold attach. + +## Operating guidance + +- Shard very large corpora across collections (attach cost is per-collection). +- Watch `/metrics`: `compass_attach_seconds_sum_millis / compass_attach_total` + is your real attach cost; `compass_refresh_reattaches_total` climbing means + compaction is outrunning refresh (raise `COMPASS_REFRESH_INTERVAL` or lower + write bursts). +- Set `COMPASS_MAX_ATTACHED` on memory-constrained workers and + `COMPASS_MAX_CONCURRENCY` in front of bursty clients. diff --git a/docs/serverless-roadmap.md b/docs/serverless-roadmap.md new file mode 100644 index 0000000..d26d38f --- /dev/null +++ b/docs/serverless-roadmap.md @@ -0,0 +1,162 @@ +# Serverless Roadmap + +> Status: Phases 0-3 SHIPPED on feat/warm-serverless (v0.4.0 candidate); Phases 4+ planned. Target: evolve Compass from a cloud-durable single-node +> engine (v0.3.0) into a fully serverless database — storage/compute separated, +> stateless workers, bounded cold starts, scale-to-zero — with **every item +> additive and open source** under Apache 2.0. Local-first, zero-config +> operation remains the default at every step; all serverless behavior is +> opt-in via config. + +## Where v0.3.0 leaves us + +Done and hardened (the foundation): + +- Object storage as source of truth: LSM of immutable UUID-keyed WAL fragments, + CAS-committed manifest, compaction with deferred GC (`storage/lsm.rs`) +- Multi-writer safety at the storage layer, proven against real S3 semantics +- Full ephemeral recovery — chunks, hierarchy, typed relations + (`collections/cloud.rs::materialize`, `rebuild_collection_from_storage`) +- Prefix scoping, O(namespaces) discovery, compensation discipline on every + partial-failure path + +Not yet serverless: + +- Serving is stateful: every node rebuilds full local indexes (Tantivy + HNSW) +- Boot rebuilds ALL collections; `materialize` holds a collection's live set in RAM +- Writes require the stateful node that has the collection attached +- Nodes never re-read the manifest, so multi-node views diverge +- One global API key; no per-namespace scoping or metering + +## Design rules (apply to every phase) + +1. **Additive or it doesn't ship.** New behavior behind config/env/roles; + `cargo run` with no configuration behaves exactly as today. +2. **The bucket is the only shared state.** Coordination primitives are the + `Storage` trait's CAS/create-only operations — no new external dependencies + (no etcd, no Redis, no Postgres). +3. **Every phase lands with**: unit tests, env-gated real-S3 integration tests + (MinIO), an adversarial review round, CHANGELOG + docs. +4. **Pluggable seams stay public**: `Storage` (backends), `VectorIndex` + (`compass-index-api`), serving modes per collection. + +--- + +## Phase 0 — Guardrails (days) → part of v0.4 + +| Item | Detail | +|---|---| +| 0.1 CI covers the cloud build | Add an `object-storage`-feature test job to CI (today CI tests the default build only) and a MinIO service container running the `s3_integration` tests on every PR. | +| 0.2 DCO | Enforce Developer Certificate of Origin sign-off in CI before external contributors arrive; keeps future licensing options open without a CLA's friction. | +| 0.3 Public tracking | This document + a GitHub milestone per phase, issues per work item. | + +## Phase 1 — Stateless write path (~2–3 wks) → v0.4 + +Writes stop requiring a node that has the collection attached. + +| Item | Detail | +|---|---| +| 1.1 Collection config in the bucket | `create_collection` writes `{ns}/collection.json` (dims, vector spaces, default space) via create-only/CAS. Recovery reads it instead of inferring specs from recovered embeddings (also fixes the "model: recovered" inference in `rebuild_collection_from_storage`). | +| 1.2 Id-block allocation | Chunk ids are minted from a local `next_id` today. Stateless writers lease id blocks via CAS on a `{ns}/id-alloc` object (e.g. 10k-id blocks); a crashed writer leaks at most one block (ids are monotonic, gaps are already fine). | +| 1.3 Writer role | `COMPASS_ROLE=writer` (or per-request): validate against the bucket-cached collection config → append WAL fragment → CAS manifest → return `{seq}`. No local index update, no collections lock. Consistency contract: **durable immediately, searchable after reader refresh** (Phase 3) or compaction. | +| 1.4 Tests | A writer node with an empty disk ingests; a reader node serves it after refresh; MinIO integration. | + +## Phase 2 — Lazy attach + streaming materialize (~2–3 wks) → v0.4 + +Cold start stops being O(all data), RAM stops being O(collection). + +| Item | Detail | +|---|---| +| 2.1 Attach-on-demand | Boot registers namespaces (already O(ns) via `list_dirs`) without rebuilding. First request to a namespace triggers attach; an LRU with a configurable budget (`COMPASS_MAX_ATTACHED` / memory target) detaches idle collections (safe — the bucket is the source of truth; detach deletes local state). | +| 2.2 Streaming materialize | `materialize` gains a sink-based variant folding segments + WAL directly into redb / Tantivy writer / mmap appends in bounded batches — no full-collection HashMap. The in-RAM variant remains for compaction (which needs the full fold anyway until Phase 5). | +| 2.3 Observability | Attach-duration histograms; `/health` reports attached/registered counts. | + +**Acceptance:** boot time independent of collection count; attaching an +N-chunk collection runs at bounded RSS. + +## Phase 3 — Manifest watch + read consistency (~2 wks) → v0.4 + +Multiple readers converge on the same view; writers' output becomes visible. + +| Item | Detail | +|---|---| +| 3.1 Refresher | Per-attached-namespace background task re-reads the manifest (compare `Version` tokens; manifests are small). Interval configurable. | +| 3.2 Incremental replay | Apply only fragments with `seq >` last-applied to the local indexes — the existing ingest/delete/relation apply logic refactored into a reusable `apply_fragment` so refresh, attach, and ingest share one code path. | +| 3.3 Read-your-writes | Writes return the manifest `seq`; queries accept optional `min_seq` (fast-path refresh or bounded wait). | +| 3.4 Tests | Two managers on one bucket: write via A, visible via B within the interval; tombstones and relations replay correctly. | + +**Phase 1–3 outcome: “warm serverless.”** Any worker attaches any namespace +on demand; writes are stateless; readers converge. Cold start is bounded but +still proportional to index size (fixed in Phase 6). + +## Phase 4 — Compactor role + leases (~1–2 wks) → v0.5 + +| Item | Detail | +|---|---| +| 4.1 Lease primitive | `{ns}/lease/compactor` object via `put_if_not_exists` with a TTL payload; expired leases are stolen via CAS. Clock-skew caveat documented (leases are long relative to plausible skew; compaction is idempotent and CAS-guarded regardless — a double-run wastes work, never corrupts). | +| 4.2 Compactor role | `COMPASS_ROLE=compactor` (same binary): scan namespaces, threshold-check, lease, run the already-storage-only `compact_storage`, release. Serving nodes' inline auto-compaction turns off when an external compactor is configured. | + +## Phase 5 — Segment format v2 (~2–3 wks) → v0.5 + +The JSON segment becomes a binary, sectioned, range-readable format. + +| Item | Detail | +|---|---| +| 5.1 Layout | Magic + version + TOC (section → offset/len), sections: chunk metadata, text, embeddings per space (contiguous f32 LE rows), relations, the serialized filter-index treemaps (the persistence code exists, currently unwired), id high-water. Zstd per section. | +| 5.2 Back-compat | `decode_segment` already falls back by version; v1 JSON segments remain readable, compaction rewrites to v2. | +| 5.3 Range reads | Attach and query paths fetch only the sections they need via `get_range` (already a true byte-range read on both backends). | + +## Phase 6 — Serve directly from object storage (~2–4 mo) → v0.6 *(the innovation epic)* + +| Item | Detail | +|---|---| +| 6.1 Vector: per-segment IVF | Compaction runs k-means per segment; centroids live in the segment TOC (tiny, RAM-cacheable per namespace), posting lists are contiguous row ranges in the embeddings section. Query: route by centroids → range-read `nprobe` cells → exact-score → merge across segments + brute-force the (small by construction) WAL tail. Filters intersect posting row-ids with the segment's treemaps. Implemented as a `VectorIndex` (`compass-index-api`) impl, pluggable next to USearch HNSW. | +| 6.2 FTS over storage | Spike: tantivy custom `Directory` over the `Storage` trait with a local block cache. Decision gate after the spike; fallback design is per-segment mini-indexes built at compaction and fetched on attach. | +| 6.3 Serving modes | Per-collection `serving_mode: attached \| stateless` (default `attached` — today's behavior). Stateless mode never rebuilds local indexes. | + +**Acceptance:** recall@10 within an agreed delta of HNSW on standard +benchmarks; p95 latency targets on cold namespaces; RAM ceiling per attached +namespace measured and documented. + +## Phase 7 — Tenancy, metering, limits (~2–3 wks, parallel with 6) → v0.6 + +| Item | Detail | +|---|---| +| 7.1 Scoped keys | Per-collection API-key scopes extending `AuthConfig`; the single global key keeps working. | +| 7.2 Usage events | Per-request metering (namespace, operation, read/write bytes, query units) emitted as structured `tracing` events with an optional export sink. OSS emits; any billing pipeline (open or closed) aggregates. | +| 7.3 Quotas | Per-key/per-namespace rate limits and quotas via tower middleware, config-driven. | + +## Phase 8 — Open control plane (~4–6 wks) → v0.7 + +An OSS reference implementation of the service layer — same repo, new crate. + +| Item | Detail | +|---|---| +| 8.1 Router | `crates/compass-router` (or `COMPASS_ROLE=router`): rendezvous-hash namespaces → workers, worker registry via storage-backed heartbeat objects (no new dependencies), request proxying with attach-on-demand, drain/failover. | +| 8.2 Scale-to-zero | Idle detach (Phase 2's LRU) + pluggable worker-lifecycle hooks; ship Kubernetes manifests/HPA examples and a compose profile as reference deployments. | +| 8.3 Ops docs | Capacity planning, S3 request-cost model, tuning guide. | + +A hosted commercial offering (billing aggregation, org management, +dashboards) can be built on top of all of this later without forking — +every technical capability above stays in the open engine. + +--- + +## Sequencing + +``` +v0.4 Phase 0 ──► Phase 1 ──► Phase 2 ──► Phase 3 (~6-8 wks) "warm serverless" +v0.5 Phase 4 ──► Phase 5 (~3-5 wks) +v0.6 Phase 6 (epic) ∥ Phase 7 (~2-4 mo) stateless serving +v0.7 Phase 8 (~4-6 wks) open control plane +``` + +## Top risks + +| Risk | Mitigation | +|---|---| +| IVF recall/latency vs HNSW | Benchmark gate in Phase 6 acceptance; HNSW attached mode remains the default until parity data exists | +| tantivy-on-object-storage feasibility | Time-boxed spike with an explicit fallback (per-segment mini-indexes) | +| Id-block allocation contention | Blocks are large (10k) and leased rarely; CAS retry loop already proven on the manifest path | +| Lease correctness under clock skew | Long TTLs, idempotent CAS-guarded compaction — worst case is wasted work | +| S3 request costs in stateless mode | Centroid/block caching, section-level range reads, request-count metrics from day one (7.2) | +| JSON→v2 segment migration | Versioned decode already shipped in v0.3.0; compaction performs the migration organically | diff --git a/scripts/e2e.sh b/scripts/e2e.sh new file mode 100755 index 0000000..e671648 --- /dev/null +++ b/scripts/e2e.sh @@ -0,0 +1,127 @@ +#!/usr/bin/env bash +# Full-surface E2E against a LIVE stack: every endpoint, every filter operator, +# every role behavior. Usage: FULL=host:port WRITER=host:port ./scripts/e2e.sh +set -u +FULL=${FULL:-localhost:4001} +WRITER=${WRITER:-localhost:4009} +pass=0; fail=0 +ok(){ echo " ✅ $1"; pass=$((pass+1)); } +bad(){ echo " ❌ $1 ($2)"; fail=$((fail+1)); } +jqn(){ python3 -c "import sys,json;d=json.load(sys.stdin);print($1)" 2>/dev/null; } +post(){ curl -s -X POST "$1" -H 'content-type: application/json' -d "$2"; } + +echo "── health + metrics ──" +[ "$(curl -s $FULL/health | jqn "d['status']")" = "ok" ] && ok health || bad health x +curl -s $FULL/metrics | grep -q compass_search_requests_total && ok metrics || bad metrics x + +echo "── collections CRUD ──" +post $FULL/collections '{"name":"e2e","embedding_dims":4}' >/dev/null +[ "$(curl -s $FULL/collections/e2e | jqn "d['name']")" = "e2e" ] && ok "create+get" || bad create x +curl -s $FULL/collections | grep -q '"e2e"' && ok list || bad list x +code=$(curl -s -o /dev/null -w '%{http_code}' -X POST $FULL/collections -H 'content-type: application/json' -d '{"name":"e2e"}') +[ "$code" -ge 400 ] && ok "duplicate create rejected" || bad dup "$code" +code=$(curl -s -o /dev/null -w '%{http_code}' -X POST $FULL/collections -H 'content-type: application/json' -d '{"name":"Bad Name!"}') +[ "$code" -ge 400 ] && ok "invalid name rejected" || bad name "$code" + +echo "── ingest: hierarchy, client refs, metadata types ──" +r=$(post $FULL/collections/e2e/ingest '{"chunks":[ + {"client_id":"src1","file_id":"v1","chunk_index":0,"doc_type":"source","text":"Premier League match Arsenal Chelsea","metadata":{"kind":"video","priority":5,"active":true,"tags":["sports","football"],"created_at":"2026-07-01T00:00:00Z"},"embeddings":{"default":[0.9,0.1,0.1,0.1]}}, + {"client_id":"seg1","file_id":"s1","chunk_index":0,"doc_type":"segment","parent_ref":"src1","group_id":"src1","text":"goal celebration minute 34","metadata":{"timerange_start_ms":2040000,"timerange_end_ms":2055000,"priority":9},"embeddings":{"default":[0.1,0.9,0.1,0.1]}}, + {"client_id":"seg2","file_id":"s2","chunk_index":0,"doc_type":"segment","parent_ref":"src1","group_id":"src1","text":"halftime interview coach","metadata":{"timerange_start_ms":2700000,"timerange_end_ms":2760000,"priority":2},"embeddings":{"default":[0.1,0.1,0.9,0.1]}}]}') +n=$(echo "$r" | jqn "d['indexed']"); seq0=$(echo "$r" | jqn "d.get('seq')") +[ "$n" = "3" ] && ok "ingest 3 (hierarchy via parent_ref)" || bad ingest "$n" +[ "$seq0" != "None" ] && ok "ingest returns seq (cloud)" || bad seq x +id_src=$(echo "$r" | jqn "d['id_map']['src1']"); id_seg1=$(echo "$r" | jqn "d['id_map']['seg1']"); id_seg2=$(echo "$r" | jqn "d['id_map']['seg2']") +r=$(post $FULL/collections/e2e/ingest '{"chunks":[{"file_id":"legacy","chunk_index":0,"text":"legacy embedding field","embedding":[0.5,0.5,0.5,0.5]}]}') +[ "$(echo "$r" | jqn "d['indexed']")" = "1" ] && ok "legacy single-embedding field" || bad legacy x +code=$(curl -s -o /dev/null -w '%{http_code}' -X POST $FULL/collections/e2e/ingest -H 'content-type: application/json' -d '{"chunks":[{"file_id":"bad","chunk_index":0,"text":"x","embeddings":{"default":[0.1,0.2]}}]}') +[ "$code" -ge 400 ] && ok "wrong-dims embedding rejected" || bad dims "$code" + +echo "── search: modes, filters, scoring, explain ──" +n=$(post $FULL/collections/e2e/search '{"query":"goal celebration","mode":"fts","top_k":5}' | jqn "len(d['results'])") +[ "$n" -ge 1 ] && ok "fts" || bad fts "$n" +n=$(post $FULL/collections/e2e/search '{"query":"","mode":"semantic","query_vector":[0.1,0.9,0.1,0.1],"top_k":1}' | jqn "d['results'][0]['chunk']['file_id']") +[ "$n" = "s1" ] && ok "semantic nearest" || bad semantic "$n" +n=$(post $FULL/collections/e2e/search '{"query":"goal","mode":"hybrid","query_vector":[0.1,0.9,0.1,0.1],"top_k":5,"score_weights":{"rrf_k":60.0,"fts_weight":2.0,"semantic_weight":0.5}}' | jqn "len(d['results'])") +[ "$n" -ge 1 ] && ok "hybrid + score_weights" || bad hybrid "$n" +n=$(post $FULL/collections/e2e/search '{"query":"","mode":"semantic","query_vector":[0.5,0.5,0.5,0.5],"top_k":10,"filters":{"kind":"video"}}' | jqn "len(d['results'])") +[ "$n" = "1" ] && ok "filter: exact string" || bad f-eq "$n" +n=$(post $FULL/collections/e2e/search '{"query":"","mode":"semantic","query_vector":[0.5,0.5,0.5,0.5],"top_k":10,"filters":{"priority":{"gte":3,"lte":10}}}' | jqn "len(d['results'])") +[ "$n" = "2" ] && ok "filter: numeric range" || bad f-range "$n" +n=$(post $FULL/collections/e2e/search '{"query":"","mode":"semantic","query_vector":[0.5,0.5,0.5,0.5],"top_k":10,"filters":{"tags":{"contains":"sports"}}}' | jqn "len(d['results'])") +[ "$n" = "1" ] && ok "filter: array contains" || bad f-contains "$n" +n=$(post $FULL/collections/e2e/search '{"query":"","mode":"semantic","query_vector":[0.5,0.5,0.5,0.5],"top_k":10,"filters":{"doc_type":{"in":["segment"]}}}' | jqn "len(d['results'])") +[ "$n" = "2" ] && ok "filter: set membership (doc_type mirror)" || bad f-in "$n" +n=$(post $FULL/collections/e2e/search '{"query":"","mode":"semantic","query_vector":[0.5,0.5,0.5,0.5],"top_k":10,"filters":{"active":true}}' | jqn "len(d['results'])") +[ "$n" = "1" ] && ok "filter: bool" || bad f-bool "$n" +r=$(post $FULL/collections/e2e/search '{"query":"match","mode":"semantic","query_vector":[0.9,0.1,0.1,0.1],"top_k":5,"filters":{"kind":"video"},"explain":true}') +[ "$(echo "$r" | jqn "d['explain']['filter']['eligible_count']")" = "1" ] && ok "explain plan" || bad explain x +n=$(post $FULL/collections/e2e/search '{"query":"goal","mode":"fts","top_k":5,"recency_preset":"mild","recency_field":"created_at"}' | jqn "len(d['results'])") +[ "$n" -ge 1 ] && ok "recency preset" || bad recency "$n" +n=$(post $FULL/collections/e2e/search '{"query":"interview","mode":"fts","top_k":5,"boosts":[{"field":"priority","gte":3,"weight":2.0}],"relationship_boost":{"parent_weight":0.3,"sibling_weight":0.1}}' | jqn "len(d['results'])") +[ "$n" -ge 1 ] && ok "boosts + relationship_boost" || bad boosts "$n" + +echo "── relations ──" +r=$(post $FULL/collections/e2e/relations "{\"relations\":[{\"source_chunk_id\":$id_seg1,\"target_chunk_id\":$id_seg2,\"relation_type\":\"follows\"}]}") +rid=$(echo "$r" | jqn "d['relations'][0]['relation_id']") +[ -n "$rid" ] && ok "create relation" || bad rel x +n=$(curl -s "$FULL/collections/e2e/chunks/$id_seg1/relations?direction=outgoing&types=follows" | jqn "d['total']") +[ "$n" = "1" ] && ok "list relations (direction+type)" || bad rel-list "$n" +st=$(curl -s "$FULL/collections/e2e/chunks/$id_seg1/relations" | jqn "d['relations'][0]['target_status']") +[ "$st" = "found" ] && ok "target_status resolution" || bad status "$st" +n=$(post $FULL/collections/e2e/search '{"query":"goal","mode":"fts","top_k":5,"include_relations":true}' | jqn "d['results'][0].get('relations') is not None") +[ "$n" = "True" ] && ok "include_relations in search" || bad inc-rel "$n" +code=$(curl -s -o /dev/null -w '%{http_code}' -X DELETE $FULL/collections/e2e/relations/$rid) +[ "$code" = "204" ] && ok "delete relation" || bad rel-del "$code" + +echo "── facets + TAMS ──" +curl -s "$FULL/collections/e2e/facets" | grep -q "video" && ok facets || bad facets x +n=$(curl -s "$FULL/collections/e2e/segments/at?asset=src1&time_ms=2050000" | jqn "len(d.get('segments',d.get('results',[])))") +[ "$n" -ge 1 ] && ok "TAMS point lookup" || bad tams "$n" + +echo "── vector spaces ──" +post $FULL/collections/e2e/vector-spaces '{"name":"wide","dims":8,"model":"test"}' >/dev/null +curl -s $FULL/collections/e2e/vector-spaces | grep -q wide && ok "add+list space" || bad vs x +r=$(post $FULL/collections/e2e/ingest '{"chunks":[{"file_id":"w","chunk_index":0,"text":"wide vec","embeddings":{"default":[0.2,0.2,0.2,0.2],"wide":[0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1]}}]}') +[ "$(echo "$r" | jqn "d['indexed']")" = "1" ] && ok "multi-space ingest" || bad ms x +n=$(post $FULL/collections/e2e/search '{"query":"","mode":"semantic","vector_space":"wide","query_vector":[0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1],"top_k":1}' | jqn "len(d['results'])") +[ "$n" = "1" ] && ok "search named space" || bad ms-search "$n" +code=$(curl -s -o /dev/null -w '%{http_code}' -X PUT $FULL/collections/e2e/default-vector-space -H 'content-type: application/json' -d '{"name":"wide"}') +[ "$code" -lt 400 ] && ok "switch default space" || bad def "$code" +curl -s -X PUT $FULL/collections/e2e/default-vector-space -H 'content-type: application/json' -d '{"name":"default"}' >/dev/null +code=$(curl -s -o /dev/null -w '%{http_code}' -X DELETE $FULL/collections/e2e/vector-spaces/wide) +[ "$code" -lt 400 ] && ok "delete space" || bad vs-del "$code" + +echo "── writer role ──" +r=$(post $WRITER/collections/e2e/ingest '{"chunks":[{"file_id":"wchunk","chunk_index":0,"text":"from the stateless writer","embeddings":{"default":[0.7,0.7,0.1,0.1]}}]}') +wseq=$(echo "$r" | jqn "d['seq']") +[ "$wseq" != "None" ] && ok "writer ingest returns seq" || bad w-ingest x +n=$(post $FULL/collections/e2e/search "{\"query\":\"\",\"mode\":\"semantic\",\"query_vector\":[0.7,0.7,0.1,0.1],\"top_k\":1,\"min_seq\":$wseq}" | jqn "d['results'][0]['chunk']['file_id']") +[ "$n" = "wchunk" ] && ok "min_seq read-your-writes across nodes" || bad ryw "$n" +code=$(curl -s -o /dev/null -w '%{http_code}' -X POST $WRITER/collections/e2e/search -H 'content-type: application/json' -d '{"query":"x"}') +[ "$code" -ge 400 ] && ok "writer refuses reads" || bad w-read "$code" +code=$(curl -s -o /dev/null -w '%{http_code}' -X POST $WRITER/collections/ghost/delete -H 'content-type: application/json' -d '{"ids":[1]}') +[ "$code" -ge 400 ] && ok "writer refuses phantom namespace" || bad w-ghost "$code" + +echo "── deletes + compact ──" +wid=$(post $FULL/collections/e2e/search '{"query":"","mode":"semantic","query_vector":[0.7,0.7,0.1,0.1],"top_k":1}' | jqn "d['results'][0]['chunk']['id']") +r=$(curl -s -X DELETE $FULL/collections/e2e/chunks/$wid) +[ "$(echo "$r" | jqn "d['deleted']")" = "1" ] && ok "delete by id (+seq $(echo "$r" | jqn "d.get('seq')"))" || bad del x +r=$(post $FULL/collections/e2e/delete '{"filters":{"kind":"video"}}') +[ "$(echo "$r" | jqn "d['deleted']")" = "1" ] && ok "delete by filter" || bad del-f x +n=$(post $FULL/collections/e2e/search '{"query":"","mode":"semantic","query_vector":[0.9,0.1,0.1,0.1],"top_k":10}' | jqn "sum(1 for r in d['results'] if r['chunk']['file_id']=='v1')") +[ "$n" = "0" ] && ok "deleted chunk masked" || bad mask "$n" +r=$(post $FULL/collections/e2e/compact '') +[ -n "$(echo "$r" | jqn "d['compacted_records']")" ] && ok compact || bad compact x +n=$(post $FULL/collections/e2e/search '{"query":"goal","mode":"fts","top_k":5}' | jqn "len(d['results'])") +[ "$n" -ge 1 ] && ok "data survives compaction" || bad post-compact "$n" + +echo "── collection delete ──" +code=$(curl -s -o /dev/null -w '%{http_code}' -X DELETE $FULL/collections/e2e) +[ "$code" -lt 400 ] && ok "delete collection" || bad coll-del "$code" +code=$(curl -s -o /dev/null -w '%{http_code}' $FULL/collections/e2e) +[ "$code" = "404" ] || [ "$(curl -s $FULL/collections/e2e)" = "null" ] && ok "collection gone" || bad gone "$code" + +echo "" +echo "E2E RESULT: $pass passed, $fail failed" +[ "$fail" = "0" ]