diff --git a/.env.example b/.env.example index 8e0f345..364a10c 100644 --- a/.env.example +++ b/.env.example @@ -59,6 +59,42 @@ RUST_LOG=compass=info # AZURE_STORAGE_CONNECTION_STRING= # AZURE_STORAGE_SAS_KEY= -# ── Telemetry (anonymous; opt out) ────────────────────────────────────────── -# COMPASS_TELEMETRY=off -# DO_NOT_TRACK=1 +# ── 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 + +# ── Serve-from-storage (cold reads) ───────────────────────────────────────── +# Serve-from-storage: semantic queries on UNATTACHED collections are answered +# directly from object storage (a few range reads, ~100s of ms) instead of +# waiting for a full index rebuild. Implies COMPASS_LAZY_ATTACH. Cloud only. +# COMPASS_COLD_SERVE=true +# Cold hits on a namespace before a background attach warms it (0 = never). +# COMPASS_WARM_AFTER=3 +# IVF clusters probed per segment per cold query (recall/latency knob). +# Default 8 reaches warm-parity recall on clustered embedding spaces; raise +# it for unstructured vector data (see docs/search-quality.md). +# COMPASS_COLD_NPROBE=8 + +# Global in-flight request cap (backpressure). Unset = effectively unlimited. +# COMPASS_MAX_CONCURRENCY=1024 + +# ── Telemetry (anonymous; OPT-IN, off by default) ─────────────────────────── +# Compass never phones home unless you set this. When on, it sends a startup +# event + daily heartbeat (random instance id, version, OS/arch, collection +# and vector counts — never document content or queries). DO_NOT_TRACK=1 is +# honored even when opted in. +# COMPASS_TELEMETRY=on diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md deleted file mode 100644 index 5aae13c..0000000 --- a/.github/ISSUE_TEMPLATE/bug_report.md +++ /dev/null @@ -1,51 +0,0 @@ ---- -name: Bug Report -about: Report a bug or unexpected behavior -title: "[BUG] " -labels: bug -assignees: '' - ---- - -## Describe the bug - -A clear and concise description of what the bug is. - -## Steps to reproduce - -1. ... -2. ... -3. ... - -## Expected behavior - -What should happen. - -## Actual behavior - -What actually happens instead. - -## Environment - -- **Compass version** (or git SHA): -- **OS and architecture**: -- **Rust version** (`rustc --version`): -- **Data directory size**: -- **Collection size** (approximate): - -## Logs - -If applicable, reproduce with debug logging: -```bash -RUST_LOG=compass=debug ./compass -``` - -Then paste relevant log output here: - -``` -[paste logs] -``` - -## Additional context - -Any other context that might be helpful. diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml index f6dbf1f..91c1c2f 100644 --- a/.github/ISSUE_TEMPLATE/config.yml +++ b/.github/ISSUE_TEMPLATE/config.yml @@ -1,5 +1,5 @@ blank_issues_enabled: true contact_links: - name: Security vulnerability - url: mailto:security@runcaptain.com + url: mailto:support@runcaptain.com about: Report security issues privately via email — do not open a public issue. diff --git a/.github/ISSUE_TEMPLATE/feature_request.md b/.github/ISSUE_TEMPLATE/feature_request.md deleted file mode 100644 index 4cf58dc..0000000 --- a/.github/ISSUE_TEMPLATE/feature_request.md +++ /dev/null @@ -1,28 +0,0 @@ ---- -name: Feature Request -about: Suggest an idea for Compass -title: "[FEATURE] " -labels: enhancement -assignees: '' - ---- - -## Description - -A clear and concise description of what you'd like to see. - -## Motivation - -Why should this feature exist? What problem does it solve? - -## Proposed solution - -Describe how you'd like the feature to work. - -## Alternatives - -Have you considered any alternative approaches? - -## Additional context - -Any other context or examples (e.g., similar features in other projects). diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 779ea3a..f46465c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -4,7 +4,8 @@ on: push: branches: [main] pull_request: - branches: [main] + # feat/** so stacked PRs (feature targeting feature) get CI too. + branches: [main, "feat/**"] env: CARGO_TERM_COLOR: always @@ -57,6 +58,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/.github/workflows/release.yml b/.github/workflows/release.yml index 79e5f95..d8ff15f 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -27,7 +27,9 @@ jobs: - name: Verify tag matches Cargo.toml run: | - CARGO_VERSION=$(grep "^version" crates/compass/Cargo.toml | head -1 | sed 's/.*"\([^"]*\)".*/\1/') + # The version lives in [workspace.package] in the ROOT manifest; + # crate manifests say `version.workspace = true`. + CARGO_VERSION=$(grep "^version" Cargo.toml | head -1 | sed 's/.*"\([^"]*\)".*/\1/') if [ "${{ steps.version.outputs.version }}" != "$CARGO_VERSION" ]; then echo "Tag version ${{ steps.version.outputs.version }} does not match Cargo.toml version $CARGO_VERSION" exit 1 diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 00db435..bb54bed 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -43,23 +43,27 @@ crates/compass/src/ mod.rs SearchMode enum + re-exports. backend.rs VectorIndex trait shim. UsearchHnswIndex (CPU) lives here. vector.rs USearch HNSW build + search + persistence (CPU primitives). - tantivy_fts.rs Full-text search via Tantivy (BM25). + tantivy_fts.rs Full-text search via Tantivy (BM25) + facet treemaps. hybrid.rs Reciprocal Rank Fusion (RRF, k=60) over FTS + semantic. + ivf.rs IVF clustering built at compaction (cold-read layout). + cold.rs Serve-from-storage query path (range reads, no attach). + filter_index.rs Roaring-treemap metadata filter index (warm pushdown). + chunk_store.rs / chunk_cache.rs redb chunk store + bounded LRU cache. + collections/ + partitions.rs Tenant-partition routing helpers. + cloud.rs Segment codec (CSEG0003), materialize, bucket config. + storage/ Storage trait + local disk + object-store backends + LSM. + metrics.rs /metrics counters. telemetry.rs: opt-in usage pings. ``` -## Vector backend abstraction +## Vector backends -All vector backends implement `compass_index_api::VectorIndex`. The default backend is `UsearchHnswIndex` (CPU, mmap-backed, disk-persistent). The opt-in GPU backend is `compass_vector_gpu::CuvsHnswIndex` (CAGRA build on GPU, HNSW search on CPU). - -Selection happens at startup in `search::backend::build_backend`, driven by the `COMPASS_BACKEND` environment variable: - -| Value | Behavior | -|-------|----------| -| `cpu` (default) | USearch on CPU. Always available. | -| `gpu` | cuVS on GPU. Requires the `gpu` feature and a CUDA-capable device. Falls back to CPU with a warning if either is missing. | -| `auto` | Probe for GPU, fall back to CPU silently if unavailable. | - -The trait is intentionally narrow: `build`, `add`, `search`, `len`, `dims`, `save`, `backend_name`. New backends should fit through this surface or extend it via a follow-up trait, not by branching on a concrete type. +The engine uses USearch HNSW directly (CPU, mmap-backed, disk-persistent) +for warm serving, plus an IVF layout inside segments (`search/ivf.rs`) for +serve-from-storage cold reads. `compass-index-api` (a narrow `VectorIndex` +trait) and `compass-vector-gpu` (cuVS) exist as standalone crates for a +future GPU integration but are NOT wired into the engine — there is no +`COMPASS_BACKEND` knob and no `gpu` feature on the `compass` crate today. ## Storage layout @@ -67,18 +71,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}` (CSEG0003 sectioned +segments: row-addressable metadata + IVF-clustered vectors; v2 readable), 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) @@ -114,7 +122,8 @@ cuVS CAGRA build on an A10G runs ~12x faster than USearch CPU build at the same 1. Create a new crate `crates/compass-vector-/`. 2. Depend on `compass-index-api` (workspace dep) and your backend library. 3. Implement `VectorIndex` (and `LoadableIndex` if loading from disk makes sense). -4. Add a `#[cfg(feature = "")]`-gated branch in `search::backend::build_backend`. +4. Wire it into the engine (there is currently no runtime backend selector — + proposing that wiring is part of such a PR; open an issue first). 5. Document the build prerequisites in `ARCHITECTURE.md` (this file). 6. Add a smoke binary under `src/bin/` that builds, queries, and prints latency. diff --git a/CHANGELOG.md b/CHANGELOG.md index 83e591a..fe3a9ae 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,45 @@ 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). +## [0.4.0] - 2026-07-04 + +### Added — serve-from-storage ("true serverless") + +- **`COMPASS_COLD_SERVE=true`**: semantic queries on collections (and tenant partitions) that are NOT attached are answered directly from object storage — a manifest read, cached centroid/TOC artifacts, and a handful of range-GETs — instead of triggering a full index rebuild. Compaction now writes IVF-clustered vector sections (`cent:`/`clu:`, k-means, unit-normalized) plus a row-addressable metadata index (`meta2`/`metaidx`) into segments (format CSEG0003; v2 segments remain readable, pre-v0.4 readers fail loudly on v3). Cold reads see the full committed state including the WAL tail and tombstones, so read-your-writes holds by construction; metadata filters apply; FTS on a cold namespace returns a clear error (inverted indexes still need an attach). Repeated cold hits (`COMPASS_WARM_AFTER`, default 3) promote a background attach so hot namespaces migrate to the fast path on their own. RAM per cold namespace is megabytes (centroids + directories), independent of collection size. + +### Added — tenant-partitioned collections + +- **`config.partition_by`**: create a collection partitioned by a metadata field (e.g. `tenant_id`) and every chunk routes to an internal per-tenant partition — a full engine namespace (own LSM, indexes, attach/evict lifecycle) behind one collection API. Searches and deletes filter by the partition field (exact → one partition; `{"in": [...]}` fans out up to 16, merged by score); chunk ids are collection-unique via the parent's CAS id allocator; partitions auto-create on first ingest (writer role included), attach on demand, are hidden from listings, and cascade-delete with the parent. This moves the scale envelope from per-collection to per-tenant: RAM and refresh cost track the HOT tenant set, so one collection can hold billions of vectors across tenants while serving on bounded memory. Not yet routed on partitioned collections (clear errors): relations, facets, TAMS lookup, vector-space CRUD. + +### 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). + +### Changed (behavior) + +- **Telemetry is now opt-in** (`COMPASS_TELEMETRY=on`); previously it defaulted on. An engine whose promise is "data never leaves your machine" should not phone home by default. + +### Scope & limitations (honest) + +- Cold serving is semantic-only: full-text (and hybrid-with-text) queries on a cold namespace return a clear error until it warms — BM25 still needs local indexes. Cold recall depends on embedding-space structure (see docs/search-quality.md); the default nprobe reaches warm parity on clustered embeddings. Cross-node convergence is periodic (refresh interval), not synchronous — use `min_seq` when you need read-your-writes (cold reads have it by construction). + ## [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/CONTRIBUTING.md b/CONTRIBUTING.md index 253ef23..a8dc67f 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -15,15 +15,24 @@ cargo run --release # serves on http://localhost:4001 Compass is a Cargo workspace. Useful invocations: +Prerequisites on Linux: `cmake`, `pkg-config`, `libssl-dev` (what CI +installs). On Windows there is a known linker clash between `esaxx-rs` and +`cxx` — use `cargo check` locally and run builds/tests in Docker or WSL. + ```bash -cargo build # builds the default member (`compass`) -cargo build -p compass-index-api # builds just the trait crate -cargo build --features gpu # adds the GPU backend (Linux + CUDA only) -cargo test --workspace # runs all tests in all crates -cargo clippy --workspace -- -D warnings # lint check (CI requires zero warnings) -cargo fmt --all --check # format check (CI requires clean diff) +cargo build # default member (`compass`) +cargo test --workspace --exclude compass-vector-gpu # all tests CI runs +cargo test -p compass --features object-storage # + S3/GCS backend tests +cargo clippy --workspace --exclude compass-vector-gpu --all-targets -- -D warnings +cargo fmt --all --check # CI requires a clean diff ``` +`compass-vector-gpu` is a standalone experimental crate (cuVS; needs CUDA +12+, CMake, a long first build) that is NOT wired into the engine yet — +every CI job excludes it, and so should you unless you're working on it. +Tests against real object storage skip cleanly unless `COMPASS_TEST_S3_BUCKET` +is set (CI runs them against MinIO). + See [`ARCHITECTURE.md`](ARCHITECTURE.md) for the module map and where to put new code. ## What we accept @@ -35,9 +44,8 @@ See [`ARCHITECTURE.md`](ARCHITECTURE.md) for the module map and where to put new - Documentation improvements, especially examples. **Out of scope (for now):** -- Storage backends other than the local filesystem. - New embedding model integrations (we plug into HuggingFace TEI / vLLM via the `embed_endpoint` config; pull requests adding new in-process embedders need a strong motivation). -- Cluster / replication features (Compass is single-node by design; horizontal scaling is via sharding behind a load balancer). +- Consensus/quorum replication. Compass scales out via object storage as the source of truth (stateless writers + serving nodes + serve-from-storage cold reads); PRs should build on that model, not introduce node-to-node coordination. If you're not sure, open an issue first and ask. @@ -45,7 +53,7 @@ If you're not sure, open an issue first and ask. - [ ] `cargo fmt --all` clean. - [ ] `cargo clippy --workspace -- -D warnings` clean. -- [ ] `cargo test --workspace` green. +- [ ] `cargo test --workspace --exclude compass-vector-gpu` green. - [ ] `CHANGELOG.md` updated under the `[Unreleased]` section. - [ ] Public API changes have rustdoc comments. - [ ] Behavior changes have a test that would have caught the regression. @@ -75,7 +83,7 @@ Open an issue with: ## Reporting security issues -Don't open a public issue. Email `founders@runcaptain.com` with the details. We'll acknowledge within two business days. +Don't open a public issue. Email `support@runcaptain.com` with the details. We'll acknowledge within two business days. ## Code of conduct diff --git a/Cargo.lock b/Cargo.lock index 1a302a3..0485e67 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -463,7 +463,7 @@ dependencies = [ [[package]] name = "compass" -version = "0.3.0" +version = "0.4.0" dependencies = [ "async-trait", "axum", @@ -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", @@ -501,7 +499,7 @@ dependencies = [ [[package]] name = "compass-index-api" -version = "0.3.0" +version = "0.4.0" dependencies = [ "serde", "thiserror 1.0.69", @@ -509,7 +507,7 @@ dependencies = [ [[package]] name = "compass-vector-gpu" -version = "0.3.0" +version = "0.4.0" dependencies = [ "compass-index-api", "cuvs", @@ -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..8b76ec4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,9 +8,10 @@ members = [ default-members = ["crates/compass"] [workspace.package] -version = "0.3.0" +version = "0.4.0" edition = "2021" -authors = ["Captain Technologies "] +authors = ["Captain Technologies "] +license = "Apache-2.0" repository = "https://github.com/runcaptain/compass" homepage = "https://runcaptain.com" rust-version = "1.88" @@ -18,6 +19,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"] } @@ -75,7 +77,7 @@ tracing = "0.1" tracing-subscriber = "0.3" # Internal crates -compass-index-api = { path = "crates/compass-index-api", version = "0.3.0" } +compass-index-api = { path = "crates/compass-index-api", version = "0.4.0" } [profile.release] # `incremental = true` on release bloats the build cache for no runtime diff --git a/README.md b/README.md index 2a4f8ab..8d5041a 100644 --- a/README.md +++ b/README.md @@ -437,7 +437,18 @@ Supported: `s3://bucket[/prefix]` (AWS S3, MinIO, Cloudflare R2 — set `COMPASS In this mode the bucket is the **source of truth**: every write lands durably in object storage first (an LSM of immutable WAL fragments + a CAS-committed manifest), and a node that boots with an empty disk discovers its collections from the bucket and rebuilds all local indexes — chunks, hierarchy, and typed relations included. Compaction (automatic past a WAL threshold, or via `POST /compact`) folds fragments into segments and physically reclaims deleted data. -Scope, honestly: reads are served from the locally rebuilt indexes (durable-via-cloud, fast-via-local) — this is not stateless multi-node serving, and cold-start recovery materializes the live set in RAM. See [CHANGELOG](CHANGELOG.md) for details. +### Serverless topologies + +With the bucket as the source of truth, nodes become disposable roles you mix per workload (full reference: [docs/deployment.md](docs/deployment.md)): + +- **Serving node** (default): full local indexes, fast reads; converges on other nodes' writes via a background manifest refresher (`COMPASS_REFRESH_INTERVAL`, default 5s). Writes return a `seq`; pass it back as `min_seq` for read-your-writes. +- **Writer node** (`COMPASS_ROLE=writer`): stateless, append-only, boots in milliseconds, refuses reads. Durable immediately; searchable on serving nodes within the refresh interval. +- **Cold serving** (`COMPASS_COLD_SERVE=true`): answers *semantic* queries on collections it has never attached, straight from object-storage range reads — first query in ~tens of ms instead of a minutes-long index rebuild. Repeated hits promote a background attach (`COMPASS_WARM_AFTER`). Full-text on a cold collection returns a clear error until it warms. Recall characteristics: [docs/search-quality.md](docs/search-quality.md). +- **Lazy attach + LRU** (`COMPASS_LAZY_ATTACH`, `COMPASS_MAX_ATTACHED`): boot registers namespaces without loading them; RAM tracks the hot set. + +### Multi-tenant partitions + +Create a collection with `"config": {"partition_by": "tenant_id"}` and every chunk routes to an internal per-tenant partition — its own indexes and attach/evict lifecycle behind one collection API. Searches and deletes filter by the partition field (exact match, or `{"in": [...]}` to fan out across up to 16 tenants); chunk ids stay collection-unique; partitions auto-create on first ingest and cascade-delete with the parent. Serving RAM tracks the hot-tenant set, not the tenant count — this also works fully offline in local mode. For local development against MinIO: @@ -566,16 +577,23 @@ 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 (unauthenticated by design, like /health — exposes collection names + counts; firewall it if that matters) ``` ## Contributing See [CONTRIBUTING.md](CONTRIBUTING.md) for development setup, PR guidelines, and commit conventions. +## Telemetry + +Telemetry is **off by default** — Compass never phones home unless you set `COMPASS_TELEMETRY=on`. When opted in, it sends a startup event and a daily heartbeat to PostHog (random instance id, version, OS/arch, collection and vector counts — never document content, queries, or metadata). `DO_NOT_TRACK=1` is honored even when opted in. + ## Security -To report a vulnerability, email **security@runcaptain.com**. See [SECURITY.md](SECURITY.md) for details. +To report a vulnerability, email **support@runcaptain.com**. See [SECURITY.md](SECURITY.md) for details. ## License diff --git a/SECURITY.md b/SECURITY.md index 673ef1c..3510741 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -2,7 +2,7 @@ ## Reporting a Vulnerability -We take the security of Compass seriously. If you discover a security vulnerability, please email founders@runcaptain.com with the following information: +We take the security of Compass seriously. If you discover a security vulnerability, please email support@runcaptain.com with the following information: 1. **Description** of the vulnerability 2. **Steps to reproduce** (if applicable) @@ -28,7 +28,7 @@ We will acknowledge your report within **two business days** and work with you t ### Model Weights -- Compass downloads model weights on first run (e.g., BGE-small via Hugging Face Hub). +Compass never downloads anything at runtime. Model weights are fetched only if you run `scripts/download-models.sh` (or `huggingface-cli`) yourself. - Verify downloaded files match expected checksums when possible. - For air-gapped deployments, pre-download and verify model weights before use. diff --git a/crates/compass-index-api/Cargo.toml b/crates/compass-index-api/Cargo.toml index 89f299d..8afefa0 100644 --- a/crates/compass-index-api/Cargo.toml +++ b/crates/compass-index-api/Cargo.toml @@ -3,6 +3,7 @@ name = "compass-index-api" version.workspace = true edition.workspace = true authors.workspace = true +license.workspace = true repository.workspace = true homepage.workspace = true rust-version.workspace = true diff --git a/crates/compass-vector-gpu/Cargo.toml b/crates/compass-vector-gpu/Cargo.toml index 3ab0b38..5b3b5d1 100644 --- a/crates/compass-vector-gpu/Cargo.toml +++ b/crates/compass-vector-gpu/Cargo.toml @@ -3,6 +3,7 @@ name = "compass-vector-gpu" version.workspace = true edition.workspace = true authors.workspace = true +license.workspace = true repository.workspace = true homepage.workspace = true rust-version.workspace = true diff --git a/crates/compass/Cargo.toml b/crates/compass/Cargo.toml index 19379ba..2e89457 100644 --- a/crates/compass/Cargo.toml +++ b/crates/compass/Cargo.toml @@ -3,6 +3,7 @@ name = "compass" version.workspace = true edition.workspace = true authors.workspace = true +license.workspace = true repository.workspace = true homepage.workspace = true rust-version.workspace = true @@ -14,26 +15,23 @@ 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"] +# Opt-in object-storage backend (S3/GCS/Azure/MinIO) via the object_store +# crate. Off by default; local-first builds skip the object_store dependency +# tree (`futures` is unconditional — the cold read path uses it). +object-storage = ["dep:object_store"] [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 } @@ -53,10 +51,8 @@ thiserror = { workspace = true } lru = { workspace = true } # Object-storage backend deps (optional, enabled by the `object-storage` feature). object_store = { workspace = true, optional = true } -futures = { workspace = true, optional = true } +futures = { workspace = true } 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..7c45a40 100644 --- a/crates/compass/src/api/mod.rs +++ b/crates/compass/src/api/mod.rs @@ -4,7 +4,9 @@ // vector space CRUD, rebuild triggers, and status checks. // // Bearer-token auth middleware is applied to all routes -// except /health. See `AuthConfig` and `auth_middleware` below. +// except /health and /metrics (both unauthenticated by design; /metrics +// exposes collection names + counts — firewall it if that matters). +// See `AuthConfig` and `auth_middleware` below. pub mod collections; pub mod delete; @@ -25,6 +27,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 +167,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..bfa8191 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,231 @@ 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, } -const SEGMENT_VERSION: u8 = 1; +/// v2 binary segment magic. v1 segments are JSON (decoded via fallback). +pub(crate) const SEG_MAGIC_V2: [u8; 8] = *b"CSEG0002"; +/// v3 adds serve-from-storage sections: row-addressable chunk metadata +/// (`meta2` + `metaidx`) and IVF-clustered vectors (`cent:`/`clu:` replace +/// `emb:` for spaces past the clustering threshold). v3 readers decode v2; +/// v2 readers FAIL LOUDLY on v3 (magic mismatch) rather than silently +/// dropping sections — do not mix pre-v0.5 readers with v0.5 writers. +pub(crate) const SEG_MAGIC_V3: [u8; 8] = *b"CSEG0003"; -/// 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). +/// Encode a segment in the v3 sectioned binary layout: +/// `[magic][u64 max_id][u32 toc_len][toc JSON][sections...]` +/// Sections: +/// `meta2` — concatenated per-chunk JSON rows (embeddings stripped); +/// each row independently parseable so cold reads can fetch +/// single chunks by byte range +/// `metaidx` — `[u64 n][n × (u64 id, u64 off, u32 len)]` sorted by id, +/// offsets into `meta2` +/// `emb:` — flat `[u32 dims][u64 n][n × (u64 id + dims×f32)]`, +/// only for spaces below the clustering threshold +/// `cent:` / `clu:` — IVF centroids + clustered vectors +/// (see search/ivf.rs) for spaces at/above the threshold; +/// vectors are stored L2-normalized +/// `rels` (JSON), `tombs` (u64 LE array), `rtombs` (JSON ids) +pub fn encode_segment_v3(seg: &Segment) -> Result, StorageError> { + let err = |e: String| StorageError::Io(format!("segment v3 encode: {e}")); + let mut sections: Vec<(String, Vec)> = Vec::new(); + + // Row-addressable metadata + index (sorted by id for range lookups). + let mut sorted: Vec<&DocumentChunk> = seg.chunks.iter().collect(); + sorted.sort_by_key(|c| c.id); + let mut meta2 = Vec::new(); + let mut metaidx = Vec::with_capacity(8 + sorted.len() * 20); + metaidx.extend_from_slice(&(sorted.len() as u64).to_le_bytes()); + let mut by_space: std::collections::BTreeMap)>> = + std::collections::BTreeMap::new(); + for c in sorted { + 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)); + } + let row = serde_json::to_vec(&m).map_err(|e| err(e.to_string()))?; + metaidx.extend_from_slice(&c.id.to_le_bytes()); + metaidx.extend_from_slice(&(meta2.len() as u64).to_le_bytes()); + metaidx.extend_from_slice(&(row.len() as u32).to_le_bytes()); + meta2.extend_from_slice(&row); + } + sections.push(("meta2".into(), meta2)); + sections.push(("metaidx".into(), metaidx)); + + for (space, rows) in by_space { + let dims = rows.first().map(|(_, v)| v.len()).unwrap_or(0); + for (_, v) in &rows { + if v.len() != dims { + return Err(err(format!("ragged dims in space '{space}'"))); + } + } + if rows.len() >= crate::search::ivf::CLUSTER_MIN_ROWS && dims > 0 { + let (cent, clu) = crate::search::ivf::build_sections(rows, dims); + sections.push((format!("cent:{space}"), cent)); + sections.push((format!("clu:{space}"), clu)); + } else { + let mut buf = Vec::with_capacity(12 + rows.len() * (8 + dims * 4)); + buf.extend_from_slice(&(dims as u32).to_le_bytes()); + buf.extend_from_slice(&(rows.len() as u64).to_le_bytes()); + for (id, v) in &rows { + 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_V3); + 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) +} + +fn decode_segment_sectioned(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(); + let mut cent_dims: 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 name == "meta2" { + // v3 row-addressable metadata: concatenated standalone JSON rows. + let mut de = serde_json::Deserializer::from_slice(body).into_iter::(); + for c in de.by_ref() { + seg.chunks.push(c.map_err(|e| err(e.to_string()))?); + } + } else if name == "metaidx" { + // Full decode doesn't need the index (meta2 rows stream in order); + // it exists for cold range reads. + } else if name.starts_with("cent:") { + let c = crate::search::ivf::parse_cent(body) + .ok_or_else(|| err(format!("bad {name} section")))?; + cent_dims.insert(name.clone(), c.dims); + } else if let Some(space) = name.strip_prefix("clu:") { + let cent_key = format!("cent:{space}"); + let dims = cent_dims + .get(¢_key) + .copied() + .ok_or_else(|| err(format!("clu:{space} without preceding cent section")))?; + for (id, v) in crate::search::ivf::parse_cluster_rows(body, dims) { + embs.entry(id).or_default().insert(space.to_string(), v); + } + } 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 (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_v3(&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 || bytes[0..8] == SEG_MAGIC_V3) { + return decode_segment_sectioned(bytes); + } if let Ok(seg) = serde_json::from_slice::(bytes) { return Ok(seg); } @@ -81,8 +386,7 @@ fn decode_segment(bytes: &[u8]) -> Result { Ok(Segment { version: 0, chunks, - relations: Vec::new(), - max_id: 0, + ..Default::default() }) } @@ -103,6 +407,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 +478,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 +710,133 @@ 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_v3(&seg).unwrap(); + assert_eq!(&bytes[0..8], b"CSEG0003"); + 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); + } + + // Past the clustering threshold the encoder emits cent:/clu: instead of + // emb:; the full decode must reconstruct every chunk's (normalized) + // embedding from the clustered layout. + #[test] + fn segment_v3_clustered_roundtrip() { + let n = crate::search::ivf::CLUSTER_MIN_ROWS + 100; + let chunks: Vec = (0..n as u64) + .map(|i| { + let mut c = chunk(i, &format!("t{i}")); + let v: Vec = (0..8).map(|d| ((i + d) % 13) as f32 + 1.0).collect(); + c.embeddings.insert("default".into(), v); + c + }) + .collect(); + let seg = Segment { + version: 2, + chunks, + relations: vec![], + max_id: n as u64, + tombstones: vec![], + relation_tombstones: vec![], + }; + let bytes = encode_segment_v3(&seg).unwrap(); + let toc_len = u32::from_le_bytes(bytes[16..20].try_into().unwrap()) as usize; + let toc = std::str::from_utf8(&bytes[20..20 + toc_len]).unwrap(); + assert!(toc.contains("cent:default"), "toc: {toc}"); + assert!(toc.contains("clu:default"), "toc: {toc}"); + assert!(!toc.contains("emb:default"), "toc: {toc}"); + + let back = decode_segment(&bytes).unwrap(); + assert_eq!(back.chunks.len(), n); + for c in &back.chunks { + let v = &c.embeddings["default"]; + assert_eq!(v.len(), 8); + let norm: f32 = v.iter().map(|x| x * x).sum::().sqrt(); + assert!((norm - 1.0).abs() < 1e-4, "clu vectors are unit-norm"); + } + } + + // 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_v3(&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/cold_serve_tests.rs b/crates/compass/src/collections/cold_serve_tests.rs new file mode 100644 index 0000000..0faedbb --- /dev/null +++ b/crates/compass/src/collections/cold_serve_tests.rs @@ -0,0 +1,404 @@ +// collections/cold_serve_tests.rs — Phase 5 serve-from-storage coverage. +// A node with cold serving on answers semantic queries on namespaces it has +// NEVER attached, straight from object-storage range reads. + +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-cold-{}-{}", + std::process::id(), + N.fetch_add(1, Ordering::Relaxed) + )) +} + +fn embed_state() -> EmbedState { + EmbedState { + bge: None, + distilled: None, + } +} + +fn mem_storage() -> Arc { + Arc::new(ObjectStoreBackend::from_store( + std::sync::Arc::new(object_store::memory::InMemory::new()), + "object-store:memory", + )) +} + +const DIMS: usize = 8; + +/// Deterministic embedding: direction depends on i % 4, plus a nudge that is +/// UNIQUE per id (ties would make self-recall assertions ambiguous). +fn vec_for(i: u64) -> Vec { + let mut v = vec![0.05f32; DIMS]; + v[(i % 4) as usize * 2] = 1.0; + v[7] = i as f32 * 1e-4; + v +} + +fn mk_chunk(i: u64, kind: &str) -> IngestChunk { + let mut metadata = HashMap::new(); + metadata.insert("kind".to_string(), MetadataValue::String(kind.to_string())); + let mut embeddings = HashMap::new(); + embeddings.insert("default".to_string(), vec_for(i)); + IngestChunk { + client_id: Some(format!("c{i}")), + file_id: format!("f{i}"), + chunk_index: 0, + page: None, + text: format!("cold document number {i}"), + metadata, + doc_type: "chunk".to_string(), + parent_id: None, + parent_ref: None, + group_id: None, + embeddings, + embedding: None, + } +} + +fn semantic_req(target: u64, top_k: usize) -> SearchRequest { + SearchRequest { + query: String::new(), + mode: "semantic".to_string(), + vector_space: None, + top_k, + query_vector: Some(vec_for(target)), + 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::default(), + min_seq: None, + } +} + +async fn cold_manager(dir: &std::path::Path, storage: Arc) -> Arc { + std::fs::create_dir_all(dir).unwrap(); + let m = CollectionManager::new_with_storage_opts(dir, storage, NodeRole::Full, true, 0, 0) + .await + .unwrap(); + m.set_cold_serve(true); + m.set_warm_after(0); // promotion off unless a test opts in + m +} + +// Core promise: a fresh node answers semantic queries on a compacted (v3, +// clustered) namespace WITHOUT attaching it — and sees WAL-tail writes and +// tombstones that landed after compaction (read-your-writes from storage). +#[tokio::test] +async fn cold_search_serves_without_attach() { + let storage = mem_storage(); + let embed = embed_state(); + + // Writer side: ingest past the clustering threshold, then compact. + let n = (crate::search::ivf::CLUSTER_MIN_ROWS + 200) as u64; + { + let dir = unique_data_dir(); + std::fs::create_dir_all(&dir).unwrap(); + let a = CollectionManager::new_with_storage(&dir, storage.clone()) + .await + .unwrap(); + a.create_collection("frozen", None, Some(DIMS), None) + .await + .unwrap(); + let mut batch = Vec::new(); + for i in 0..n { + batch.push(mk_chunk(i, if i % 2 == 0 { "even" } else { "odd" })); + if batch.len() == 1000 { + a.ingest("frozen", std::mem::take(&mut batch), &embed) + .await + .unwrap(); + } + } + if !batch.is_empty() { + a.ingest("frozen", batch, &embed).await.unwrap(); + } + let live = a.compact_collection("frozen").await.unwrap(); + assert_eq!(live, n); + // Post-compaction writes + a delete stay in the WAL tail. + a.ingest("frozen", vec![mk_chunk(n, "tail")], &embed) + .await + .unwrap(); + a.delete_chunks("frozen", &[2]).await.unwrap(); + let _ = std::fs::remove_dir_all(&dir); + } + + // Cold node: never attaches, still answers. + let dir_b = unique_data_dir(); + let b = cold_manager(&dir_b, storage.clone()).await; + + // Exact-vector self-recall: the target chunk must be the top hit. + let target = 1234u64; + let (results, _, _, _) = b + .search("frozen", &semantic_req(target, 5), &embed) + .await + .unwrap(); + assert!(!results.is_empty(), "cold search returned nothing"); + assert_eq!( + results[0].0.file_id, + format!("f{target}"), + "self-recall: exact stored vector must rank first" + ); + assert_eq!(results[0].2, "semantic-cold"); + + // The namespace is still NOT attached (that's the whole point). + assert!( + !b.collections.read().await.contains_key("frozen"), + "cold search must not attach" + ); + + // Tail write is visible; tombstoned chunk is not. + let (results, _, _, _) = b + .search("frozen", &semantic_req(n, 5), &embed) + .await + .unwrap(); + assert!( + results.iter().any(|r| r.0.file_id == format!("f{n}")), + "post-compaction tail write must be cold-visible" + ); + let (results, _, _, _) = b + .search("frozen", &semantic_req(2, 20), &embed) + .await + .unwrap(); + assert!( + results.iter().all(|r| r.0.file_id != "f2"), + "tombstoned chunk leaked into cold results" + ); + + // Metadata filters apply cold. + let mut req = semantic_req(target, 10); + req.filters.insert( + "kind".to_string(), + FilterValue::Exact(MetadataValue::String("even".to_string())), + ); + let (results, _, _, _) = b.search("frozen", &req, &embed).await.unwrap(); + assert!(!results.is_empty()); + for r in &results { + assert_eq!( + r.0.metadata.get("kind"), + Some(&MetadataValue::String("even".to_string())) + ); + } + + // FTS stays honest: clear error, not silent emptiness. + let mut req = semantic_req(0, 5); + req.mode = "fts".to_string(); + req.query = "cold".to_string(); + req.query_vector = None; + let err = b.search("frozen", &req, &embed).await.unwrap_err(); + assert!(err.to_string().contains("cold"), "{err}"); + + let _ = std::fs::remove_dir_all(&dir_b); +} + +// Cold serving composes with tenant partitions: a partition namespace is +// cold-served through the same router, tenant isolation intact. +#[tokio::test] +async fn cold_search_composes_with_partitions() { + let storage = mem_storage(); + let embed = embed_state(); + { + let dir = unique_data_dir(); + std::fs::create_dir_all(&dir).unwrap(); + let a = CollectionManager::new_with_storage(&dir, storage.clone()) + .await + .unwrap(); + a.create_collection( + "mt", + None, + Some(DIMS), + Some(CollectionConfig { + embed_model: "bge-small".to_string(), + partition_by: Some("tenant".to_string()), + }), + ) + .await + .unwrap(); + let mut chunks = Vec::new(); + for i in 0..40u64 { + let mut c = mk_chunk(i, "x"); + c.metadata.insert( + "tenant".to_string(), + MetadataValue::String(if i % 2 == 0 { "acme" } else { "globex" }.to_string()), + ); + chunks.push(c); + } + a.ingest("mt", chunks, &embed).await.unwrap(); + let _ = std::fs::remove_dir_all(&dir); + } + + let dir_b = unique_data_dir(); + let b = cold_manager(&dir_b, storage.clone()).await; + let mut req = semantic_req(4, 10); // id 4 is acme (even) + req.filters.insert( + "tenant".to_string(), + FilterValue::Exact(MetadataValue::String("acme".to_string())), + ); + let (results, _, _, _) = b.search("mt", &req, &embed).await.unwrap(); + assert!(!results.is_empty()); + for r in &results { + assert_eq!( + r.0.metadata.get("tenant"), + Some(&MetadataValue::String("acme".to_string())), + "tenant isolation must hold on the cold path" + ); + } + assert!( + !b.collections.read().await.contains_key("mt--part--acme"), + "partition must be cold-served, not attached" + ); + let _ = std::fs::remove_dir_all(&dir_b); +} + +// Repeated cold hits promote a background attach; once attached, queries +// take the hot path. +#[tokio::test] +async fn cold_hits_promote_background_attach() { + let storage = mem_storage(); + let embed = embed_state(); + { + let dir = unique_data_dir(); + std::fs::create_dir_all(&dir).unwrap(); + let a = CollectionManager::new_with_storage(&dir, storage.clone()) + .await + .unwrap(); + a.create_collection("warmup", None, Some(DIMS), None) + .await + .unwrap(); + a.ingest( + "warmup", + (0..20u64).map(|i| mk_chunk(i, "x")).collect(), + &embed, + ) + .await + .unwrap(); + let _ = std::fs::remove_dir_all(&dir); + } + + let dir_b = unique_data_dir(); + let b = cold_manager(&dir_b, storage.clone()).await; + b.set_warm_after(2); + for _ in 0..2 { + let (results, _, _, _) = b + .search("warmup", &semantic_req(3, 3), &embed) + .await + .unwrap(); + assert!(!results.is_empty()); + } + // The promotion attach runs in the background; poll briefly. + let mut attached = false; + for _ in 0..100 { + if b.collections.read().await.contains_key("warmup") { + attached = true; + break; + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + assert!(attached, "warm promotion never attached the namespace"); + // Post-promotion searches run the hot path (engine != semantic-cold). + let (results, _, _, _) = b + .search("warmup", &semantic_req(3, 3), &embed) + .await + .unwrap(); + assert!(!results.is_empty()); + assert_ne!(results[0].2, "semantic-cold"); + let _ = std::fs::remove_dir_all(&dir_b); +} + +// H1 regression: a LAZY (or cold-serve) node must ingest into a partitioned +// collection for a brand-new tenant WITHOUT the parent ever being attached — +// the partition template comes from the bucket config. +#[tokio::test] +async fn lazy_node_ingests_new_tenant_without_attaching_parent() { + let storage = mem_storage(); + let embed = embed_state(); + { + let dir = unique_data_dir(); + std::fs::create_dir_all(&dir).unwrap(); + let a = CollectionManager::new_with_storage(&dir, storage.clone()) + .await + .unwrap(); + a.create_collection( + "lz", + None, + Some(DIMS), + Some(CollectionConfig { + embed_model: "bge-small".to_string(), + partition_by: Some("tenant".to_string()), + }), + ) + .await + .unwrap(); + let _ = std::fs::remove_dir_all(&dir); + } + + let dir_b = unique_data_dir(); + let b = cold_manager(&dir_b, storage.clone()).await; // lazy + cold serve + let mut c = mk_chunk(0, "x"); + c.metadata.insert( + "tenant".to_string(), + MetadataValue::String("fresh-tenant".to_string()), + ); + let (n, _, _) = b + .ingest("lz", vec![c], &embed) + .await + .expect("lazy node must route partitioned ingest from bucket config"); + assert_eq!(n, 1); + assert!( + !b.collections.read().await.contains_key("lz"), + "the parent must not have been attached to serve the ingest" + ); + // And the data is queryable through the partition filter. + let mut req = semantic_req(0, 3); + req.filters.insert( + "tenant".to_string(), + FilterValue::Exact(MetadataValue::String("fresh-tenant".to_string())), + ); + let (results, _, _, _) = b.search("lz", &req, &embed).await.unwrap(); + assert_eq!(results.len(), 1); + let _ = std::fs::remove_dir_all(&dir_b); +} + +// Live-stack repro: the cold node boots BEFORE the collection exists (empty +// registry), another node creates + ingests, cold must still answer. +#[tokio::test] +async fn cold_serves_collection_created_after_boot() { + let storage = mem_storage(); + let embed = embed_state(); + let dir_b = unique_data_dir(); + let b = cold_manager(&dir_b, storage.clone()).await; // boots on empty bucket + + let dir_a = unique_data_dir(); + std::fs::create_dir_all(&dir_a).unwrap(); + let a = CollectionManager::new_with_storage(&dir_a, storage.clone()) + .await + .unwrap(); + a.create_collection("late", None, Some(DIMS), None) + .await + .unwrap(); + a.ingest("late", vec![mk_chunk(1, "x")], &embed) + .await + .unwrap(); + + let (results, _, _, _) = b + .search("late", &semantic_req(1, 3), &embed) + .await + .expect("cold node must serve a collection created after its boot"); + assert_eq!(results.len(), 1); + assert_eq!(results[0].2, "semantic-cold"); + let _ = std::fs::remove_dir_all(&dir_a); + let _ = std::fs::remove_dir_all(&dir_b); +} 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..a1feb36 100644 --- a/crates/compass/src/collections/mod.rs +++ b/crates/compass/src/collections/mod.rs @@ -10,6 +10,13 @@ // search with full scoring pipeline, vector space CRUD, background rebuild jobs. pub mod cloud; +#[cfg(all(test, feature = "object-storage"))] +mod cold_serve_tests; +#[cfg(all(test, feature = "object-storage"))] +mod partition_cloud_tests; +#[cfg(test)] +mod partition_tests; +pub mod partitions; pub mod rebuild; pub mod relation_store; pub mod relationships; @@ -18,6 +25,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 +63,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 +143,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 +164,64 @@ 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, + /// Serve-from-storage (COMPASS_COLD_SERVE): semantic queries on + /// UNATTACHED namespaces are answered with object-storage range reads + /// instead of triggering an attach. Cloud + lazy mode only. + cold_serve: std::sync::atomic::AtomicBool, + /// Cold queries on a namespace before a background attach is kicked off + /// (COMPASS_WARM_AFTER; 0 = never warm automatically). + warm_after: std::sync::atomic::AtomicU32, + /// Cold-hit counters per namespace (drives warm promotion). + cold_hits: std::sync::Mutex>, + /// Cached cold-read artifacts, keyed by "{ns}/{segment_id}". Segments are + /// immutable, so entries never go stale; bounded by simple clearing. + cold_segments: tokio::sync::RwLock>>, + /// Weak self-reference for background tasks spawned from &self methods + /// (warm promotion). Set once right after construction. + self_weak: std::sync::OnceLock>, +} + +/// 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 +235,63 @@ 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 cold = std::env::var("COMPASS_COLD_SERVE") + .map(|v| v == "true" || v == "1") + .unwrap_or(false); + // Cold serving implies lazy attach: its whole point is answering + // queries WITHOUT attaching, so eager boot-time attach is senseless. + let lazy = cold + || 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); + let manager = Self::new_with_storage_opts( + data_dir, + storage, + role, + lazy, + max_attached, + refresh_interval_secs, + ) + .await?; + if cold { + manager.set_cold_serve(true); + } + if let Ok(Some(n)) = std::env::var("COMPASS_WARM_AFTER").map(|v| v.parse().ok()) { + manager.set_warm_after(n); + } + Ok(manager) + } + + /// 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 +299,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,7 +312,27 @@ 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, + cold_serve: std::sync::atomic::AtomicBool::new(false), + warm_after: std::sync::atomic::AtomicU32::new(3), + cold_hits: std::sync::Mutex::new(HashMap::new()), + cold_segments: tokio::sync::RwLock::new(HashMap::new()), + self_weak: std::sync::OnceLock::new(), }); + let _ = manager.self_weak.set(Arc::downgrade(&manager)); + + // 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)?; @@ -166,6 +362,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 +385,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 +427,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 +461,6 @@ impl CollectionManager { key_to_chunk_id: Vec::new(), mmap_vectors: None, vectors: Vec::new(), - dims: space_config.dims, }), ); } @@ -254,16 +479,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 +513,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 +558,277 @@ impl CollectionManager { embedding_dims: Option, config: Option, ) -> Result> { - 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()); + // The partition separator is reserved: user collections must not + // squat on internal partition namespaces. + if partitions::is_partition_ns(name) { + return Err(format!( + "Collection name '{name}' contains the reserved partition separator '{}'", + partitions::PART_SEP + ) + .into()); + } + if let Some(cfg) = &config { + if let Some(field) = &cfg.partition_by { + if field.is_empty() { + return Err("partition_by must name a metadata field".into()); + } + } } + self.create_collection_inner(name, vector_spaces, embedding_dims, config) + .await + } - // 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(), - }, + /// Shared create path. Partition namespaces (containing [`partitions::PART_SEP`]) + /// may only be created internally by the ingest router. + async fn create_collection_inner( + &self, + name: &str, + vector_spaces: Option>, + 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(), ); - m - }); + } + validate_name_segment(name, "Collection")?; - let default_space = spaces.keys().next().cloned(); - let dims = spaces.values().next().map(|s| s.dims).unwrap_or(384); + let collection = { + let mut collections = self.collections.write().await; + if collections.contains_key(name) { + return Err(format!("Collection '{}' already exists", name).into()); + } - 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 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, + }; - store::save_metadata(&self.data_dir, &collection)?; + store::save_metadata(&self.data_dir, &collection)?; - // Build empty FTS index - let tantivy_dir = store::tantivy_dir(&self.data_dir, name); - let fts = tantivy_fts::build_index(&tantivy_dir, &[], 0)?; + // Build empty FTS index + let tantivy_dir = store::tantivy_dir(&self.data_dir, name); + let fts = tantivy_fts::build_index(&tantivy_dir, &[])?; - // 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, - }), - ); - } + // 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(), + }), + ); + } - // 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)?; + // 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(), + }; + + collections.insert(name.to_string(), loaded); + collection + }; // write lock released — never hold it across S3 round-trips. + + // 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. + // Partition namespaces mint from the PARENT's allocator + // (collection-unique ids) and are never seeded themselves. + if let Err(e) = if partitions::is_partition_ns(name) { + Ok(()) + } else { + 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(_)) => { + if partitions::is_partition_ns(name) { + // A racing writer bootstrapped this partition's + // manifest between our config write and here. That is + // the expected create race for partitions — the + // namespace (config + manifest) is exactly what we + // wanted; adopt it. (Destroying the config here left + // a config-less namespace that cold serving and warm + // attach then disagreed about.) + rollback_local().await; + return Err(format!( + "Collection '{}' already exists in object storage", + name + ) + .into()); + } + // 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(), - }; + // Partitioned parents allocate chunk ids from a shared CAS allocator + // in LOCAL mode too (partitions must never mint colliding ids). This + // is a single JSON file under the collection dir — no WAL/manifest. + if !self.cloud_mode + && collection.config.partition_by.is_some() + && !partitions::is_partition_ns(name) + { + if let Err(e) = crate::storage::id_alloc::seed(self.storage.as_ref(), name, 0).await { + // Roll back: a partitioned parent without an allocator could + // never ingest (the migrate path needs a manifest that local + // mode deliberately doesn't have). + self.collections.write().await.remove(name); + let _ = store::delete_collection_data(&self.data_dir, name); + return Err(format!("id allocator seed failed: {e}").into()); + } + } - 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() + }; + // Partition namespaces are internal — the parent represents them. + out.retain(|c| !partitions::is_partition_ns(&c.name)); + 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 partitions::is_partition_ns(&name) { + continue; + } + 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 +837,77 @@ 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()); + // Partitioned parent: cascade over every partition namespace FIRST, + // so a failure mid-cascade leaves the parent (and the retry path) + // intact. Partitions are discovered from all sources — attached map, + // lazy registry, local dirs, and the bucket (a writer node may have + // created partitions this node never saw). + if !partitions::is_partition_ns(name) { + let prefix = format!("{name}{}", partitions::PART_SEP); + let mut parts: std::collections::HashSet = std::collections::HashSet::new(); + { + let collections = self.collections.read().await; + parts.extend( + collections + .keys() + .filter(|k| k.starts_with(&prefix)) + .cloned(), + ); + } + parts.extend( + self.registered + .read() + .await + .iter() + .filter(|k| k.starts_with(&prefix)) + .cloned(), + ); + if let Ok(entries) = std::fs::read_dir(&self.data_dir) { + for e in entries.flatten() { + if let Some(n) = e.file_name().to_str() { + if n.starts_with(&prefix) { + parts.insert(n.to_string()); + } + } + } + } + if self.cloud_mode { + if let Ok(all) = crate::storage::lsm::list_namespaces(self.storage.as_ref()).await { + parts.extend(all.into_iter().filter(|n| n.starts_with(&prefix))); + } + } + for part in parts { + Box::pin(self.delete_collection(&part)).await?; + } + } + // 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). @@ -446,6 +930,8 @@ impl CollectionManager { dims: usize, model: &str, ) -> Result<(), Box> { + self.reject_if_partitioned(collection_name, "vector-space changes") + .await?; // Validate `space_name` before it touches the filesystem. The name is // interpolated into on-disk paths (`{space_name}.bin`, `.index`, // `.keymap`), so an unconstrained value like `../../tmp/pwn` could @@ -453,36 +939,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(()) } @@ -492,20 +1011,48 @@ impl CollectionManager { collection_name: &str, space_name: &str, ) -> Result<(), Box> { + self.reject_if_partitioned(collection_name, "vector-space changes") + .await?; // Same path-traversal guard as add_vector_space — the name flows into // `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 +1072,82 @@ 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))?; + self.reject_if_partitioned(collection_name, "vector-space changes") + .await?; + 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(); @@ -582,3457 +1180,3311 @@ impl CollectionManager { store::vectors_dir(&self.data_dir, collection_name) } - // ── Ingest ─────────────────────────────────────────────────────────── - - /// Ingest chunks with batch parent resolution, named embeddings, and relationships. - pub async fn ingest( - &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))?; + // ── Serve-from-storage (Phase 5) ───────────────────────────────────── - let count = ingest_chunks.len(); + pub fn set_cold_serve(&self, on: bool) { + self.cold_serve + .store(on, std::sync::atomic::Ordering::Relaxed); + } - // 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); + /// Cold hits before background warm promotion (0 disables promotion). + pub fn set_warm_after(&self, n: u32) { + self.warm_after + .store(n, std::sync::atomic::Ordering::Relaxed); + } - 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); - } - } + fn cold_serve(&self) -> bool { + self.cold_serve.load(std::sync::atomic::Ordering::Relaxed) + } - // 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() - .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, - ); + /// Cold semantic search: manifest read (freshness anchor) → cached + /// per-segment artifacts → cluster probes → tail brute-force → hydrate. + /// FTS needs an inverted index and is not cold-servable — callers get a + /// clear error steering them to semantic mode (or a warmed node). + #[allow(clippy::type_complexity)] + async fn search_cold( + &self, + ns: &str, + req: &SearchRequest, + embed_state: &EmbedState, + ) -> Result< + ( + Vec<( + DocumentChunk, + f32, + String, + Option>, + Option>, + )>, + usize, + u64, + Option, + ), + Box, + > { + let start = std::time::Instant::now(); + crate::metrics::inc(&crate::metrics::COLD_SEARCHES_TOTAL); + // The cold path is pure vector search. Anything that needs local + // indexes or the scoring pipeline is REJECTED, not silently ignored — + // identical requests must never return materially different rankings + // cold vs. warm without a signal. ("hybrid" with an empty text query + // degenerates to semantic legitimately and is allowed.) + if req.mode == "fts" || (req.mode == "hybrid" && !req.query.is_empty()) { + return Err(format!( + "collection '{ns}' is cold (not attached): '{}' search needs local \ + indexes. Use mode \"semantic\", or query again after the namespace warms.", + req.mode + ) + .into()); + } + if req.recency.is_some() + || req.recency_preset.is_some() + || !req.boosts.is_empty() + || req.relationship_boost.is_some() + || req.include_relations + { + return Err(format!( + "collection '{ns}' is cold (not attached): recency/boosts/relationship \ + options need the warm scoring pipeline. Drop them, or query again after \ + the namespace warms." + ) + .into()); + } - // Phase 3: Build DocumentChunks and collect embeddings per vector space - let default_space = loaded - .metadata - .default_vector_space + // The bucket config names the default space. A namespace can hold + // data WITHOUT a config (older create paths could destroy the config + // after a race) — warm attach accepts manifest-exists, so cold must + // too, falling back to the requested/"default" space. + let default_space = match self.bucket_config(ns, false).await { + Ok(cfg) => cfg.default_vector_space.clone(), + Err(e) if e.downcast_ref::().is_some() => { + if !self.storage.exists(&format!("{ns}/manifest")).await? { + return Err(e); + } + None + } + Err(e) => return Err(e), + }; + let space = req + .vector_space .clone() - .unwrap_or_else(|| "default".into()); - let mut chunks: Vec = Vec::with_capacity(count); - // space_name -> Vec<(chunk_id, embedding)> - let mut space_vectors: HashMap)>> = HashMap::new(); - // Deferred relationship additions (applied after the S3 append, so we can - // release the collections lock during network I/O). - let mut rel_adds: Vec<(u64, Option, Option)> = Vec::with_capacity(count); + .or(default_space) + .unwrap_or_else(|| "default".to_string()); + + let query_vec: Vec = match &req.query_vector { + Some(v) => v.clone(), + None => embed_state.embed_query(&req.query).map_err(|e| { + format!("cold search needs a query_vector or a loaded embed model: {e}") + })?, + }; - for (i, ic) in ingest_chunks.into_iter().enumerate() { - let id = assigned_ids[i]; - let (parent_id, group_id) = resolved[i].clone(); + let (manifest, _) = crate::storage::lsm::read_manifest(self.storage.as_ref(), ns).await?; + if let Some(min_seq) = req.min_seq { + // Cold reads see everything committed to the manifest, so + // read-your-writes holds by construction — only a seq beyond the + // write history is unsatisfiable. + if min_seq >= manifest.next_seq { + return Err(format!( + "min_seq {} is beyond the collection's write history ({})", + min_seq, manifest.next_seq + ) + .into()); + } + } - // Collect named embeddings - let mut embeddings = ic.embeddings; - // Legacy: single embedding -> map to default space - if let Some(emb) = ic.embedding { - if !embeddings.contains_key(&default_space) { - embeddings.insert(default_space.clone(), emb); - } + // Cached artifacts per segment (immutable → cache by id). + let mut segments = Vec::with_capacity(manifest.segments.len()); + for sref in &manifest.segments { + let key = format!("{ns}/{}", sref.id); + if let Some(cs) = self.cold_segments.read().await.get(&key).cloned() { + segments.push(cs); + continue; } - // If no embeddings provided at all, compute using the built-in - // embedder — but ONLY if its output matches the default space's - // dims (a 384-dim BGE vector in a 4-dim space would corrupt the - // vector file). Chunks without a usable embedding stay FTS-only. - if embeddings.is_empty() { - if let Ok(emb) = embed_state.embed_query(&ic.text) { - let expected = loaded - .metadata - .vector_spaces - .get(&default_space) - .map(|c| c.dims) - .unwrap_or(loaded.metadata.embedding_dims); - if emb.len() == expected { - embeddings.insert(default_space.clone(), emb); - } - } + let cs = Arc::new( + crate::search::cold::ColdSegment::open(self.storage.as_ref(), ns, &sref.id).await?, + ); + let mut cache = self.cold_segments.write().await; + if cache.len() >= 1024 { + cache.clear(); // crude bound; entries rebuild in a few reads } + cache.insert(key, cs.clone()); + segments.push(cs); + } - // Validate USER-provided embedding lengths against each space's - // configured dims BEFORE anything is written. One wrong-length - // vector would silently corrupt the mmap vector file in release - // builds (offsets shift for every vector after it). - for (space_name, vec) in &embeddings { - let expected = loaded - .metadata - .vector_spaces - .get(space_name) - .map(|c| c.dims) - .unwrap_or(loaded.metadata.embedding_dims); - if vec.len() != expected { - return Err(format!( - "chunk {i}: embedding for vector space '{space_name}' has {} dims, \ - expected {expected}", - vec.len() - ) - .into()); - } - } + let nprobe = std::env::var("COMPASS_COLD_NPROBE") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(crate::search::cold::DEFAULT_NPROBE); + let hits = crate::search::cold::search( + self.storage.as_ref(), + ns, + &segments, + &manifest, + &space, + &query_vec, + req.top_k, + nprobe, + &req.filters, + ) + .await?; - // Store embeddings by vector space for batch index building - for (space_name, vec) in &embeddings { - space_vectors - .entry(space_name.clone()) - .or_default() - .push((id, vec.clone())); - } + self.maybe_warm(ns); - let chunk = DocumentChunk { - id, - collection: collection_name.to_string(), - file_id: ic.file_id, - chunk_index: ic.chunk_index, - page: ic.page, - text: ic.text, - metadata: ic.metadata, - doc_type: ic.doc_type, - parent_id, - group_id: group_id.clone(), - embeddings, - embedding: None, // v2 uses named embeddings - }; + let took_us = start.elapsed().as_micros() as u64; + let total = hits.len(); + let explain = req.explain.then(|| ExplainPlan { + filter: FilterExplain { + eligible_count: total as u64, + universe_count: 0, // unknown without attaching — cold reads don't scan + selectivity: 0.0, + }, + ann: AnnExplain { + engine: "cold-ivf".to_string(), + candidates_inspected: None, + ef_search_used: 0, + }, + }); + let results = hits + .into_iter() + .map(|(c, score)| (c, score, "semantic-cold".to_string(), None, None)) + .collect(); + Ok((results, total, took_us, explain)) + } - // Defer applying relationships / chunk map to `loaded` until AFTER - // the S3 append (so we can drop the lock during network I/O, #2). - rel_adds.push((id, parent_id, group_id)); - chunks.push(chunk); + /// Count a cold hit; at the warm threshold, spawn a background attach so + /// a repeatedly-queried namespace migrates to the fast path on its own. + fn maybe_warm(&self, ns: &str) { + let after = self.warm_after.load(std::sync::atomic::Ordering::Relaxed); + if after == 0 { + return; + } + let fire = { + let mut map = self.cold_hits.lock().unwrap(); + let e = map.entry(ns.to_string()).or_insert(0); + *e += 1; + // Reset at the threshold so a namespace that got promoted and + // later LRU-evicted can warm AGAIN after `after` fresh cold hits + // (an == latch would seal shut forever after the first firing). + if *e >= after { + *e = 0; + true + } else { + false + } + }; + if fire { + if let Some(m) = self.self_weak.get().and_then(|w| w.upgrade()) { + let ns = ns.to_string(); + tokio::spawn(async move { + crate::metrics::inc(&crate::metrics::WARM_PROMOTIONS_TOTAL); + tracing::info!("cold namespace '{ns}' hit warm threshold; attaching"); + if let Err(e) = m.ensure_attached(&ns).await { + tracing::warn!("warm promotion attach for '{ns}' failed: {e}"); + // Reset so a later burst can retry. + m.cold_hits.lock().unwrap().remove(&ns); + } + }); + } } + } - // Release the collections write lock BEFORE the S3 network round-trip, so - // a slow S3 call doesn't stall every other collection (#2). Nothing local - // has been mutated yet (chunks/relationships were deferred into local Vecs - // above), so there is no state to roll back if the append fails. - drop(collections); + // ── Tenant partitions (Phase 6) ────────────────────────────────────── - // Phase 3a (cloud): DURABLE S3 WAL append FIRST, before any local commit - // (fixes F14 split-brain — a failed append leaves nothing local, clean retry). + /// The partition field of a collection, or None for normal collections + /// and partition namespaces themselves. Reads attached metadata when + /// present, else the (cached) bucket config — deliberately WITHOUT + /// attaching: cold-served and lazy namespaces must be routable from + /// config alone. + async fn partition_field( + &self, + name: &str, + ) -> Result, Box> { + if partitions::is_partition_ns(name) { + return Ok(None); + } + if let Some(loaded) = self.collections.read().await.get(name) { + return Ok(loaded.metadata.config.partition_by.clone()); + } if self.cloud_mode { - 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, ingest not applied: {e}"))?; - tracing::info!( - "Cloud ingest: WAL fragment seq={} ({} chunks) durable for '{}'", - seq, - records, - collection_name - ); - maybe_auto_compact( - self.storage.clone(), - collection_name.to_string(), - self.compacting.clone(), + // Unattached: the bucket config answers without an attach. A + // MISSING config means "not partitioned" (the caller's own lookup + // produces the not-found) — but a transient storage error must + // PROPAGATE: swallowing it would reclassify a partitioned + // collection as unpartitioned and misroute tenant writes into the + // parent namespace, where no partition-routed search looks. + return Ok( + match cloud::read_bucket_config(self.storage.as_ref(), name).await? { + Some(cfg) => cfg.config.partition_by, + None => None, + }, ); } + Ok(None) + } - // Re-acquire the write lock and apply local state (durable S3 record, if - // any, already written). Ids were pre-assigned from a monotonic counter, - // 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, + /// Make sure a partition namespace exists and is servable, creating it on + /// first sight (inheriting the parent's vector spaces + embed model). + /// Racing creators and partitions created by writer nodes resolve via the + /// create path's own already-exists handling. + async fn ensure_partition( + &self, + parent: &str, + pval: &str, + ) -> Result> { + let ns = partitions::partition_ns(parent, pval); + if self.collections.read().await.contains_key(&ns) { + return Ok(ns); + } + if self.cloud_mode && self.registered.read().await.contains(&ns) { + return Ok(ns); // lazy attach loads it at the entry point + } + // Parent template: attached metadata when present, else the bucket + // config — a lazy/cold-serve node routes partitioned ingest without + // ever attaching the parent, so requiring attachment here would break + // first ingest of a new tenant on exactly those nodes. + let attached_template = { + let collections = self.collections.read().await; + collections.get(parent).map(|p| { + ( + p.metadata.vector_spaces.clone(), + p.metadata.config.embed_model.clone(), + ) + }) + }; + let (spaces, embed_model) = match attached_template { + Some(t) => t, + None if self.cloud_mode => { + let cfg = self.bucket_config(parent, false).await?; + (cfg.vector_spaces.clone(), cfg.config.embed_model.clone()) + } 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 - // manifest referencing only our orphan fragment) yields nothing. - if self.cloud_mode { - if let Err(te) = crate::storage::lsm::append_tombstone( - self.storage.as_ref(), - collection_name, - &assigned_ids, - ) - .await - { - tracing::error!( - "compensation (deleted-in-gap): S3 tombstone for '{}' failed: {}", - collection_name, - te - ); + return Err(not_found(format_args!("Collection '{}' not found", parent))); + } + }; + let config = CollectionConfig { + embed_model, + partition_by: None, + }; + match self + .create_collection_inner(&ns, Some(spaces), None, Some(config)) + .await + { + Ok(_) => Ok(ns), + // Lost a create race (local map, bucket config, or bucket data) — + // the partition exists; ensure_attached at the entry point loads it. + Err(e) if e.to_string().contains("already") => Ok(ns), + Err(e) => Err(e), + } + } + + /// Ingest into a partitioned collection: one delegated ingest per touched + /// partition. `seq` is passed through when exactly one partition was + /// touched; multi-partition batches return None (each partition has its + /// own manifest and thus its own seq domain). + async fn ingest_partitioned( + &self, + parent: &str, + field: &str, + ingest_chunks: Vec, + embed_state: &EmbedState, + ) -> Result<(usize, HashMap, Option), Box> + { + let groups = partitions::group_by_partition(field, ingest_chunks)?; + let multi = groups.len() > 1; + let mut total = 0usize; + let mut id_map = HashMap::new(); + let mut last_seq = None; + for (pval, group) in groups { + let ns = self.ensure_partition(parent, &pval).await?; + let (n, ids, seq) = Box::pin(self.ingest(&ns, group, embed_state)).await?; + total += n; + id_map.extend(ids); + last_seq = seq; + } + Ok((total, id_map, if multi { None } else { last_seq })) + } + + /// Search a partitioned collection: route to the partitions named by the + /// partition-field filter, merge by score, truncate to top_k. Partitions + /// that do not exist yet contribute zero results (a tenant with no data + /// is empty, not an error). + #[allow(clippy::type_complexity)] + async fn search_partitioned( + &self, + parent: &str, + field: &str, + req: &SearchRequest, + embed_state: &EmbedState, + ) -> Result< + ( + Vec<( + DocumentChunk, + f32, + String, + Option>, + Option>, + )>, + usize, + u64, + Option, + ), + Box, + > { + let pvals = partitions::partition_values_from_filters(field, &req.filters)?; + if req.min_seq.is_some() && pvals.len() > 1 { + return Err("min_seq applies to a single partition's write history; \ + filter to one partition value when using it" + .into()); + } + let mut merged = Vec::new(); + let mut total = 0usize; + let mut took = 0u64; + let mut explain = None; + for pval in pvals { + let ns = partitions::partition_ns(parent, &pval); + match Box::pin(self.search(&ns, req, embed_state)).await { + Ok((results, t, us, ex)) => { + merged.extend(results); + total += t; + took += us; + if explain.is_none() { + explain = ex; } } - return Err(format!("Collection '{}' not found", collection_name).into()); + Err(e) if e.downcast_ref::().is_some() => continue, + Err(e) => return Err(e), } - }; - // 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 - // tombstone (below) — otherwise a partial local commit + orphan S3 - // fragment would resurrect/duplicate the batch on a cold restart (F2). - let commit_result = Self::apply_ingest_commit( - &self.data_dir, - collection_name, - loaded, - rel_adds, - &chunks, - space_vectors, - count, - ); - if let Err(e) = commit_result { - if self.cloud_mode { - // Compensate for the durable S3 fragment whose local commit - // failed. Tombstone the ids in THREE places so they can never - // resurface, on any restart path: - // 1. local redb tombstones table — survives `load_collection` - // rehydrating from redb on a persistent-disk node (the - // normal deployment). Without this, the chunk sits in redb - // (insert_batch may have succeeded before FTS/HNSW failed) - // and would be pulled back into RAM on restart. - // 2. the in-RAM tombstone set — masks it at query time now. - // 3. an S3 tombstone — so a fresh-disk rebuild-from-manifest - // also drops it. - if let Err(te) = loaded.chunk_store.tombstone_batch(&assigned_ids) { - tracing::error!( - "compensation: failed to write local tombstones for '{}': {} \ - (chunk may need manual delete)", - collection_name, - te - ); - } - for id in &assigned_ids { - loaded.tombstones.insert(*id); - loaded.chunks.remove(id); - } - 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(), - collection_name, - &assigned_ids, - ) - .await - { - // Flaky S3: the orphan fragment now has no S3 tombstone. Local - // tombstones (redb + RAM) still mask it on THIS node; surface - // loudly so a fresh-disk rebuild risk is visible to operators. - tracing::error!( - "compensation: durable S3 tombstone for '{}' FAILED: {} — \ - orphan fragment may resurrect on a fresh-disk rebuild; \ - run POST /collections/{}/compact once S3 is healthy", - collection_name, - te, - collection_name - ); - } - } - return Err(e); } - - tracing::info!("Ingested {} chunks into '{}'", count, collection_name); - - Ok((count, client_id_map)) + merged.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + merged.truncate(req.top_k); + Ok((merged, total, took, explain)) } - /// Apply an ingest batch's local state (chunk map, redb, FTS, HNSW, metadata, - /// relationships, filter index). All-or-caller-compensates: any `?` failure - /// leaves partial local state, which the caller undoes + tombstones in cloud - /// mode. Synchronous (no `.await`) — the S3 write already happened. - #[allow(clippy::too_many_arguments)] - fn apply_ingest_commit( - data_dir: &Path, - collection_name: &str, - loaded: &mut LoadedCollection, - rel_adds: Vec<(u64, Option, Option)>, - chunks: &[DocumentChunk], - space_vectors: HashMap)>>, - count: usize, - ) -> Result<(), Box> { - for (id, parent_id, group_id) in rel_adds { - loaded.relationships.add(id, parent_id, group_id); + /// Writer-side partition bootstrap: create-only bucket objects for a + /// partition namespace (config copy + empty manifest). Idempotent — racing + /// writers and serving nodes all converge on the first writer's objects. + /// No per-partition id allocator is seeded (ids mint from the parent's). + async fn ensure_partition_ns_cloud( + &self, + parent_cfg: &cloud::BucketConfig, + pval: &str, + ) -> Result> { + let ns = partitions::partition_ns(&parent_cfg.name, pval); + if self.bucket_configs.read().await.contains_key(&ns) { + return Ok(ns); } - for chunk in chunks { - loaded.chunks.insert(chunk.id, chunk.clone()); + // First sight of this partition: re-validate the PARENT with a fresh + // read before creating bucket objects. A writer's cached parent + // config outlives a cascade delete — bootstrapping from it would + // resurrect the collection as an orphan namespace. + let parent_cfg = self.bucket_config(&parent_cfg.name, true).await?; + let mut part_cfg = parent_cfg.clone(); + part_cfg.name = ns.clone(); + part_cfg.config.partition_by = None; + match cloud::write_bucket_config_if_absent(self.storage.as_ref(), &ns, &part_cfg).await { + Ok(()) | Err(crate::storage::StorageError::AlreadyExists(_)) => {} + Err(e) => return Err(format!("partition config write failed: {e}").into()), } + match crate::storage::lsm::init_namespace(self.storage.as_ref(), &ns).await { + Ok(()) | Err(crate::storage::StorageError::AlreadyExists(_)) => {} + Err(e) => return Err(format!("partition manifest init failed: {e}").into()), + } + Ok(ns) + } - // Phase 3b: Persist chunks to the disk-backed store BEFORE updating - // FTS/HNSW. If this write fails we error out before any index commits, - // so we never end up with a Tantivy or HNSW index referencing chunks - // that don't exist on disk. redb writes are atomic per batch. - let to_persist: Vec<(u64, DocumentChunk)> = - chunks.iter().map(|c| (c.id, c.clone())).collect(); - loaded.chunk_store.insert_batch(&to_persist)?; - - // Phase 4: Update Tantivy FTS index - let tantivy_dir = store::tantivy_dir(data_dir, collection_name); - loaded.fts = tantivy_fts::build_index(&tantivy_dir, chunks, loaded.metadata.chunk_count)?; - - // Phase 5: Update each vector space's HNSW index - let vectors_dir = store::vectors_dir(data_dir, collection_name); - for (space_name, new_vecs) in space_vectors { - // Same fallback the ingest-time validation uses, so a space unknown - // to metadata can't validate against one dims and build with another. - let dims = loaded - .metadata - .vector_spaces - .get(&space_name) - .map(|c| c.dims) - .unwrap_or(loaded.metadata.embedding_dims); + /// Typed fence for operations not yet routed on partitioned collections. + /// Writer nodes consult the bucket config (they hold no local metadata); + /// without this a writer would durably append e.g. relations into the + /// parent namespace, which no serving node ever materializes. + async fn reject_if_partitioned( + &self, + name: &str, + what: &str, + ) -> Result<(), Box> { + if self.is_partitioned_any_role(name).await? { + return Err(format!( + "{what} is not supported on a partitioned collection yet \ + (collection '{name}' is partitioned)" + ) + .into()); + } + Ok(()) + } - let index_path = vectors_dir.join(format!("{}.index", space_name)); - let vecs_path = vectors_dir.join(format!("{}.bin", space_name)); + /// Role-aware "is this collection partitioned?": serving nodes read local + /// metadata (attaching if needed); writers consult the bucket config. + async fn is_partitioned_any_role( + &self, + name: &str, + ) -> Result> { + if partitions::is_partition_ns(name) { + return Ok(false); + } + if self.role == NodeRole::Writer { + return match self.bucket_config(name, false).await { + Ok(c) => Ok(c.config.partition_by.is_some()), + // Unknown collection: not partitioned (downstream 404s). + Err(e) if e.downcast_ref::().is_some() => Ok(false), + // Transient storage errors PROPAGATE — treating them as + // "not partitioned" would let a writer append tombstones + // into the parent namespace no serving node materializes. + Err(e) => Err(e), + }; + } + Ok(self.partition_field(name).await?.is_some()) + } - // Check if we can do incremental add (existing index + mmap vectors) - let existing = loaded.vector_spaces.get(&space_name); - let can_incremental = existing.map(|e| e.mmap_vectors.is_some()).unwrap_or(false); + // ── Ingest ─────────────────────────────────────────────────────────── - if can_incremental { - // Incremental path: append to mmap file, add to HNSW, save - let arc = loaded.vector_spaces.remove(&space_name).unwrap(); - let Ok(mut vs) = Arc::try_unwrap(arc) else { - // Another thread holds a reference — fall back to full rebuild - let existing = loaded.vector_spaces.get(&space_name); - let mut all_ids: Vec = existing - .map(|e| e.key_to_chunk_id.clone()) - .unwrap_or_default(); - let mut all_vecs: Vec> = existing - .and_then(|e| e.mmap_vectors.as_ref()) - .map(|m| m.to_vecs()) - .unwrap_or_default(); - for (cid, vec) in new_vecs { - all_ids.push(cid); - all_vecs.push(vec); - } - let vs = vector::build_vector_index( - &index_path, - &vecs_path, - &all_ids, - &all_vecs, - dims, - )?; - loaded.vector_spaces.insert(space_name, Arc::new(vs)); - continue; + /// Ingest chunks with batch parent resolution, named embeddings, and relationships. + /// 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; + // Partitions mint from the PARENT's allocator: chunk ids stay unique + // across the whole partitioned collection. + let ns = partitions::alloc_ns(ns); + 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()), + } + } - // 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 - // semantic search on it until restart. - let result = (|| -> Result<(), Box> { - // Append new vectors to mmap file - if let Some(ref mut mmap) = vs.mmap_vectors { - mmap.append(&new_vecs)?; - } - - // Extend the key mapping and persist it IMMEDIATELY after - // the mmap append, before any HNSW work — so the two files - // never desync on disk (a stale keymap makes search fabricate - // chunk ids from raw key indexes after restart). - let base_key = vs.key_to_chunk_id.len(); - for (cid, _) in &new_vecs { - vs.key_to_chunk_id.push(*cid); - } - 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()) - let total = vs.key_to_chunk_id.len(); - if total >= 1000 && (vs.index.is_none() || index_path.exists()) { - 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()); - 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))?; + /// 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(); } - index - .save(index_path_str) - .map_err(|e| format!("Failed to save index: {}", e))?; - vs.index = Some(index); } - Ok(()) - })(); - // Space goes back in whatever happened; a partial update is - // recoverable (caller compensates the batch), a vanished space - // is a silent outage. - loaded.vector_spaces.insert(space_name, Arc::new(vs)); - result?; - } else { - // Full rebuild path (first ingest or legacy data) - let mut all_ids: Vec = existing - .map(|e| e.key_to_chunk_id.clone()) - .unwrap_or_default(); - let mut all_vecs: Vec> = - existing.map(|e| e.vectors.clone()).unwrap_or_default(); - - for (cid, vec) in new_vecs { - all_ids.push(cid); - all_vecs.push(vec); + 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 + ))) } - - let vs = - vector::build_vector_index(&index_path, &vecs_path, &all_ids, &all_vecs, dims)?; - loaded.vector_spaces.insert(space_name, Arc::new(vs)); } } - - // Phase 6: Save metadata + relationships, then rebuild the filter index. - // Persist the advanced next_id high-water mark so ids are never reused - // even if the local chunk store is later empty on restart. - loaded.metadata.chunk_count += count as u64; - loaded.metadata.next_id = loaded.next_id; - 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); - Ok(()) } - // ── Search ─────────────────────────────────────────────────────────── + /// 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) + } - /// Search with full scoring pipeline: retrieve (filter-aware) → score → return. - /// - /// The filter is applied INSIDE the HNSW walk via USearch's filter - /// callback when set. Recall does not collapse on selective filters. - /// See `docs/v0.4-filter-aware-ann.md`. - pub async fn search( + /// 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, - req: &SearchRequest, + ingest_chunks: Vec, embed_state: &EmbedState, - ) -> Result< - ( - Vec<( - DocumentChunk, - f32, - String, - Option>, - Option>, - )>, - usize, - u64, - Option, - ), - Box, - > { - 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 mode = SearchMode::from_str_param(&req.mode); - let rerank_k = req.top_k * 3; // fetch extra candidates for scoring + ) -> 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?; + + // Partitioned collection: group by the partition field, make each + // partition's bucket objects exist (idempotent create-only writes — + // a writer may see a tenant before any serving node does), delegate. + if let Some(field) = cfg.config.partition_by.clone() { + let groups = partitions::group_by_partition(&field, ingest_chunks)?; + let multi = groups.len() > 1; + let mut total = 0usize; + let mut id_map = HashMap::new(); + let mut last_seq = None; + for (pval, group) in groups { + let ns = self.ensure_partition_ns_cloud(&cfg, &pval).await?; + let (n, ids, seq) = + Box::pin(self.ingest_stateless(&ns, group, embed_state)).await?; + total += n; + id_map.extend(ids); + last_seq = seq; + } + return Ok((total, id_map, if multi { None } else { last_seq })); + } - // Determine which vector space to use - let space_name = req - .vector_space - .as_deref() - .or(loaded.metadata.default_vector_space.as_deref()) - .unwrap_or("default"); - - // ── Step 0: Compile filter, resolve eligible bitmap ────────────── - // FilterExpr::compile is cheap. eligible() is a roaring intersection - // across the predicate bitmaps; sub-millisecond at any realistic size. - // The bitmap routes through both FTS and semantic retrieval below. - let filter_expr = FilterExpr::compile(&req.filters); - let eligible = loaded.filter_index.eligible(&filter_expr); - let universe_count = loaded.filter_index.len(); - let selectivity_val = selectivity(&eligible, universe_count); - let filter_active = !filter_expr.is_empty(); - - // Engine / candidates-inspected / ef metadata, captured for /explain. - let mut explain_engine: Option<&'static str> = None; - let mut explain_candidates: Option = None; + // 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); + } - // ── 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)?; - // 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. - if filter_active { - raw.into_iter() - .filter(|(id, _)| eligible.contains(*id)) - .collect() - } else { - raw + // 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); + } } - } else { - Vec::new() - }; + 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, + ); - let semantic_results = if matches!(mode, SearchMode::Semantic | SearchMode::Hybrid) { - if let Some(vs) = loaded.vector_spaces.get(space_name) { - let query_vec_opt: Option> = req - .query_vector - .clone() - .or_else(|| embed_state.embed_query(&req.query).ok()); - if let Some(query_vec) = query_vec_opt { - let vs_clone = vs.clone(); - if filter_active { - // Filter-aware path: USearch's filter callback prunes - // ineligible nodes during the HNSW walk. No over-fetch, - // no post-filter recall collapse. - let eligible_clone = eligible.clone(); - let (vr, explain) = tokio::task::spawn_blocking(move || { - vector::search_vectors_filtered( - &query_vec, - &vs_clone, - rerank_k, - &eligible_clone, - ) - }) - .await - .unwrap_or_else(|_| (Vec::new(), vector::FilteredSearchExplain::default())); - explain_engine = Some(if explain.used_hnsw { - "hnsw" + 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 { - "brute_force" - }); - if explain.used_hnsw { - explain_candidates = Some(explain.candidates_inspected); + tracing::warn!( + "writer ingest: built-in embedder produces {} dims but space \ + '{}' expects {} — chunk {} will be FTS-only", + emb.len(), + default_space, + expected, + i + ); } - vr.iter().map(|r| (r.chunk_id, r.score)).collect::>() - } else { - // No filter: skip the predicate-callback overhead and - // use the existing unfiltered HNSW path. - let vr = tokio::task::spawn_blocking(move || { - vector::search_vectors(&query_vec, &vs_clone, rerank_k) - }) - .await - .unwrap_or_default(); - explain_engine = Some("hnsw"); - vr.iter().map(|r| (r.chunk_id, r.score)).collect::>() } - } else { - Vec::new() } - } else { - Vec::new() + 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, + }); } - } else { - Vec::new() + Ok((chunks, client_id_map)) }; - - // ── Step 2: Merge via RRF (for hybrid) or use single-mode results ── - let mut candidates: Vec = match mode { - SearchMode::Hybrid if !fts_results.is_empty() || !semantic_results.is_empty() => { - let (rrf_k, fts_w, sem_w) = match &req.score_weights { - Some(sw) => ( - sw.rrf_k as f32, - sw.fts_weight as f32, - sw.semantic_weight as f32, - ), - None => (60.0, 1.0, 1.0), - }; - let merged = hybrid::merge_rrf( - &fts_results, - &semantic_results, - rerank_k, - rrf_k, - fts_w, - sem_w, - ); - merged - .iter() - .map(|r| ScoredCandidate { - chunk_id: r.chunk_id, - base_score: r.rrf_score, - final_score: r.rrf_score, - source: r.source.as_str().to_string(), - }) - .collect() + 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)? } - SearchMode::Fts => fts_results - .iter() - .map(|(id, score)| ScoredCandidate { - chunk_id: *id, - base_score: *score, - final_score: *score, - source: "fts".to_string(), - }) - .collect(), - SearchMode::Semantic => semantic_results - .iter() - .map(|(id, score)| ScoredCandidate { - chunk_id: *id, - base_score: *score, - final_score: *score, - source: "semantic".to_string(), - }) - .collect(), - _ => Vec::new(), }; - // ── Step 3: Filter is already applied (filter-aware retrieval). ─ - // The bitmap pushdown happens inside both FTS post-filter and - // USearch's filter callback above, so we no longer need a post-merge - // `retain`. Kept as an assertion in debug builds to catch invariant - // drift if a new retrieval path bypasses the eligibility check. - debug_assert!( - !filter_active || candidates.iter().all(|c| eligible.contains(c.chunk_id)), - "filter-aware retrieval produced a candidate outside the eligible bitmap" + // 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))) + } - // ── Step 3b: Drop soft-deleted (tombstoned) chunks ────────────── - // The HNSW/FTS indexes still physically contain deleted ids until the - // next rebuild/compaction, so we filter them out here. Cheap O(1) set - // membership per candidate. - if !loaded.tombstones.is_empty() { - candidates.retain(|c| !loaded.tombstones.contains(&c.chunk_id)); + 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; } - // ── Step 4: Apply scoring pipeline ────────────────────────────── - // Resolve recency preset into a full config (explicit `recency` wins) - let recency_config = req.recency.clone().or_else(|| { - req.recency_preset.as_deref().and_then(|preset| { - req.recency_field - .as_deref() - .map(|field| RecencyConfig::from_preset(preset, field.to_string())) - .flatten() - }) - }); - - let has_scoring = - 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())) - }) - .collect(); - - let candidate_ids: Vec = candidates.iter().map(|c| c.chunk_id).collect(); - let (parent_ids, sibling_map) = loaded.relationships.build_scoring_maps(&candidate_ids); - - scoring::apply_scoring_pipeline( - &mut candidates, - &chunk_metadata, - &parent_ids, - &sibling_map, - &recency_config, - &req.boosts, - &req.relationship_boost, - ); + // Partitioned collection: group by the partition field and delegate + // each group to its partition namespace (partitions.rs). + if let Some(field) = self.partition_field(collection_name).await? { + return self + .ingest_partitioned(collection_name, &field, ingest_chunks, embed_state) + .await; } - // ── Step 5: Truncate to top_k and build response ──────────────── - candidates.truncate(req.top_k); - let total = candidates.len(); + 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. + // + // Partition namespaces use the block allocator in LOCAL mode too: all + // partitions of one collection mint from the PARENT's allocator, so a + // per-partition next_id counter would collide across siblings. + let cloud_ids: Option> = + if (self.cloud_mode || partitions::is_partition_ns(collection_name)) && count > 0 { + Some(self.take_ids_cloud(collection_name, count).await?) + } else { + None + }; - // ── Step 5a: Parent metadata enrichment for segment hits ───────── - // For each segment hit with a parent_id, inline the parent's top-level - // metadata so callers (typically AI agents) avoid a second round-trip - // to fetch source-level attributes. Parents are deduplicated: N - // 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 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)) + })?; - // ── Step 5b: Relation enrichment (opt-in) ─────────────────────── - // When include_relations is set, fetch each hit's edges in ONE batched, - // on-demand read from the disk-backed relation store (never resident in - // RAM). target_status is resolved against the in-memory chunk cache: - // "found" if the target chunk exists locally, else "missing". - let mut relations_by_chunk: HashMap> = HashMap::new(); - if req.include_relations { - let types = req.relation_types.as_deref(); - relations_by_chunk = loaded.relation_store.for_chunks( - &candidate_chunk_ids, - req.relation_direction, - types, - )?; - 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) - { - "found".to_string() - } else { - "missing".to_string() - }; + // 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); } } - let hits: Vec<( - DocumentChunk, - f32, - String, - Option>, - Option>, - )> = candidates + // 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() - .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(), - ) - } else { - None - }; - ( - chunk.clone(), - c.final_score, - c.source.clone(), - parent_metadata, - relations, - ) - }) - }) + .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, + ); - let took_us = start.elapsed().as_micros() as u64; - - // ── Step 6: Build /explain plan if requested ──────────────────── - let explain_plan = if req.explain { - Some(ExplainPlan { - filter: FilterExplain { - eligible_count: eligible.len(), - universe_count, - selectivity: selectivity_val, - }, - ann: AnnExplain { - engine: explain_engine.unwrap_or("none").to_string(), - candidates_inspected: explain_candidates, - ef_search_used: hnsw_ef_search_default(), - }, - }) - } else { - None - }; + // Phase 3: Build DocumentChunks and collect embeddings per vector space + let default_space = loaded + .metadata + .default_vector_space + .clone() + .unwrap_or_else(|| "default".into()); + let mut chunks: Vec = Vec::with_capacity(count); + // space_name -> Vec<(chunk_id, embedding)> + let mut space_vectors: HashMap)>> = HashMap::new(); + // Deferred relationship additions (applied after the S3 append, so we can + // release the collections lock during network I/O). + let mut rel_adds: Vec<(u64, Option, Option)> = Vec::with_capacity(count); - Ok((hits, total, took_us, explain_plan)) - } + for (i, ic) in ingest_chunks.into_iter().enumerate() { + let id = assigned_ids[i]; + let (parent_id, group_id) = resolved[i].clone(); - // ── Chunk Relations ────────────────────────────────────────────────── + // Collect named embeddings + let mut embeddings = ic.embeddings; + // Legacy: single embedding -> map to default space + if let Some(emb) = ic.embedding { + if !embeddings.contains_key(&default_space) { + embeddings.insert(default_space.clone(), emb); + } + } + // If no embeddings provided at all, compute using the built-in + // embedder — but ONLY if its output matches the default space's + // dims (a 384-dim BGE vector in a 4-dim space would corrupt the + // vector file). Chunks without a usable embedding stay FTS-only. + if embeddings.is_empty() { + if let Ok(emb) = embed_state.embed_query(&ic.text) { + let expected = loaded + .metadata + .vector_spaces + .get(&default_space) + .map(|c| c.dims) + .unwrap_or(loaded.metadata.embedding_dims); + if emb.len() == expected { + embeddings.insert(default_space.clone(), emb); + } + } + } - /// Create a batch of chunk relations. The server mints a UUIDv4 - /// `relation_id` and stamps `created_at` for each. Self-relations - /// (`source == target`) are rejected. Returns the created edges with their - /// assigned ids. - pub async fn create_relations( - &self, - collection_name: &str, - new: Vec, - ) -> Result, Box> { - // 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))?; - 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()); + // Validate USER-provided embedding lengths against each space's + // configured dims BEFORE anything is written. One wrong-length + // vector would silently corrupt the mmap vector file in release + // builds (offsets shift for every vector after it). + for (space_name, vec) in &embeddings { + let expected = loaded + .metadata + .vector_spaces + .get(space_name) + .map(|c| c.dims) + .unwrap_or(loaded.metadata.embedding_dims); + if vec.len() != expected { + return Err(format!( + "chunk {i}: embedding for vector space '{space_name}' has {} dims, \ + expected {expected}", + vec.len() + ) + .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: if loaded.chunks.contains_key(&r.target_chunk_id) - && !loaded.tombstones.contains(&r.target_chunk_id) - { - "found".to_string() - } else { - "missing".to_string() - }, - metadata: r.metadata, - created_at: now, - }); } - } // read lock released before S3 I/O. - // Phase 2 (NO lock): durable S3 relation-upsert FIRST (S3-first ordering). - 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( + // Store embeddings by vector space for batch index building + for (space_name, vec) in &embeddings { + space_vectors + .entry(space_name.clone()) + .or_default() + .push((id, vec.clone())); + } + + let chunk = DocumentChunk { + id, + collection: collection_name.to_string(), + file_id: ic.file_id, + chunk_index: ic.chunk_index, + page: ic.page, + text: ic.text, + metadata: ic.metadata, + doc_type: ic.doc_type, + parent_id, + group_id: group_id.clone(), + embeddings, + embedding: None, // v2 uses named embeddings + }; + + // Defer applying relationships / chunk map to `loaded` until AFTER + // the S3 append (so we can drop the lock during network I/O, #2). + rel_adds.push((id, parent_id, group_id)); + chunks.push(chunk); + } + + // Release the collections write lock BEFORE the S3 network round-trip, so + // a slow S3 call doesn't stall every other collection (#2). Nothing local + // has been mutated yet (chunks/relationships were deferred into local Vecs + // above), so there is no state to roll back if the append fails. + drop(collections); + + // 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; + let seq = crate::storage::lsm::append_fragment( self.storage.as_ref(), collection_name, bytes::Bytes::from(payload), records, ) .await - .map_err(|e| format!("cloud relation-upsert append failed: {e}"))?; + .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, + records, + collection_name + ); + maybe_auto_compact( + self.storage.clone(), + collection_name.to_string(), + self.compacting.clone(), + ); } - // Phase 3 (read lock): apply locally (durable S3 record already written). - // On EITHER failure mode — collection deleted in the lock gap, or a - // local insert error — compensate with a RelationDelete for the minted - // 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()), + // Re-acquire the write lock and apply local state (durable S3 record, if + // any, already written). Ids were pre-assigned from a monotonic counter, + // 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.last_used + .store(next_lru_tick(), std::sync::atomic::Ordering::Relaxed); + l } - }; - if let Err(e) = apply_result { - if self.cloud_mode && !built.is_empty() { - let ids: Vec = built.iter().map(|r| r.relation_id.clone()).collect(); - if let Err(te) = crate::storage::lsm::append_relation_delete( - self.storage.as_ref(), - collection_name, - &ids, - ) - .await + None => { + // 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() { - tracing::error!( - "compensation: relation-delete for '{}' failed: {} — orphan \ - relation fragment may resurrect on rebuild", + 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( + self.storage.as_ref(), collection_name, - te - ); + &assigned_ids, + ) + .await + { + tracing::error!( + "compensation (deleted-in-gap): S3 tombstone for '{}' failed: {}", + collection_name, + te + ); + } } + return Err(not_found(format_args!( + "Collection \'{}\' not found", + collection_name + ))); } - return Err(e); - } - Ok(built) - } - - /// Delete a relation by id. Returns true if it existed. - pub async fn delete_relation( - &self, - collection_name: &str, - relation_id: &str, - ) -> Result> { - // 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()); + }; + // 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)); } } - // 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. - if self.cloud_mode { - 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}"))?; - } - - // 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) - } - - /// List a single chunk's relations, with `target_status` resolved against - /// the current chunk set. - pub async fn get_chunk_relations( - &self, - collection_name: &str, - chunk_id: u64, - direction: RelationDirection, - types: Option<&[String]>, - ) -> 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 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) - { - "found".to_string() - } else { - "missing".to_string() - }; - } - Ok(edges) - } - - // ── Delete (soft-delete via tombstones) ────────────────────────────── - - /// Soft-delete a set of chunk ids. Tombstones them (so they immediately - /// vanish from search results), persists the tombstones, prunes incident - /// relations, and — in object-storage mode — appends a tombstone WAL - /// 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); + // 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 + // tombstone (below) — otherwise a partial local commit + orphan S3 + // fragment would resurrect/duplicate the batch on a cold restart (F2). + let commit_result = Self::apply_ingest_commit( + &self.data_dir, + collection_name, + loaded, + rel_adds, + &chunks, + 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; + } } - - // 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) + if let Err(e) = commit_result { + if self.cloud_mode { + // Compensate for the durable S3 fragment whose local commit + // failed. Tombstone the ids in THREE places so they can never + // resurface, on any restart path: + // 1. local redb tombstones table — survives `load_collection` + // rehydrating from redb on a persistent-disk node (the + // normal deployment). Without this, the chunk sits in redb + // (insert_batch may have succeeded before FTS/HNSW failed) + // and would be pulled back into RAM on restart. + // 2. the in-RAM tombstone set — masks it at query time now. + // 3. an S3 tombstone — so a fresh-disk rebuild-from-manifest + // also drops it. + if let Err(te) = loaded.chunk_store.tombstone_batch(&assigned_ids) { + tracing::error!( + "compensation: failed to write local tombstones for '{}': {} \ + (chunk may need manual delete)", + collection_name, + te + ); + } + for id in &assigned_ids { + loaded.tombstones.insert(*id); + if let Ok(Some(c)) = loaded.chunk_store.get(*id) { + loaded.filter_index.remove(*id, &filter_meta(&c)); + } + } + drop(collections); + if let Err(te) = crate::storage::lsm::append_tombstone( + self.storage.as_ref(), + collection_name, + &assigned_ids, + ) .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 { - loaded.tombstones.insert(*id); - } - // Persist the corrected live count IMMEDIATELY after the tombstones — - // before the fallible relation pruning — so a pruning error can't leave - // 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)?; - // 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); - // 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 { - let edges = loaded - .relation_store - .for_chunk(id, RelationDirection::Both, None)?; - for e in edges { - loaded.relation_store.delete(&e.relation_id)?; + { + // Flaky S3: the orphan fragment now has no S3 tombstone. Local + // tombstones (redb + RAM) still mask it on THIS node; surface + // loudly so a fresh-disk rebuild risk is visible to operators. + tracing::error!( + "compensation: durable S3 tombstone for '{}' FAILED: {} — \ + orphan fragment may resurrect on a fresh-disk rebuild; \ + run POST /collections/{}/compact once S3 is healthy", + collection_name, + te, + collection_name + ); + } } + return Err(e); } - tracing::info!( - "Deleted {} chunk(s) from '{}' (tombstoned{})", - apply.len(), - collection_name, - if self.cloud_mode { - " + WAL tombstone" - } else { - "" - } - ); - Ok(apply.len()) + tracing::info!("Ingested {} chunks into '{}'", count, collection_name); + + Ok((count, client_id_map, 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, + /// Apply an ingest batch's local state (chunk map, redb, FTS, HNSW, metadata, + /// relationships, filter index). All-or-caller-compensates: any `?` failure + /// leaves partial local state, which the caller undoes + tombstones in cloud + /// mode. Synchronous (no `.await`) — the S3 write already happened. + #[allow(clippy::too_many_arguments)] + fn apply_ingest_commit( + 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, + rel_adds: Vec<(u64, Option, Option)>, + chunks: &[DocumentChunk], + space_vectors: HashMap)>>, + count: usize, + ) -> Result<(), Box> { + for (id, parent_id, group_id) in rel_adds { + loaded.relationships.add(id, parent_id, group_id); } - self.delete_chunks(collection_name, &ids).await - } - // ── Cloud compaction (object-storage mode) ─────────────────────────── + // Phase 3b: Persist chunks to the disk-backed store BEFORE updating + // FTS/HNSW. If this write fails we error out before any index commits, + // so we never end up with a Tantivy or HNSW index referencing chunks + // that don't exist on disk. redb writes are atomic per batch. + let to_persist: Vec<(u64, DocumentChunk)> = + chunks.iter().map(|c| (c.id, c.clone())).collect(); + loaded.chunk_store.insert_batch(&to_persist)?; - /// 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); - } - // 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()); - } - } + // 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); + let mut new_fts = tantivy_fts::build_index(&tantivy_dir, chunks)?; + new_fts.facet_bitsets.absorb(&loaded.fts.facet_bitsets); + loaded.fts = new_fts; - Ok(compact_storage(self.storage.as_ref(), collection_name).await?) - } + // Phase 5: Update each vector space's HNSW index + let vectors_dir = store::vectors_dir(data_dir, collection_name); + for (space_name, new_vecs) in space_vectors { + // Same fallback the ingest-time validation uses, so a space unknown + // to metadata can't validate against one dims and build with another. + let dims = loaded + .metadata + .vector_spaces + .get(&space_name) + .map(|c| c.dims) + .unwrap_or(loaded.metadata.embedding_dims); - /// Rebuild a collection's LOCAL indexes from its object-storage manifest - /// (materialize segments + WAL fragments → chunks → local redb/Tantivy/HNSW/ - /// filter index). This is cold-start recovery: an ephemeral node with an - /// empty local disk reconstructs the collection entirely from S3. Returns the - /// number of live chunks recovered. Cloud mode only. - pub async fn rebuild_collection_from_storage( - &self, - collection_name: &str, - ) -> Result> { - validate_name_segment(collection_name, "Collection")?; - let (manifest, _) = - crate::storage::lsm::read_manifest(self.storage.as_ref(), collection_name).await?; - let materialized = - cloud::materialize(self.storage.as_ref(), collection_name, &manifest).await?; - let chunks: Vec = materialized.chunks.values().cloned().collect(); - let live_count = chunks.len(); + let index_path = vectors_dir.join(format!("{}.index", space_name)); + let vecs_path = vectors_dir.join(format!("{}.bin", space_name)); - // 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(), - }); - } - } - 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 - // `+ if live_count > 0` form reused the highest id when every chunk was - // deleted. - let next_id = if materialized.max_id > 0 || live_count > 0 { - materialized.max_id + 1 - } else { - 0 - }; + // Check if we can do incremental add (existing index + mmap vectors) + let existing = loaded.vector_spaces.get(&space_name); + let can_incremental = existing.map(|e| e.mmap_vectors.is_some()).unwrap_or(false); - let metadata = Collection { - name: collection_name.to_string(), - created_at: Utc::now(), - 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(), - }; - store::save_metadata(&self.data_dir, &metadata)?; + if can_incremental { + // Incremental path: append to mmap file, add to HNSW, save + let arc = loaded.vector_spaces.remove(&space_name).unwrap(); + let Ok(mut vs) = Arc::try_unwrap(arc) else { + // Another thread holds a reference — fall back to full rebuild + let existing = loaded.vector_spaces.get(&space_name); + let mut all_ids: Vec = existing + .map(|e| e.key_to_chunk_id.clone()) + .unwrap_or_default(); + let mut all_vecs: Vec> = existing + .and_then(|e| e.mmap_vectors.as_ref()) + .map(|m| m.to_vecs()) + .unwrap_or_default(); + for (cid, vec) in new_vecs { + all_ids.push(cid); + all_vecs.push(vec); + } + let vs = vector::build_vector_index( + &index_path, + &vecs_path, + &all_ids, + &all_vecs, + dims, + )?; + loaded.vector_spaces.insert(space_name, Arc::new(vs)); + continue; + }; - // 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 - // local redb (e.g. a stale chunk left by a partially-committed ingest that - // was later compensated with a tombstone) must not survive. Wipe first. - let chunks_db = store::chunks_db_path(&self.data_dir, collection_name); - if let Some(parent) = chunks_db.parent() { - std::fs::create_dir_all(parent)?; - } - let _ = std::fs::remove_file(&chunks_db); - let chunk_store = 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)?; + // 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 + // semantic search on it until restart. + let result = (|| -> Result<(), Box> { + // Append new vectors to mmap file + if let Some(ref mut mmap) = vs.mmap_vectors { + mmap.append(&new_vecs)?; + } - // FTS index. - let tantivy_dir = store::tantivy_dir(&self.data_dir, collection_name); - let fts = tantivy_fts::build_index(&tantivy_dir, &chunks, 0)?; + // Extend the key mapping and persist it IMMEDIATELY after + // the mmap append, before any HNSW work — so the two files + // never desync on disk (a stale keymap makes search fabricate + // chunk ids from raw key indexes after restart). + let base_key = vs.key_to_chunk_id.len(); + for (cid, _) in &new_vecs { + vs.key_to_chunk_id.push(*cid); + } + let map_path = index_path.with_extension("keymap"); + vector::save_key_map(&map_path, &vs.key_to_chunk_id)?; - // Vector spaces (HNSW) from embeddings. - let vectors_dir = store::vectors_dir(&self.data_dir, collection_name); - let mut vs_map: HashMap> = HashMap::new(); - for (space, cfg) in &vector_spaces { - let mut ids = Vec::new(); - let mut vecs = Vec::new(); - for c in &chunks { - if let Some(emb) = c.embeddings.get(space) { - ids.push(c.id); - vecs.push(emb.clone()); + // 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 { + let index_path_str = index_path + .to_str() + .ok_or("USearch index path is not valid UTF-8")?; + 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))?; + 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))?; + } + *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); } - } - let index_path = vectors_dir.join(format!("{}.index", space)); - let vecs_path = vectors_dir.join(format!("{}.bin", space)); - let vs = vector::build_vector_index(&index_path, &vecs_path, &ids, &vecs, cfg.dims)?; - vs_map.insert(space.clone(), Arc::new(vs)); - } - - // Relationships + filter index from the recovered chunks. - let mut relationships = RelationshipStore::new(); - 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()); + loaded.vector_spaces.insert(space_name, Arc::new(vs)); + result?; + } else { + // Full rebuild path (first ingest or legacy data) + let mut all_ids: Vec = existing + .map(|e| e.key_to_chunk_id.clone()) + .unwrap_or_default(); + let mut all_vecs: Vec> = + existing.map(|e| e.vectors.clone()).unwrap_or_default(); - // Reconstruct the typed-relation store from the materialized relations - // (recovered from the S3 WAL/segments) — so relations survive a cold - // restart, not just chunks. - let relations_db = store::relations_db_path(&self.data_dir, collection_name); - let _ = std::fs::remove_file(&relations_db); // start clean, then repopulate - let relation_store = RelationStore::open(&relations_db)?; - let recovered_relations: Vec = - materialized.relations.values().cloned().collect(); - if !recovered_relations.is_empty() { - relation_store.insert_batch(&recovered_relations)?; + for (cid, vec) in new_vecs { + all_ids.push(cid); + all_vecs.push(vec); + } + + let vs = + vector::build_vector_index(&index_path, &vecs_path, &all_ids, &all_vecs, dims)?; + loaded.vector_spaces.insert(space_name, Arc::new(vs)); + } } - let loaded = LoadedCollection { - metadata, - fts, - vector_spaces: vs_map, - relationships, - chunks: chunk_map, - chunk_store, - relation_store, - tombstones: std::collections::HashSet::new(), - next_id, - filter_index, - }; - let mut collections = self.collections.write().await; - collections.insert(collection_name.to_string(), loaded); - Ok(live_count) + // Phase 6: Save metadata + relationships, then rebuild the filter index. + // Persist the advanced next_id high-water mark so ids are never reused + // even if the local chunk store is later empty on restart. + loaded.metadata.chunk_count += count as u64; + loaded.metadata.next_id = loaded.next_id; + 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)?; + // 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(()) } - /// Get facet counts for a collection. - pub async fn get_facets( + // ── Search ─────────────────────────────────────────────────────────── + + /// Search with full scoring pipeline: retrieve (filter-aware) → score → return. + /// + /// The filter is applied INSIDE the HNSW walk via USearch's filter + /// callback when set. Recall does not collapse on selective filters. + /// See `docs/v0.4-filter-aware-ann.md`. + pub async fn search( &self, collection_name: &str, - query: &str, - fields: &[String], + req: &SearchRequest, + embed_state: &EmbedState, ) -> Result< - (HashMap>, u64), + ( + Vec<( + DocumentChunk, + f32, + String, + Option>, + Option>, + )>, + usize, + u64, + Option, + ), Box, > { - 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) - } + 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); - /// Get all chunk texts and IDs for rebuild jobs. - pub async fn get_all_chunk_data( - &self, - 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))?; + // Partitioned collection: route to the partitions named by the filter. + if let Some(field) = self.partition_field(collection_name).await? { + return self + .search_partitioned(collection_name, &field, req, embed_state) + .await; + } - let mut texts = Vec::new(); - let mut ids = Vec::new(); - for (&id, chunk) in &loaded.chunks { - ids.push(id); - texts.push(chunk.text.clone()); + // Serve-from-storage: an UNATTACHED namespace answers semantic + // queries with a handful of object-storage range reads — no attach, + // no index rebuild. Repeated cold hits promote a background attach. + if self.cold_serve() + && self.cloud_mode + && !self.collections.read().await.contains_key(collection_name) + { + return self.search_cold(collection_name, req, embed_state).await; } - Ok((texts, ids)) - } + self.ensure_attached(collection_name).await?; - /// Temporal point/range lookup for TAMS-style segments. - /// - /// Returns every chunk where `doc_type == "segment"`, `group_id == Some(asset)`, - /// and the time window matches the requested query. If no time params are - /// provided, returns all segments for the asset (enumeration mode). - /// - /// Time unit: all parameters and the `timerange_start_ms` / `timerange_end_ms` - /// metadata fields are in integer milliseconds. Results are sorted ascending - /// by `timerange_start_ms` for stable ordering. - pub async fn segments_at( - &self, - collection_name: &str, - asset: &str, - time_ms: Option, - time_start_ms: Option, - time_end_ms: Option, - ) -> Result, Box> { + // 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 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)) + })?; + 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 + + // Determine which vector space to use + let space_name = req + .vector_space + .as_deref() + .or(loaded.metadata.default_vector_space.as_deref()) + .unwrap_or("default"); + + // ── Step 0: Compile filter, resolve eligible bitmap ────────────── + // FilterExpr::compile is cheap. eligible() is a roaring intersection + // across the predicate bitmaps; sub-millisecond at any realistic size. + // The bitmap routes through both FTS and semantic retrieval below. + let filter_expr = FilterExpr::compile(&req.filters); + let eligible = loaded.filter_index.eligible(&filter_expr); + let universe_count = loaded.filter_index.len(); + let selectivity_val = selectivity(&eligible, universe_count); + let filter_active = !filter_expr.is_empty(); + + // Engine / candidates-inspected / ef metadata, captured for /explain. + let mut explain_engine: Option<&'static str> = None; + let mut explain_candidates: Option = None; + + // ── 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, 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. + if filter_active { + raw.into_iter() + .filter(|(id, _)| eligible.contains(*id)) + .collect() + } else { + raw + } + } else { + Vec::new() + }; + + let semantic_results = if matches!(mode, SearchMode::Semantic | SearchMode::Hybrid) { + if let Some(vs) = loaded.vector_spaces.get(space_name) { + let query_vec_opt: Option> = req + .query_vector + .clone() + .or_else(|| embed_state.embed_query(&req.query).ok()); + if let Some(query_vec) = query_vec_opt { + let vs_clone = vs.clone(); + if filter_active { + // Filter-aware path: USearch's filter callback prunes + // ineligible nodes during the HNSW walk. No over-fetch, + // no post-filter recall collapse. + let eligible_clone = eligible.clone(); + let (vr, explain) = tokio::task::spawn_blocking(move || { + vector::search_vectors_filtered( + &query_vec, + &vs_clone, + rerank_k, + &eligible_clone, + ) + }) + .await + .unwrap_or_else(|_| (Vec::new(), vector::FilteredSearchExplain::default())); + explain_engine = Some(if explain.used_hnsw { + "hnsw" + } else { + "brute_force" + }); + if explain.used_hnsw { + explain_candidates = Some(explain.candidates_inspected); + } + vr.iter().map(|r| (r.chunk_id, r.score)).collect::>() + } else { + // No filter: skip the predicate-callback overhead and + // use the existing unfiltered HNSW path. + let vr = tokio::task::spawn_blocking(move || { + vector::search_vectors(&query_vec, &vs_clone, rerank_k) + }) + .await + .unwrap_or_default(); + explain_engine = Some("hnsw"); + vr.iter().map(|r| (r.chunk_id, r.score)).collect::>() + } + } else { + Vec::new() + } + } else { + Vec::new() + } + } else { + Vec::new() + }; + + // ── Step 2: Merge via RRF (for hybrid) or use single-mode results ── + let mut candidates: Vec = match mode { + SearchMode::Hybrid if !fts_results.is_empty() || !semantic_results.is_empty() => { + let (rrf_k, fts_w, sem_w) = match &req.score_weights { + Some(sw) => ( + sw.rrf_k as f32, + sw.fts_weight as f32, + sw.semantic_weight as f32, + ), + None => (60.0, 1.0, 1.0), + }; + let merged = hybrid::merge_rrf( + &fts_results, + &semantic_results, + rerank_k, + rrf_k, + fts_w, + sem_w, + ); + merged + .iter() + .map(|r| ScoredCandidate { + chunk_id: r.chunk_id, + base_score: r.rrf_score, + final_score: r.rrf_score, + source: r.source.as_str().to_string(), + }) + .collect() + } + SearchMode::Fts => fts_results + .iter() + .map(|(id, score)| ScoredCandidate { + chunk_id: *id, + base_score: *score, + final_score: *score, + source: "fts".to_string(), + }) + .collect(), + SearchMode::Semantic => semantic_results + .iter() + .map(|(id, score)| ScoredCandidate { + chunk_id: *id, + base_score: *score, + final_score: *score, + source: "semantic".to_string(), + }) + .collect(), + _ => Vec::new(), + }; + + // ── Step 3: Filter is already applied (filter-aware retrieval). ─ + // The bitmap pushdown happens inside both FTS post-filter and + // USearch's filter callback above, so we no longer need a post-merge + // `retain`. Kept as an assertion in debug builds to catch invariant + // drift if a new retrieval path bypasses the eligibility check. + debug_assert!( + !filter_active || candidates.iter().all(|c| eligible.contains(c.chunk_id)), + "filter-aware retrieval produced a candidate outside the eligible bitmap" + ); - // Sort ascending by timerange_start_ms. Segments missing the metadata - // sort to the end (f64::INFINITY) instead of position 0, so callers - // don't see malformed data masquerading as the earliest segment. - // `total_cmp` is NaN-safe and deterministic (NaN sorts after Infinity). - results.sort_by(|a, b| { - let ta = a - .metadata - .get("timerange_start_ms") - .and_then(MetadataValue::as_f64) - .unwrap_or(f64::INFINITY); - let tb = b - .metadata - .get("timerange_start_ms") - .and_then(MetadataValue::as_f64) - .unwrap_or(f64::INFINITY); - ta.total_cmp(&tb) + // ── Step 3b: Drop soft-deleted (tombstoned) chunks ────────────── + // The HNSW/FTS indexes still physically contain deleted ids until the + // next rebuild/compaction, so we filter them out here. Cheap O(1) set + // membership per candidate. + if !loaded.tombstones.is_empty() { + candidates.retain(|c| !loaded.tombstones.contains(&c.chunk_id)); + } + + // ── Step 4: Apply scoring pipeline ────────────────────────────── + // Resolve recency preset into a full config (explicit `recency` wins) + let recency_config = req.recency.clone().or_else(|| { + req.recency_preset.as_deref().and_then(|preset| { + req.recency_field + .as_deref() + .map(|field| RecencyConfig::from_preset(preset, field.to_string())) + .flatten() + }) }); - Ok(results) - } -} + let has_scoring = + recency_config.is_some() || !req.boosts.is_empty() || req.relationship_boost.is_some(); -/// Whether a segment chunk's [timerange_start_ms, timerange_end_ms] window -/// matches the requested time query. All values are in integer milliseconds. -/// -/// - If no time params are provided, returns true (enumeration mode). -/// - If `time_ms` is set, returns true when -/// `timerange_start_ms <= time_ms <= timerange_end_ms`. -/// - If `time_start_ms` and/or `time_end_ms` are set, returns true when the -/// segment's window overlaps the query range. Missing bounds default to -/// ±infinity. -/// - Segments missing `timerange_start_ms` or `timerange_end_ms` are excluded -/// when any time filter is set. -/// - Instants (zero-duration events) are stored as segments where -/// `timerange_start_ms == timerange_end_ms`. A point query at that exact -/// millisecond matches the instant; range queries that overlap that -/// millisecond also match. -pub(crate) fn segment_in_time_window( - chunk: &DocumentChunk, - time_ms: Option, - time_start_ms: Option, - time_end_ms: Option, -) -> bool { - let no_filter = time_ms.is_none() && time_start_ms.is_none() && time_end_ms.is_none(); - if no_filter { - return true; - } - let ts = chunk - .metadata - .get("timerange_start_ms") - .and_then(MetadataValue::as_f64); - let te = chunk - .metadata - .get("timerange_end_ms") - .and_then(MetadataValue::as_f64); - let (s, e) = match (ts, te) { - (Some(s), Some(e)) => (s, e), - _ => return false, - }; - if let Some(t) = time_ms { - return s <= t && t <= e; - } - let lo = time_start_ms.unwrap_or(f64::NEG_INFINITY); - let hi = time_end_ms.unwrap_or(f64::INFINITY); - s <= hi && e >= lo -} + if has_scoring && !candidates.is_empty() { + 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(); -#[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, - } - } + let candidate_ids: Vec = candidates.iter().map(|c| c.chunk_id).collect(); + let (parent_ids, sibling_map) = loaded.relationships.build_scoring_maps(&candidate_ids); - /// 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) - } + scoring::apply_scoring_pipeline( + &mut candidates, + &chunk_metadata, + &parent_ids, + &sibling_map, + &recency_config, + &req.boosts, + &req.relationship_boost, + ); + } - #[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)); - } + // ── Step 5: Truncate to top_k and build response ──────────────── + candidates.truncate(req.top_k); + let total = candidates.len(); - #[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)); - } + // ── Step 5a: Parent metadata enrichment for segment hits ───────── + // For each segment hit with a parent_id, inline the parent's top-level + // metadata so callers (typically AI agents) avoid a second round-trip + // to fetch source-level attributes. Parents are deduplicated: N + // 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.chunk_store); - #[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)); - } + // ── Step 5b: Relation enrichment (opt-in) ─────────────────────── + // When include_relations is set, fetch each hit's edges in ONE batched, + // on-demand read from the disk-backed relation store (never resident in + // RAM). target_status is resolved against the in-memory chunk cache: + // "found" if the target chunk exists locally, else "missing". + let mut relations_by_chunk: HashMap> = HashMap::new(); + if req.include_relations { + let types = req.relation_types.as_deref(); + relations_by_chunk = loaded.relation_store.for_chunks( + &candidate_chunk_ids, + req.relation_direction, + types, + )?; + for edges in relations_by_chunk.values_mut() { + for edge in edges.iter_mut() { + edge.target_status = if loaded.filter_index.contains(edge.target_chunk_id) { + "found".to_string() + } else { + "missing".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))); - } + let hits: Vec<( + DocumentChunk, + f32, + String, + Option>, + Option>, + )> = candidates + .iter() + .filter_map(|c| { + 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, + ) + }) + }) + .collect(); - #[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))); - } + let took_us = start.elapsed().as_micros() as u64; - #[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))); - } + // ── Step 6: Build /explain plan if requested ──────────────────── + let explain_plan = if req.explain { + Some(ExplainPlan { + filter: FilterExplain { + eligible_count: eligible.len(), + universe_count, + selectivity: selectivity_val, + }, + ann: AnnExplain { + engine: explain_engine.unwrap_or("none").to_string(), + candidates_inspected: explain_candidates, + ef_search_used: hnsw_ef_search_default(), + }, + }) + } else { + None + }; - #[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)); + Ok((hits, total, took_us, explain_plan)) } - #[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) - )); - } -} + // ── Chunk Relations ────────────────────────────────────────────────── -/// 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; + /// Create a batch of chunk relations. The server mints a UUIDv4 + /// `relation_id` and stamps `created_at` for each. Self-relations + /// (`source == target`) are rejected. Returns the created edges with their + /// assigned ids. + pub async fn create_relations( + &self, + collection_name: &str, + new: Vec, + ) -> Result, Box> { + self.reject_if_partitioned(collection_name, "creating relations") + .await?; + // 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(|| { + 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()); + } + 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: if loaded.filter_index.contains(r.target_chunk_id) { + "found".to_string() + } else { + "missing".to_string() + }, + metadata: r.metadata, + created_at: now, + }); + } + } // read lock released before S3 I/O. -/// 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 -/// manifest to reference only it. Runs the CAS-retry loop so it converges -/// against concurrent appends. Safe to spawn detached — compaction never mutates -/// the local indexes (they already hold the data). -pub(crate) async fn compact_storage( - storage: &dyn Storage, - ns: &str, -) -> Result { - const MAX_RETRIES: u32 = 10; - 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); + // 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; + let seq = 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}"))?; + appended_seq = Some(seq); } - 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, - &version, - &manifest, - bytes::Bytes::from(segment_bytes), - records, - ) - .await - { - Ok(()) => { - tracing::info!( - "Compacted '{}': {} live records in one segment", - ns, - records - ); - return Ok(records); + // Phase 3 (read lock): apply locally (durable S3 record already written). + // On EITHER failure mode — collection deleted in the lock gap, or a + // local insert error — compensate with a RelationDelete for the minted + // 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 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 + ))), } - Err(crate::storage::StorageError::VersionConflict { .. }) => continue, - Err(e) => return Err(e), + }; + if let Err(e) = apply_result { + if self.cloud_mode && !built.is_empty() { + let ids: Vec = built.iter().map(|r| r.relation_id.clone()).collect(); + if let Err(te) = crate::storage::lsm::append_relation_delete( + self.storage.as_ref(), + collection_name, + &ids, + ) + .await + { + tracing::error!( + "compensation: relation-delete for '{}' failed: {} — orphan \ + relation fragment may resurrect on rebuild", + collection_name, + te + ); + } + } + return Err(e); } + Ok(built) } - Err(crate::storage::StorageError::Io( - "compaction failed after max CAS retries (persistent contention)".into(), - )) -} -/// Spawn a detached background compaction if the cloud collection's uncompacted -/// fragment count is over the threshold. Best-effort: logs and moves on. Called -/// after cloud ingest/delete so deletes actually reclaim space over time. -fn maybe_auto_compact( - storage: Arc, - ns: String, - inflight: Arc>>, -) { - // Single-flight: if a compaction for this ns is already running, skip. This - // stops concurrent triggers from each writing (and, on CAS loss, leaking) a - // full segment. - { - let mut set = inflight.lock().unwrap_or_else(|e| e.into_inner()); - if !set.insert(ns.clone()) { - return; // already compacting this ns + /// Delete a relation by id. Returns true if it existed. + pub async fn delete_relation( + &self, + collection_name: &str, + relation_id: &str, + ) -> Result> { + self.reject_if_partitioned(collection_name, "deleting relations") + .await?; + // 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); } - } - tokio::spawn(async move { - let result = async { - let (manifest, _) = crate::storage::lsm::read_manifest(storage.as_ref(), &ns).await?; - let uncompacted = manifest.uncompacted().count(); - if uncompacted >= AUTO_COMPACT_FRAGMENT_THRESHOLD { - tracing::info!("Auto-compacting '{}' ({} fragments)", ns, uncompacted); - compact_storage(storage.as_ref(), &ns).await?; + 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(not_found(format_args!( + "Collection \'{}\' not found", + collection_name + ))); } - Ok::<(), crate::storage::StorageError>(()) } - .await; - if let Err(e) = result { - tracing::warn!("auto-compaction of '{}' failed: {}", ns, e); + + // 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 { + 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); } - // Clear the in-flight flag so a later trigger can run. - inflight - .lock() - .unwrap_or_else(|e| e.into_inner()) - .remove(&ns); - }); -} -/// Build the filter index from the chunk map, EXCLUDING tombstoned ids — so -/// `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; + // 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 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); + 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 } - idx.finalize(); - idx -} -/// Build a deduplicated cache of parent chunk metadata for a set of candidate -/// chunk ids. Used by `search()` to enrich segment hits with their parent's -/// top-level metadata without paying for repeated lookups when multiple -/// segments share the same parent. -/// -/// N segments pointing at the same parent trigger exactly one HashMap lookup. -/// Candidates that are not segments, or are segments without a `parent_id`, -/// contribute nothing to the cache. -/// -/// Orphan parents (segment has a `parent_id` but the parent chunk is not in -/// `chunks`) are NOT inserted into the cache. This means `parent_metadata_for` -/// returns `None` for them, which lets callers distinguish "no parent at all" -/// from "parent exists with empty metadata." -pub(crate) fn build_parent_metadata_cache( - candidate_chunk_ids: &[u64], - chunks: &HashMap, -) -> HashMap> { - let mut cache: HashMap> = HashMap::new(); - for cid in candidate_chunk_ids { - let Some(chunk) = chunks.get(cid) else { - continue; - }; - if chunk.doc_type != "segment" { - continue; - } - let Some(pid) = chunk.parent_id else { - continue; - }; - if cache.contains_key(&pid) { - continue; + /// List a single chunk's relations, with `target_status` resolved against + /// the current chunk set. + pub async fn get_chunk_relations( + &self, + collection_name: &str, + chunk_id: u64, + direction: RelationDirection, + types: Option<&[String]>, + ) -> Result, Box> { + self.reject_if_partitioned(collection_name, "listing relations") + .await?; + if self.role == NodeRole::Writer { + return Err("this node runs in writer role and does not serve queries".into()); } - // 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) { - cache.insert(pid, parent.metadata.clone()); + self.ensure_attached(collection_name).await?; + 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); + let mut edges = loaded + .relation_store + .for_chunk(chunk_id, direction, types)?; + for edge in edges.iter_mut() { + edge.target_status = if loaded.filter_index.contains(edge.target_chunk_id) { + "found".to_string() + } else { + "missing".to_string() + }; } + Ok(edges) } - cache -} -/// Look up parent metadata for a given chunk from a pre-built cache. -/// -/// Returns `None` when the chunk is not a segment or has no `parent_id`. -/// Returns `Some(metadata)` (possibly empty) when the chunk is a segment -/// whose `parent_id` was included in the cache. -pub(crate) fn parent_metadata_for( - chunk: &DocumentChunk, - cache: &HashMap>, -) -> Option> { - if chunk.doc_type != "segment" { - return None; + // ── Delete (soft-delete via tombstones) ────────────────────────────── + + /// Soft-delete a set of chunk ids. Tombstones them (so they immediately + /// vanish from search results), persists the tombstones, prunes incident + /// relations, and — in object-storage mode — appends a tombstone WAL + /// 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. + /// 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 — + // before the fallible relation pruning — so a pruning error can't leave + // chunk_count permanently overstated. + let removed = apply.len() as u64; + loaded.metadata.chunk_count = loaded.metadata.chunk_count.saturating_sub(removed); + store::save_metadata(data_dir, &loaded.metadata)?; + // Keep the filter index in step with the tombstones so `eligible` / + // 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 { + let edges = loaded + .relation_store + .for_chunk(id, RelationDirection::Both, None)?; + for e in edges { + loaded.relation_store.delete(&e.relation_id)?; + } + } + Ok(()) } - chunk.parent_id.and_then(|pid| cache.get(&pid).cloned()) -} -#[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, + /// 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, + 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(()) + } } } - fn into_map(chunks: Vec) -> HashMap { - chunks.into_iter().map(|c| (c.id, c)).collect() + /// 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() } - #[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())) + /// 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, + ns: &str, + ) -> Result<(), Box> { + // Partition namespaces attach on demand EVEN in non-lazy cloud mode: + // partitions appear dynamically (a writer node can mint one at any + // time), so "everything attached at boot" can never hold for them. + let dynamic_partition = self.cloud_mode && partitions::is_partition_ns(ns); + if !self.lazy_attach && !dynamic_partition { + 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, ); - } - - #[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" + tracing::info!( + "Attached '{}' on demand ({} chunks in {:.2}s)", + ns, + n, + start.elapsed().as_secs_f64() ); - 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())) - ); - } + self.maybe_evict_lru(ns).await; + Ok(()) } - #[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, + /// 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. In NON-lazy mode only + // dynamic partition namespaces may be evicted — a normal collection + // evicted there could never re-attach (ensure_attached early-returns + // for non-partitions when lazy attach is off). + 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) + .filter(|(name, _)| self.lazy_attach || partitions::is_partition_ns(name)) + .min_by_key(|(_, l)| l.last_used.load(std::sync::atomic::Ordering::Relaxed)) + .map(|(name, _)| name.clone()) + } }; - 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()); + 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); } -} -#[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). + /// 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 { - 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 + 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 + ); + } + } } - // 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(); + 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 + }; - 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" - ); + 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); + } - // Cleanup - let _ = std::fs::remove_dir_all(&data_dir); + // 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) } - #[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(); + /// 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); + } + } - // 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(); + /// 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); + } } + } - // 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, + pub async fn delete_chunks( + &self, + collection_name: &str, + ids: &[u64], + ) -> Result<(usize, Option), Box> { + crate::metrics::inc(&crate::metrics::DELETE_REQUESTS_TOTAL); + // Ids alone don't say which partition holds them (a fan-out probe of + // every partition would be unbounded) — partitioned collections + // delete via POST /delete with the partition filter. This applies to + // writer nodes too: a tombstone appended to the PARENT namespace + // would never be materialized by any serving node. + if self.is_partitioned_any_role(collection_name).await? { + let hint = if self.role == NodeRole::Writer { + "route the delete through a serving node (writers cannot resolve \ + partition-scoped filters)" + } else { + "delete via filters (POST .../delete with the partition field), \ + not bare ids" + }; + return Err(format!("collection '{collection_name}' is partitioned: {hint}").into()); + } + // 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(), + partitions::alloc_ns(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 - .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 - ); + .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?; - let _ = std::fs::remove_dir_all(&data_dir); - } -} + // 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)); + } -#[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:?}" + // 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)) } - #[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()); + /// 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(), + ); + } + // Partitioned collection: route to the partitions named by the filter + // (same routing rule as search; NotFound partitions delete nothing). + if let Some(field) = self.partition_field(collection_name).await? { + let pvals = partitions::partition_values_from_filters(&field, filters)?; + let multi = pvals.len() > 1; + let mut total = 0usize; + let mut last_seq = None; + for pval in pvals { + let ns = partitions::partition_ns(collection_name, &pval); + match Box::pin(self.delete_by_filter(&ns, filters)).await { + Ok((n, seq)) => { + total += n; + last_seq = seq; + } + Err(e) if e.downcast_ref::().is_some() => continue, + Err(e) => return Err(e), + } + } + return Ok((total, if multi { None } else { last_seq })); + } + 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 } -} -#[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(); + // ── Cloud compaction (object-storage mode) ─────────────────────────── - // 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)); + /// 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 + ))); + } } - 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, + Ok(compact_storage(self.storage.as_ref(), collection_name).await?) + } + + /// Rebuild a collection's LOCAL indexes from its object-storage manifest + /// (materialize segments + WAL fragments → chunks → local redb/Tantivy/HNSW/ + /// filter index). This is cold-start recovery: an ephemeral node with an + /// empty local disk reconstructs the collection entirely from S3. Returns the + /// number of live chunks recovered. Cloud mode only. + pub async fn rebuild_collection_from_storage( + &self, + collection_name: &str, + ) -> Result> { + validate_name_segment(collection_name, "Collection")?; + let (manifest, _) = + crate::storage::lsm::read_manifest(self.storage.as_ref(), collection_name).await?; + let materialized = + cloud::materialize(self.storage.as_ref(), collection_name, &manifest).await?; + let chunks: Vec = materialized.chunks.values().cloned().collect(); + let live_count = chunks.len(); + + // 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(), + ) + } + }; + // 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 + // `+ if live_count > 0` form reused the highest id when every chunk was + // deleted. + let next_id = if materialized.max_id > 0 || live_count > 0 { + materialized.max_id + 1 + } else { + 0 }; - 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 metadata = Collection { + name: collection_name.to_string(), + 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: coll_config, + applied_seq: manifest.next_seq, }; - 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(), - }, - ], + 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 - .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()); + { + if !matches!(e, crate::storage::StorageError::AlreadyExists(_)) { + tracing::warn!( + "bucket config back-fill for '{}' failed: {}", + collection_name, + e + ); + } + } + } - // 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, - }; + // 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 + // local redb (e.g. a stale chunk left by a partially-committed ingest that + // was later compensated with a tombstone) must not survive. Wipe first. + let chunks_db = store::chunks_db_path(&self.data_dir, collection_name); + if let Some(parent) = chunks_db.parent() { + std::fs::create_dir_all(parent)?; + } + let _ = std::fs::remove_file(&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)?; - // 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())); + // FTS index. + let tantivy_dir = store::tantivy_dir(&self.data_dir, collection_name); + let fts = tantivy_fts::build_index(&tantivy_dir, &chunks)?; - // 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); + // Vector spaces (HNSW) from embeddings. + let vectors_dir = store::vectors_dir(&self.data_dir, collection_name); + let mut vs_map: HashMap> = HashMap::new(); + for (space, cfg) in &vector_spaces { + let mut ids = Vec::new(); + let mut vecs = Vec::new(); + for c in &chunks { + if let Some(emb) = c.embeddings.get(space) { + ids.push(c.id); + vecs.push(emb.clone()); + } + } + let index_path = vectors_dir.join(format!("{}.index", space)); + let vecs_path = vectors_dir.join(format!("{}.bin", space)); + let vs = vector::build_vector_index(&index_path, &vecs_path, &ids, &vecs, cfg.dims)?; + vs_map.insert(space.clone(), Arc::new(vs)); + } - let _ = std::fs::remove_dir_all(&data_dir); - } + // Relationships + filter index from the recovered chunks. + let mut relationships = RelationshipStore::new(); + for c in &chunks { + relationships.add(c.id, c.parent_id, c.group_id.clone()); + } + // 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)); + } - #[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(); + // Reconstruct the typed-relation store from the materialized relations + // (recovered from the S3 WAL/segments) — so relations survive a cold + // restart, not just chunks. + let relations_db = store::relations_db_path(&self.data_dir, collection_name); + let _ = std::fs::remove_file(&relations_db); // start clean, then repopulate + let relation_store = RelationStore::open(&relations_db)?; + let recovered_relations: Vec = + materialized.relations.values().cloned().collect(); + if !recovered_relations.is_empty() { + relation_store.insert_batch(&recovered_relations)?; + } - { - 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" - ); + 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, + chunk_store, + relation_store, + tombstones: std::collections::HashSet::new(), + next_id, + filter_index, + }; + let mut collections = self.collections.write().await; + collections.insert(collection_name.to_string(), loaded); + Ok(live_count) + } - // 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; + /// Get facet counts for a collection. + pub async fn get_facets( + &self, + collection_name: &str, + query: &str, + fields: &[String], + ) -> Result< + (HashMap>, u64), + Box, + > { + if self.role == NodeRole::Writer { + return Err("this node runs in writer role and does not serve queries".into()); } + self.reject_if_partitioned(collection_name, "facet counting") + .await?; + self.ensure_attached(collection_name).await?; + 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); + tantivy_fts::get_facets(&loaded.fts, query, fields, loaded.filter_index.universe()) + } - // 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() - ); - } + /// Get all chunk texts and IDs for rebuild jobs. + pub async fn get_all_chunk_data( + &self, + collection_name: &str, + ) -> Result<(Vec, Vec), Box> { + 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 _ = std::fs::remove_dir_all(&data_dir); + let mut texts = Vec::new(); + let mut ids = Vec::new(); + loaded.chunk_store.for_each(|id, chunk| { + if !loaded.tombstones.contains(&id) { + ids.push(id); + texts.push(chunk.text); + } + })?; + Ok((texts, ids)) } - // 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(); + /// Temporal point/range lookup for TAMS-style segments. + /// + /// Returns every chunk where `doc_type == "segment"`, `group_id == Some(asset)`, + /// and the time window matches the requested query. If no time params are + /// provided, returns all segments for the asset (enumeration mode). + /// + /// Time unit: all parameters and the `timerange_start_ms` / `timerange_end_ms` + /// metadata fields are in integer milliseconds. Results are sorted ascending + /// by `timerange_start_ms` for stable ordering. + pub async fn segments_at( + &self, + collection_name: &str, + asset: &str, + time_ms: Option, + time_start_ms: Option, + time_end_ms: Option, + ) -> Result, Box> { + self.reject_if_partitioned(collection_name, "temporal segment lookup") + .await?; + let collections = self.collections.read().await; + let loaded = collections.get(collection_name).ok_or_else(|| { + not_found(format_args!("Collection \'{}\' not found", collection_name)) + })?; - // Delete chunk 0 — both edges (as source and as target) must be pruned. - manager.delete_chunks("delrel", &[0]).await.unwrap(); + 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(); - 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(); + // Sort ascending by timerange_start_ms. Segments missing the metadata + // sort to the end (f64::INFINITY) instead of position 0, so callers + // don't see malformed data masquerading as the earliest segment. + // `total_cmp` is NaN-safe and deterministic (NaN sorts after Infinity). + results.sort_by(|a, b| { + let ta = a + .metadata + .get("timerange_start_ms") + .and_then(MetadataValue::as_f64) + .unwrap_or(f64::INFINITY); + let tb = b + .metadata + .get("timerange_start_ms") + .and_then(MetadataValue::as_f64) + .unwrap_or(f64::INFINITY); + ta.total_cmp(&tb) + }); - // 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); + Ok(results) } } -#[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(); +/// Whether a segment chunk's [timerange_start_ms, timerange_end_ms] window +/// matches the requested time query. All values are in integer milliseconds. +/// +/// - If no time params are provided, returns true (enumeration mode). +/// - If `time_ms` is set, returns true when +/// `timerange_start_ms <= time_ms <= timerange_end_ms`. +/// - If `time_start_ms` and/or `time_end_ms` are set, returns true when the +/// segment's window overlaps the query range. Missing bounds default to +/// ±infinity. +/// - Segments missing `timerange_start_ms` or `timerange_end_ms` are excluded +/// when any time filter is set. +/// - Instants (zero-duration events) are stored as segments where +/// `timerange_start_ms == timerange_end_ms`. A point query at that exact +/// millisecond matches the instant; range queries that overlap that +/// millisecond also match. +pub(crate) fn segment_in_time_window( + chunk: &DocumentChunk, + time_ms: Option, + time_start_ms: Option, + time_end_ms: Option, +) -> bool { + let no_filter = time_ms.is_none() && time_start_ms.is_none() && time_end_ms.is_none(); + if no_filter { + return true; + } + let ts = chunk + .metadata + .get("timerange_start_ms") + .and_then(MetadataValue::as_f64); + let te = chunk + .metadata + .get("timerange_end_ms") + .and_then(MetadataValue::as_f64); + let (s, e) = match (ts, te) { + (Some(s), Some(e)) => (s, e), + _ => return false, + }; + if let Some(t) = time_ms { + return s <= t && t <= e; + } + let lo = time_start_ms.unwrap_or(f64::NEG_INFINITY); + let hi = time_end_ms.unwrap_or(f64::INFINITY); + s <= hi && e >= lo +} - // 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); +#[cfg(test)] +mod segments_at_tests; + +/// 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); + +impl std::fmt::Display for NotFound { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(&self.0) } +} - #[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); +impl std::error::Error for NotFound {} - // 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]); +fn not_found(what: impl std::fmt::Display) -> Box { + Box::new(NotFound(what.to_string())) +} - let _ = std::fs::remove_dir_all(&data_dir); - } +/// 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; - // 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()); +/// 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 +/// manifest to reference only it. Runs the CAS-retry loop so it converges +/// against concurrent appends. Safe to spawn detached — compaction never mutates +/// the local indexes (they already hold the data). +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; - 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, + // 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?; + 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_v3(&segment)?; + match crate::storage::lsm::append_segment( + storage, + ns, + &version, + &manifest, + bytes::Bytes::from(bytes), + records, + folded_through, ) .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); + { + 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), + } } - // #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(); + // 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, + &version, + &manifest, + bytes::Bytes::from(segment_bytes), + records, + ) + .await { - 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(); + Ok(()) => { + tracing::info!("Merged '{}' segments: {} live records", ns, records); + return Ok(records); + } + Err(crate::storage::StorageError::VersionConflict { .. }) => continue, + Err(e) => return Err(e), } + } + Err(crate::storage::StorageError::Io( + "compaction failed after max CAS retries (persistent contention)".into(), + )) +} - // 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; +/// Spawn a detached background compaction if the cloud collection's uncompacted +/// fragment count is over the threshold. Best-effort: logs and moves on. Called +/// after cloud ingest/delete so deletes actually reclaim space over time. +fn maybe_auto_compact( + storage: Arc, + ns: String, + inflight: Arc>>, +) { + // Single-flight: if a compaction for this ns is already running, skip. This + // stops concurrent triggers from each writing (and, on CAS loss, leaking) a + // full segment. + { + let mut set = inflight.lock().unwrap_or_else(|e| e.into_inner()); + if !set.insert(ns.clone()) { + return; // already compacting this ns + } + } + tokio::spawn(async move { + let result = async { + let (manifest, _) = crate::storage::lsm::read_manifest(storage.as_ref(), &ns).await?; + let uncompacted = manifest.uncompacted().count(); + if uncompacted >= AUTO_COMPACT_FRAGMENT_THRESHOLD { + tracing::info!("Auto-compacting '{}' ({} fragments)", ns, uncompacted); + compact_storage(storage.as_ref(), &ns).await?; } - tokio::time::sleep(std::time::Duration::from_millis(100)).await; + Ok::<(), crate::storage::StorageError>(()) } - 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(); + .await; + if let Err(e) = result { + tracing::warn!("auto-compaction of '{}' failed: {}", ns, e); } - // 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(); + // Clear the in-flight flag so a later trigger can run. + inflight + .lock() + .unwrap_or_else(|e| e.into_inner()) + .remove(&ns); + }); +} - // 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(); +/// Build the filter index from the chunk map, EXCLUDING tombstoned ids — so +/// `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. +/// 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) +} - // 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); +/// 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 +} - let _ = std::fs::remove_dir_all(&data_dir); +/// Build a deduplicated cache of parent chunk metadata for a set of candidate +/// chunk ids. Used by `search()` to enrich segment hits with their parent's +/// top-level metadata without paying for repeated lookups when multiple +/// segments share the same parent. +/// +/// N segments pointing at the same parent trigger exactly one HashMap lookup. +/// Candidates that are not segments, or are segments without a `parent_id`, +/// contribute nothing to the cache. +/// +/// Orphan parents (segment has a `parent_id` but the parent chunk is not in +/// `chunks`) are NOT inserted into the cache. This means `parent_metadata_for` +/// returns `None` for them, which lets callers distinguish "no parent at all" +/// from "parent exists with empty metadata." +pub(crate) fn build_parent_metadata_cache( + candidate_chunk_ids: &[u64], + chunks: &ChunkCache, +) -> HashMap> { + let mut cache: HashMap> = HashMap::new(); + for cid in candidate_chunk_ids { + let Ok(Some(chunk)) = chunks.get(*cid) else { + continue; + }; + if chunk.doc_type != "segment" { + continue; + } + let Some(pid) = chunk.parent_id else { + continue; + }; + if cache.contains_key(&pid) { + continue; + } + // Only cache parents that actually exist. Missing parents stay out + // of the cache so `parent_metadata_for` returns None for them. + if let Ok(Some(parent)) = chunks.get(pid) { + cache.insert(pid, parent.metadata.clone()); + } } + cache +} - // 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(); +/// Look up parent metadata for a given chunk from a pre-built cache. +/// +/// Returns `None` when the chunk is not a segment or has no `parent_id`. +/// Returns `Some(metadata)` (possibly empty) when the chunk is a segment +/// whose `parent_id` was included in the cache. +pub(crate) fn parent_metadata_for( + chunk: &DocumentChunk, + cache: &HashMap>, +) -> Option> { + if chunk.doc_type != "segment" { + return None; + } + chunk.parent_id.and_then(|pid| cache.get(&pid).cloned()) +} - 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(); - } +#[cfg(test)] +mod parent_metadata_tests; - // 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(); +#[cfg(test)] +mod persistence_tests; - // 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" - ); +#[cfg(test)] +mod validate_name_segment_tests; - 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 - ); +#[cfg(test)] +mod filter_aware_search_tests; - let _ = std::fs::remove_dir_all(&data_dir); - } -} +#[cfg(all(test, feature = "object-storage"))] +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/partition_cloud_tests.rs b/crates/compass/src/collections/partition_cloud_tests.rs new file mode 100644 index 0000000..49f2884 --- /dev/null +++ b/crates/compass/src/collections/partition_cloud_tests.rs @@ -0,0 +1,272 @@ +// collections/partition_cloud_tests.rs — Phase 6 cloud-mode coverage: writer +// routing, cross-node partition discovery, and cold rebuild. Uses the +// in-memory object_store backend (same code path as real S3). + +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-partition-cloud-{}-{}", + std::process::id(), + N.fetch_add(1, Ordering::Relaxed) + )) +} + +fn embed_state() -> EmbedState { + EmbedState { + bge: None, + distilled: None, + } +} + +fn mem_storage() -> Arc { + Arc::new(ObjectStoreBackend::from_store( + std::sync::Arc::new(object_store::memory::InMemory::new()), + "object-store:memory", + )) +} + +fn tenant_chunk(tenant: &str, file: &str, text: &str, vec: [f32; 4]) -> IngestChunk { + let mut metadata = HashMap::new(); + metadata.insert( + "tenant".to_string(), + MetadataValue::String(tenant.to_string()), + ); + let mut embeddings = HashMap::new(); + embeddings.insert("default".to_string(), vec.to_vec()); + IngestChunk { + client_id: None, + file_id: file.to_string(), + chunk_index: 0, + page: None, + text: text.to_string(), + metadata, + doc_type: "chunk".to_string(), + parent_id: None, + parent_ref: None, + group_id: None, + embeddings, + embedding: None, + } +} + +fn partitioned_config() -> Option { + Some(CollectionConfig { + embed_model: "bge-small".to_string(), + partition_by: Some("tenant".to_string()), + }) +} + +fn search_req(query: &str, tenant: &str) -> SearchRequest { + let mut filters = HashMap::new(); + filters.insert( + "tenant".to_string(), + FilterValue::Exact(MetadataValue::String(tenant.to_string())), + ); + SearchRequest { + query: query.to_string(), + mode: "fts".to_string(), + vector_space: None, + top_k: 10, + query_vector: None, + filters, + 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::default(), + min_seq: None, + } +} + +// A stateless writer routes partitioned ingest into per-partition namespaces +// it bootstraps itself; a serving node that booted BEFORE those partitions +// existed attaches them on demand and serves the data. +#[tokio::test] +async fn writer_partitioned_ingest_visible_on_serving_node() { + let storage = mem_storage(); + let embed = embed_state(); + + // Serving node boots first and creates the (empty) partitioned parent. + let serve_dir = unique_data_dir(); + std::fs::create_dir_all(&serve_dir).unwrap(); + let serving = CollectionManager::new_with_storage(&serve_dir, storage.clone()) + .await + .unwrap(); + serving + .create_collection("wp", None, Some(4), partitioned_config()) + .await + .unwrap(); + + // Writer node ingests for two tenants the serving node has never seen. + let writer_dir = unique_data_dir(); + std::fs::create_dir_all(&writer_dir).unwrap(); + let writer = CollectionManager::new_with_storage_opts( + &writer_dir, + storage.clone(), + NodeRole::Writer, + false, + usize::MAX, + 0, + ) + .await + .unwrap(); + let (n, id_map, _) = writer + .ingest( + "wp", + vec![ + { + let mut c = + tenant_chunk("acme", "a1", "durable acme fact", [0.9, 0.1, 0.0, 0.0]); + c.client_id = Some("a1".to_string()); + c + }, + { + let mut c = + tenant_chunk("globex", "g1", "durable globex fact", [0.1, 0.9, 0.0, 0.0]); + c.client_id = Some("g1".to_string()); + c + }, + ], + &embed, + ) + .await + .unwrap(); + assert_eq!(n, 2); + let ids: std::collections::HashSet = id_map.values().copied().collect(); + assert_eq!(ids.len(), 2, "writer ids must be collection-unique"); + + // The serving node attaches the new partitions on first query. + let (results, _, _, _) = serving + .search("wp", &search_req("durable", "acme"), &embed) + .await + .unwrap(); + assert_eq!(results.len(), 1, "acme partition attaches on demand"); + assert_eq!(results[0].0.file_id, "a1"); + let (results, _, _, _) = serving + .search("wp", &search_req("durable", "globex"), &embed) + .await + .unwrap(); + assert_eq!(results.len(), 1); + assert_eq!(results[0].0.file_id, "g1"); + + // Writer refuses partitioned operations that would black-hole data. + let err = writer.delete_chunks("wp", &[0]).await.unwrap_err(); + assert!( + err.to_string() + .contains("route the delete through a serving node"), + "{err}" + ); + let err = writer + .create_relations( + "wp", + vec![CreateRelation { + source_chunk_id: 0, + target_chunk_id: 1, + target_document_id: None, + relation_type: "cites".to_string(), + metadata: HashMap::new(), + }], + ) + .await + .unwrap_err(); + assert!( + err.to_string().contains("not supported on a partitioned"), + "{err}" + ); + + let _ = std::fs::remove_dir_all(&serve_dir); + let _ = std::fs::remove_dir_all(&writer_dir); +} + +// A brand-new node with an empty disk recovers a partitioned collection — +// parent config, every partition's data, and the shared id allocator — from +// the bucket alone. +#[tokio::test] +async fn partitioned_collection_cold_rebuild_from_bucket() { + let storage = mem_storage(); + let embed = embed_state(); + + { + let dir_a = unique_data_dir(); + std::fs::create_dir_all(&dir_a).unwrap(); + let a = CollectionManager::new_with_storage(&dir_a, storage.clone()) + .await + .unwrap(); + a.create_collection("cold", None, Some(4), partitioned_config()) + .await + .unwrap(); + a.ingest( + "cold", + vec![ + tenant_chunk("acme", "a1", "cold acme", [0.9, 0.1, 0.0, 0.0]), + tenant_chunk("globex", "g1", "cold globex", [0.1, 0.9, 0.0, 0.0]), + ], + &embed, + ) + .await + .unwrap(); + let _ = std::fs::remove_dir_all(&dir_a); + } // node A gone, local disk gone; only the bucket remains. + + let dir_b = unique_data_dir(); + std::fs::create_dir_all(&dir_b).unwrap(); + let b = CollectionManager::new_with_storage(&dir_b, storage.clone()) + .await + .unwrap(); + + // Parent is listed (partitions hidden), routing metadata survived. + let listed = b.list_collections().await; + let names: Vec<&str> = listed.iter().map(|c| c.name.as_str()).collect(); + assert!(names.contains(&"cold"), "{names:?}"); + assert!(!names.iter().any(|n| n.contains(partitions::PART_SEP))); + + let (results, _, _, _) = b + .search("cold", &search_req("cold", "acme"), &embed) + .await + .unwrap(); + assert_eq!(results.len(), 1); + assert_eq!(results[0].0.file_id, "a1"); + + // New ingest keeps minting collection-unique ids from the recovered + // allocator (never reuses the pre-rebuild ids). + let (_, id_map, _) = b + .ingest( + "cold", + vec![{ + let mut c = tenant_chunk("acme", "a2", "post-rebuild", [0.7, 0.3, 0.0, 0.0]); + c.client_id = Some("a2".to_string()); + c + }], + &embed, + ) + .await + .unwrap(); + let new_id = *id_map.values().next().unwrap(); + assert!( + new_id >= 2, + "rebuilt node must not reuse ids (got {new_id})" + ); + + // Cascade delete purges parent + partitions from the bucket. + b.delete_collection("cold").await.unwrap(); + let remaining = crate::storage::lsm::list_namespaces(storage.as_ref()) + .await + .unwrap(); + assert!( + remaining.iter().all(|n| !n.starts_with("cold")), + "bucket still has: {remaining:?}" + ); + + let _ = std::fs::remove_dir_all(&dir_b); +} diff --git a/crates/compass/src/collections/partition_tests.rs b/crates/compass/src/collections/partition_tests.rs new file mode 100644 index 0000000..0b86039 --- /dev/null +++ b/crates/compass/src/collections/partition_tests.rs @@ -0,0 +1,405 @@ +// collections/partition_tests.rs — tenant-partitioned collections (Phase 6). +// Child module of `collections`: `super::*` sees the parent's private items. + +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); + std::env::temp_dir().join(format!( + "compass-partition-{}-{}", + std::process::id(), + N.fetch_add(1, Ordering::Relaxed) + )) +} + +fn embed_state() -> EmbedState { + EmbedState { + bge: None, + distilled: None, + } +} + +fn tenant_chunk(tenant: &str, file: &str, text: &str, vec: [f32; 4]) -> IngestChunk { + let mut metadata = HashMap::new(); + metadata.insert( + "tenant".to_string(), + MetadataValue::String(tenant.to_string()), + ); + let mut embeddings = HashMap::new(); + embeddings.insert("default".to_string(), vec.to_vec()); + IngestChunk { + client_id: None, + file_id: file.to_string(), + chunk_index: 0, + page: None, + text: text.to_string(), + metadata, + doc_type: "chunk".to_string(), + parent_id: None, + parent_ref: None, + group_id: None, + embeddings, + embedding: None, + } +} + +fn partitioned_config() -> Option { + Some(CollectionConfig { + embed_model: "bge-small".to_string(), + partition_by: Some("tenant".to_string()), + }) +} + +fn tenant_filter(tenant: &str) -> HashMap { + let mut f = HashMap::new(); + f.insert( + "tenant".to_string(), + FilterValue::Exact(MetadataValue::String(tenant.to_string())), + ); + f +} + +fn search_req(query: &str, vec: Option<[f32; 4]>, tenant: &str) -> SearchRequest { + SearchRequest { + query: query.to_string(), + mode: if vec.is_some() { + "semantic".to_string() + } else { + "fts".to_string() + }, + vector_space: None, + top_k: 10, + query_vector: vec.map(|v| v.to_vec()), + filters: tenant_filter(tenant), + 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::default(), + min_seq: None, + } +} + +async fn local_manager(data_dir: &std::path::Path) -> std::sync::Arc { + std::fs::create_dir_all(data_dir).unwrap(); + CollectionManager::new(data_dir).await.unwrap() +} + +// ── Local mode ────────────────────────────────────────────────────────── + +#[tokio::test] +async fn partitioned_ingest_routes_and_isolates() { + let data_dir = unique_data_dir(); + let manager = local_manager(&data_dir).await; + let embed = embed_state(); + manager + .create_collection("multi", None, Some(4), partitioned_config()) + .await + .unwrap(); + + let (n, _, _) = manager + .ingest( + "multi", + vec![ + tenant_chunk("acme", "a1", "acme secret report", [0.9, 0.1, 0.0, 0.0]), + tenant_chunk("acme", "a2", "acme quarterly numbers", [0.8, 0.2, 0.0, 0.0]), + tenant_chunk("globex", "g1", "globex secret memo", [0.1, 0.9, 0.0, 0.0]), + ], + &embed, + ) + .await + .unwrap(); + assert_eq!(n, 3); + + // Tenant-scoped search sees ONLY that tenant's chunks — even for a query + // term both tenants share. + let (results, _, _, _) = manager + .search("multi", &search_req("secret", None, "acme"), &embed) + .await + .unwrap(); + assert_eq!( + results.len(), + 1, + "acme must see exactly its own 'secret' hit" + ); + assert_eq!(results[0].0.file_id, "a1"); + + let (results, _, _, _) = manager + .search("multi", &search_req("secret", None, "globex"), &embed) + .await + .unwrap(); + assert_eq!(results.len(), 1); + assert_eq!(results[0].0.file_id, "g1"); + + // A tenant that never ingested is empty, not an error. + let (results, total, _, _) = manager + .search("multi", &search_req("secret", None, "initech"), &embed) + .await + .unwrap(); + assert!(results.is_empty()); + assert_eq!(total, 0); + + // Unfiltered search on a partitioned collection is a clear error. + let mut req = search_req("secret", None, "acme"); + req.filters.clear(); + let err = manager.search("multi", &req, &embed).await.unwrap_err(); + assert!(err.to_string().contains("partitioned by 'tenant'"), "{err}"); + + let _ = std::fs::remove_dir_all(&data_dir); +} + +#[tokio::test] +async fn partition_ids_are_collection_unique() { + let data_dir = unique_data_dir(); + let manager = local_manager(&data_dir).await; + let embed = embed_state(); + manager + .create_collection("uniq", None, Some(4), partitioned_config()) + .await + .unwrap(); + + // Interleave ingests across tenants; every assigned id must be distinct + // (partitions share the parent's id allocator). + let mut all_ids = Vec::new(); + for round in 0..3 { + for tenant in ["t-a", "t-b", "t-c"] { + let (_, id_map, _) = manager + .ingest( + "uniq", + vec![{ + let mut c = tenant_chunk( + tenant, + &format!("{tenant}-{round}"), + "payload", + [0.5, 0.5, 0.0, 0.0], + ); + c.client_id = Some(format!("{tenant}-{round}")); + c + }], + &embed, + ) + .await + .unwrap(); + all_ids.extend(id_map.values().copied()); + } + } + assert_eq!(all_ids.len(), 9); + let unique: std::collections::HashSet = all_ids.iter().copied().collect(); + assert_eq!(unique.len(), 9, "ids must never collide across partitions"); + + let _ = std::fs::remove_dir_all(&data_dir); +} + +#[tokio::test] +async fn partitioned_delete_routes_by_filter_and_rejects_bare_ids() { + let data_dir = unique_data_dir(); + let manager = local_manager(&data_dir).await; + let embed = embed_state(); + manager + .create_collection("deltest", None, Some(4), partitioned_config()) + .await + .unwrap(); + manager + .ingest( + "deltest", + vec![ + tenant_chunk("acme", "a1", "doomed", [0.9, 0.1, 0.0, 0.0]), + tenant_chunk("globex", "g1", "doomed", [0.1, 0.9, 0.0, 0.0]), + ], + &embed, + ) + .await + .unwrap(); + + // Bare ids are ambiguous across partitions — rejected with guidance. + let err = manager.delete_chunks("deltest", &[0]).await.unwrap_err(); + assert!(err.to_string().contains("delete via filters"), "{err}"); + + // Filter-scoped delete removes acme's chunk only. + let (n, _) = manager + .delete_by_filter("deltest", &tenant_filter("acme")) + .await + .unwrap(); + assert_eq!(n, 1); + let (results, _, _, _) = manager + .search("deltest", &search_req("doomed", None, "acme"), &embed) + .await + .unwrap(); + assert!(results.is_empty(), "acme's chunk is gone"); + let (results, _, _, _) = manager + .search("deltest", &search_req("doomed", None, "globex"), &embed) + .await + .unwrap(); + assert_eq!(results.len(), 1, "globex is untouched"); + + let _ = std::fs::remove_dir_all(&data_dir); +} + +#[tokio::test] +async fn partitions_hidden_from_listing_and_cascade_deleted() { + let data_dir = unique_data_dir(); + let manager = local_manager(&data_dir).await; + let embed = embed_state(); + manager + .create_collection("casc", None, Some(4), partitioned_config()) + .await + .unwrap(); + manager + .ingest( + "casc", + vec![ + tenant_chunk("acme", "a1", "x", [0.9, 0.1, 0.0, 0.0]), + tenant_chunk("globex", "g1", "y", [0.1, 0.9, 0.0, 0.0]), + ], + &embed, + ) + .await + .unwrap(); + + // Listing shows the parent only — partition namespaces are internal. + let listed = manager.list_collections().await; + let names: Vec<&str> = listed.iter().map(|c| c.name.as_str()).collect(); + assert!(names.contains(&"casc")); + assert!( + !names.iter().any(|n| n.contains(partitions::PART_SEP)), + "partition namespaces must not be listed: {names:?}" + ); + + // Deleting the parent removes every partition's data on disk. + manager.delete_collection("casc").await.unwrap(); + let leftovers: Vec = std::fs::read_dir(&data_dir) + .map(|rd| { + rd.flatten() + .filter_map(|e| e.file_name().to_str().map(String::from)) + .filter(|n| n.starts_with("casc")) + .collect() + }) + .unwrap_or_default(); + assert!( + leftovers.is_empty(), + "cascade left dirs behind: {leftovers:?}" + ); + + let _ = std::fs::remove_dir_all(&data_dir); +} + +#[tokio::test] +async fn partitioned_fences_and_validation() { + let data_dir = unique_data_dir(); + let manager = local_manager(&data_dir).await; + let embed = embed_state(); + manager + .create_collection("fenced", None, Some(4), partitioned_config()) + .await + .unwrap(); + manager + .ingest( + "fenced", + vec![tenant_chunk("acme", "a1", "x", [0.9, 0.1, 0.0, 0.0])], + &embed, + ) + .await + .unwrap(); + + // Reserved separator in user collection names. + let err = manager + .create_collection("evil--part--x", None, Some(4), None) + .await + .unwrap_err(); + assert!(err.to_string().contains("reserved partition separator")); + + // Chunks missing the partition field are rejected with the field name. + let mut bad = tenant_chunk("acme", "b", "x", [0.1, 0.1, 0.0, 0.0]); + bad.metadata.clear(); + let err = manager + .ingest("fenced", vec![bad], &embed) + .await + .unwrap_err(); + assert!(err.to_string().contains("missing partition field 'tenant'")); + + // Unrouted operations fail closed with a clear message. + let err = manager + .get_facets("fenced", "", &["tenant".to_string()]) + .await + .unwrap_err(); + assert!(err.to_string().contains("not supported on a partitioned")); + let err = manager + .create_relations( + "fenced", + vec![CreateRelation { + source_chunk_id: 0, + target_chunk_id: 1, + target_document_id: None, + relation_type: "cites".to_string(), + metadata: HashMap::new(), + }], + ) + .await + .unwrap_err(); + assert!(err.to_string().contains("not supported on a partitioned")); + let err = manager + .add_vector_space("fenced", "extra", 8, "model") + .await + .unwrap_err(); + assert!(err.to_string().contains("not supported on a partitioned")); + + let _ = std::fs::remove_dir_all(&data_dir); +} + +#[tokio::test] +async fn partitioned_collection_survives_restart() { + let data_dir = unique_data_dir(); + let embed = embed_state(); + { + let manager = local_manager(&data_dir).await; + manager + .create_collection("persist", None, Some(4), partitioned_config()) + .await + .unwrap(); + manager + .ingest( + "persist", + vec![ + tenant_chunk("acme", "a1", "durable acme", [0.9, 0.1, 0.0, 0.0]), + tenant_chunk("globex", "g1", "durable globex", [0.1, 0.9, 0.0, 0.0]), + ], + &embed, + ) + .await + .unwrap(); + } + + let manager = CollectionManager::new(&data_dir).await.unwrap(); + // Routing metadata survives: tenant-scoped search still works, ids keep + // minting from the shared allocator without collision. + let (results, _, _, _) = manager + .search("persist", &search_req("durable", None, "acme"), &embed) + .await + .unwrap(); + assert_eq!(results.len(), 1); + assert_eq!(results[0].0.file_id, "a1"); + let (_, id_map, _) = manager + .ingest( + "persist", + vec![{ + let mut c = tenant_chunk("acme", "a2", "post-restart", [0.7, 0.3, 0.0, 0.0]); + c.client_id = Some("a2".to_string()); + c + }], + &embed, + ) + .await + .unwrap(); + let new_id = *id_map.values().next().unwrap(); + assert!(new_id >= 2, "restart must not reuse ids (got {new_id})"); + + let _ = std::fs::remove_dir_all(&data_dir); +} diff --git a/crates/compass/src/collections/partitions.rs b/crates/compass/src/collections/partitions.rs new file mode 100644 index 0000000..2be819b --- /dev/null +++ b/crates/compass/src/collections/partitions.rs @@ -0,0 +1,216 @@ +// collections/partitions.rs — tenant-partitioned collections (roadmap Phase 6). +// +// Design: a partition IS a full internal collection. Each partition gets its +// own namespace — LSM manifest, WAL, segments, Tantivy dir, vector files, +// attach/evict lifecycle — by reusing the existing per-collection engine +// wholesale. Phase 6 is a ROUTER at the manager entry points, not a new +// engine: +// +// - `create` with `config.partition_by = "tenant_id"` marks the parent. +// The parent namespace holds config + the shared id allocator; chunk data +// lives only in partitions. +// - Ingest groups chunks by `metadata[partition_by]` and delegates each +// group to the partition's namespace, auto-creating it on first sight. +// - Search/deletes require a filter on the partition field and route to +// exactly the named partitions (set membership fans out, capped). +// - Chunk ids are COLLECTION-unique: every partition claims id blocks from +// the PARENT's allocator (in local mode too — the allocator is just a +// CAS-updated file; this does not create WAL/manifest objects). +// +// Why this shape: per-tenant cost isolation falls out of the existing +// machinery. A query for tenant T attaches T's partition only; LRU eviction +// and refresh scale with the HOT tenant set, not the collection. The +// per-namespace scale envelope now bounds the largest TENANT, not the +// collection, which is what makes a billion-vector multi-tenant collection +// servable on bounded RAM. +// +// Scope fences (MVP, enforced with clear errors): relations, facets, TAMS +// temporal lookup, and vector-space CRUD are not yet routed for partitioned +// collections; partition keys must be kebab-case strings; the partition +// field is immutable after create. + +use crate::models::{FilterValue, IngestChunk, MetadataValue}; +use std::collections::HashMap; + +/// Separator between parent collection name and partition value in the +/// internal namespace. User-facing collection names must not contain it +/// (enforced at create); it is otherwise valid kebab-case, so every existing +/// storage/path rule accepts partition namespaces unchanged. +pub const PART_SEP: &str = "--part--"; + +/// Max partitions a single set-membership search may fan out to. +pub const MAX_SEARCH_FANOUT: usize = 16; + +/// Internal namespace for one partition of a parent collection. +pub fn partition_ns(parent: &str, pval: &str) -> String { + format!("{parent}{PART_SEP}{pval}") +} + +/// Is this namespace a partition (vs a user-facing collection)? +pub fn is_partition_ns(ns: &str) -> bool { + ns.contains(PART_SEP) +} + +/// The parent collection of a partition namespace, or None for a normal one. +pub fn parent_of(ns: &str) -> Option<&str> { + ns.split_once(PART_SEP).map(|(parent, _)| parent) +} + +/// The namespace whose id allocator a collection mints from: partitions share +/// the PARENT's allocator so chunk ids are unique across the whole collection +/// (delete-by-id and search results would otherwise be ambiguous). +pub fn alloc_ns(ns: &str) -> &str { + parent_of(ns).unwrap_or(ns) +} + +/// Partition values become path/namespace segments — hold them to the same +/// kebab-case rule as collection names, and bound the length so a hostile +/// value can't manufacture absurd object keys. +pub fn validate_partition_value(v: &str) -> Result<(), Box> { + if v.is_empty() || v.len() > 64 || !v.chars().all(|c| c.is_ascii_alphanumeric() || c == '-') { + return Err(format!( + "partition value '{v}' is invalid: use letters, digits, and hyphens (max 64 chars)" + ) + .into()); + } + if v.contains(PART_SEP) { + return Err(format!("partition value '{v}' must not contain '{PART_SEP}'").into()); + } + Ok(()) +} + +/// Group an ingest batch by its partition value, validating that every chunk +/// carries a usable `metadata[field]` string. +pub fn group_by_partition( + field: &str, + chunks: Vec, +) -> Result>, Box> { + let mut groups: HashMap> = HashMap::new(); + for chunk in chunks { + let pval = match chunk.metadata.get(field) { + Some(MetadataValue::String(s)) => s.clone(), + Some(_) => { + return Err(format!( + "chunk '{}': partition field '{field}' must be a string", + chunk.file_id + ) + .into()); + } + None => { + return Err(format!( + "chunk '{}' is missing partition field '{field}' \ + (this collection is partitioned by it)", + chunk.file_id + ) + .into()); + } + }; + validate_partition_value(&pval)?; + groups.entry(pval).or_default().push(chunk); + } + Ok(groups) +} + +/// Resolve which partition values a filtered request targets. Exact match +/// routes to one partition; set membership (`in`) fans out (capped). Anything +/// else is an error — a partitioned collection cannot be scanned blind. +pub fn partition_values_from_filters( + field: &str, + filters: &HashMap, +) -> Result, Box> { + let missing = || -> Box { + format!( + "this collection is partitioned by '{field}': include filters.{field} \ + (exact value, or {{\"in\": [...]}} for up to {MAX_SEARCH_FANOUT} partitions)" + ) + .into() + }; + let values = match filters.get(field) { + Some(FilterValue::Exact(MetadataValue::String(s))) => vec![s.clone()], + Some(FilterValue::Condition(cond)) => match &cond.in_values { + Some(vs) if !vs.is_empty() => vs.clone(), + _ => return Err(missing()), + }, + Some(_) => return Err(missing()), + None => return Err(missing()), + }; + if values.len() > MAX_SEARCH_FANOUT { + return Err(format!( + "filters.{field} names {} partitions; max fan-out is {MAX_SEARCH_FANOUT}", + values.len() + ) + .into()); + } + for v in &values { + validate_partition_value(v)?; + } + Ok(values) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::models::FilterCondition; + + #[test] + fn namespace_roundtrip() { + let ns = partition_ns("videos", "acme"); + assert_eq!(ns, "videos--part--acme"); + assert!(is_partition_ns(&ns)); + assert_eq!(parent_of(&ns), Some("videos")); + assert_eq!(alloc_ns(&ns), "videos"); + assert!(!is_partition_ns("videos")); + assert_eq!(parent_of("videos"), None); + assert_eq!(alloc_ns("videos"), "videos"); + } + + #[test] + fn partition_value_rules() { + assert!(validate_partition_value("acme-01").is_ok()); + assert!(validate_partition_value("").is_err()); + assert!(validate_partition_value("has space").is_err()); + assert!(validate_partition_value("a--part--b").is_err()); + assert!(validate_partition_value(&"x".repeat(65)).is_err()); + } + + #[test] + fn filter_routing() { + let field = "tenant"; + let mut f = HashMap::new(); + assert!(partition_values_from_filters(field, &f).is_err()); + + f.insert( + field.to_string(), + FilterValue::Exact(MetadataValue::String("acme".into())), + ); + assert_eq!( + partition_values_from_filters(field, &f).unwrap(), + vec!["acme".to_string()] + ); + + f.insert( + field.to_string(), + FilterValue::Condition(FilterCondition { + gte: None, + lte: None, + contains: None, + in_values: Some(vec!["a".into(), "b".into()]), + }), + ); + assert_eq!(partition_values_from_filters(field, &f).unwrap().len(), 2); + + let too_many: Vec = (0..MAX_SEARCH_FANOUT + 1) + .map(|i| format!("t{i}")) + .collect(); + f.insert( + field.to_string(), + FilterValue::Condition(FilterCondition { + gte: None, + lte: None, + contains: None, + in_values: Some(too_many), + }), + ); + assert!(partition_values_from_filters(field, &f).is_err()); + } +} 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..f559891 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; @@ -105,7 +102,7 @@ async fn main() -> Result<(), Box> { .allow_methods(Any) .allow_headers(Any); - // Anonymous telemetry — opt out with COMPASS_TELEMETRY=off or DO_NOT_TRACK=1 + // Anonymous telemetry — OPT-IN ONLY (COMPASS_TELEMETRY=on); off by default telemetry::spawn_telemetry(data_dir.clone(), app_state.manager.clone()); // Bearer-token auth via COMPASS_API_KEY. When unset, auth is disabled. diff --git a/crates/compass/src/metrics.rs b/crates/compass/src/metrics.rs new file mode 100644 index 0000000..96f9ae8 --- /dev/null +++ b/crates/compass/src/metrics.rs @@ -0,0 +1,57 @@ +//! 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, + COLD_SEARCHES_TOTAL, + WARM_PROMOTIONS_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..7f49c75 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 { @@ -147,6 +155,11 @@ fn default_dims() -> usize { pub struct CollectionConfig { #[serde(default = "default_embed_model")] pub embed_model: String, + /// Tenant-partitioned collections: the metadata field whose (string) + /// value routes each chunk to its own internal partition namespace. + /// Immutable after create. None = normal single-namespace collection. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub partition_by: Option, } fn default_embed_model() -> String { @@ -262,6 +275,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 +369,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 +636,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/cold.rs b/crates/compass/src/search/cold.rs new file mode 100644 index 0000000..979b31d --- /dev/null +++ b/crates/compass/src/search/cold.rs @@ -0,0 +1,762 @@ +// search/cold.rs — serve-from-storage: answer semantic queries on a +// collection that is NOT attached, with a handful of object-storage range +// reads instead of a full index rebuild. +// +// Query flow (per namespace): +// 1. GET manifest (freshness anchor: everything committed is visible, so +// cold reads satisfy read-your-writes by construction) +// 2. per segment (immutable → artifacts cached by segment id): +// header+TOC, `cent:` centroids, `tombs`, `metaidx` — all small +// 3. rank clusters by centroid dot-product, range-GET the top `nprobe` +// clusters, score their (unit-norm) vectors against the query +// 4. brute-force the uncompacted WAL tail (bounded by the auto-compact +// threshold) and apply tombstones; newest version of an id wins +// 5. hydrate the top candidates' chunk JSON by byte range via `metaidx`, +// apply metadata filters, return +// +// RAM cost per cold namespace: centroids + directories + tombstones + the +// metadata index — megabytes, independent of collection size. + +use crate::models::{DocumentChunk, FilterValue, MetadataValue}; +use crate::search::filter_pushdown::{FilterExpr, Predicate}; +use crate::search::ivf; +use crate::storage::{lsm, Storage, StorageError}; +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +/// Clusters probed per segment per query (env-tunable via the manager). +pub const DEFAULT_NPROBE: usize = 8; +/// Overfetch factor before deduping down to top_k. +const OVERFETCH: usize = 4; +/// Additional overfetch multiplier when metadata filters are present (cold +/// filtering is post-selection; see the recall-contract note in search()). +const FILTER_OVERFETCH: usize = 8; +/// Header bytes fetched optimistically (magic + max_id + toc_len + TOC). +const HEADER_PROBE: u64 = 16 * 1024; +/// metaidx at or below this size is fetched whole; larger ones page in +/// blocks on demand. +const METAIDX_FULL_MAX: u64 = 8 * 1024 * 1024; +const METAIDX_BLOCK_ROWS: usize = 2048; + +type BoxErr = Box; + +/// Cached, immutable cold-read artifacts for ONE segment object. +pub struct ColdSegment { + ns: String, + segment_id: String, + /// section name -> (absolute byte offset in the object, length) + sections: HashMap, + /// space -> parsed centroids + cluster directory + cents: HashMap, + /// ids tombstoned BY this segment (they apply to OLDER segments only — + /// see the generation rule in `search`) + tombstones: HashSet, + metaidx: MetaIdx, +} + +enum MetaIdx { + /// Whole index resident: sorted (id, off, len) rows. + Full(Vec<(u64, u64, u32)>), + /// Sparse anchors (first id of each block) + block byte range info; blocks + /// are fetched on demand per query (not cached — queries touch few). + Paged { + anchors: Vec, // first id of block i + n_rows: u64, + idx_offset: u64, // absolute offset of the first row (after the count) + }, +} + +fn seg_key(ns: &str, id: &str) -> String { + format!("{ns}/segments/{id}") +} + +async fn get_range( + storage: &dyn Storage, + key: &str, + start: u64, + len: u64, +) -> Result { + storage.get_range(key, start..start + len).await +} + +impl ColdSegment { + /// Build the cached artifacts with a few small reads. Total fetched: + /// TOC + centroids + tombstones + (metaidx or its anchors). + pub async fn open(storage: &dyn Storage, ns: &str, segment_id: &str) -> Result { + Self::open_with_limits(storage, ns, segment_id, METAIDX_FULL_MAX).await + } + + /// `open` with an explicit full-fetch threshold — lets tests exercise the + /// paged metadata-index path without a 400k-chunk segment. + async fn open_with_limits( + storage: &dyn Storage, + ns: &str, + segment_id: &str, + metaidx_full_max: u64, + ) -> Result { + let key = seg_key(ns, segment_id); + // Header + TOC (optimistic single read; re-read if the TOC is huge). + let head = storage.get_range(&key, 0..HEADER_PROBE).await?; + if head.len() < 20 { + return Err(format!("segment {segment_id}: truncated header").into()); + } + // Only v3 segments carry the cold-read sections (metaidx/meta2 and + // clusters). v2 would brute-force its whole flat section and then + // drop every hit at hydration — reject loudly instead. + if head[0..8] != crate::collections::cloud::SEG_MAGIC_V3 { + return Err(format!( + "segment {segment_id} predates the cold-servable format; \ + run POST /collections/:name/compact once to upgrade it" + ) + .into()); + } + let toc_len = u32::from_le_bytes(head[16..20].try_into().unwrap()) as u64; + let toc_bytes = if 20 + toc_len <= head.len() as u64 { + head.slice(20..(20 + toc_len) as usize) + } else { + get_range(storage, &key, 20, toc_len).await? + }; + let toc: Vec<(String, u64)> = serde_json::from_slice(&toc_bytes) + .map_err(|e| format!("segment {segment_id}: bad TOC: {e}"))?; + let mut sections = HashMap::new(); + let mut pos = 20 + toc_len; + for (name, len) in toc { + sections.insert(name, (pos, len)); + pos += len; + } + + // Centroids for every clustered space (small — cache them all). + let mut cents = HashMap::new(); + for (name, &(off, len)) in §ions { + if let Some(space) = name.strip_prefix("cent:") { + let body = get_range(storage, &key, off, len).await?; + let c = ivf::parse_cent(&body) + .ok_or_else(|| format!("segment {segment_id}: bad {name}"))?; + cents.insert(space.to_string(), c); + } + } + + // Tombstones (u64 LE array; bounded by deletes-per-fold). + let mut tombstones = HashSet::new(); + if let Some(&(off, len)) = sections.get("tombs") { + if len > 0 { + let body = get_range(storage, &key, off, len).await?; + for c in body.chunks_exact(8) { + tombstones.insert(u64::from_le_bytes(c.try_into().unwrap())); + } + } + } + + // Metadata index: whole if small, paged anchors otherwise. + let metaidx = match sections.get("metaidx") { + Some(&(off, len)) if len > 8 => { + if len <= metaidx_full_max { + let body = get_range(storage, &key, off, len).await?; + let n = u64::from_le_bytes(body[0..8].try_into().unwrap()) as usize; + if body.len() < 8 + n * 20 { + return Err(format!( + "segment {segment_id}: truncated metaidx ({n} rows declared)" + ) + .into()); + } + let mut rows = Vec::with_capacity(n); + for i in 0..n { + let p = 8 + i * 20; + rows.push(( + u64::from_le_bytes(body[p..p + 8].try_into().unwrap()), + u64::from_le_bytes(body[p + 8..p + 16].try_into().unwrap()), + u32::from_le_bytes(body[p + 16..p + 20].try_into().unwrap()), + )); + } + MetaIdx::Full(rows) + } else { + // Big index: fetch it whole ONCE at open (bounded by + // index size, ~20B/row) but keep only per-block anchor + // ids resident; lookups page 20B×2048 blocks on demand. + let body = get_range(storage, &key, off, len).await?; + let n = u64::from_le_bytes(body[0..8].try_into().unwrap()); + let mut anchors = Vec::new(); + let mut i = 0u64; + while i < n { + let p = (8 + i * 20) as usize; + anchors.push(u64::from_le_bytes(body[p..p + 8].try_into().unwrap())); + i += METAIDX_BLOCK_ROWS as u64; + } + MetaIdx::Paged { + anchors, + n_rows: n, + idx_offset: off + 8, + } + } + } + _ => MetaIdx::Full(Vec::new()), + }; + + Ok(Self { + ns: ns.to_string(), + segment_id: segment_id.to_string(), + sections, + cents, + tombstones, + metaidx, + }) + } + + /// Rank this segment's clusters for `q` (unit-norm) and return the top + /// `nprobe` cluster byte ranges to fetch. + fn probe_plan(&self, space: &str, q: &[f32], nprobe: usize) -> Vec<(u64, u64)> { + let Some(cent) = self.cents.get(space) else { + return Vec::new(); + }; + let Some(&(clu_off, _)) = self.sections.get(&format!("clu:{space}")) else { + return Vec::new(); + }; + let mut ranked: Vec<(usize, f32)> = cent + .centroids + .iter() + .enumerate() + .filter(|(i, _)| cent.dir[*i].count > 0) + .map(|(i, c)| (i, ivf::dot(c, q))) + .collect(); + ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + ranked + .into_iter() + .take(nprobe) + .map(|(i, _)| { + let d = cent.dir[i]; + (clu_off + d.offset, d.len) + }) + .collect() + } + + fn dims_for(&self, space: &str) -> Option { + self.cents.get(space).map(|c| c.dims) + } + + /// The flat `emb:` byte range (segments below the clustering + /// threshold) — brute-forced whole. + fn flat_range(&self, space: &str) -> Option<(u64, u64)> { + self.sections.get(&format!("emb:{space}")).copied() + } + + /// Look up the meta2 byte ranges for a set of ids. + async fn meta_ranges( + &self, + storage: &dyn Storage, + ids: &[u64], + ) -> Result, BoxErr> { + let Some(&(meta_off, _)) = self.sections.get("meta2") else { + return Ok(Vec::new()); + }; + let key = seg_key(&self.ns, &self.segment_id); + let mut out = Vec::new(); + match &self.metaidx { + MetaIdx::Full(rows) => { + for &id in ids { + if let Ok(i) = rows.binary_search_by_key(&id, |r| r.0) { + let (rid, off, len) = rows[i]; + out.push((rid, meta_off + off, len)); + } + } + } + MetaIdx::Paged { + anchors, + n_rows, + idx_offset, + } => { + // Group wanted ids by block, fetch each needed block once. + let mut by_block: HashMap> = HashMap::new(); + for &id in ids { + let block = match anchors.binary_search(&id) { + Ok(i) => i, + Err(0) => continue, // below the first anchor: absent + Err(i) => i - 1, + }; + by_block.entry(block).or_default().push(id); + } + for (block, wanted) in by_block { + let start_row = (block * METAIDX_BLOCK_ROWS) as u64; + let rows_here = (*n_rows - start_row).min(METAIDX_BLOCK_ROWS as u64); + let body = + get_range(storage, &key, idx_offset + start_row * 20, rows_here * 20) + .await?; + for r in body.chunks_exact(20) { + let rid = u64::from_le_bytes(r[0..8].try_into().unwrap()); + if wanted.contains(&rid) { + out.push(( + rid, + meta_off + u64::from_le_bytes(r[8..16].try_into().unwrap()), + u32::from_le_bytes(r[16..20].try_into().unwrap()), + )); + } + } + } + } + } + Ok(out) + } +} + +/// One scored candidate before hydration. `generation` orders duplicates of +/// the same id: higher wins (segments in manifest order, tail above all). +struct Candidate { + id: u64, + score: f32, + generation: usize, + /// Tail candidates already carry their chunk. + chunk: Option, + segment: Option, // index into segments, for hydration +} + +/// A cold semantic search over one namespace. `segments` are the cached +/// artifacts in manifest order; the WAL tail is read fresh per query. +#[allow(clippy::too_many_arguments)] +pub async fn search( + storage: &dyn Storage, + ns: &str, + segments: &[Arc], + manifest: &lsm::Manifest, + space: &str, + query: &[f32], + top_k: usize, + nprobe: usize, + filters: &HashMap, +) -> Result, BoxErr> { + let mut q = query.to_vec(); + ivf::normalize(&mut q); + + // Tombstone semantics must match materialize(): a segment's carried + // tombstones apply only to OLDER segments (its own chunks are written + // after them, and a NEWER segment's re-ingest of the same id must + // survive). So a candidate from generation g dies only to a tombstone + // from generation > g. A single flat union would permanently suppress + // re-ingested chunks that warm search serves. + let tomb_of = |generation: usize| -> &HashSet { &segments[generation].tombstones }; + let killed_by_newer = |id: u64, generation: usize| -> bool { + ((generation + 1)..segments.len()).any(|j| tomb_of(j).contains(&id)) + }; + // Tail tombstones (replayed in seq order below) are the newest + // generation of all: they kill any segment candidate. + let mut tail_dead: HashSet = HashSet::new(); + + // WAL tail: bounded by the auto-compaction threshold in healthy + // operation. A pathologically long tail (compaction disabled/failing) + // would make this a full-dataset materialization per query — refuse + // loudly instead of degrading into that silently. + let tail_len = manifest.uncompacted().count(); + if tail_len > 2 * crate::collections::AUTO_COMPACT_FRAGMENT_THRESHOLD { + return Err(format!( + "namespace '{ns}' has {tail_len} uncompacted WAL fragments — too many to \ + cold-serve. Run POST /collections/:name/compact (or check why \ + auto-compaction is not running), then retry." + ) + .into()); + } + let tail = lsm::read_uncompacted_fragments(storage, ns, manifest).await?; + let mut tail_chunks: HashMap = HashMap::new(); + for (fref, payload) in &tail { + match fref.kind { + lsm::FragmentKind::Data => { + let chunks: Vec = serde_json::from_slice(payload) + .map_err(|e| format!("tail fragment decode: {e}"))?; + for c in chunks { + tail_dead.remove(&c.id); // re-ingest after delete resurrects + tail_chunks.insert(c.id, c); + } + } + lsm::FragmentKind::Tombstone => { + let ids: Vec = serde_json::from_slice(payload) + .map_err(|e| format!("tail tombstone decode: {e}"))?; + for id in ids { + tail_dead.insert(id); + tail_chunks.remove(&id); + } + } + _ => {} + } + } + + // Filters are applied POST-candidate-selection on the cold path (there + // is no roaring index to push down without attaching), so a selective + // filter needs a deeper candidate pool. Recall contract: cold filtered + // queries can under-return when matches are rarer than ~1/FILTER_OVERFETCH + // of the probed neighborhoods; warm search has no such limit. + let overfetch = if filters.is_empty() { + OVERFETCH + } else { + OVERFETCH * FILTER_OVERFETCH + }; + let want = (top_k * overfetch).max(top_k).min(512); + let mut candidates: Vec = Vec::new(); + + // Segment candidates: probe clusters (or brute-force flat sections). + for (gen, seg) in segments.iter().enumerate() { + let key = seg_key(ns, &seg.segment_id); + let mut ranges = seg.probe_plan(space, &q, nprobe); + let dims = match seg.dims_for(space) { + Some(d) => d, + None => match seg.flat_range(space) { + Some((off, len)) if len >= 12 => { + // Flat section: [u32 dims][u64 n][rows] — brute force it. + let head = get_range(storage, &key, off, 12).await?; + let dims = u32::from_le_bytes(head[0..4].try_into().unwrap()) as usize; + ranges = vec![(off + 12, len - 12)]; + dims + } + _ => continue, // space absent in this segment + }, + }; + if dims != q.len() { + return Err(format!( + "query has {} dims but segment space '{space}' has {dims}", + q.len() + ) + .into()); + } + // Fetch probed ranges concurrently. + let bodies = futures::future::try_join_all( + ranges + .iter() + .map(|&(off, len)| get_range(storage, &key, off, len)), + ) + .await?; + for body in bodies { + for (id, v) in ivf::parse_cluster_rows(&body, dims) { + if tail_dead.contains(&id) + || tail_chunks.contains_key(&id) + || killed_by_newer(id, gen) + { + continue; + } + // Flat sections store raw vectors; clustered store unit-norm. + // Normalizing again is idempotent for the latter. + let mut v = v; + ivf::normalize(&mut v); + candidates.push(Candidate { + id, + score: ivf::dot(&v, &q), + generation: gen, + chunk: None, + segment: Some(gen), + }); + } + } + } + + // Tail candidates: brute-force the fresh writes. + let tail_gen = segments.len(); + for (id, c) in &tail_chunks { + if let Some(emb) = c.embeddings.get(space) { + let mut v = emb.clone(); + ivf::normalize(&mut v); + candidates.push(Candidate { + id: *id, + score: ivf::dot(&v, &q), + generation: tail_gen, + chunk: Some(c.clone()), + segment: None, + }); + } + } + + // Dedupe by id, newest generation wins; then keep the global top `want`. + candidates.sort_by(|a, b| a.id.cmp(&b.id).then(b.generation.cmp(&a.generation))); + candidates.dedup_by_key(|c| c.id); + candidates.sort_by(|a, b| { + b.score + .partial_cmp(&a.score) + .unwrap_or(std::cmp::Ordering::Equal) + }); + candidates.truncate(want); + + // Hydrate: group segment candidates per segment, batch the meta lookups. + let expr = FilterExpr::compile(filters); + let mut hydrated: Vec<(DocumentChunk, f32)> = Vec::new(); + let mut by_seg: HashMap> = HashMap::new(); + let mut scores: HashMap = HashMap::new(); + for c in &candidates { + scores.insert(c.id, c.score); + match (&c.chunk, c.segment) { + (Some(ch), _) => { + if eval_filters(&expr, ch) { + hydrated.push((ch.clone(), c.score)); + } + } + (None, Some(seg_i)) => by_seg.entry(seg_i).or_default().push(c.id), + _ => {} + } + } + for (seg_i, ids) in by_seg { + let seg = &segments[seg_i]; + let key = seg_key(ns, &seg.segment_id); + let mut ranges = seg.meta_ranges(storage, &ids).await?; + // Coalesce adjacent-ish rows into fewer GETs. + ranges.sort_by_key(|r| r.1); + let mut batches: Vec<(u64, u64, Vec<(u64, u64, u32)>)> = Vec::new(); + for r in ranges { + match batches.last_mut() { + Some((_start, end, rows)) if r.1 <= *end + 64 * 1024 => { + *end = (*end).max(r.1 + r.2 as u64); + rows.push(r); + } + _ => batches.push((r.1, r.1 + r.2 as u64, vec![r])), + } + } + for (start, end, rows) in batches { + let body = get_range(storage, &key, start, end - start).await?; + for (id, off, len) in rows { + let lo = (off - start) as usize; + let chunk: DocumentChunk = serde_json::from_slice(&body[lo..lo + len as usize]) + .map_err(|e| format!("meta2 row decode (id {id}): {e}"))?; + if eval_filters(&expr, &chunk) { + hydrated.push((chunk, scores.get(&id).copied().unwrap_or(0.0))); + } + } + } + } + + hydrated.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + hydrated.truncate(top_k); + Ok(hydrated) +} + +/// Metadata filter evaluation for cold hits (the roaring FilterIndex only +/// exists for attached collections). Mirrors FilterIndex::eligible semantics, +/// including the doc_type-as-metadata rule. +fn eval_filters(expr: &FilterExpr, chunk: &DocumentChunk) -> bool { + if expr.is_empty() { + return true; + } + let get = |field: &str| -> Option { + if field == "doc_type" { + Some(MetadataValue::String(chunk.doc_type.clone())) + } else { + chunk.metadata.get(field).cloned() + } + }; + expr.predicates.iter().all(|p| match p { + Predicate::Eq { field, value } => get(field).as_ref() == Some(value), + Predicate::Range { field, gte, lte } => match 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, + }, + // Parity with FilterIndex: `contains` matches STRING LISTS only (the + // warm index populates string_list_contains from StringList values) — + // matching bare strings here would make cold return hits warm never + // would. + Predicate::Contains { field, value } => match get(field) { + Some(MetadataValue::StringList(xs)) => xs.iter().any(|x| x == value), + _ => false, + }, + Predicate::In { field, values } => match get(field) { + Some(MetadataValue::String(s)) => values.contains(&s), + _ => false, + }, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::storage::local::LocalDiskStorage; + + fn storage(name: &str) -> (std::path::PathBuf, Arc) { + use std::sync::atomic::{AtomicU64, Ordering}; + static N: AtomicU64 = AtomicU64::new(0); + let root = std::env::temp_dir().join(format!( + "compass_cold_unit_{}_{}_{}", + name, + std::process::id(), + N.fetch_add(1, Ordering::Relaxed) + )); + let _ = std::fs::remove_dir_all(&root); + let s: Arc = Arc::new(LocalDiskStorage::new(&root).unwrap()); + (root, s) + } + + fn seg_with_chunks(ids: &[u64], tombstones: &[u64]) -> Vec { + use crate::models::DocumentChunk; + let chunks: Vec = ids + .iter() + .map(|&i| { + let mut c = DocumentChunk { + id: i, + collection: "ns".into(), + file_id: format!("f{i}"), + chunk_index: 0, + page: None, + text: format!("t{i}"), + metadata: Default::default(), + doc_type: "chunk".into(), + parent_id: None, + group_id: None, + embeddings: Default::default(), + embedding: None, + }; + c.embeddings + .insert("default".into(), vec![i as f32, 1.0, 0.0, 0.0]); + c + }) + .collect(); + crate::collections::cloud::encode_segment_v3(&crate::collections::cloud::Segment { + version: 2, + chunks, + relations: vec![], + max_id: ids.iter().copied().max().unwrap_or(0), + tombstones: tombstones.to_vec(), + relation_tombstones: vec![], + }) + .unwrap() + } + + // C1 regression: an OLDER segment's carried tombstone must not suppress + // the same id re-ingested into a NEWER segment (materialize parity). + #[tokio::test] + async fn newer_segment_survives_older_tombstone() { + let (root, s) = storage("gen"); + // seg A (gen 0): chunk 1 live, carries tombstone for id 5. + s.put( + "ns/segments/a", + bytes::Bytes::from(seg_with_chunks(&[1], &[5])), + ) + .await + .unwrap(); + // seg B (gen 1): id 5 re-ingested. + s.put( + "ns/segments/b", + bytes::Bytes::from(seg_with_chunks(&[5], &[])), + ) + .await + .unwrap(); + let manifest = lsm::Manifest { + segments: vec![ + lsm::SegmentRef { + id: "a".into(), + records: 1, + }, + lsm::SegmentRef { + id: "b".into(), + records: 1, + }, + ], + ..Default::default() + }; + let segs = vec![ + Arc::new(ColdSegment::open(s.as_ref(), "ns", "a").await.unwrap()), + Arc::new(ColdSegment::open(s.as_ref(), "ns", "b").await.unwrap()), + ]; + let hits = search( + s.as_ref(), + "ns", + &segs, + &manifest, + "default", + &[5.0, 1.0, 0.0, 0.0], + 10, + DEFAULT_NPROBE, + &Default::default(), + ) + .await + .unwrap(); + assert!( + hits.iter().any(|(c, _)| c.id == 5), + "id 5 lives in the NEWER segment; the older tombstone must not kill it" + ); + // And the reverse still holds: a NEWER segment's tombstone kills an + // OLDER segment's chunk. + s.put( + "ns/segments/c", + bytes::Bytes::from(seg_with_chunks(&[9], &[1])), + ) + .await + .unwrap(); + let manifest2 = lsm::Manifest { + segments: vec![ + lsm::SegmentRef { + id: "a".into(), + records: 1, + }, + lsm::SegmentRef { + id: "c".into(), + records: 1, + }, + ], + ..Default::default() + }; + let segs2 = vec![ + segs[0].clone(), + Arc::new(ColdSegment::open(s.as_ref(), "ns", "c").await.unwrap()), + ]; + let hits = search( + s.as_ref(), + "ns", + &segs2, + &manifest2, + "default", + &[1.0, 1.0, 0.0, 0.0], + 10, + DEFAULT_NPROBE, + &Default::default(), + ) + .await + .unwrap(); + assert!( + hits.iter().all(|(c, _)| c.id != 1), + "newer segment's tombstone must kill the older chunk" + ); + let _ = std::fs::remove_dir_all(&root); + } + + // Paged metadata-index path: force it with a tiny full-fetch threshold + // and verify hydration still resolves every candidate. + #[tokio::test] + async fn paged_metaidx_hydrates() { + let (root, s) = storage("paged"); + let ids: Vec = (0..50).collect(); + s.put( + "ns/segments/p", + bytes::Bytes::from(seg_with_chunks(&ids, &[])), + ) + .await + .unwrap(); + let seg = ColdSegment::open_with_limits(s.as_ref(), "ns", "p", 16) + .await + .unwrap(); + assert!( + matches!(seg.metaidx, MetaIdx::Paged { .. }), + "tiny threshold must force the paged variant" + ); + let ranges = seg + .meta_ranges(s.as_ref(), &[0, 7, 49, 999_999]) + .await + .unwrap(); + let found: std::collections::HashSet = ranges.iter().map(|r| r.0).collect(); + assert_eq!( + found, + [0u64, 7, 49].into_iter().collect(), + "paged lookups must resolve present ids and skip absent ones" + ); + let _ = std::fs::remove_dir_all(&root); + } + + // Pre-v3 segments are rejected loudly — v2 has no metaidx, so cold + // serving it would read the whole flat section and then drop every hit. + #[tokio::test] + async fn v2_segment_rejected_with_upgrade_hint() { + let (root, s) = storage("v2"); + let mut bytes = Vec::new(); + bytes.extend_from_slice(&crate::collections::cloud::SEG_MAGIC_V2); + bytes.extend_from_slice(&0u64.to_le_bytes()); + bytes.extend_from_slice(&2u32.to_le_bytes()); + bytes.extend_from_slice(b"[]"); + s.put("ns/segments/old", bytes::Bytes::from(bytes)) + .await + .unwrap(); + let err = match ColdSegment::open(s.as_ref(), "ns", "old").await { + Err(e) => e, + Ok(_) => panic!("v2 segment must be rejected"), + }; + assert!(err.to_string().contains("compact"), "{err}"); + let _ = std::fs::remove_dir_all(&root); + } +} 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..fc0c03b 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,21 @@ 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 + // in filter_index.rs tests (FilterIndex::eligible, the warm evaluator) + // and cold.rs (eval_filters, its cold-path mirror). 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 +104,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/ivf.rs b/crates/compass/src/search/ivf.rs new file mode 100644 index 0000000..8c3d84c --- /dev/null +++ b/crates/compass/src/search/ivf.rs @@ -0,0 +1,314 @@ +// search/ivf.rs — IVF (inverted-file) clustering for serve-from-storage. +// +// Compaction k-means-clusters each vector space and writes two sections into +// the segment: `cent:` (tiny: centroids + a cluster directory) and +// `clu:` (the vectors, grouped by cluster). A cold query then needs +// only: the centroids (cached, a few hundred KB), and range-GETs of the +// `nprobe` nearest clusters — instead of materializing the whole segment. +// That is what turns "attach = rebuild everything" into "query = a handful +// of small reads": the difference between warm and true serverless. +// +// Section formats (all little-endian): +// cent: = [u32 k][u32 dims] +// [k × dims × f32 centroids] +// [k × (u64 offset, u64 len, u32 count)] cluster directory, +// offsets relative to the START of clu:'s body +// clu: = concatenation of clusters, each [count × (u64 id, dims×f32)] +// +// Vectors are stored L2-NORMALIZED in `clu` so scoring is a plain dot +// product (cosine == dot on unit vectors); centroids are means of normalized +// vectors, re-normalized. + +/// Below this many rows a space is stored as the flat `emb:` section and cold +/// queries brute-force it — clustering tiny sets costs more than it saves. +pub const CLUSTER_MIN_ROWS: usize = 5_000; + +/// Cap on k-means training sample: training cost is O(sample × k × dims); +/// assignment of ALL rows is a single pass afterwards. +const TRAIN_SAMPLE_MAX: usize = 20_000; +const KMEANS_ITERS: usize = 8; + +/// Number of clusters for n rows: sqrt(n), clamped. At 5M rows and 384 dims +/// this keeps a cluster ~3.5MB — a few parallel range-GETs per query. +pub fn cluster_count(n: usize) -> usize { + ((n as f64).sqrt() as usize).clamp(16, 4096) +} + +#[inline] +pub fn normalize(v: &mut [f32]) { + let norm: f32 = v.iter().map(|x| x * x).sum::().sqrt(); + if norm > 0.0 { + for x in v.iter_mut() { + *x /= norm; + } + } +} + +#[inline] +pub fn dot(a: &[f32], b: &[f32]) -> f32 { + a.iter().zip(b).map(|(x, y)| x * y).sum() +} + +/// K-means (Lloyd) over normalized vectors, trained on a deterministic +/// sample. Returns (centroids, assignment of EVERY input row). +/// Deterministic: seeded by row order, no RNG. +pub fn kmeans(rows: &[(u64, Vec)], dims: usize, k: usize) -> (Vec>, Vec) { + let n = rows.len(); + let k = k.min(n).max(1); + + // Deterministic training sample: evenly-strided rows. + let stride = (n / TRAIN_SAMPLE_MAX).max(1); + let sample: Vec<&[f32]> = rows + .iter() + .step_by(stride) + .map(|(_, v)| v.as_slice()) + .collect(); + + // Init: evenly-strided sample points as seeds. + let seed_stride = (sample.len() / k).max(1); + let mut centroids: Vec> = sample + .iter() + .step_by(seed_stride) + .take(k) + .map(|v| v.to_vec()) + .collect(); + while centroids.len() < k { + centroids.push(centroids[centroids.len() % sample.len().max(1)].clone()); + } + + let nearest = |cents: &[Vec], v: &[f32]| -> usize { + let mut best = 0usize; + let mut best_d = f32::MIN; + for (i, c) in cents.iter().enumerate() { + let d = dot(c, v); // unit vectors: max dot == min angle + if d > best_d { + best_d = d; + best = i; + } + } + best + }; + + for _ in 0..KMEANS_ITERS { + let mut sums = vec![vec![0f32; dims]; k]; + let mut counts = vec![0usize; k]; + for v in &sample { + let c = nearest(¢roids, v); + for (s, x) in sums[c].iter_mut().zip(v.iter()) { + *s += x; + } + counts[c] += 1; + } + for (i, (sum, cnt)) in sums.iter_mut().zip(counts.iter()).enumerate() { + if *cnt > 0 { + for x in sum.iter_mut() { + *x /= *cnt as f32; + } + normalize(sum); + centroids[i] = std::mem::take(sum); + } + // Empty cluster: keep the old centroid (harmless; directory entry + // just ends up with count 0). + } + } + + let assignment: Vec = rows.iter().map(|(_, v)| nearest(¢roids, v)).collect(); + (centroids, assignment) +} + +/// Directory entry for one cluster inside `clu:`. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct ClusterRef { + pub offset: u64, + pub len: u64, + pub count: u32, +} + +/// Build the `cent:` and `clu:` section bodies for one space. +/// Input rows may be un-normalized; they are normalized in place here. +pub fn build_sections(mut rows: Vec<(u64, Vec)>, dims: usize) -> (Vec, Vec) { + for (_, v) in rows.iter_mut() { + normalize(v); + } + let k = cluster_count(rows.len()); + let (centroids, assignment) = kmeans(&rows, dims, k); + + // Group row indexes by cluster, then lay clusters out contiguously. + let mut by_cluster: Vec> = vec![Vec::new(); k]; + for (row_idx, &c) in assignment.iter().enumerate() { + by_cluster[c].push(row_idx); + } + + let row_size = 8 + dims * 4; + let mut clu = Vec::with_capacity(rows.len() * row_size); + let mut dir: Vec = Vec::with_capacity(k); + for members in &by_cluster { + let offset = clu.len() as u64; + for &ri in members { + let (id, v) = &rows[ri]; + clu.extend_from_slice(&id.to_le_bytes()); + for x in v { + clu.extend_from_slice(&x.to_le_bytes()); + } + } + dir.push(ClusterRef { + offset, + len: (members.len() * row_size) as u64, + count: members.len() as u32, + }); + } + + let mut cent = Vec::with_capacity(8 + k * dims * 4 + k * 20); + cent.extend_from_slice(&(k as u32).to_le_bytes()); + cent.extend_from_slice(&(dims as u32).to_le_bytes()); + for c in ¢roids { + for x in c { + cent.extend_from_slice(&x.to_le_bytes()); + } + } + for d in &dir { + cent.extend_from_slice(&d.offset.to_le_bytes()); + cent.extend_from_slice(&d.len.to_le_bytes()); + cent.extend_from_slice(&d.count.to_le_bytes()); + } + (cent, clu) +} + +/// Parsed `cent:` section. +pub struct Centroids { + pub dims: usize, + pub centroids: Vec>, + pub dir: Vec, +} + +pub fn parse_cent(body: &[u8]) -> Option { + if body.len() < 8 { + return None; + } + let k = u32::from_le_bytes(body[0..4].try_into().ok()?) as usize; + let dims = u32::from_le_bytes(body[4..8].try_into().ok()?) as usize; + let cent_bytes = k.checked_mul(dims)?.checked_mul(4)?; + let dir_bytes = k.checked_mul(20)?; + if body.len() < 8 + cent_bytes + dir_bytes { + return None; + } + let mut centroids = Vec::with_capacity(k); + let mut pos = 8; + for _ in 0..k { + let mut v = Vec::with_capacity(dims); + for _ in 0..dims { + v.push(f32::from_le_bytes(body[pos..pos + 4].try_into().ok()?)); + pos += 4; + } + centroids.push(v); + } + let mut dir = Vec::with_capacity(k); + for _ in 0..k { + let offset = u64::from_le_bytes(body[pos..pos + 8].try_into().ok()?); + let len = u64::from_le_bytes(body[pos + 8..pos + 16].try_into().ok()?); + let count = u32::from_le_bytes(body[pos + 16..pos + 20].try_into().ok()?); + pos += 20; + dir.push(ClusterRef { offset, len, count }); + } + Some(Centroids { + dims, + centroids, + dir, + }) +} + +/// Iterate `(id, vector)` rows out of a cluster blob. +pub fn parse_cluster_rows(body: &[u8], dims: usize) -> impl Iterator)> + '_ { + let row = 8 + dims * 4; + body.chunks_exact(row).map(move |r| { + let id = u64::from_le_bytes(r[0..8].try_into().unwrap()); + let mut v = Vec::with_capacity(dims); + for d in 0..dims { + let o = 8 + d * 4; + v.push(f32::from_le_bytes(r[o..o + 4].try_into().unwrap())); + } + (id, v) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn synthetic(n: usize, dims: usize) -> Vec<(u64, Vec)> { + // Deterministic, mildly clustered data: 8 anchor directions + noise. + (0..n) + .map(|i| { + let anchor = i % 8; + let v: Vec = (0..dims) + .map(|d| { + let base = if d % 8 == anchor { 1.0 } else { 0.1 }; + base + ((i * 31 + d * 17) % 97) as f32 / 970.0 + }) + .collect(); + (i as u64, v) + }) + .collect() + } + + #[test] + fn sections_roundtrip_and_cover_all_rows() { + let dims = 16; + let rows = synthetic(1000, dims); + let (cent, clu) = build_sections(rows.clone(), dims); + let parsed = parse_cent(¢).expect("cent parses"); + assert_eq!(parsed.dims, dims); + let total: u32 = parsed.dir.iter().map(|d| d.count).sum(); + assert_eq!(total as usize, rows.len(), "every row lands in a cluster"); + // Every directory range decodes to exactly `count` rows and all ids + // survive. + let mut seen = std::collections::HashSet::new(); + for d in &parsed.dir { + let body = &clu[d.offset as usize..(d.offset + d.len) as usize]; + let rows: Vec<_> = parse_cluster_rows(body, dims).collect(); + assert_eq!(rows.len(), d.count as usize); + for (id, v) in rows { + assert_eq!(v.len(), dims); + assert!(seen.insert(id), "id {id} duplicated across clusters"); + } + } + assert_eq!(seen.len(), 1000); + } + + #[test] + fn nearest_cluster_probe_finds_exact_vector() { + // Self-recall: probing the nearest clusters for a vector that IS in + // the index must find it with a modest nprobe. + let dims = 16; + let rows = synthetic(2000, dims); + let (cent, clu) = build_sections(rows.clone(), dims); + let parsed = parse_cent(¢).unwrap(); + + let mut hits = 0; + let probes = 4; + for probe_i in (0..2000).step_by(97) { + let mut q = rows[probe_i].1.clone(); + normalize(&mut q); + let mut ranked: Vec<(usize, f32)> = parsed + .centroids + .iter() + .enumerate() + .map(|(i, c)| (i, dot(c, &q))) + .collect(); + ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); + let found = ranked.iter().take(probes).any(|(ci, _)| { + let d = parsed.dir[*ci]; + let body = &clu[d.offset as usize..(d.offset + d.len) as usize]; + parse_cluster_rows(body, dims).any(|(id, _)| id == rows[probe_i].0) + }); + if found { + hits += 1; + } + } + let total = (0..2000).step_by(97).count(); + assert!( + hits * 10 >= total * 9, + "self-recall with nprobe={probes}: {hits}/{total}" + ); + } +} diff --git a/crates/compass/src/search/mod.rs b/crates/compass/src/search/mod.rs index 0a302d4..37b47a8 100644 --- a/crates/compass/src/search/mod.rs +++ b/crates/compass/src/search/mod.rs @@ -5,33 +5,19 @@ // 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. +pub mod cold; #[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 ivf; 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/local.rs b/crates/compass/src/storage/local.rs index a0412da..b645fd3 100644 --- a/crates/compass/src/storage/local.rs +++ b/crates/compass/src/storage/local.rs @@ -132,7 +132,11 @@ impl Storage for LocalDiskStorage { .metadata() .map_err(|e| StorageError::Io(e.to_string()))? .len(); - if range.start > range.end || range.end > size { + // Contract parity with the object-store backend: a range end past + // the object is CLAMPED (S3/GCS Range semantics), not an error — + // cold reads probe fixed-size headers on objects of unknown length. + let range = range.start..range.end.min(size); + if range.start > range.end { return Err(StorageError::InvalidRange { start: range.start, end: range.end, @@ -361,11 +365,9 @@ mod tests { let s = store("range"); s.put("k", Bytes::from_static(b"0123456789")).await.unwrap(); assert_eq!(&s.get_range("k", 2..5).await.unwrap()[..], b"234"); - // Out-of-bounds range errors. - assert!(matches!( - s.get_range("k", 5..100).await, - Err(StorageError::InvalidRange { .. }) - )); + // End past EOF is CLAMPED (object-store Range semantics — cold reads + // probe fixed-size headers on objects of unknown length). + assert_eq!(&s.get_range("k", 5..100).await.unwrap()[..], b"56789"); } // Boundary cases for the seek-based get_range (each a plausible off-by-one). @@ -387,9 +389,12 @@ mod tests { #[allow(clippy::reversed_empty_ranges)] let reversed = s.get_range("k", 6..3).await; assert!(matches!(reversed, Err(StorageError::InvalidRange { .. }))); - // end past EOF → error. + // end past EOF → clamped to the object (matches S3/GCS semantics). + assert_eq!(&s.get_range("k", 8..11).await.unwrap()[..], b"89"); + // start past EOF stays an error (object_store 416 parity): only the + // END is clamped. assert!(matches!( - s.get_range("k", 8..11).await, + s.get_range("k", 20..30).await, Err(StorageError::InvalidRange { .. }) )); } 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..b3f1bb5 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, @@ -101,8 +108,10 @@ pub trait Storage: Send + Sync { /// Whole-object read. async fn get(&self, key: &str) -> Result; - /// Range read — fetch only `range` bytes of the object. The primitive that - /// makes large segments servable without loading the whole object. + /// Range read — the serve-from-storage primitive (segment TOCs point at + /// byte ranges; cold queries fetch only the sections they need). A range + /// end past the object is CLAMPED, never an error (S3/GCS semantics; + /// LocalDiskStorage matches). async fn get_range(&self, key: &str, range: Range) -> Result; /// Read the object together with its current version, for a CAS cycle. @@ -162,6 +171,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/crates/compass/src/telemetry.rs b/crates/compass/src/telemetry.rs index 9d87b97..cc26504 100644 --- a/crates/compass/src/telemetry.rs +++ b/crates/compass/src/telemetry.rs @@ -15,15 +15,23 @@ const POSTHOG_API_KEY: &str = "phc_BFvsmH5rpe8GqJ8zwfqhH9jGAdZMXcNZhEao8mnDEd3X" const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(24 * 60 * 60); const STARTUP_DELAY: Duration = Duration::from_secs(60); -/// Returns true if telemetry is enabled (default). +/// Returns true if telemetry is enabled. DEFAULT OFF: Compass's core promise +/// is "data never leaves your machine" — an engine pitched on privacy must +/// not phone home unless the operator explicitly opts in +/// (COMPASS_TELEMETRY=on). DO_NOT_TRACK is honored even when opted in. pub fn is_enabled() -> bool { - if let Ok(v) = std::env::var("COMPASS_TELEMETRY") { - return !matches!(v.to_lowercase().as_str(), "off" | "false" | "0" | "no"); - } if let Ok(v) = std::env::var("DO_NOT_TRACK") { - return !matches!(v.as_str(), "1" | "true"); + if matches!(v.as_str(), "1" | "true") { + return false; + } } - true + matches!( + std::env::var("COMPASS_TELEMETRY") + .unwrap_or_default() + .to_lowercase() + .as_str(), + "on" | "true" | "1" | "yes" + ) } /// Persistent instance ID — generated once, stored in data_dir/instance_id. @@ -100,7 +108,7 @@ pub fn spawn_telemetry( let instance_id = get_or_create_instance_id(&data_dir); tracing::info!( - "Anonymous telemetry enabled (instance: {}). Set COMPASS_TELEMETRY=off to disable.", + "Anonymous telemetry enabled by explicit opt-in (instance: {}). Unset COMPASS_TELEMETRY to disable.", &instance_id[..8] ); diff --git a/docs/deployment.md b/docs/deployment.md new file mode 100644 index 0000000..a6fe8b8 --- /dev/null +++ b/docs/deployment.md @@ -0,0 +1,84 @@ +# Deployment topologies + +Compass runs in three shapes. All of them are the same binary; the shape is +chosen entirely by environment variables. + +## 1. Local single node (default) + +Zero config. Local disk is the source of truth; everything is embedded. + +```bash +./compass # or: docker run -p 4001:4001 -v ./data:/app/data compass +``` + +- No cloud credentials, no telemetry, no network calls. +- Tenant partitions (`partition_by`) work fully in this mode. +- Backup = copy `DATA_DIR`. + +## 2. Cloud: serving nodes + stateless writers + +Object storage is the source of truth; nodes are disposable. + +```bash +# Serving node(s): full local indexes, fast reads, background convergence +COMPASS_STORAGE=s3://bucket AWS_ACCESS_KEY_ID=… AWS_SECRET_ACCESS_KEY=… ./compass + +# Writer node(s): stateless append-only ingest, boots in milliseconds +COMPASS_ROLE=writer COMPASS_STORAGE=s3://bucket … ./compass +``` + +- Writers validate against the bucket config, mint chunk ids from a CAS + allocator (never collide with anyone), append one WAL fragment, return a + `seq`. Durable immediately; searchable on serving nodes within + `COMPASS_REFRESH_INTERVAL` (default 5s). +- Read-your-writes: pass a write's `seq` as `min_seq` on search. +- A serving node that loses its disk rebuilds every collection from the + bucket on boot. Kill -9 is a supported operation. +- Optional: `COMPASS_LAZY_ATTACH=true` + `COMPASS_MAX_ATTACHED=N` bound RAM + to the hot collection set (LRU detach; re-attach on demand). + +Upgrade caveat: do not run pre-v0.4 and v0.4 writers against one bucket; +old readers fail loudly on the v0.4 segment format rather than mis-reading. + +## 3. Cloud: cold serving (serverless reads) + +```bash +COMPASS_COLD_SERVE=true COMPASS_STORAGE=s3://bucket … ./compass +``` + +- Boots in <1s regardless of how much data the bucket holds; RAM starts at + ~tens of MB. +- Semantic queries on collections the node has NEVER attached are answered + from object-storage range reads (manifest → cached centroids → a few + cluster reads → byte-range hydration). Freshness is read-your-writes by + construction — every cold query reads the live manifest. +- `COMPASS_WARM_AFTER` (default 3) cold hits promote a background attach: + cold → warm → hot automatically. +- Honest limits: cold is semantic-only (FTS errors until the namespace + warms); recall on unstructured vector spaces needs a higher + `COMPASS_COLD_NPROBE` (see [search-quality.md](search-quality.md)); + scoring options (recency/boosts/relations) are rejected cold rather than + silently ignored. + +## Multi-tenant collections (any topology) + +```bash +curl -X POST :4001/collections -d '{"name":"app","embedding_dims":384, + "config":{"partition_by":"tenant_id"}}' +``` + +Every chunk routes to an internal per-tenant partition by +`metadata.tenant_id`. Searches/deletes must filter on the partition field +(exact, or `{"in":[…]}` for ≤16 tenants). Ids are collection-unique; +partitions auto-create on first ingest (writer nodes included), hide from +listings, cascade-delete with the parent. Serving cost tracks the HOT tenant +set — 50 or 200 tenants boot identically. + +## What the fleet does NOT give you (yet) + +- Tenant-affinity routing between nodes: put a proxy in front and hash a + tenant header to a node, or every node will warm every hot tenant. +- Per-tenant auth: `COMPASS_API_KEY` is one key for the whole node; tenant + scoping is the caller's responsibility today. +- `/metrics` and `/health` are unauthenticated by design; firewall them if + collection names/counts are sensitive. diff --git a/docs/scale-envelope.md b/docs/scale-envelope.md new file mode 100644 index 0000000..47e5480 --- /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 is reached via tenant partitions + serve-from-storage: the per-NAMESPACE envelope above bounds the largest tenant, not the collection, and cold reads serve unattached namespaces from object storage (see docs/search-quality.md for the recall contract). + 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/search-quality.md b/docs/search-quality.md new file mode 100644 index 0000000..9f65373 --- /dev/null +++ b/docs/search-quality.md @@ -0,0 +1,51 @@ +# Search quality — measured recall & latency + +Method: 20k docs × 64 dims, 200 held-out queries (perturbed documents), +ground truth = exact cosine top-10 (numpy). Two datasets: **structured** +(32 topic clusters + mild noise — the shape real text/image embeddings have) +and **adversarial uniform** (noise-dominated, nearly structureless — the +worst case for any ANN index). Latency measured over HTTP against Docker + +MinIO on a laptop; treat relative numbers, not absolutes. + +## Warm path (attached: HNSW, ef_search=128) + +| dataset | recall@10 | p50 | p95 | +|---|---|---|---| +| structured | 1.000 | 1.0ms | 1.1ms | +| structured, filtered (50% selectivity) | 1.000 | 1.0ms | 1.2ms | +| adversarial uniform | 0.895 | 1.2ms | 1.3ms | +| adversarial, filtered | 0.962 | 1.4ms | 1.7ms | + +Full-text (BM25): exact-token top-1 50/50; topical precision@10 = 1.000. + +## Cold path (serve-from-storage: IVF over object storage) + +`COMPASS_COLD_NPROBE` clusters probed per segment (default 8): + +| dataset | nprobe | recall@10 | p50 | +|---|---|---|---| +| structured | 4 | 0.947 | 13ms | +| structured | **8 (default)** | **1.000** | 13ms | +| structured | 16 | 1.000 | 14ms | +| adversarial uniform | 4 | 0.269 | 12ms | +| adversarial uniform | 8 | 0.409 | 13ms | +| adversarial uniform | 16 | 0.590 | 14ms | +| adversarial uniform | k (exhaustive) | 1.000 | 30ms | + +## The honest contract + +- On **clustered embedding spaces** — which is what real embedding models + produce — cold recall reaches warm parity at the default nprobe, at + ~10× warm latency (a handful of object-storage range reads). +- On **unstructured/uniform vector spaces**, IVF recall drops steeply (this + is inherent to inverted-file indexes, not a Compass bug — the exhaustive + row proves the pipeline is exact). If your vectors are random-ish + (hashes, uncalibrated projections), raise `COMPASS_COLD_NPROBE` + aggressively or rely on warm serving (`COMPASS_WARM_AFTER` promotes hot + namespaces automatically). +- Cold filtered queries apply filters post-selection with an 8× deeper + candidate pool; extremely selective filters (≪1% match rate) can + under-return on the cold path — warm search has no such limit. + +Reproduce: `scratchpad` eval scripts live in the PR discussion; the harness +is ~100 lines of numpy + HTTP and pins seeds. diff --git a/docs/serverless-roadmap.md b/docs/serverless-roadmap.md new file mode 100644 index 0000000..e7c6fbf --- /dev/null +++ b/docs/serverless-roadmap.md @@ -0,0 +1,162 @@ +# Serverless Roadmap + +> Status: warm serverless (stateless writers, refresh, lazy attach), tenant partitions, and serve-from-storage cold reads are SHIPPED on the v0.4 branches (phase numbers below predate the final split). Remaining: routing/affinity hooks, per-tenant auth binding, cold FTS, intra-tenant sharding. 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/docs/v0.4-filter-aware-ann.md b/docs/v0.4-filter-aware-ann.md index 31c3efd..fe41d9b 100644 --- a/docs/v0.4-filter-aware-ann.md +++ b/docs/v0.4-filter-aware-ann.md @@ -2,7 +2,6 @@ Status: done. Decision: ship Path A (USearch native `filtered_search`). -Sister design doc: [docs/v0.4-vision.md](./v0.4-vision.md). ## Verdict diff --git a/scripts/e2e.sh b/scripts/e2e.sh new file mode 100755 index 0000000..d1d4ee4 --- /dev/null +++ b/scripts/e2e.sh @@ -0,0 +1,172 @@ +#!/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} +COLD=${COLD:-} # optional: a COMPASS_COLD_SERVE node against the same bucket +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 "── tenant partitions ──" +post $FULL/collections '{"name":"mt","embedding_dims":4,"config":{"partition_by":"tenant"}}' >/dev/null +r=$(post $FULL/collections/mt/ingest '{"chunks":[ + {"client_id":"p1","file_id":"p1","chunk_index":0,"doc_type":"chunk","text":"shared secret alpha","metadata":{"tenant":"acme"},"embeddings":{"default":[0.9,0.1,0.1,0.1]}}, + {"client_id":"p2","file_id":"p2","chunk_index":0,"doc_type":"chunk","text":"shared secret beta","metadata":{"tenant":"globex"},"embeddings":{"default":[0.1,0.9,0.1,0.1]}}]}') +[ "$(echo "$r" | jqn "d['indexed']")" = "2" ] && ok "partitioned ingest routes" || bad p-ingest x +n=$(post $FULL/collections/mt/search '{"query":"secret","mode":"fts","top_k":10,"filters":{"tenant":"acme"}}' | jqn "len(d['results'])") +f1=$(post $FULL/collections/mt/search '{"query":"secret","mode":"fts","top_k":10,"filters":{"tenant":"acme"}}' | jqn "d['results'][0]['chunk']['file_id']") +[ "$n" = "1" ] && [ "$f1" = "p1" ] && ok "tenant isolation (acme sees only its hit)" || bad p-iso "$n/$f1" +code=$(curl -s -o /dev/null -w '%{http_code}' -X POST $FULL/collections/mt/search -H 'content-type: application/json' -d '{"query":"secret","mode":"fts"}') +[ "$code" -ge 400 ] && ok "unfiltered partitioned search rejected" || bad p-nofilter "$code" +n=$(post $FULL/collections/mt/search '{"query":"secret","mode":"fts","top_k":10,"filters":{"tenant":{"in":["acme","globex"]}}}' | jqn "len(d['results'])") +[ "$n" = "2" ] && ok "set-membership fan-out merges tenants" || bad p-fanout "$n" +r=$(post $WRITER/collections/mt/ingest '{"chunks":[{"client_id":"p3","file_id":"p3","chunk_index":0,"doc_type":"chunk","text":"writer minted tenant","metadata":{"tenant":"initech"},"embeddings":{"default":[0.1,0.1,0.9,0.1]}}]}') +[ "$(echo "$r" | jqn "d['indexed']")" = "1" ] && ok "writer partitioned ingest" || bad p-writer x +sleep 1 +n=$(post $FULL/collections/mt/search '{"query":"minted","mode":"fts","top_k":5,"filters":{"tenant":"initech"}}' | jqn "len(d['results'])") +[ "$n" = "1" ] && ok "writer-minted partition attaches on serving node" || bad p-attach "$n" +r=$(post $FULL/collections/mt/delete '{"filters":{"tenant":"acme"}}') +[ "$(echo "$r" | jqn "d['deleted']")" = "1" ] && ok "partition-scoped delete" || bad p-del x +code=$(curl -s -o /dev/null -w '%{http_code}' -X DELETE $FULL/collections/mt) +[ "$code" -lt 400 ] && ok "partitioned collection cascade delete" || bad p-casc "$code" +curl -s $FULL/collections | grep -q "part--" && bad "partitions hidden from listing" leak || ok "partitions hidden from listing" + +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" + +if [ -n "$COLD" ]; then +echo "── serve-from-storage (cold node) ──" +post $FULL/collections '{"name":"icy","embedding_dims":4}' >/dev/null +post $FULL/collections/icy/ingest '{"chunks":[ + {"file_id":"i1","chunk_index":0,"doc_type":"chunk","text":"glacier core","metadata":{"kind":"ice"},"embeddings":{"default":[0.9,0.1,0.1,0.1]}}, + {"file_id":"i2","chunk_index":0,"doc_type":"chunk","text":"magma core","metadata":{"kind":"fire"},"embeddings":{"default":[0.1,0.9,0.1,0.1]}}]}' >/dev/null +r=$(post $COLD/collections/icy/search '{"query":"","mode":"semantic","top_k":3,"query_vector":[0.9,0.1,0.1,0.1]}') +f1=$(echo "$r" | jqn "d['results'][0]['chunk']['file_id']") +[ "$f1" = "i1" ] && ok "cold node answers without attach" || bad cold-search "$f1" +n=$(post $COLD/collections/icy/search '{"query":"","mode":"semantic","top_k":3,"query_vector":[0.9,0.1,0.1,0.1],"filters":{"kind":"fire"}}' | jqn "d['results'][0]['chunk']['file_id']") +[ "$n" = "i2" ] && ok "cold filters apply" || bad cold-filter "$n" +code=$(curl -s -o /dev/null -w '%{http_code}' -X POST $COLD/collections/icy/search -H 'content-type: application/json' -d '{"query":"glacier","mode":"fts"}') +[ "$code" -ge 400 ] && ok "cold FTS rejected with guidance" || bad cold-fts "$code" +wseq2=$(post $WRITER/collections/icy/ingest '{"chunks":[{"file_id":"i3","chunk_index":0,"doc_type":"chunk","text":"fresh tail","metadata":{"kind":"new"},"embeddings":{"default":[0.1,0.1,0.9,0.1]}}]}' | jqn "d['seq']") +f3=$(post $COLD/collections/icy/search '{"query":"","mode":"semantic","top_k":1,"query_vector":[0.1,0.1,0.9,0.1]}' | jqn "d['results'][0]['chunk']['file_id']") +[ "$f3" = "i3" ] && ok "cold read-your-writes (writer tail visible instantly, seq $wseq2)" || bad cold-ryw "$f3" +curl -s $COLD/metrics | grep -q "compass_cold_searches_total [1-9]" && ok "cold metrics counting" || bad cold-metrics x +curl -s -o /dev/null -X DELETE $FULL/collections/icy +fi + +echo "" +echo "E2E RESULT: $pass passed, $fail failed" +[ "$fail" = "0" ]