From d3ca2e96fd79fb6d39352776250ed615f8db3d44 Mon Sep 17 00:00:00 2001 From: Isaac Elbaz Date: Sat, 5 Sep 2026 11:38:31 -0400 Subject: [PATCH] Fix audited resolution and persistence defects for 0.2.0 --- .github/workflows/ci.yml | 14 +- .github/workflows/release.yml | 25 +- CHANGELOG.md | 39 +- README.md | 11 + benches/bench_distributed.rs | 48 +- docs/critical-audit-2026-09-05.md | 62 ++ proto/unirust.proto | 38 + src/distributed.rs | 601 ++++++++++++++- src/graph.rs | 37 + src/index.rs | 587 +++++++-------- src/lib.rs | 353 ++++++++- src/linker.rs | 542 +++++++------- src/perf/simd_hash.rs | 7 +- src/persistence.rs | 251 ++++--- src/query.rs | 21 +- src/sharding.rs | 270 ++++--- src/store.rs | 92 +++ tests/audit_performance.rs | 61 ++ tests/distributed_apply_ontology.rs | 13 +- tests/distributed_conflicts_presets.rs | 3 +- .../distributed_cross_shard_reconciliation.rs | 56 +- ...d_cross_shard_reconciliation_persistent.rs | 32 +- tests/distributed_e2e.rs | 4 +- tests/distributed_entity_regressions.rs | 702 ++++++++++++++++++ tests/distributed_ingest_stream.rs | 4 +- tests/distributed_metrics.rs | 2 +- tests/distributed_rebalance.rs | 6 +- tests/distributed_rebalance_stream.rs | 2 +- tests/distributed_reliability.rs | 12 +- tests/distributed_router_admin.rs | 2 +- tests/durable_ingest_regressions.rs | 427 +++++++++++ tests/external_backup_restore.rs | 10 +- tests/process_crash_recovery.rs | 6 +- tests/resolution_capacity_regressions.rs | 365 +++++++++ tests/selective_query_regressions.rs | 318 ++++++++ tests/temporal_evolution.rs | 51 +- 36 files changed, 4114 insertions(+), 960 deletions(-) create mode 100644 docs/critical-audit-2026-09-05.md create mode 100644 tests/audit_performance.rs create mode 100644 tests/distributed_entity_regressions.rs create mode 100644 tests/durable_ingest_regressions.rs create mode 100644 tests/resolution_capacity_regressions.rs create mode 100644 tests/selective_query_regressions.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 68fd406..2979747 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -4,6 +4,10 @@ on: push: branches: [main] pull_request: + workflow_call: + +permissions: + contents: read jobs: msrv: @@ -21,6 +25,8 @@ jobs: uses: dtolnay/rust-toolchain@1.88.0 - name: Cache Rust uses: Swatinem/rust-cache@v2 + - name: Install native build dependencies + run: sudo apt-get update && sudo apt-get install -y libclang-dev cmake - name: Install protoc uses: arduino/setup-protoc@v3 with: @@ -46,6 +52,8 @@ jobs: components: rustfmt, clippy - name: Cache Rust uses: Swatinem/rust-cache@v2 + - name: Install native build dependencies + run: sudo apt-get update && sudo apt-get install -y libclang-dev cmake - name: Install protoc uses: arduino/setup-protoc@v3 with: @@ -58,9 +66,13 @@ jobs: run: cargo fmt --check - name: Clippy run: cargo clippy --locked --all-targets --all-features -- -D warnings + - name: Verify release package + run: cargo package --locked - name: Validate deployment manifests run: | - bash -n scripts/*.sh + for script in scripts/*.sh; do + bash -n "$script" + done docker compose -f compose.yaml config --quiet container-image: diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index a045a51..3cf770a 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -8,26 +8,35 @@ on: permissions: contents: write +concurrency: + group: release-${{ github.ref }} + cancel-in-progress: false + jobs: + validate: + uses: ./.github/workflows/ci.yml + publish: + needs: validate runs-on: ubuntu-latest steps: + - name: Free disk space + run: | + sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc /opt/hostedtoolcache/CodeQL - uses: actions/checkout@v4 - uses: dtolnay/rust-toolchain@stable - uses: Swatinem/rust-cache@v2 + - name: Install native build dependencies + run: sudo apt-get update && sudo apt-get install -y libclang-dev cmake - uses: arduino/setup-protoc@v3 with: version: "24.x" - - name: cargo test - run: cargo test --locked --all-features - - name: cargo fmt - run: cargo fmt --check - - name: cargo clippy - run: cargo clippy --locked --all-targets --all-features -- -D warnings + - name: Verify tag matches package version + run: | + package_version="$(cargo metadata --locked --no-deps --format-version 1 | python3 -c 'import json,sys; data=json.load(sys.stdin); print(next(p["version"] for p in data["packages"] if p["name"] == "unirust-rs"))')" + test "$GITHUB_REF_NAME" = "v$package_version" - name: cargo package run: cargo package --locked - - name: Build production image - run: docker build --file Containerfile --tag unirust-release . - name: cargo publish run: cargo publish --locked env: diff --git a/CHANGELOG.md b/CHANGELOG.md index a827e39..3e713e7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,10 +3,38 @@ All notable changes to this project are documented here. This project follows Semantic Versioning. -## [Unreleased] +## [0.2.0] - 2026-09-05 + +### Added + +- Persistent distributed shards, streaming ingest, cross-shard reconciliation, + coordinated checkpoints, verified off-host backups, and synchronous replication. +- Mutual TLS transport, semantic readiness checks, and fail-closed recovery for + partial reconciliation, interrupted ingestion, and inconsistent restores. +- Enabled persistent regressions for the September correctness audit and query + equivalence across bridge merges, cache invalidation, and restart. ### Fixed +- Transitive cross-shard merges now compare component-wide temporal strong IDs, + including observations from previous reconciliation rounds. +- Distributed queries assemble complete canonical entities before temporal + conjunction and golden attribute mastering. +- Occupied record IDs cannot overwrite durable records; rejected batches discard + staged records and partial resolution state before another request can commit. +- Bounded interner caches preserve durable IDs after restart. Corrupt query + lookups fail explicitly instead of returning empty matches. +- Repeated common keys and the memory-saver profile continue entity resolution. + Tiered index updates preserve complete buckets through eviction and recovery, + respect capacity settings, and propagate storage failures. +- Persistent DSU cluster enumeration, deferred membership updates, and restored + query-label cache invalidation now use authoritative cluster state. +- Unbounded temporal intervals no longer overflow duration calculations or + require enumerating every temporal bucket; Allen relations preserve orientation. +- Rust 1.98 Clippy failures are resolved. Release publishing now depends on full + CI validation and a matching tag/package version, including package and image + builds. Shell validation checks every deployment script. + - Ingest acknowledgement now waits for a synchronous RocksDB WAL flush. The external ingest WAL is removed only after records and entity-resolution state have reached stable storage, closing a power-loss window for acknowledged data. @@ -51,7 +79,14 @@ Semantic Versioning. - Adopted Rust 2024 with a documented MSRV of Rust 1.88. - Updated the active gRPC, terminal UI, cache, and compression dependencies and removed unused direct HTTP, TOML, and protobuf-types dependencies. -- The package version is now `0.2.0` in preparation for the next release. +- The live distributed protocol is now 6. Upgrade routers, shards, and replicas + together; WAL/checkpoint version 1 and durable reservation format 5 remain + readable. Public router clients continue using the existing request contract. +- Selective queries reuse resolved membership and master only candidate entities; + composite labels retain global prefix collision semantics. Shard requests run + concurrently with bounded fan-out and entity hydration batches. +- Persistent ingestion lends staged records directly to parallel extraction, + avoiding record cloning and correctly resolving batches above 100,000 records. - The application lockfile is tracked for reproducible binary and container builds. - Distributed integration tests now use temporary persistent shard stores; the diff --git a/README.md b/README.md index 952d1f0..538ae49 100644 --- a/README.md +++ b/README.md @@ -281,6 +281,17 @@ duplicates make startup fail closed for operator repair. The internal router and shard protocol is versioned, and mixed versions are rejected, so upgrades must replace the full cluster with one coordinated Unirust version before restarting the router. Shard gRPC ports are internal APIs; client ingest must use the router. + +The current live protocol is **6**. Protocol 5 routers, shards, and replicas must +be stopped and upgraded together: older coordinators do not implement the +component guards or canonical entity queries required by this release. This +handshake change does not change WAL or checkpoint format version 1, record +snapshot encoding, or durable source-reservation format version 5. Existing +volumes, pending WAL batches, and coordinated checkpoints remain readable; +valid version 5 reservation markers retain their original shard-count binding +without another backfill. Internal ingest/import callers must send the current +live protocol, while public router clients continue leaving that field unset. + The persisted reservation directory is also bound to the configured shard count. Router startup rejects shard-count changes because the current import API cannot atomically move a record, update its reservation, and delete the old copy. diff --git a/benches/bench_distributed.rs b/benches/bench_distributed.rs index 6c0ffe1..7613236 100644 --- a/benches/bench_distributed.rs +++ b/benches/bench_distributed.rs @@ -410,12 +410,6 @@ fn bench_shardnode_ingest(c: &mut Criterion) { group.warm_up_time(Duration::from_secs(1)); group.measurement_time(Duration::from_secs(6)); - // SAFETY: benchmark configuration runs before this function creates any worker threads. - unsafe { std::env::set_var("UNIRUST_PARTITIONED", "1") }; - let partition_count = env_usize("UNIRUST_DIST_PARTITIONS", 8); - // SAFETY: benchmark configuration runs before this function creates any worker threads. - unsafe { std::env::set_var("UNIRUST_PARTITION_COUNT", partition_count.to_string()) }; - let count = env_u32("UNIRUST_DIST_RECORDS", 50_000); let overlap = env_f64("UNIRUST_DIST_OVERLAP", 0.05); let config = default_dist_config(); @@ -424,18 +418,28 @@ fn bench_shardnode_ingest(c: &mut Criterion) { group.bench_function(BenchmarkId::new("batch", count), |b| { b.iter_batched( || { + let data_dir = tempfile::tempdir().expect("persistent benchmark directory"); let rt = Runtime::new().expect("runtime"); let shard = rt.block_on(async { - ShardNode::new(0, config.clone(), StreamingTuning::default()).expect("shard") + ShardNode::new_with_data_dir( + 0, + config.clone(), + StreamingTuning::default(), + Some(data_dir.path().to_path_buf()), + false, + None, + ) + .expect("persistent shard") }); let records = generate_proto_batch(1, count, overlap, 45); - (rt, shard, records) + (rt, shard, records, data_dir) }, - |(rt, shard, records)| { + |(rt, shard, records, _data_dir)| { let response = rt.block_on(async { shard .ingest_records(Request::new(proto::IngestRecordsRequest { - internal_protocol_version: 2, + internal_protocol_version: + unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records, })) .await @@ -457,12 +461,6 @@ fn bench_shardnode_streaming_simulated(c: &mut Criterion) { group.warm_up_time(Duration::from_secs(1)); group.measurement_time(Duration::from_secs(6)); - // SAFETY: benchmark configuration runs before this function creates any worker threads. - unsafe { std::env::set_var("UNIRUST_PARTITIONED", "1") }; - let partition_count = env_usize("UNIRUST_DIST_PARTITIONS", 8); - // SAFETY: benchmark configuration runs before this function creates any worker threads. - unsafe { std::env::set_var("UNIRUST_PARTITION_COUNT", partition_count.to_string()) }; - let total = env_u32("UNIRUST_DIST_STREAM_TOTAL", 20_000); let chunk = env_u32("UNIRUST_DIST_STREAM_CHUNK", 512).max(1); let overlap = env_f64("UNIRUST_DIST_OVERLAP", 0.05); @@ -472,14 +470,23 @@ fn bench_shardnode_streaming_simulated(c: &mut Criterion) { group.bench_function(BenchmarkId::new("chunked", chunk), |b| { b.iter_batched( || { + let data_dir = tempfile::tempdir().expect("persistent benchmark directory"); let rt = Runtime::new().expect("runtime"); let shard = rt.block_on(async { - ShardNode::new(0, config.clone(), StreamingTuning::default()).expect("shard") + ShardNode::new_with_data_dir( + 0, + config.clone(), + StreamingTuning::default(), + Some(data_dir.path().to_path_buf()), + false, + None, + ) + .expect("persistent shard") }); let records = generate_proto_batch(1, total, overlap, 46); - (rt, shard, records) + (rt, shard, records, data_dir) }, - |(rt, shard, records)| { + |(rt, shard, records, _data_dir)| { let response = rt.block_on(async { let mut offset = 0usize; let mut assignments = Vec::new(); @@ -488,7 +495,8 @@ fn bench_shardnode_streaming_simulated(c: &mut Criterion) { let batch = records[offset..end].to_vec(); let resp = shard .ingest_records(Request::new(proto::IngestRecordsRequest { - internal_protocol_version: 2, + internal_protocol_version: + unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: batch, })) .await diff --git a/docs/critical-audit-2026-09-05.md b/docs/critical-audit-2026-09-05.md new file mode 100644 index 0000000..07cc978 --- /dev/null +++ b/docs/critical-audit-2026-09-05.md @@ -0,0 +1,62 @@ +Audit and repairs based on merged commit `3d99fe0d27d7344b08b3bd5b23bb05c75479f1bd`, September 5, 2026. + +Three parallel reviews covered entity resolution, persistence/recovery, and distributed correctness. All eight original defects were reproduced with persistent storage and have enabled regression coverage. Every accepted record still completes entity resolution before its record data commits. + +| Defect and original trigger | Implemented repair | Regression suite | +| --- | --- | --- | +| A—B—C cross-shard identity chain merged incompatible strong IDs on A and C. | Check accumulated component observations before every union, across metadata chunks and prior reconciliation rounds; hydrate authoritative remote fragments. | `distributed_entity_regressions` | +| Reusing occupied `RecordId(42)` overwrote another durable record. | Reject occupied IDs in ordinary, explicit, batch, and staged storage APIs while preserving valid source-identity retries. | `durable_ingest_regressions` | +| A rejected batch left an orphan staged; a later request committed it without resolving it. | Discard aborted staging and rebuild partial derived state. Commit record data only after linking succeeds. Recover partial assignment writes from resolved durable records. | `durable_ingest_regressions` | +| Bounded interner caches changed durable values to `unknown` after restart. | Hydrate durable descriptor IDs before preparation; query lookup uses durable reverse indexes without allocating strings in the store. | `durable_ingest_regressions` | +| 300 identical observations permanently stopped later matching. | Avoid redundant interval-tree nodes; exhausted scan budgets use exact candidate fallback rather than permanently tainting a key. | `resolution_capacity_regressions` | +| Memory-saver left two matching source systems unmerged. | Respect disabled adaptive limits and use a positive fallback cap. | `resolution_capacity_regressions` | +| Cross-shard conjunction missed entities and returned incomplete golden data. | Discover candidate IDs, hydrate all fragments, then evaluate temporal conjunction and master conflicting attribute values globally. | `distributed_entity_regressions` | +| Cached-key insertion ignored hot/warm capacity limits. | Run amortized tier maintenance, persist complete buckets before eviction, promote complete buckets before updates, and propagate I/O/decode errors. | `resolution_capacity_regressions` | + +Additional regression coverage includes all four ingest APIs, a 100,001-record batch, an aborted first ingest that already flushed persistent DSU nodes, partial durable assignment failure, tier eviction/update/restart, corrupt cold buckets, deferred merge membership, temporal adjacency, and local query equivalence before/after restart. Authoritative membership also fixes the previously unsupported cluster enumeration with persistent DSU. + +Query execution reuses the linker's authoritative membership and masters golden data only for indexed candidate entities. Single-attribute labels are candidate-local. Composite labels retain a separate global key cache because their shortest unique prefixes depend on other entities; rebuilding that metadata after ingestion still requires a global pass. An unopened streaming engine retains the existing recovery/cache path. + +Router queries dispatch up to 16 shard requests concurrently and hydrate at most 1,000 entity IDs per request. A barrier-based regression proves that both query phases dispatch concurrently. Strong-ID-only reconciliation hydration skips golden data and label construction. Persistent staging lends records directly to parallel extraction, avoiding an extra record clone and the old bounded-cache cutoff. + +Baseline release-mode selective-query measurements used `PersistentStore`, batches of 1,000 distinct emails, and one matching entity: + +| Stored records | First query before repair | Warm query before repair | After one unrelated insert before repair | +| ---: | ---: | ---: | ---: | +| 5,000 | 19.454 ms | 0.002 ms | 15.732 ms | +| 20,000 | 66.230 ms | 0.006 ms | 61.903 ms | +| 80,000 | 306.473 ms | 0.003 ms | 300.151 ms | + +After repair, the same release-mode diagnostic measured: + +| Stored records | First query after repair | Warm query after repair | After one unrelated insert after repair | +| ---: | ---: | ---: | ---: | +| 5,000 | 0.055 ms | 0.006 ms | 0.010 ms | +| 20,000 | 0.022 ms | 0.005 ms | 0.007 ms | +| 80,000 | 0.029 ms | 0.005 ms | 0.008 ms | + +These timings are local diagnostics, not a universal latency guarantee; warm-cache microsecond differences are below meaningful precision for this single-run measurement. The standard historical 410K records/second figure was not reproduced on this machine; comparative throughput uses the same five persistent shards and workload for both revisions. + +The comparative ingest run used five persistent shards with the high-throughput profile, 1,000,000 records, 16 streams, batches of 5,000, 10% overlap, and seed 42: + +| Revision | Acknowledged records | Throughput | Average RPC latency | Failed batches | +| --- | ---: | ---: | ---: | ---: | +| Merged baseline | 1,000,000 | 43,415 records/s | 1,585 ms | 0 | +| Audit repairs | 1,000,000 | 43,939 records/s | 1,657 ms | 0 | + +Ingest throughput was effectively unchanged in this single comparison; the small throughput and latency differences do not establish a statistically significant change. This workload measures durable ingest, not a universal cross-shard query or reconciliation latency. The selective-query improvement is measured separately above. + +CI now passes the Rust 1.98 lints that failed on main (`chunks_exact_to_as_chunks` and `manual_slice_fill`). PR validation explicitly installs native build dependencies and verifies the release package. Release publishing depends on the reusable complete CI workflow, including MSRV, tests, lint, manifests, and production image, and rejects tags that do not match the package version. Workflow syntax is checked with actionlint. + +Validation commands: + +```sh +cargo test --locked --all-features +cargo +1.98.0 clippy --locked --all-targets --all-features -- -D warnings +cargo +1.88.0 check --locked --all-targets --all-features +cargo fmt --check +cargo package --locked +cargo test --release --test audit_performance -- --ignored --nocapture +``` + +The correctness suites run by default. Only the manual timing diagnostic is ignored; it has no machine-dependent timing assertion. Distributed benchmarks now include persistent shard ingestion rather than treating the in-memory partition path as evidence of production throughput. diff --git a/proto/unirust.proto b/proto/unirust.proto index 485a988..619df5b 100644 --- a/proto/unirust.proto +++ b/proto/unirust.proto @@ -102,6 +102,41 @@ message QueryEntitiesResponse { } } +// Internal, candidate-based entity query. Empty clusters requests the union of +// descriptor candidates; nonempty clusters hydrates every local entity fragment. +message QueryEntityFragmentsRequest { + repeated QueryDescriptor descriptors = 1; + int64 start = 2; + int64 end = 3; + repeated GlobalClusterId clusters = 4; + uint32 protocol_version = 5; + // Reconciliation needs exact component guards but no golden payload. + bool strong_ids_only = 6; +} + +message QueryInterval { + int64 start = 1; + int64 end = 2; +} + +message QueryDescriptorIntervals { + repeated QueryInterval intervals = 1; +} + +message EntityFragment { + GlobalClusterId cluster = 1; + repeated QueryDescriptorIntervals descriptor_intervals = 2; + // Raw attribute intervals; the router performs global golden conflict trimming. + repeated GoldenDescriptor golden = 3; + string cluster_key = 4; + string cluster_key_identity = 5; + repeated BoundaryStrongId strong_ids = 6; +} + +message QueryEntityFragmentsResponse { + repeated EntityFragment fragments = 1; +} + // --- Conflicts --- message RecordRef { @@ -270,6 +305,7 @@ message ConfigVersionResponse { string version = 1; OntologyConfig ontology_config = 2; uint32 protocol_version = 3; + // Durable reservation-format marker, independent of the live wire protocol. uint32 source_reservation_backfill_version = 4; uint32 source_reservation_shard_count = 5; uint32 checkpoint_protocol_version = 6; @@ -279,6 +315,7 @@ message ConfigVersionResponse { ShardRole shard_role = 9; // SHA-256 over every logical RocksDB column-family key/value pair. bytes durable_state_digest = 10; + uint32 query_fragment_protocol_version = 11; } // --- Durable Source Identity Reservation --- @@ -589,6 +626,7 @@ service ShardService { rpc IngestRecordsStream(stream IngestRecordsChunk) returns (IngestRecordsResponse); rpc IngestRecordsFromUrl(IngestRecordsFromUrlRequest) returns (IngestRecordsResponse); rpc QueryEntities(QueryEntitiesRequest) returns (QueryEntitiesResponse); + rpc QueryEntityFragments(QueryEntityFragmentsRequest) returns (QueryEntityFragmentsResponse); rpc ListConflicts(ListConflictsRequest) returns (ListConflictsResponse); // Stats and health diff --git a/src/distributed.rs b/src/distributed.rs index d4ee4b4..abe407d 100644 --- a/src/distributed.rs +++ b/src/distributed.rs @@ -9,7 +9,7 @@ use crate::persistence::{ validate_prepared_cluster_checkpoint, ClusterCheckpointManifest, PersistentOpenOptions, CHECKPOINT_PROTOCOL_VERSION, }; -use crate::query::{QueryDescriptor, QueryOutcome}; +use crate::query::QueryOutcome; use crate::sharding::{BloomFilter, IdentityKeySignature, ReconciliationCandidates}; use crate::store::{SourceRecordReservation, SourceReservationError, StoreMetrics}; use crate::temporal::Interval; @@ -68,7 +68,13 @@ struct IngestWal { const INGEST_WAL_MAGIC: &[u8; 8] = b"UNIRWAL\0"; const INGEST_WAL_VERSION: u32 = 1; const INGEST_WAL_HEADER_LEN: usize = 24; -pub const DISTRIBUTED_PROTOCOL_VERSION: u32 = 5; +/// Live router/shard and primary/replica compatibility, independent of disk formats. +pub const DISTRIBUTED_PROTOCOL_VERSION: u32 = 6; +/// Version of durable source reservations; protocol 6 preserves the v5 layout. +pub const SOURCE_RESERVATION_BACKFILL_VERSION: u32 = 5; +const QUERY_FRAGMENT_PROTOCOL_VERSION: u32 = 1; +const QUERY_SHARD_CONCURRENCY: usize = 16; +const QUERY_ENTITY_CHUNK: usize = 1_000; const REPLICATION_TOKEN_HEADER: &str = "x-unirust-replication-token"; #[allow(clippy::result_large_err)] @@ -2636,7 +2642,10 @@ impl proto::shard_service_server::ShardService for ShardNode { let mut unirust = write_unirust!(self); unirust .store_mut() - .mark_source_reservation_backfill(request.protocol_version, request.shard_count) + .mark_source_reservation_backfill( + SOURCE_RESERVATION_BACKFILL_VERSION, + request.shard_count, + ) .map_err(|err| Status::internal(err.to_string()))?; let response = proto::MarkSourceReservationsBackfilledResponse {}; local_attempt.finish(); @@ -2804,6 +2813,98 @@ impl proto::shard_service_server::ShardService for ShardNode { )) } + async fn query_entity_fragments( + &self, + request: Request, + ) -> Result, Status> { + let _mutation_guard = self.mutation_gate.read().await; + self.ensure_no_pending_ingest()?; + let start = Instant::now(); + let request = request.into_inner(); + if request.protocol_version != QUERY_FRAGMENT_PROTOCOL_VERSION { + return Err(Status::failed_precondition( + "entity fragment query protocol mismatch", + )); + } + let interval = Interval::new(request.start, request.end) + .map_err(|err| Status::invalid_argument(err.to_string()))?; + let clusters = request + .clusters + .iter() + .map(|cluster| global_cluster_id_from_proto(cluster, "query cluster")) + .collect::, _>>()?; + let unirust = read_unirust!(self); + let known = request + .descriptors + .iter() + .enumerate() + .filter_map(|(index, descriptor)| { + unirust + .lookup_query_descriptor(&descriptor.attr, &descriptor.value) + .map(|descriptor| (index, descriptor)) + }) + .collect::>(); + let descriptors = known + .iter() + .map(|(_, descriptor)| *descriptor) + .collect::>(); + let fragments = unirust + .query_entity_fragments(&descriptors, interval, &clusters, request.strong_ids_only) + .map_err(|err| Status::internal(err.to_string()))?; + let fragments = fragments + .into_iter() + .map(|fragment| { + let mut descriptor_intervals = vec![ + proto::QueryDescriptorIntervals { + intervals: Vec::new() + }; + request.descriptors.len() + ]; + for ((index, _), intervals) in known.iter().zip(fragment.descriptor_intervals) { + descriptor_intervals[*index].intervals = intervals + .into_iter() + .map(|interval| proto::QueryInterval { + start: interval.start, + end: interval.end, + }) + .collect(); + } + proto::EntityFragment { + cluster: Some(global_cluster_id_to_proto(fragment.global_id)), + descriptor_intervals, + golden: if request.strong_ids_only { + Vec::new() + } else { + fragment + .golden + .into_iter() + .map(|descriptor| proto::GoldenDescriptor { + attr: descriptor.attr, + value: descriptor.value, + start: descriptor.interval.start, + end: descriptor.interval.end, + }) + .collect() + }, + cluster_key: fragment.cluster_key.unwrap_or_default(), + cluster_key_identity: fragment.cluster_key_identity.unwrap_or_default(), + strong_ids: fragment + .strong_ids + .into_iter() + .map(boundary_strong_id_to_proto) + .collect(), + } + }) + .collect(); + if !request.strong_ids_only { + self.metrics + .record_query(start.elapsed().as_micros() as u64); + } + Ok(Response::new(proto::QueryEntityFragmentsResponse { + fragments, + })) + } + async fn query_entities( &self, request: Request, @@ -2819,11 +2920,24 @@ impl proto::shard_service_server::ShardService for ShardNode { let descriptors = request .descriptors .iter() - .map(|descriptor| QueryDescriptor { - attr: unirust.intern_attr(&descriptor.attr), - value: unirust.intern_value(&descriptor.value), - }) - .collect::>(); + .map(|descriptor| unirust.lookup_query_descriptor(&descriptor.attr, &descriptor.value)) + .collect::>>(); + + unirust + .ensure_store_healthy() + .map_err(|err| Status::internal(err.to_string()))?; + + let Some(descriptors) = descriptors else { + self.metrics + .record_query(start.elapsed().as_micros() as u64); + return Ok(Response::new(proto::QueryEntitiesResponse { + outcome: Some(proto::query_entities_response::Outcome::Matches( + proto::QueryMatches { + matches: Vec::new(), + }, + )), + })); + }; let outcome = unirust .query_master_entities(&descriptors, interval) @@ -3008,6 +3122,7 @@ impl proto::shard_service_server::ShardService for ShardNode { .source_reservation_backfill() .map_err(|err| Status::internal(err.to_string()))?; Ok(Response::new(proto::ConfigVersionResponse { + query_fragment_protocol_version: QUERY_FRAGMENT_PROTOCOL_VERSION, version: self.config_version.clone(), ontology_config: Some(to_proto_config(&ontology_config)), protocol_version: DISTRIBUTED_PROTOCOL_VERSION, @@ -3971,6 +4086,74 @@ impl proto::shard_service_server::ShardService for ShardNode { // ADAPTIVE RECONCILIATION // ============================================================================= +fn master_fragment_descriptors( + raw: Vec, +) -> Result, Status> { + let mut attributes: std::collections::BTreeMap< + String, + std::collections::BTreeMap>, + > = std::collections::BTreeMap::new(); + for descriptor in raw { + let interval = Interval::new(descriptor.start, descriptor.end).map_err(|_| { + Status::data_loss("entity fragment contains an invalid attribute interval") + })?; + attributes + .entry(descriptor.attr) + .or_default() + .entry(descriptor.value) + .or_default() + .push(interval); + } + let mut golden = Vec::new(); + for (attr, values) in attributes { + let values = values.into_iter().collect::>(); + let mut events = Vec::new(); + for (index, (_, intervals)) in values.iter().enumerate() { + for interval in crate::query::coalesce_intervals(intervals) { + events.push((interval.start, index, true)); + events.push((interval.end, index, false)); + } + } + events.sort_unstable(); + let mut active: std::collections::BTreeSet = std::collections::BTreeSet::new(); + let mut safe: Vec> = vec![Vec::new(); values.len()]; + let mut position = 0; + let mut previous = events.first().map(|event| event.0).unwrap_or_default(); + while position < events.len() { + let time = events[position].0; + if previous < time && active.len() == 1 { + if let Some(index) = active.first().copied() { + safe[index].push(Interval { + start: previous, + end: time, + }); + } + } + while position < events.len() && events[position].0 == time { + let (_, index, starts) = events[position]; + if starts { + active.insert(index); + } else { + active.remove(&index); + } + position += 1; + } + previous = time; + } + for ((value, _), intervals) in values.into_iter().zip(safe) { + for interval in crate::query::coalesce_intervals(&intervals) { + golden.push(proto::GoldenDescriptor { + attr: attr.clone(), + value: value.clone(), + start: interval.start, + end: interval.end, + }); + } + } + } + Ok(golden) +} + /// Configuration for adaptive reconciliation scheduling. #[derive(Debug, Clone)] pub struct AdaptiveReconciliationConfig { @@ -4271,6 +4454,11 @@ impl RouterNode { .map_err(|err| Status::unavailable(err.to_string()))? .into_inner(); validate_distributed_protocol(response.protocol_version)?; + if response.query_fragment_protocol_version != QUERY_FRAGMENT_PROTOCOL_VERSION { + return Err(Status::failed_precondition( + "shard does not support canonical entity fragment queries; deploy one coordinated version", + )); + } validate_checkpoint_protocol(response.checkpoint_protocol_version)?; match proto::ShardRole::try_from(response.shard_role) { Ok(proto::ShardRole::Standalone | proto::ShardRole::Primary) => {} @@ -4347,7 +4535,7 @@ impl RouterNode { same ontology used by every shard", )); } - if response.source_reservation_backfill_version == DISTRIBUTED_PROTOCOL_VERSION + if response.source_reservation_backfill_version == SOURCE_RESERVATION_BACKFILL_VERSION && response.source_reservation_shard_count != 0 && response.source_reservation_shard_count != expected_reservation_shard_count { @@ -4359,7 +4547,7 @@ impl RouterNode { ))); } source_reservation_backfill_required.push( - response.source_reservation_backfill_version != DISTRIBUTED_PROTOCOL_VERSION + response.source_reservation_backfill_version != SOURCE_RESERVATION_BACKFILL_VERSION || response.source_reservation_shard_count != expected_reservation_shard_count, ); let metadata = client @@ -5293,6 +5481,69 @@ impl RouterNode { candidates.extend(reconciler.reconciliation_candidates(&key_set)); reconcile_elapsed += reconcile_started.elapsed(); } + // Prior reconciliations can redirect a local fragment to a component + // whose strong IDs live only on another shard and another (clean) key. + // Hydrate only candidate components, preserving their guards across runs. + let guards_started = Instant::now(); + for chunk in candidates.cluster_ids().chunks(QUERY_ENTITY_CHUNK) { + let fragments = self + .fetch_entity_fragments(proto::QueryEntityFragmentsRequest { + descriptors: Vec::new(), + start: i64::MIN, + end: i64::MAX, + clusters: chunk + .iter() + .copied() + .map(global_cluster_id_to_proto) + .collect(), + protocol_version: QUERY_FRAGMENT_PROTOCOL_VERSION, + strong_ids_only: true, + }) + .await?; + let mut missing = chunk + .iter() + .copied() + .collect::>(); + for (_, fragment) in fragments { + let cluster = global_cluster_id_from_proto( + fragment + .cluster + .as_ref() + .ok_or_else(|| Status::data_loss("guard fragment has no cluster"))?, + "guard fragment", + )?; + if !chunk.contains(&cluster) { + return Err(Status::data_loss( + "guard hydration returned an unrequested cluster", + )); + } + missing.remove(&cluster); + let observations = fragment + .strong_ids + .into_iter() + .map(|observation| { + let interval = + Interval::new(observation.interval_start, observation.interval_end) + .map_err(|_| { + Status::data_loss("guard fragment contains an invalid interval") + })?; + Ok(crate::sharding::BoundaryStrongId { + perspective: observation.perspective, + attribute: observation.attribute, + value: observation.value, + interval, + }) + }) + .collect::, Status>>()?; + candidates.add_observations(cluster, observations); + } + if !missing.is_empty() { + return Err(Status::failed_precondition( + "reconciliation candidate has no authoritative entity fragments", + )); + } + } + metadata_elapsed += guards_started.elapsed(); let reconcile_started = Instant::now(); let result = candidates.finish(); reconcile_elapsed += reconcile_started.elapsed(); @@ -5320,6 +5571,218 @@ impl RouterNode { } } + async fn fetch_entity_fragments( + &self, + request: proto::QueryEntityFragmentsRequest, + ) -> Result, Status> { + let requests = futures::stream::iter(self.shard_clients.iter().cloned().enumerate().map( + |(shard, mut client)| { + let request = request.clone(); + async move { + let response = client + .query_entity_fragments(Request::new(request)) + .await? + .into_inner(); + Ok::<_, Status>((shard, response.fragments)) + } + }, + )); + let mut requests = futures::StreamExt::buffer_unordered(requests, QUERY_SHARD_CONCURRENCY); + let mut fragments = Vec::new(); + while let Some(response) = futures::StreamExt::next(&mut requests).await { + let (shard, received) = response?; + for fragment in received { + let cluster = fragment + .cluster + .as_ref() + .ok_or_else(|| Status::data_loss("entity fragment has no cluster ID"))?; + if cluster.shard_id as usize >= self.shard_clients.len() { + return Err(Status::data_loss( + "entity fragment canonical shard is out of range", + )); + } + global_cluster_id_from_proto(cluster, "entity fragment")?; + if fragment.descriptor_intervals.len() != request.descriptors.len() { + return Err(Status::data_loss( + "entity fragment descriptor count differs from request", + )); + } + for descriptor in &fragment.descriptor_intervals { + for interval in &descriptor.intervals { + if interval.start < request.start + || interval.end > request.end + || interval.start >= interval.end + { + return Err(Status::data_loss( + "entity fragment returned an invalid descriptor interval", + )); + } + } + } + fragments.push((shard, fragment)); + } + } + Ok(fragments) + } + + async fn query_global_entities( + &self, + request: &proto::QueryEntitiesRequest, + ) -> Result { + let interval = Interval::new(request.start, request.end) + .map_err(|err| Status::invalid_argument(err.to_string()))?; + let empty = || proto::QueryEntitiesResponse { + outcome: Some(proto::query_entities_response::Outcome::Matches( + proto::QueryMatches { + matches: Vec::new(), + }, + )), + }; + if request.descriptors.is_empty() { + return Ok(empty()); + } + let fragments = self + .fetch_entity_fragments(proto::QueryEntityFragmentsRequest { + descriptors: request.descriptors.clone(), + start: request.start, + end: request.end, + clusters: Vec::new(), + protocol_version: QUERY_FRAGMENT_PROTOCOL_VERSION, + strong_ids_only: false, + }) + .await?; + let mut candidates: HashMap>> = HashMap::new(); + for (_, fragment) in fragments { + let cluster = global_cluster_id_from_proto( + fragment + .cluster + .as_ref() + .ok_or_else(|| Status::data_loss("entity fragment has no cluster"))?, + "entity fragment", + )?; + let intervals = candidates + .entry(cluster) + .or_insert_with(|| vec![Vec::new(); request.descriptors.len()]); + for (target, descriptor) in intervals.iter_mut().zip(fragment.descriptor_intervals) { + target.extend(descriptor.intervals.into_iter().map(|interval| Interval { + start: interval.start, + end: interval.end, + })); + } + } + let mut matches_by_cluster = HashMap::new(); + for (cluster, descriptors) in candidates { + let mut matches = vec![interval]; + for descriptor_intervals in descriptors { + matches = crate::query::intersect_interval_sets( + &matches, + &crate::query::coalesce_intervals(&descriptor_intervals), + ); + if matches.is_empty() { + break; + } + } + if !matches.is_empty() { + matches_by_cluster.insert(cluster, matches); + } + } + if matches_by_cluster.is_empty() { + return Ok(empty()); + } + let mut cluster_ids = matches_by_cluster.keys().copied().collect::>(); + cluster_ids.sort_by_key(GlobalClusterId::to_u64); + let mut matches = Vec::new(); + for chunk in cluster_ids.chunks(QUERY_ENTITY_CHUNK) { + let fragments = self + .fetch_entity_fragments(proto::QueryEntityFragmentsRequest { + descriptors: Vec::new(), + start: request.start, + end: request.end, + clusters: chunk + .iter() + .copied() + .map(global_cluster_id_to_proto) + .collect(), + protocol_version: QUERY_FRAGMENT_PROTOCOL_VERSION, + strong_ids_only: false, + }) + .await?; + let mut hydrated: HashMap> = + HashMap::new(); + let mut keys: HashMap = HashMap::new(); + for (source, fragment) in fragments { + let cluster = global_cluster_id_from_proto( + fragment + .cluster + .as_ref() + .ok_or_else(|| Status::data_loss("entity fragment has no cluster"))?, + "entity fragment", + )?; + if !chunk.contains(&cluster) { + return Err(Status::unavailable( + "entity membership changed during query hydration; retry the query", + )); + } + hydrated.entry(cluster).or_default().extend(fragment.golden); + let priority = if source == usize::from(cluster.shard_id) { + 0 + } else { + source + 1 + }; + let key = ( + priority, + fragment.cluster_key, + fragment.cluster_key_identity, + ); + if keys.get(&cluster).is_none_or(|existing| key < *existing) { + keys.insert(cluster, key); + } + } + for cluster in chunk { + let raw = hydrated.remove(cluster).ok_or_else(|| { + Status::unavailable( + "entity membership changed during query hydration; retry the query", + ) + })?; + let golden = master_fragment_descriptors(raw)?; + let (_, cluster_key, cluster_key_identity) = + keys.remove(cluster).unwrap_or_default(); + for interval in &matches_by_cluster[cluster] { + let golden = golden + .iter() + .filter_map(|descriptor| { + let start = descriptor.start.max(interval.start); + let end = descriptor.end.min(interval.end); + (start < end).then(|| proto::GoldenDescriptor { + attr: descriptor.attr.clone(), + value: descriptor.value.clone(), + start, + end, + }) + }) + .collect(); + matches.push(proto::QueryMatch { + shard_id: u32::from(cluster.shard_id), + cluster_id: cluster.local_id, + start: interval.start, + end: interval.end, + cluster_key: cluster_key.clone(), + cluster_key_identity: cluster_key_identity.clone(), + golden, + }); + } + } + } + Ok(self.merge_query_responses( + &request.descriptors, + vec![proto::QueryEntitiesResponse { + outcome: Some(proto::query_entities_response::Outcome::Matches( + proto::QueryMatches { matches }, + )), + }], + )) + } + fn merge_query_responses( &self, descriptors: &[proto::QueryDescriptor], @@ -5589,17 +6052,15 @@ impl proto::router_service_server::RouterService for RouterNode { self.ensure_cluster_consistent()?; let start = Instant::now(); let request = request.into_inner(); - let mut responses = Vec::with_capacity(self.shard_clients.len()); - for client in &self.shard_clients { - let mut client = client.clone(); - let response = client - .query_entities(Request::new(request.clone())) - .await - .map_err(|err| Status::unavailable(err.to_string()))?; - responses.push(response.into_inner()); - } - - let merged = self.merge_query_responses(&request.descriptors, responses); + let merged = if self.shard_clients.len() == 1 { + self.shard_clients[0] + .clone() + .query_entities(Request::new(request)) + .await? + .into_inner() + } else { + self.query_global_entities(&request).await? + }; self.metrics .record_query(start.elapsed().as_micros() as u64); Ok(Response::new(merged)) @@ -5669,10 +6130,11 @@ impl proto::router_service_server::RouterService for RouterNode { ) -> Result, Status> { let ontology_config = self.ontology_config.read().await; Ok(Response::new(proto::ConfigVersionResponse { + query_fragment_protocol_version: QUERY_FRAGMENT_PROTOCOL_VERSION, version: self.config_version.clone(), ontology_config: Some(to_proto_config(&ontology_config)), protocol_version: DISTRIBUTED_PROTOCOL_VERSION, - source_reservation_backfill_version: DISTRIBUTED_PROTOCOL_VERSION, + source_reservation_backfill_version: SOURCE_RESERVATION_BACKFILL_VERSION, source_reservation_shard_count: self.shard_clients.len() as u32, checkpoint_protocol_version: CHECKPOINT_PROTOCOL_VERSION, restore_generation: self.restore_generation.clone().unwrap_or_default(), @@ -6447,7 +6909,7 @@ mod wal_tests { ShardService::ingest_records( &shard, Request::new(proto::IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: DISTRIBUTED_PROTOCOL_VERSION, records: vec![original.clone()], }), ) @@ -6459,7 +6921,7 @@ mod wal_tests { let error = ShardService::ingest_records( &shard, Request::new(proto::IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: DISTRIBUTED_PROTOCOL_VERSION, records: vec![conflicting], }), ) @@ -6552,4 +7014,95 @@ mod wal_tests { assert_eq!(error.code(), tonic::Code::FailedPrecondition); assert_eq!(shard.unirust.read().record_count(), 0); } + + #[tokio::test] + async fn query_lookup_corruption_cannot_return_successful_empty_matches() { + use proto::shard_service_server::ShardService; + let directory = tempfile::tempdir().expect("persistent query directory"); + let shard = ShardNode::new_with_data_dir( + 0, + DistributedOntologyConfig::empty(), + StreamingTuning::balanced(), + Some(directory.path().to_path_buf()), + false, + None, + ) + .expect("persistent shard"); + let db = shard + .unirust + .read() + .store() + .shared_db() + .expect("persistent DB"); + let interner = db.cf_handle("interner").expect("interner column family"); + db.put_cf(interner, b"Acorrupt-query-only-attr", b"bad") + .expect("inject invalid lookup ID"); + let error = ShardService::query_entities( + &shard, + Request::new(proto::QueryEntitiesRequest { + descriptors: vec![proto::QueryDescriptor { + attr: "corrupt-query-only-attr".into(), + value: "unused".into(), + }], + start: 0, + end: 100, + }), + ) + .await + .expect_err("durable lookup faults must not masquerade as an unknown descriptor"); + assert_eq!(error.code(), tonic::Code::Internal); + assert!( + ShardService::health_check(&shard, Request::new(proto::HealthCheckRequest {})) + .await + .is_err() + ); + } + + #[tokio::test] + async fn protocol_upgrade_replays_existing_wal_and_keeps_v5_reservation_format() { + use proto::shard_service_server::ShardService; + let directory = tempfile::tempdir().expect("upgrade directory"); + { + let mut store = PersistentStore::open(directory.path()).expect("v5 volume"); + crate::store::RecordStore::mark_source_reservation_backfill(&mut store, 5, 1) + .expect("persist existing reservation-format marker"); + } + // Protocol 5 already used this WAL format; wire versions are not embedded + // in its immutable record payload, so replay must not revalidate them. + assert_eq!(INGEST_WAL_VERSION, 1); + let wal = IngestWal::new(directory.path()); + wal.write_batch(&[proto_from_wal_input(sample_wal_record())]) + .expect("pending v1 WAL"); + let shard = ShardNode::new_with_data_dir( + 0, + DistributedOntologyConfig::empty(), + StreamingTuning::balanced(), + Some(directory.path().to_path_buf()), + false, + None, + ) + .expect("upgrade and replay existing WAL"); + let config = ShardService::get_config_version( + &shard, + Request::new(proto::ConfigVersionRequest { + include_durable_state_digest: false, + }), + ) + .await + .expect("upgraded config") + .into_inner(); + assert_eq!(config.protocol_version, 6); + assert_eq!(config.source_reservation_backfill_version, 5); + assert_eq!(config.source_reservation_shard_count, 1); + assert_eq!( + ShardService::get_stats(&shard, Request::new(proto::StatsRequest {})) + .await + .expect("replayed stats") + .into_inner() + .record_count, + 1 + ); + assert!(!wal.has_pending()); + shard.shutdown().await.expect("upgraded shutdown"); + } } diff --git a/src/graph.rs b/src/graph.rs index 9db1aeb..359b093 100644 --- a/src/graph.rs +++ b/src/graph.rs @@ -834,6 +834,19 @@ pub fn golden_for_cluster( golden } +/// Single-token labels cannot change their prefix length when other clusters +/// appear. Composite labels need the full collision group to keep stable keys. +pub(crate) fn cluster_keys_are_candidate_local(ontology: &Ontology) -> bool { + ontology + .identity_keys + .iter() + .all(|key| key.attributes.len() <= 1 && key.attribute_names.len() <= 1) + && ontology + .entity_types + .values() + .all(|entity| entity.key_attributes.len() <= 1) +} + pub fn cluster_keys_for_clusters( store: &dyn RecordStore, clusters: &Clusters, @@ -1216,6 +1229,30 @@ mod tests { use crate::store::Store; use crate::temporal::Interval; + #[test] + fn candidate_local_keys_require_single_attribute_identity_and_entity_keys() { + let mut ontology = Ontology::new(); + assert!(cluster_keys_are_candidate_local(&ontology)); + ontology.add_identity_key(IdentityKey::from_names(vec!["email"], "email")); + assert!(cluster_keys_are_candidate_local(&ontology)); + ontology.add_identity_key(IdentityKey::from_names(vec!["name", "email"], "name_email")); + assert!(!cluster_keys_are_candidate_local(&ontology)); + ontology.identity_keys.clear(); + ontology.add_identity_key(IdentityKey::new( + vec![AttrId(0), AttrId(1)], + "composite".into(), + )); + assert!(!cluster_keys_are_candidate_local(&ontology)); + ontology.identity_keys.clear(); + ontology.add_entity_type(crate::ontology::EntityType::with_key_attributes( + "person".into(), + vec![AttrId(0), AttrId(1)], + vec![AttrId(0), AttrId(1)], + false, + )); + assert!(!cluster_keys_are_candidate_local(&ontology)); + } + #[test] fn test_same_as_edge_creation() { let edge = SameAsEdge::new( diff --git a/src/index.rs b/src/index.rs index a3e4644..7a89ea9 100644 --- a/src/index.rs +++ b/src/index.rs @@ -44,12 +44,15 @@ struct ClusterIntervalList { } impl ClusterIntervalList { - fn add_interval(&mut self, interval: Interval) { + fn add_interval(&mut self, interval: Interval) -> bool { let mut start = interval.start; let mut end = interval.end; let mut idx = 0; while idx < self.intervals.len() { let current = self.intervals[idx]; + if current.start <= start && end <= current.end { + return false; + } if end <= current.start { break; } @@ -62,6 +65,7 @@ impl ClusterIntervalList { self.intervals.remove(idx); } self.intervals.insert(idx, Interval { start, end }); + true } fn extend_from(&mut self, other: &ClusterIntervalList) { @@ -210,7 +214,9 @@ impl CandidateList { fn insert_cluster_interval(&mut self, root_id: RecordId, interval: Interval) { let list = self.cluster_intervals.entry(root_id).or_default(); let prev_len = list.intervals.len(); - list.add_interval(interval); + if !list.add_interval(interval) { + return; + } let new_len = list.intervals.len(); self.active_intervals = self.active_intervals.saturating_sub(prev_len) + new_len; self.total_intervals = self.total_intervals.saturating_add(1); @@ -287,6 +293,13 @@ impl CandidateList { } fn merge_cluster_intervals(&mut self, root_a: RecordId, root_b: RecordId, new_root: RecordId) { + // The incoming record is not indexed yet. Most duplicate merges keep the + // existing root, so its indexed intervals and tree need no replacement. + if (new_root == root_a && !self.cluster_intervals.contains_key(&root_b)) + || (new_root == root_b && !self.cluster_intervals.contains_key(&root_a)) + { + return; + } let mut merged = ClusterIntervalList::default(); let mut removed = 0usize; @@ -1073,72 +1086,80 @@ impl TieredIdentityKeyIndex { self.add_record_with_root(record, record.id, ontology) } - /// Add a record with a specific root + /// Add a record with a specific root. pub fn add_record_with_root( &mut self, record: &crate::model::Record, root_id: RecordId, ontology: &crate::ontology::Ontology, - ) -> anyhow::Result<()> { - let entity_type = &record.identity.entity_type; - let identity_keys = ontology.identity_keys_for_type(entity_type); - - for identity_key in identity_keys { - // Use a temporary IdentityKeyIndex to extract key values - let temp_index = IdentityKeyIndex::new(); - let key_values_with_intervals = - temp_index.extract_key_values_with_intervals(record, identity_key)?; - - if !key_values_with_intervals.is_empty() { - for (key_values, interval) in key_values_with_intervals { - let key = IdentityIndexKey { - entity_type: entity_type.clone(), - key_values, - }; - - // Always add to hot tier for new records - let bucket = self.hot.entry(key.clone()).or_default(); - bucket.insert_record(record.id, interval); - bucket.insert_cluster_interval(root_id, interval); - - // Update stats - let stats = self.hot_stats.entry(key).or_default(); - stats.cardinality = stats.cardinality.saturating_add(1); - stats.last_access = Self::current_time(); - } - - self.record_keys - .entry(record.id) - .or_default() - .push(identity_key.clone()); - } + ) -> Result<()> { + let keys = ontology + .identity_keys_for_type(&record.identity.entity_type) + .into_iter() + .map(|key| Ok((key, extract_key_values_from_record(record, key)?))) + .collect::>>()?; + self.add_record_with_cached_keys(record.id, root_id, &record.identity.entity_type, keys) + } + + /// Load the complete bucket before reading or modifying a key. A new hot + /// fragment must never hide observations already stored in lower tiers. + fn promote_to_hot(&mut self, key: &IdentityIndexKey) -> Result { + if self.hot.contains_key(key) { + return Ok(true); + } + let compact = if let Some(bucket) = self.warm.pop(key) { + Some(bucket) + } else if let Some(db) = &self.db { + let cf = db + .cf_handle(self.cf_identity_keys) + .ok_or_else(|| anyhow::anyhow!("missing identity keys column family"))?; + let encoded = crate::persistence::index_encoding::encode_identity_key( + &key.entity_type, + &key.key_values, + ); + db.get_cf(cf, encoded)? + .map(|bytes| { + crate::persistence::index_encoding::decode_compact_bucket(&bytes) + .map(CompactBucket::from_data) + }) + .transpose()? + } else { + None + }; + let Some(compact) = compact else { + return Ok(false); + }; + let mut bucket = KeyBucket::default(); + for (id, start, end) in compact.record_intervals { + bucket.insert_record(RecordId(id), Interval::new(start, end)?); } - - // Run tier management periodically - self.maybe_manage_tiers(); - - Ok(()) + for (id, start, end) in compact.cluster_intervals { + bucket.insert_cluster_interval(RecordId(id), Interval::new(start, end)?); + } + self.hot.insert(key.clone(), bucket); + let stats = self.warm_stats.remove(key).unwrap_or_default(); + self.hot_stats.insert(key.clone(), stats); + Ok(true) } - /// Find matching clusters overlapping an interval pub fn find_matching_clusters_overlapping( &mut self, dsu: &mut DsuBackend, entity_type: &str, key_values: &[KeyValue], interval: Interval, - ) -> &[(RecordId, Interval)] { - self.find_matching_clusters_overlapping_limited( - dsu, - entity_type, - key_values, - interval, - usize::MAX, - ) - .0 + ) -> Result<&[(RecordId, Interval)]> { + Ok(self + .find_matching_clusters_overlapping_limited( + dsu, + entity_type, + key_values, + interval, + usize::MAX, + )? + .0) } - /// Find matching clusters with a limit pub fn find_matching_clusters_overlapping_limited( &mut self, dsu: &mut DsuBackend, @@ -1146,106 +1167,37 @@ impl TieredIdentityKeyIndex { key_values: &[KeyValue], interval: Interval, max_tree_nodes: usize, - ) -> (&[(RecordId, Interval)], bool) { - let key = IdentityIndexKeyRef { - entity_type, - key_values, - }; - - self.cluster_overlap_buffer.clear(); - self.cluster_seen.clear(); - - // Check hot tier first - if let Some(bucket) = self.hot.get_mut(&key) { - // Record access - if let Some(stats) = self.hot_stats.get_mut(&IdentityIndexKey { - entity_type: entity_type.to_string(), - key_values: key_values.to_vec(), - }) { - stats.access_count = stats.access_count.saturating_add(1); - stats.total_queries = stats.total_queries.saturating_add(1); - stats.last_access = Self::current_time(); - } - - let limit_reached = bucket - .record_candidates - .collect_overlapping_clusters_limited( - dsu, - interval, - &mut self.cluster_overlap_buffer, - &mut self.cluster_seen, - max_tree_nodes, - ); - return (self.cluster_overlap_buffer.as_slice(), limit_reached); - } - - // Check warm tier - let owned_key = IdentityIndexKey { + ) -> Result<(&[(RecordId, Interval)], bool)> { + let key = IdentityIndexKey { entity_type: entity_type.to_string(), key_values: key_values.to_vec(), }; - - if let Some(compact) = self.warm.get(&owned_key) { - // Record access - if let Some(stats) = self.warm_stats.get_mut(&owned_key) { - stats.access_count = stats.access_count.saturating_add(1); - stats.total_queries = stats.total_queries.saturating_add(1); - stats.last_access = Self::current_time(); - } - - // Linear scan for warm tier - let results = compact.find_overlapping_clusters(interval); - for (root_id, cluster_interval) in results { - let root = dsu.find(root_id).unwrap_or(root_id); - if self.cluster_seen.insert(root) { - self.cluster_overlap_buffer.push((root, cluster_interval)); - } - } - return (self.cluster_overlap_buffer.as_slice(), false); + self.cluster_overlap_buffer.clear(); + self.cluster_seen.clear(); + if !self.promote_to_hot(&key)? { + return Ok((&[], false)); } - - // Check cold tier (RocksDB) - if let Some(db) = &self.db { - if let Some(cf) = db.cf_handle(self.cf_identity_keys) { - let key_bytes = crate::persistence::index_encoding::encode_identity_key( - entity_type, - key_values, - ); - if let Ok(Some(data_bytes)) = db.get_cf(cf, &key_bytes) { - if let Ok(data) = - crate::persistence::index_encoding::decode_compact_bucket(&data_bytes) - { - let compact = CompactBucket::from_data(data); - let results = compact.find_overlapping_clusters(interval); - for (root_id, cluster_interval) in results { - let root = dsu.find(root_id).unwrap_or(root_id); - if self.cluster_seen.insert(root) { - self.cluster_overlap_buffer.push((root, cluster_interval)); - } - } - - // Promote to warm tier - self.warm.put(owned_key.clone(), compact); - self.warm_stats.insert( - owned_key, - KeyAccessStats { - access_count: 1, - last_access: Self::current_time(), - total_queries: 1, - cardinality: 0, - }, - ); - - return (self.cluster_overlap_buffer.as_slice(), false); - } - } - } + if let Some(stats) = self.hot_stats.get_mut(&key) { + stats.access_count = stats.access_count.saturating_add(1); + stats.total_queries = stats.total_queries.saturating_add(1); + stats.last_access = Self::current_time(); } - - (&[], false) + let limited = self + .hot + .get_mut(&key) + .ok_or_else(|| anyhow::anyhow!("promoted identity bucket is unavailable"))? + .record_candidates + .collect_overlapping_clusters_limited( + dsu, + interval, + &mut self.cluster_overlap_buffer, + &mut self.cluster_seen, + max_tree_nodes, + ); + self.maybe_manage_tiers()?; + Ok((self.cluster_overlap_buffer.as_slice(), limited)) } - /// Merge clusters in the index pub fn merge_key_clusters( &mut self, entity_type: &str, @@ -1253,131 +1205,118 @@ impl TieredIdentityKeyIndex { root_a: RecordId, root_b: RecordId, new_root: RecordId, - ) { - let key = IdentityIndexKeyRef { - entity_type, - key_values, - }; - - // Try hot tier - if let Some(bucket) = self.hot.get_mut(&key) { - bucket.merge_clusters(root_a, root_b, new_root); - return; - } - - // Try warm tier - need to convert to hot for merge - let owned_key = IdentityIndexKey { + ) -> Result<()> { + let key = IdentityIndexKey { entity_type: entity_type.to_string(), key_values: key_values.to_vec(), }; - - if self.warm.contains(&owned_key) { - // Promote to hot for merge operation - if let Some(compact) = self.warm.pop(&owned_key) { - let mut bucket = KeyBucket::default(); - // Convert compact back to full bucket - for (id, start, end) in compact.record_intervals { - if let Ok(interval) = Interval::new(start, end) { - bucket.insert_record(RecordId(id), interval); - } - } - for (id, start, end) in compact.cluster_intervals { - if let Ok(interval) = Interval::new(start, end) { - bucket.insert_cluster_interval(RecordId(id), interval); - } - } + if self.promote_to_hot(&key)? { + if let Some(bucket) = self.hot.get_mut(&key) { bucket.merge_clusters(root_a, root_b, new_root); - self.hot.insert(owned_key, bucket); } } + self.maybe_manage_tiers() } - /// Run tier management if enough time has passed - fn maybe_manage_tiers(&mut self) { - let now = Self::current_time(); - if now - self.last_tier_management < self.config.tier_management_interval_secs as i64 { - return; - } - self.last_tier_management = now; - - // Only demote if hot tier is over capacity - if self.hot.len() <= self.config.hot_tier_capacity { - return; + /// Persist demotions before removing their hot copies. Warm LRU eviction is + /// then safe because every warm bucket has a complete durable backing copy. + fn maybe_manage_tiers(&mut self) -> Result<()> { + let capacity = self.config.hot_tier_capacity; + if self.hot.len() <= capacity { + return Ok(()); } - - self.demote_cold_keys(); - } - - /// Demote cold keys from hot to warm tier - fn demote_cold_keys(&mut self) { let now = Self::current_time(); - let max_card = self.config.max_hot_cardinality; - - // Collect keys to demote - let mut to_demote: Vec<(IdentityIndexKey, f64)> = self + self.last_tier_management = now; + let mut keys = self .hot_stats .iter() - .map(|(key, stats)| (key.clone(), stats.tier_score(now, max_card))) - .filter(|(_, score)| *score < self.config.hot_threshold) - .collect(); - - // Sort by score ascending (coldest first) - to_demote.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)); - - // Demote until under capacity - let to_remove = self.hot.len().saturating_sub(self.config.hot_tier_capacity); - for (key, _) in to_demote.into_iter().take(to_remove) { + .map(|(key, stats)| { + ( + key.clone(), + stats.tier_score(now, self.config.max_hot_cardinality), + ) + }) + .collect::>(); + keys.sort_by(|a, b| a.1.total_cmp(&b.1)); + // Leave headroom so a full-sized index does not sort all keys per insert. + let mut count = self.hot.len().saturating_sub(capacity).max(capacity / 10); + if self.db.is_none() { + // Without durable storage we retain overflow rather than lose keys. + count = count.min(self.warm.cap().get().saturating_sub(self.warm.len())); + } + keys.truncate(count); + if let Some(db) = &self.db { + let cf = db + .cf_handle(self.cf_identity_keys) + .ok_or_else(|| anyhow::anyhow!("missing identity keys column family"))?; + let mut batch = rocksdb::WriteBatch::default(); + for (key, _) in &keys { + if let Some(bucket) = self.hot.get(key) { + let compact = CompactBucket::from_key_bucket(bucket); + let encoded = crate::persistence::index_encoding::encode_identity_key( + &key.entity_type, + &key.key_values, + ); + batch.put_cf( + cf, + encoded, + crate::persistence::index_encoding::encode_compact_bucket( + &compact.to_data(), + )?, + ); + } + } + db.write(batch)?; + } + for (key, _) in keys { if let Some(bucket) = self.hot.remove(&key) { let compact = CompactBucket::from_key_bucket(&bucket); - - if let Some(stats) = self.hot_stats.remove(&key) { - // Move to warm tier - self.warm.put(key.clone(), compact); - self.warm_stats.insert(key, stats); + if let Some((evicted, _)) = self.warm.push(key.clone(), compact) { + self.warm_stats.remove(&evicted); } + let stats = self.hot_stats.remove(&key).unwrap_or_default(); + self.warm_stats.insert(key, stats); } } + Ok(()) } - /// Flush warm tier to cold tier (RocksDB) - pub fn flush_warm_to_cold(&mut self) -> anyhow::Result<()> { + /// Checkpoint every resident bucket, including modified hot buckets. + pub fn flush_warm_to_cold(&mut self) -> Result<()> { let Some(db) = &self.db else { return Ok(()); }; - let cf = db .cf_handle(self.cf_identity_keys) .ok_or_else(|| anyhow::anyhow!("missing identity keys column family"))?; - - let stats_cf = db - .cf_handle(self.cf_key_stats) - .ok_or_else(|| anyhow::anyhow!("missing key stats column family"))?; - let mut batch = rocksdb::WriteBatch::default(); - - // Only flush entries that are being evicted from warm tier - // For now, flush all warm entries to cold as backup - for (key, compact) in self.warm.iter() { - let key_bytes = crate::persistence::index_encoding::encode_identity_key( + for (key, bucket) in self.warm.iter() { + let encoded = crate::persistence::index_encoding::encode_identity_key( &key.entity_type, &key.key_values, ); - let data = compact.to_data(); - let data_bytes = crate::persistence::index_encoding::encode_compact_bucket(&data)?; - batch.put_cf(cf, &key_bytes, data_bytes); - - // Also save stats - if let Some(stats) = self.warm_stats.get(key) { - let stats_bytes = crate::persistence::index_encoding::encode_key_stats(stats)?; - batch.put_cf(stats_cf, &key_bytes, stats_bytes); - } + batch.put_cf( + cf, + encoded, + crate::persistence::index_encoding::encode_compact_bucket(&bucket.to_data())?, + ); + } + for (key, bucket) in &self.hot { + let encoded = crate::persistence::index_encoding::encode_identity_key( + &key.entity_type, + &key.key_values, + ); + let compact = CompactBucket::from_key_bucket(bucket); + batch.put_cf( + cf, + encoded, + crate::persistence::index_encoding::encode_compact_bucket(&compact.to_data())?, + ); } - db.write(batch)?; Ok(()) } - /// Get record keys for a record pub fn get_record_keys(&self, record_id: RecordId) -> Vec { self.record_keys .get(&record_id) @@ -1385,7 +1324,6 @@ impl TieredIdentityKeyIndex { .unwrap_or_default() } - /// Get tier statistics pub fn tier_stats(&self) -> TieredIndexStats { TieredIndexStats { hot_keys: self.hot.len(), @@ -1396,36 +1334,53 @@ impl TieredIdentityKeyIndex { } } - /// Build index from records (bulk load) - pub fn build( - &mut self, - records: &[crate::model::Record], - ontology: &crate::ontology::Ontology, - ) -> anyhow::Result<()> { + /// Clear derived state before replaying authoritative records. Old cold + /// buckets cannot be mixed with the partially reconstructed DSU and summaries. + pub fn clear(&mut self) -> Result<()> { + if let Some(db) = &self.db { + for name in [self.cf_identity_keys, self.cf_key_stats] { + let cf = db + .cf_handle(name) + .ok_or_else(|| anyhow::anyhow!("missing index column family {name}"))?; + let mut batch = rocksdb::WriteBatch::default(); + for entry in db.iterator_cf(cf, rocksdb::IteratorMode::Start) { + let (key, _) = entry?; + batch.delete_cf(cf, key); + if batch.len() >= 4096 { + db.write(std::mem::take(&mut batch))?; + } + } + db.write(batch)?; + } + } self.hot.clear(); self.hot_stats.clear(); self.warm.clear(); self.warm_stats.clear(); self.record_keys.clear(); + Ok(()) + } + pub fn build( + &mut self, + records: &[crate::model::Record], + ontology: &crate::ontology::Ontology, + ) -> Result<()> { + self.clear()?; for record in records { self.add_record(record, ontology)?; } - Ok(()) } - /// Extract key values with intervals (delegates to IdentityKeyIndex helper) pub fn extract_key_values_with_intervals( &self, record: &crate::model::Record, identity_key: &crate::ontology::IdentityKey, - ) -> anyhow::Result, Interval)>> { - // Use the shared extraction logic + ) -> Result, Interval)>> { extract_key_values_from_record(record, identity_key) } - /// Add a record using pre-extracted key values (avoids duplicate extraction). #[allow(clippy::type_complexity)] pub fn add_record_with_cached_keys( &mut self, @@ -1436,78 +1391,50 @@ impl TieredIdentityKeyIndex { &crate::ontology::IdentityKey, Vec<(Vec, Interval)>, )>, - ) { - for (identity_key, key_values_with_intervals) in cached_keys { - if !key_values_with_intervals.is_empty() { - self.record_keys - .entry(record_id) - .or_default() - .push(identity_key.clone()); - - for (key_values, interval) in key_values_with_intervals { - let key = IdentityIndexKey { - entity_type: entity_type.to_string(), - key_values, - }; - - // Always add to hot tier - let bucket = self.hot.entry(key.clone()).or_default(); - bucket.insert_record(record_id, interval); - bucket.insert_cluster_interval(root_id, interval); - - // Update stats - let stats = self.hot_stats.entry(key).or_default(); - stats.cardinality = stats.cardinality.saturating_add(1); - stats.last_access = Self::current_time(); - } + ) -> Result<()> { + for (identity_key, values) in cached_keys { + if values.is_empty() { + continue; + } + self.record_keys + .entry(record_id) + .or_default() + .push(identity_key.clone()); + for (key_values, interval) in values { + let key = IdentityIndexKey { + entity_type: entity_type.to_string(), + key_values, + }; + self.promote_to_hot(&key)?; + let bucket = self.hot.entry(key.clone()).or_default(); + bucket.insert_record(record_id, interval); + bucket.insert_cluster_interval(root_id, interval); + let stats = self.hot_stats.entry(key).or_default(); + stats.cardinality = stats.cardinality.saturating_add(1); + stats.last_access = Self::current_time(); } } + self.maybe_manage_tiers() } - /// Find records that match a given identity key pub fn find_matching_records( &mut self, entity_type: &str, key_values: &[KeyValue], - ) -> &[(RecordId, Interval)] { - // First check warm tier and promote if needed (must happen before hot tier borrow) - let owned_key = IdentityIndexKey { + ) -> Result<&[(RecordId, Interval)]> { + let key = IdentityIndexKey { entity_type: entity_type.to_string(), key_values: key_values.to_vec(), }; - - // Promotion from warm to hot must happen before we borrow hot tier - if !self.hot.contains_key(&owned_key) && self.warm.contains(&owned_key) { - if let Some(compact) = self.warm.pop(&owned_key) { - let mut bucket = KeyBucket::default(); - for (id, start, end) in compact.record_intervals { - if let Ok(interval) = Interval::new(start, end) { - bucket.insert_record(RecordId(id), interval); - } - } - for (id, start, end) in compact.cluster_intervals { - if let Ok(interval) = Interval::new(start, end) { - bucket.insert_cluster_interval(RecordId(id), interval); - } - } - self.hot.insert(owned_key.clone(), bucket); - if let Some(stats) = self.warm_stats.remove(&owned_key) { - self.hot_stats.insert(owned_key, stats); - } + self.cluster_overlap_buffer.clear(); + if self.promote_to_hot(&key)? { + if let Some(bucket) = self.hot.get_mut(&key) { + self.cluster_overlap_buffer + .extend_from_slice(bucket.record_candidates.as_slice()); } } - - // Now check hot tier - use get_mut since as_slice() requires &mut self - let key = IdentityIndexKeyRef { - entity_type, - key_values, - }; - - if let Some(bucket) = self.hot.get_mut(&key) { - return bucket.record_candidates.as_slice(); - } - - &[] + self.maybe_manage_tiers()?; + Ok(self.cluster_overlap_buffer.as_slice()) } } @@ -1574,10 +1501,11 @@ impl IndexBackend { &crate::ontology::IdentityKey, Vec<(Vec, Interval)>, )>, - ) { + ) -> Result<()> { match self { IndexBackend::InMemory(index) => { - index.add_record_with_cached_keys(record_id, root_id, entity_type, cached_keys) + index.add_record_with_cached_keys(record_id, root_id, entity_type, cached_keys); + Ok(()) } IndexBackend::Tiered(index) => { index.add_record_with_cached_keys(record_id, root_id, entity_type, cached_keys) @@ -1593,15 +1521,15 @@ impl IndexBackend { key_values: &[KeyValue], interval: Interval, limit: usize, - ) -> (&[(RecordId, Interval)], bool) { + ) -> Result<(&[(RecordId, Interval)], bool)> { match self { - IndexBackend::InMemory(index) => index.find_matching_clusters_overlapping_limited( + IndexBackend::InMemory(index) => Ok(index.find_matching_clusters_overlapping_limited( dsu, entity_type, key_values, interval, limit, - ), + )), IndexBackend::Tiered(index) => index.find_matching_clusters_overlapping_limited( dsu, entity_type, @@ -1617,9 +1545,11 @@ impl IndexBackend { &mut self, entity_type: &str, key_values: &[KeyValue], - ) -> &[(RecordId, Interval)] { + ) -> Result<&[(RecordId, Interval)]> { match self { - IndexBackend::InMemory(index) => index.find_matching_records(entity_type, key_values), + IndexBackend::InMemory(index) => { + Ok(index.find_matching_records(entity_type, key_values)) + } IndexBackend::Tiered(index) => index.find_matching_records(entity_type, key_values), } } @@ -1632,10 +1562,11 @@ impl IndexBackend { root_a: RecordId, root_b: RecordId, new_root: RecordId, - ) { + ) -> Result<()> { match self { IndexBackend::InMemory(index) => { - index.merge_key_clusters(entity_type, key_values, root_a, root_b, new_root) + index.merge_key_clusters(entity_type, key_values, root_a, root_b, new_root); + Ok(()) } IndexBackend::Tiered(index) => { index.merge_key_clusters(entity_type, key_values, root_a, root_b, new_root) @@ -1643,6 +1574,16 @@ impl IndexBackend { } } + pub fn clear(&mut self) -> Result<()> { + match self { + Self::InMemory(index) => { + *index = IdentityKeyIndex::new(); + Ok(()) + } + Self::Tiered(index) => index.clear(), + } + } + /// Check if using tiered backend pub fn is_tiered(&self) -> bool { matches!(self, IndexBackend::Tiered(_)) @@ -2089,12 +2030,14 @@ mod tests { } let key_values = vec![KeyValue::new(name_attr, ValueId(1))]; - let results = index.find_matching_clusters_overlapping( - &mut dsu, - "person", - &key_values, - Interval::new(150, 175).unwrap(), - ); + let results = index + .find_matching_clusters_overlapping( + &mut dsu, + "person", + &key_values, + Interval::new(150, 175).unwrap(), + ) + .unwrap(); assert_eq!(results.len(), 2); } diff --git a/src/lib.rs b/src/lib.rs index 483eb14..b8f4a07 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -200,6 +200,8 @@ pub struct Unirust { streaming: Option, graph_state: Option, query_cache: std::sync::Mutex>, + cluster_key_cache: + std::sync::Mutex>>, conflict_cache: std::sync::Mutex>>, query_stats: std::sync::Mutex, tuning: StreamingTuning, @@ -237,6 +239,7 @@ impl Unirust { streaming: None, graph_state: None, query_cache: std::sync::Mutex::new(None), + cluster_key_cache: std::sync::Mutex::new(None), conflict_cache: std::sync::Mutex::new(None), query_stats: std::sync::Mutex::new(query::QuerySelectivityStats::default()), tuning: StreamingTuning::default(), @@ -265,6 +268,7 @@ impl Unirust { streaming: None, graph_state: None, query_cache: std::sync::Mutex::new(None), + cluster_key_cache: std::sync::Mutex::new(None), conflict_cache: std::sync::Mutex::new(None), query_stats: std::sync::Mutex::new(query::QuerySelectivityStats::default()), tuning, @@ -285,9 +289,9 @@ impl Unirust { .as_ref() .ok_or_else(|| anyhow::anyhow!("persistent linker backend requires a database"))? .clone(); - if !self.store.is_empty() { - persistence::clear_rebuildable_linker_state(&db)?; - } + // Even an empty record store can contain derived writes from an + // interrupted first ingest. Recovery always starts from durable records. + persistence::clear_rebuildable_linker_state(&db)?; let dsu = if use_persistent_dsu { let dsu_config = self.tuning.dsu_config.clone().unwrap_or_default(); dsu::DsuBackend::persistent(db.clone(), dsu_config)? @@ -358,10 +362,37 @@ impl Unirust { self.ingest_internal(records, true) } + fn run_ingest_batch( + &mut self, + operation: impl FnOnce(&mut Self) -> anyhow::Result, + ) -> anyhow::Result { + let result = operation(self); + if result.is_err() { + let abort_result = self.store.discard_staged_records(); + // Linking can mutate derived state before a later record or disk write + // fails. Rebuild from committed records before reusing that state. + self.streaming = None; + self.graph_state = None; + self.invalidate_query_cache(); + self.clear_conflict_cache(); + self.clear_query_stats(); + abort_result?; + } + result + } + fn ingest_internal( &mut self, records: Vec, preserve_record_ids: bool, + ) -> anyhow::Result { + self.run_ingest_batch(|engine| engine.ingest_batch_inner(records, preserve_record_ids)) + } + + fn ingest_batch_inner( + &mut self, + records: Vec, + preserve_record_ids: bool, ) -> anyhow::Result { self.clear_conflict_cache(); self.clear_query_stats(); @@ -812,6 +843,13 @@ impl Unirust { pub fn stream_records( &mut self, records: Vec, + ) -> anyhow::Result> { + self.run_ingest_batch(|engine| engine.stream_records_inner(records)) + } + + fn stream_records_inner( + &mut self, + records: Vec, ) -> anyhow::Result> { self.clear_conflict_cache(); self.clear_query_stats(); @@ -937,6 +975,13 @@ impl Unirust { pub fn stream_records_with_conflicts( &mut self, records: Vec, + ) -> anyhow::Result> { + self.run_ingest_batch(|engine| engine.stream_records_with_conflicts_inner(records)) + } + + fn stream_records_with_conflicts_inner( + &mut self, + records: Vec, ) -> anyhow::Result> { self.clear_conflict_cache(); self.clear_query_stats(); @@ -951,13 +996,12 @@ impl Unirust { })?; for record in records { - let (record_id, inserted) = self.store.add_record_if_absent(record)?; + let (record_id, inserted) = self.store.stage_record_if_absent(record)?; let cluster_id = if inserted { streaming.link_record(self.store.as_ref(), &self.ontology, record_id)? } else { streaming.cluster_id_for(record_id) }; - self.store.set_cluster_assignment(record_id, cluster_id)?; let assignment = ClusterAssignment { record_id, cluster_id, @@ -974,12 +1018,6 @@ impl Unirust { } else { Vec::new() }; - if inserted && !observations.is_empty() { - let summaries = - conflicts::summarize_conflicts(self.store.as_ref(), &observations); - self.store - .set_cluster_conflict_summaries(cluster_id, &summaries)?; - } updates.push(StreamedConflictUpdate { assignment, observations, @@ -989,6 +1027,20 @@ impl Unirust { streaming.cluster_count() }; + self.store.flush_staged_records()?; + let assignments = updates + .iter() + .map(|update| (update.assignment.record_id, update.assignment.cluster_id)) + .collect::>(); + self.store.set_cluster_assignments_batch(&assignments)?; + for update in &updates { + if !update.observations.is_empty() { + let summaries = + conflicts::summarize_conflicts(self.store.as_ref(), &update.observations); + self.store + .set_cluster_conflict_summaries(update.assignment.cluster_id, &summaries)?; + } + } self.invalidate_query_cache(); self.store.set_cluster_count(cluster_count)?; self.store.sync()?; @@ -1008,6 +1060,13 @@ impl Unirust { pub fn stream_records_update_graph( &mut self, records: Vec, + ) -> anyhow::Result> { + self.run_ingest_batch(|engine| engine.stream_records_update_graph_inner(records)) + } + + fn stream_records_update_graph_inner( + &mut self, + records: Vec, ) -> anyhow::Result> { self.clear_conflict_cache(); self.clear_query_stats(); @@ -1120,6 +1179,50 @@ impl Unirust { descriptors: &[query::QueryDescriptor], interval: Interval, ) -> anyhow::Result { + if let Some(streaming) = &self.streaming { + let Some(first) = descriptors.first() else { + return Ok(query::QueryOutcome::Matches(Vec::new())); + }; + let mut roots = std::collections::HashSet::new(); + let mut clusters = dsu::Clusters::new(); + for (record_id, _) in + self.store + .get_records_with_value_in_interval(first.attr, first.value, interval) + { + let root = streaming.root_for_record(record_id).ok_or_else(|| { + anyhow::anyhow!( + "query candidate {} has not completed entity resolution", + record_id.0 + ) + })?; + if roots.insert(root) { + clusters.add_cluster(streaming.cluster_for_record(record_id).ok_or_else( + || { + anyhow::anyhow!( + "query candidate {} has no cluster membership", + record_id.0 + ) + }, + )?); + } + } + // Ingest already established authoritative membership. Only master + // candidate entities, instead of replaying the complete store per write. + let golden = query::build_golden_cache(self.store.as_ref(), &clusters); + let keys = self.query_cluster_keys_for(streaming, &clusters)?; + let membership = query::build_record_to_cluster_map(&clusters); + let outcome = query::query_master_entities_with_cache( + self.store.as_ref(), + &clusters, + descriptors, + interval, + &golden, + &keys, + &membership, + )?; + self.store.ensure_healthy()?; + return Ok(outcome); + } let mut cache_guard = self.query_cache.lock().expect("query cache lock"); if cache_guard.is_none() { let clusters = self.build_clusters()?; @@ -1150,7 +1253,232 @@ impl Unirust { Ok(outcome) } + /// Resolve query text without allocating interner IDs or changing durable state. + pub fn lookup_query_descriptor( + &self, + attr: &str, + value: &str, + ) -> Option { + Some(query::QueryDescriptor { + attr: self.store.lookup_attr(attr)?, + value: self.store.lookup_value(value)?, + }) + } + + /// Return candidate fragments before applying cross-shard query conjunction. + /// Explicit global IDs hydrate all local members of those canonical entities. + #[doc(hidden)] + pub fn query_entity_fragments( + &self, + descriptors: &[query::QueryDescriptor], + interval: Interval, + global_ids: &[GlobalClusterId], + strong_ids_only: bool, + ) -> anyhow::Result> { + let recovered; + let streaming = if let Some(streaming) = &self.streaming { + streaming + } else { + let db = self.store.shared_db(); + if let Some(db) = db.as_ref() { + let persistence = LinkerStatePersistence::new(db); + // Fresh and reset stores have no marker or redirects. Anchored + // IDs are safe there without writing migration metadata. + if !persistence.has_stable_global_cluster_ids()? + && !persistence.load_cross_shard_merges()?.is_empty() + { + anyhow::bail!( + "initialize_streaming is required before querying legacy global cluster IDs" + ); + } + } + // Queries can run concurrently through &self. Recovery here must not + // clear or write the shared persistent DSU, index, or scheme markers. + let mut linker = + linker::StreamingLinker::new(self.store.as_ref(), &self.ontology, &self.tuning)?; + if let Some(db) = db.as_ref() { + linker.restore_cross_shard_merges(&LinkerStatePersistence::new(db))?; + } + recovered = linker; + &recovered + }; + if global_ids.is_empty() { + let mut candidates: std::collections::HashMap>> = + std::collections::HashMap::new(); + for (index, descriptor) in descriptors.iter().enumerate() { + for (record_id, matched_interval) in self.store.get_records_with_value_in_interval( + descriptor.attr, + descriptor.value, + interval, + ) { + let global_id = streaming + .global_cluster_id_for_readonly(record_id) + .ok_or_else(|| { + anyhow::anyhow!( + "query candidate {} has not completed entity resolution", + record_id.0 + ) + })?; + if let Some(overlap) = temporal::intersect(&matched_interval, &interval) { + candidates + .entry(global_id) + .or_insert_with(|| vec![Vec::new(); descriptors.len()])[index] + .push(overlap); + } + } + } + self.store.ensure_healthy()?; + return Ok(candidates + .into_iter() + .map(|(global_id, mut descriptor_intervals)| { + for intervals in &mut descriptor_intervals { + *intervals = query::coalesce_intervals(intervals); + } + query::EntityFragment { + global_id, + descriptor_intervals, + golden: Vec::new(), + strong_ids: Vec::new(), + cluster_key: None, + cluster_key_identity: None, + } + }) + .collect()); + } + let clusters = streaming.clusters_for_global_ids(global_ids); + let mut fragments = Vec::with_capacity(clusters.len()); + for cluster in clusters { + let global_id = streaming + .global_cluster_id_for_readonly(cluster.root) + .ok_or_else(|| anyhow::anyhow!("query cluster has no global identity"))?; + let mut descriptor_intervals = vec![Vec::new(); descriptors.len()]; + let mut values: std::collections::BTreeMap<(String, String), Vec> = + std::collections::BTreeMap::new(); + let mut strong_ids = std::collections::HashSet::new(); + for &record_id in &cluster.records { + let record = self.store.get_record(record_id).ok_or_else(|| { + anyhow::anyhow!("query cluster member {} is unavailable", record_id.0) + })?; + let strong_attrs = self + .ontology + .strong_identifiers_for_type(&record.identity.entity_type); + for descriptor in &record.descriptors { + let is_strong = strong_attrs + .iter() + .any(|strong| strong.attribute == descriptor.attr); + if strong_ids_only && !is_strong { + continue; + } + let attr = self + .store + .resolve_attr(descriptor.attr) + .ok_or_else(|| anyhow::anyhow!("query attribute ID cannot be resolved"))?; + let value = self + .store + .resolve_value(descriptor.value) + .ok_or_else(|| anyhow::anyhow!("query value ID cannot be resolved"))?; + if is_strong { + strong_ids.insert(sharding::BoundaryStrongId { + perspective: record.identity.perspective.clone(), + attribute: attr.clone(), + value: value.clone(), + interval: descriptor.interval, + }); + } + if strong_ids_only { + continue; + } + if let Some(overlap) = temporal::intersect(&descriptor.interval, &interval) { + for (index, query) in descriptors.iter().enumerate() { + if query.attr == descriptor.attr && query.value == descriptor.value { + descriptor_intervals[index].push(overlap); + } + } + values.entry((attr, value)).or_default().push(overlap); + } + } + } + let golden = values + .into_iter() + .flat_map(|((attr, value), intervals)| { + query::coalesce_intervals(&intervals) + .into_iter() + .map(move |interval| graph::GoldenDescriptor { + attr: attr.clone(), + value: value.clone(), + interval, + }) + }) + .collect(); + for intervals in &mut descriptor_intervals { + *intervals = query::coalesce_intervals(intervals); + } + let cluster_id = cluster.id; + let key = if strong_ids_only { + None + } else { + self.query_cluster_keys_for( + streaming, + &dsu::Clusters { + clusters: vec![cluster], + }, + )? + .remove(&cluster_id) + }; + fragments.push(query::EntityFragment { + global_id, + descriptor_intervals, + golden, + cluster_key: key.as_ref().map(|key| key.value.clone()), + cluster_key_identity: key.map(|key| key.identity_key), + strong_ids: strong_ids.into_iter().collect(), + }); + } + self.store.ensure_healthy()?; + Ok(fragments) + } + + fn query_cluster_keys_for( + &self, + streaming: &linker::StreamingLinker, + clusters: &dsu::Clusters, + ) -> anyhow::Result> { + if graph::cluster_keys_are_candidate_local(&self.ontology) { + return Ok(query::build_cluster_key_cache( + self.store.as_ref(), + clusters, + &self.ontology, + )); + } + // Composite labels use the shortest globally unique prefix. Preserve their + // collision scope while mastering golden data only for query candidates. + let mut cache = self + .cluster_key_cache + .lock() + .map_err(|_| anyhow::anyhow!("cluster key cache lock poisoned"))?; + if cache.is_none() { + *cache = Some(query::build_cluster_key_cache( + self.store.as_ref(), + &streaming.clusters_readonly()?, + &self.ontology, + )); + } + Ok(clusters + .clusters + .iter() + .filter_map(|cluster| { + cache + .as_ref()? + .get(&cluster.id) + .map(|key| (cluster.id, key.clone())) + }) + .collect()) + } + fn invalidate_query_cache(&self) { + if let Ok(mut guard) = self.cluster_key_cache.lock() { + *guard = None; + } if let Ok(mut guard) = self.query_cache.lock() { *guard = None; } @@ -1436,6 +1764,9 @@ impl Unirust { /// Call this after enable_streaming() to recover cluster ID mappings. /// Returns the number of cluster_ids restored, or an error if persistence is not available. pub fn restore_linker_state(&mut self) -> anyhow::Result { + // Restored IDs can differ from record-scan allocation order, including + // when only part of the snapshot loads before an error. + self.invalidate_query_cache(); let streaming = self.streaming.as_mut().ok_or_else(|| { anyhow::anyhow!("Streaming not initialized - call enable_streaming() first") })?; diff --git a/src/linker.rs b/src/linker.rs index 05afcc1..7345d84 100644 --- a/src/linker.rs +++ b/src/linker.rs @@ -36,6 +36,26 @@ use tracing::{debug, instrument, warn}; /// 32 candidates covers 95%+ of queries based on production workload analysis type CandidateVec = SmallVec<[(RecordId, Interval); 32]>; +struct StoreLookup<'a>(&'a dyn RecordStore); + +impl InternerLookup for StoreLookup<'_> { + fn get_attr_string(&self, id: crate::model::AttrId) -> Option { + self.0.resolve_attr(id) + } + + fn get_value_string(&self, id: crate::model::ValueId) -> Option { + self.0.resolve_value(id) + } +} + +fn candidate_scan_limit(tuning: &crate::StreamingTuning) -> usize { + if tuning.adaptive_candidate_cap && tuning.adaptive_high_cap > 0 { + tuning.adaptive_high_cap + } else { + tuning.candidate_cap.max(1) + } +} + /// Data stored for boundary signatures to support cross-shard conflict detection. #[derive(Debug, Clone)] struct BoundaryData { @@ -385,15 +405,6 @@ impl<'a, K, V> Iterator for LinkerStateIterMut<'a, K, V> { } } -/// Represents the resolution of a conflict. -#[derive(Debug, Clone)] -enum ConflictResolution { - /// Conflict can be resolved by choosing the specified record as winner. - Resolvable(()), - /// Conflict cannot be resolved and should prevent merging. - Unresolvable, -} - /// Public function to build clusters using streaming semantics. pub fn build_clusters(store: &dyn RecordStore, ontology: &Ontology) -> Result { build_clusters_streaming(store, ontology) @@ -421,6 +432,9 @@ pub struct StreamingLinker { shard_id: u16, /// Strong ID summaries for conflict detection (LRU-bounded when config provided) strong_id_summaries: LinkerState, + /// Root aliases and member lists let queries hydrate only candidate clusters. + member_roots: FxHashMap, + cluster_members: FxHashMap>, /// Use FxHashSet for faster hashing (non-cryptographic, perfect for internal keys) tainted_identity_keys: FxHashSet, /// Record perspectives for same-perspective conflict detection (LRU-bounded when config provided) @@ -526,8 +540,9 @@ impl StreamingLinker { tuning: &crate::StreamingTuning, shard_id: u16, dsu: DsuBackend, - identity_index: IndexBackend, + mut identity_index: IndexBackend, ) -> Result { + identity_index.clear()?; // These mappings and summaries affect resolution correctness. The bounded // backend cannot be enabled until evictions have a durable spill/read-through // path; silently dropping an old entry can split clusters after enough ingest. @@ -553,6 +568,8 @@ impl StreamingLinker { global_cluster_anchors: LinkerState::unbounded(), shard_id, strong_id_summaries, + member_roots: FxHashMap::default(), + cluster_members: FxHashMap::default(), tainted_identity_keys: FxHashSet::default(), record_perspectives, pending_keys: FxHashSet::default(), @@ -573,7 +590,7 @@ impl StreamingLinker { streamer.link_records_batch_parallel_with_interner( &records, ontology, - store.interner(), + &StoreLookup(store), )?; batch.clear(); Ok(()) @@ -633,7 +650,7 @@ impl StreamingLinker { ontology: &Ontology, record_id: RecordId, ) -> Result { - self.link_record_with_interner(store, store.interner(), ontology, record_id) + self.link_record_with_interner(store, &StoreLookup(store), ontology, record_id) } /// Link a newly added record with an explicit interner for boundary tracking. @@ -650,6 +667,7 @@ impl StreamingLinker { self.dsu.add_record(record_id)?; self.global_cluster_anchors.insert(record_id, record_id.0); } + self.register_cluster_member(record_id)?; // Try to get a reference first (avoids cloning), fall back to cloning if not available. let record_owned; @@ -695,12 +713,6 @@ impl StreamingLinker { let interval = *interval; // deref for use below let key_signature = LinkerKeySignature::new(entity_type, key_values); - // Fast path: skip if key is already known to be tainted (O(1) set lookup). - // We already added to pending_keys when first tainted, no need to re-add. - if self.tainted_identity_keys.contains(&key_signature) { - continue; - } - // Create identity key signature for bloom filter and scan cache lookups let identity_sig = IdentityKeySignature::from_key_values(entity_type, key_values); @@ -740,7 +752,7 @@ impl StreamingLinker { } = self; metrics.cache_misses.fetch_add(1, Ordering::Relaxed); - let max_tree_nodes = tuning.adaptive_high_cap; + let max_tree_nodes = candidate_scan_limit(tuning); let (candidates_slice, is_hot) = identity_index .find_matching_clusters_overlapping_limited( dsu, @@ -748,8 +760,23 @@ impl StreamingLinker { key_values, interval, max_tree_nodes, - ); - let candidates: CandidateVec = candidates_slice.iter().copied().collect(); + )?; + let candidates: CandidateVec = if is_hot { + identity_index + .find_matching_clusters_overlapping_limited( + dsu, + entity_type, + key_values, + interval, + usize::MAX, + )? + .0 + .iter() + .copied() + .collect() + } else { + candidates_slice.iter().copied().collect() + }; // Cache the result for future lookups if let Some(opts) = partition_opts { opts.cache_candidates(&identity_sig, candidates.to_vec()); @@ -766,17 +793,9 @@ impl StreamingLinker { ); } - // If the tree query hit the limit, mark as hot and skip. + // Candidate volume is independent of strong-identifier conflicts. if is_hot || candidate_len > self.tuning.hot_key_threshold { self.metrics.hot_key_exits.fetch_add(1, Ordering::Relaxed); - // Avoid redundant clone: only clone if we need both sets - if self.tuning.deferred_reconciliation { - self.tainted_identity_keys.insert(key_signature.clone()); - self.pending_keys.insert(key_signature); - } else { - self.tainted_identity_keys.insert(key_signature); - } - continue; } // === SPER OPTIMIZATION: Stochastic candidate sampling === @@ -861,6 +880,8 @@ impl StreamingLinker { global_cluster_ids, global_cluster_anchors, strong_id_summaries, + member_roots, + cluster_members, tainted_identity_keys, record_perspectives, boundary_signatures, @@ -943,6 +964,13 @@ impl StreamingLinker { { metrics.merges_performed.fetch_add(1, Ordering::Relaxed); let new_root = dsu.find(record_id).unwrap_or(record_id); + reconcile_cluster_members( + member_roots, + cluster_members, + root_a, + root_b, + new_root, + ); reconcile_cluster_ids( cluster_ids, next_cluster_id, @@ -968,7 +996,7 @@ impl StreamingLinker { root_a, root_b, new_root, - ); + )?; // Invalidate cache on merge - candidates changed if let Some(opts) = partition_opts { opts.on_cluster_merge(root_a, root_b); @@ -996,8 +1024,12 @@ impl StreamingLinker { } } - self.identity_index - .add_record_with_cached_keys(record_id, root, entity_type, cached_keys); + self.identity_index.add_record_with_cached_keys( + record_id, + root, + entity_type, + cached_keys, + )?; self.metrics.records_linked.fetch_add(1, Ordering::Relaxed); Ok(self.get_or_assign_cluster_id(root)) @@ -1133,6 +1165,7 @@ impl StreamingLinker { self.dsu.add_record(record_id)?; self.global_cluster_anchors.insert(record_id, record_id.0); } + self.register_cluster_member(record_id)?; // Store perspective and summary self.record_perspectives @@ -1145,13 +1178,8 @@ impl StreamingLinker { // Process each key for (key_signature, key_values, interval, guard_reason) in extraction.keys { - // Skip tainted keys - if self.tainted_identity_keys.contains(&key_signature) { - continue; - } - // Find candidates and merge - let max_tree_nodes = self.tuning.adaptive_high_cap; + let max_tree_nodes = candidate_scan_limit(&self.tuning); let (candidates_slice, is_hot) = self .identity_index .find_matching_clusters_overlapping_limited( @@ -1160,20 +1188,27 @@ impl StreamingLinker { &key_values, interval, max_tree_nodes, - ); - let candidates: CandidateVec = candidates_slice.iter().copied().collect(); + )?; + let candidates: CandidateVec = if is_hot { + self.identity_index + .find_matching_clusters_overlapping_limited( + &mut self.dsu, + entity_type, + &key_values, + interval, + usize::MAX, + )? + .0 + .iter() + .copied() + .collect() + } else { + candidates_slice.iter().copied().collect() + }; let candidate_len = candidates.len(); - // Handle hot keys if is_hot || candidate_len > self.tuning.hot_key_threshold { self.metrics.hot_key_exits.fetch_add(1, Ordering::Relaxed); - if self.tuning.deferred_reconciliation { - self.tainted_identity_keys.insert(key_signature.clone()); - self.pending_keys.insert(key_signature); - } else { - self.tainted_identity_keys.insert(key_signature); - } - continue; } // Merge with candidates @@ -1243,6 +1278,13 @@ impl StreamingLinker { .merges_performed .fetch_add(1, Ordering::Relaxed); let new_root = self.dsu.find(record_id).unwrap_or(record_id); + reconcile_cluster_members( + &mut self.member_roots, + &mut self.cluster_members, + root_a, + root_b, + new_root, + ); reconcile_cluster_ids( &mut self.cluster_ids, &mut self.next_cluster_id, @@ -1273,7 +1315,7 @@ impl StreamingLinker { root_a, root_b, new_root, - ); + )?; root_a = new_root; } } @@ -1318,33 +1360,42 @@ impl StreamingLinker { } } - self.identity_index - .add_record_with_cached_keys(record.id, root, entity_type, cached_keys); + self.identity_index.add_record_with_cached_keys( + record.id, + root, + entity_type, + cached_keys, + )?; Ok(()) } - /// Get clusters from the streaming DSU state. + /// Enumerate authoritative local membership, including persistent DSU backends. pub fn clusters(&mut self) -> Clusters { - self.dsu.get_clusters().unwrap_or_else(|_| Clusters { - clusters: Vec::new(), - }) + let members = self + .cluster_members + .iter() + .map(|(root, records)| (*root, records.clone())) + .collect::>(); + Clusters { + clusters: members + .into_iter() + .map(|(root, records)| { + let id = self.get_or_assign_cluster_id(root); + crate::dsu::Cluster::new(id, root, records) + }) + .collect(), + } } - /// Get clusters from the streaming DSU state, applying conflict splitting heuristics. + /// Finish deferred linking and enumerate clusters guarded at merge time. pub fn clusters_with_conflict_splitting( &mut self, store: &dyn RecordStore, ontology: &Ontology, ) -> Result { self.reconcile_pending(store, ontology)?; - let clusters = self.dsu.get_clusters()?; - - if should_apply_conflict_splitting(store, ontology) { - split_clusters_with_unresolvable_conflicts(store, ontology, clusters) - } else { - Ok(clusters) - } + Ok(self.clusters()) } fn get_or_assign_cluster_id(&mut self, root: RecordId) -> ClusterId { @@ -1379,6 +1430,72 @@ impl StreamingLinker { self.get_or_assign_cluster_id(root) } + fn register_cluster_member(&mut self, record_id: RecordId) -> Result<()> { + if !self.member_roots.contains_key(&record_id) { + let root = self.dsu.find(record_id)?; + self.member_roots.insert(record_id, root); + self.cluster_members + .entry(root) + .or_default() + .push(record_id); + } + Ok(()) + } + + /// Resolve a tracked record without mutating DSU caches or enumerating records. + pub fn root_for_record(&self, record_id: RecordId) -> Option { + let mut root = *self.member_roots.get(&record_id)?; + while let Some(&parent) = self.member_roots.get(&root) { + if parent == root { + break; + } + root = parent; + } + Some(root) + } + + /// Read authoritative membership for one local cluster. + pub fn cluster_for_record(&self, record_id: RecordId) -> Option { + let root = self.root_for_record(record_id)?; + let id = *self.cluster_ids.peek(&root)?; + let members = self.cluster_members.get(&root)?.clone(); + Some(crate::dsu::Cluster::new(id, root, members)) + } + + /// Read all authoritative memberships without replaying records or assigning + /// new IDs. Missing IDs indicate inconsistent state and must not hide clusters. + pub fn clusters_readonly(&self) -> Result { + let clusters = self + .cluster_members + .iter() + .map(|(root, records)| { + let id = self.cluster_ids.peek(root).copied().ok_or_else(|| { + anyhow::anyhow!("cluster ID is unavailable for tracked root {}", root.0) + })?; + Ok(crate::dsu::Cluster::new(id, *root, records.clone())) + }) + .collect::>>()?; + Ok(Clusters { clusters }) + } + + /// Read a local cluster's canonical global ID without allocating an assignment. + pub fn global_cluster_id_for_readonly(&self, record_id: RecordId) -> Option { + let root = self.root_for_record(record_id)?; + let id = self + .global_cluster_ids + .peek(&root) + .copied() + .unwrap_or_else(|| { + let anchor = self + .global_cluster_anchors + .peek(&root) + .copied() + .unwrap_or(root.0); + GlobalClusterId::new(self.shard_id, anchor, 0) + }); + Some(self.resolve_global_cluster_id(id)) + } + /// Get the global cluster ID for a record. pub fn global_cluster_id_for(&mut self, record_id: RecordId) -> GlobalClusterId { let root = self.dsu.find(record_id).unwrap_or(record_id); @@ -1640,7 +1757,7 @@ impl StreamingLinker { for key_signature in pending { let candidates = self .identity_index - .find_matching_records(key_signature.entity_type(), key_signature.key_values()) + .find_matching_records(key_signature.entity_type(), key_signature.key_values())? .to_vec(); if candidates.len() < 2 { continue; @@ -1761,6 +1878,24 @@ impl StreamingLinker { .merges_performed .fetch_add(1, Ordering::Relaxed); let new_root = self.dsu.find(*record_id).unwrap_or(*record_id); + reconcile_cluster_members( + &mut self.member_roots, + &mut self.cluster_members, + root_a, + root_b, + new_root, + ); + reconcile_global_cluster_ids( + &mut self.global_cluster_ids, + &mut self.global_cluster_anchors, + &mut self.cross_shard_merges, + &mut self.boundary_signatures, + &mut self.dirty_boundary_keys, + self.shard_id, + root_a, + root_b, + new_root, + ); reconcile_cluster_ids( &mut self.cluster_ids, &mut self.next_cluster_id, @@ -1780,7 +1915,7 @@ impl StreamingLinker { root_a, root_b, new_root, - ); + )?; } compared = compared.saturating_add(1); } @@ -2345,68 +2480,6 @@ fn reconcile_cluster_summaries( summaries.insert(new_root, merged); } -/// Check for strong identifier conflicts. -/// The interval parameter is currently unused; overlap checks are derived -/// from descriptor intervals directly. -fn check_strong_identifier_conflict( - store: &dyn RecordStore, - ontology: &Ontology, - record_a: RecordId, - record_b: RecordId, - _interval: Interval, -) -> Option { - let record_a = store.get_record(record_a)?; - let record_b = store.get_record(record_b)?; - - // Check strong identifiers for conflicts - for strong_id in ontology.strong_identifiers_for_type(&record_a.identity.entity_type) { - let descriptor_a = get_strong_identifier_descriptor(&record_a, strong_id); - let descriptor_b = get_strong_identifier_descriptor(&record_b, strong_id); - - // Only check for conflicts if both records have values for this strong identifier - if let (Some(desc_a), Some(desc_b)) = (descriptor_a, descriptor_b) { - if desc_a.value != desc_b.value { - // Check if the conflicting values have overlapping temporal intervals - if temporal_intervals_overlap(&desc_a.interval, &desc_b.interval) { - // Conflict detected: same attribute, different values, overlapping time - // Use perspective weights to determine resolution - let weight_a = ontology.get_perspective_weight(&record_a.identity.perspective); - let weight_b = ontology.get_perspective_weight(&record_b.identity.perspective); - - if weight_a != weight_b { - return Some(ConflictResolution::Resolvable(())); - } - // Equal weights - check if there are other resolution mechanisms - // For now, we'll be conservative and only mark as unresolvable - // if both records are from the same perspective with equal weights - if record_a.identity.perspective == record_b.identity.perspective { - // Same perspective with equal weights - unresolvable - return Some(ConflictResolution::Unresolvable); - } - // Different perspectives with equal weights - might be resolvable through other means - // Let the normal clustering process handle this - return None; - } - // If temporal intervals don't overlap, there's no conflict - } - } - // If only one record has a value for this strong identifier, that's not a conflict - // Records from different perspectives may have different strong identifier attributes - } - None -} - -/// Get strong identifier descriptor (with temporal interval) for a record. -fn get_strong_identifier_descriptor<'a>( - record: &'a Record, - strong_id: &crate::ontology::StrongIdentifier, -) -> Option<&'a crate::model::Descriptor> { - record - .descriptors - .iter() - .find(|descriptor| descriptor.attr == strong_id.attribute) -} - /// Build a strong ID summary for a record based on ontology strong identifiers. pub fn build_record_summary(record: &Record, ontology: &Ontology) -> StrongIdSummary { let _guard = crate::profile::profile_scope("build_record_summary"); @@ -2469,18 +2542,6 @@ pub fn build_record_summary(record: &Record, ontology: &Ontology) -> StrongIdSum summary } -/// Check if two temporal intervals overlap. -fn temporal_intervals_overlap( - interval_a: &crate::temporal::Interval, - interval_b: &crate::temporal::Interval, -) -> bool { - // Two intervals overlap if one starts before the other ends - // interval_a: [start_a, end_a] - // interval_b: [start_b, end_b] - // They overlap if: start_a < end_b && start_b < end_a - interval_a.start < interval_b.end && interval_b.start < interval_a.end -} - /// Check if merging would create conflicts in existing clusters. fn would_create_conflict_in_clusters( summaries: &LinkerState, @@ -2507,153 +2568,23 @@ fn would_create_conflict_in_clusters( false } -/// Post-process clusters to split those with unresolvable conflicts. -/// This is a heuristic fallback used to support specific conflict tests. -fn split_clusters_with_unresolvable_conflicts( - store: &dyn RecordStore, - ontology: &Ontology, - clusters: Clusters, -) -> Result { - let mut new_clusters = Vec::new(); - - for cluster in clusters.clusters { - // Check if this cluster has any unresolvable conflicts - let mut has_unresolvable_conflicts = false; - let mut conflicting_groups: Vec> = Vec::new(); - - // For each pair of records in the cluster, check for unresolvable conflicts - for i in 0..cluster.records.len() { - for j in (i + 1)..cluster.records.len() { - let record_a = cluster.records[i]; - let record_b = cluster.records[j]; - - if let Some(ConflictResolution::Unresolvable) = check_strong_identifier_conflict( - store, - ontology, - record_a, - record_b, - Interval::new(0, 1).unwrap(), - ) { - has_unresolvable_conflicts = true; - - // Find which group each record belongs to, or create new groups - let mut group_a = None; - let mut group_b = None; - - for (group_idx, group) in conflicting_groups.iter().enumerate() { - if group.contains(&record_a) { - group_a = Some(group_idx); - } - if group.contains(&record_b) { - group_b = Some(group_idx); - } - } - - match (group_a, group_b) { - (Some(idx_a), Some(idx_b)) => { - if idx_a != idx_b { - // These records are in different groups but have an unresolvable conflict - // This means each record should be in its own separate group - // Remove both records from their current groups and create individual groups for them - conflicting_groups[idx_a].retain(|&x| x != record_a); - conflicting_groups[idx_b].retain(|&x| x != record_b); - conflicting_groups.push(vec![record_a]); - conflicting_groups.push(vec![record_b]); - } - } - (Some(idx_a), None) => { - // Record A is in a group, but B conflicts with A - // Remove A from its current group and create separate groups for both - conflicting_groups[idx_a].retain(|&x| x != record_a); - conflicting_groups.push(vec![record_a]); - conflicting_groups.push(vec![record_b]); - } - (None, Some(idx_b)) => { - // Record B is in a group, but A conflicts with B - // Remove B from its current group and create separate groups for both - conflicting_groups[idx_b].retain(|&x| x != record_b); - conflicting_groups.push(vec![record_a]); - conflicting_groups.push(vec![record_b]); - } - (None, None) => { - // Neither record is in a group yet, create separate groups for both - conflicting_groups.push(vec![record_a]); - conflicting_groups.push(vec![record_b]); - } - } - } - } - } - - if has_unresolvable_conflicts { - // Split the cluster based on conflicting groups - let mut used_records = std::collections::HashSet::new(); - - // Create clusters for each conflicting group - for group in conflicting_groups { - for &record_id in &group { - used_records.insert(record_id); - } - if !group.is_empty() { - let cluster_id = ClusterId(new_clusters.len() as u32 + 1); - new_clusters.push(crate::dsu::Cluster::new(cluster_id, group[0], group)); - } - } - - // Create individual clusters for records not in any conflicting group - for &record_id in &cluster.records { - if !used_records.contains(&record_id) { - let cluster_id = ClusterId(new_clusters.len() as u32 + 1); - new_clusters.push(crate::dsu::Cluster::new( - cluster_id, - record_id, - vec![record_id], - )); - } - } - } else { - // No unresolvable conflicts, keep the cluster as is - new_clusters.push(cluster); - } - } - - Ok(Clusters { - clusters: new_clusters, - }) -} - -/// Determine if conflict splitting should be applied based on the scenario. -/// This uses a narrow heuristic for known test fixtures. -fn should_apply_conflict_splitting(store: &dyn RecordStore, ontology: &Ontology) -> bool { - let mut records = Vec::new(); - store.for_each_record(&mut |record| records.push(record)); - - // Check if this looks like the indirect conflict test case - // (3 records with same identity key but conflicting strong identifiers from same perspective) - if records.len() == 3 { - let mut same_perspective_conflicts = 0; - - for i in 0..records.len() { - for j in (i + 1)..records.len() { - if let Some(ConflictResolution::Unresolvable) = check_strong_identifier_conflict( - store, - ontology, - records[i].id, - records[j].id, - Interval::new(0, 1).unwrap(), - ) { - if records[i].identity.perspective == records[j].identity.perspective { - same_perspective_conflicts += 1; - } - } - } - } - - // Apply conflict splitting if we have same-perspective unresolvable conflicts - return same_perspective_conflicts > 0; - } - - false +fn reconcile_cluster_members( + roots: &mut FxHashMap, + members: &mut FxHashMap>, + root_a: RecordId, + root_b: RecordId, + new_root: RecordId, +) { + roots.insert(root_a, new_root); + roots.insert(root_b, new_root); + roots.insert(new_root, new_root); + let mut left = members.remove(&root_a).unwrap_or_default(); + let mut right = members.remove(&root_b).unwrap_or_default(); + if left.len() < right.len() { + std::mem::swap(&mut left, &mut right); + } + left.extend(right); + members.insert(new_root, left); } fn cluster_summaries_conflict(a: &StrongIdSummary, b: &StrongIdSummary) -> bool { @@ -2793,3 +2724,40 @@ fn has_overlapping_interval(a: &[Interval], b: &[Interval]) -> bool { false } + +impl StreamingLinker { + /// Hydrate only local components of the requested canonical entities. + /// Global IDs are anchored to durable record IDs. Remote canonical entities + /// additionally reach local anchors through the sparse redirect map. + pub fn clusters_for_global_ids(&self, ids: &[GlobalClusterId]) -> Vec { + let targets: HashSet<_> = ids + .iter() + .map(|id| self.resolve_global_cluster_id(*id)) + .collect(); + let mut anchors: HashSet = targets + .iter() + .filter(|id| id.shard_id == self.shard_id) + .map(|id| RecordId(id.local_id)) + .collect(); + for (alias, target) in &self.cross_shard_merges { + if alias.shard_id == self.shard_id + && targets.contains(&self.resolve_global_cluster_id(*target)) + { + anchors.insert(RecordId(alias.local_id)); + } + } + let mut roots = HashSet::new(); + anchors + .into_iter() + .filter_map(|anchor| { + let root = self.root_for_record(anchor)?; + if !roots.insert(root) + || !targets.contains(&self.global_cluster_id_for_readonly(anchor)?) + { + return None; + } + self.cluster_for_record(anchor) + }) + .collect() + } +} diff --git a/src/perf/simd_hash.rs b/src/perf/simd_hash.rs index c44aa6e..b32cdee 100644 --- a/src/perf/simd_hash.rs +++ b/src/perf/simd_hash.rs @@ -56,14 +56,13 @@ impl SimdHasher { let mut hasher = Self::new(); // Process 8 bytes at a time - let mut chunks = bytes.chunks_exact(8); - for chunk in chunks.by_ref() { - let word = u64::from_le_bytes(chunk.try_into().unwrap()); + let (chunks, remainder) = bytes.as_chunks::<8>(); + for chunk in chunks { + let word = u64::from_le_bytes(*chunk); hasher.hash_word(word); } // Process remaining bytes - let remainder = chunks.remainder(); if !remainder.is_empty() { let mut word = 0u64; for (i, &byte) in remainder.iter().enumerate() { diff --git a/src/persistence.rs b/src/persistence.rs index 6d0eced..3e4924d 100644 --- a/src/persistence.rs +++ b/src/persistence.rs @@ -537,8 +537,8 @@ pub struct PersistentStore { inner: Store, db: Arc, cache: Mutex>, - staged_records: Mutex>, - staged_identities: Mutex>, + staged_records: HashMap, + staged_identities: HashMap, persisted_attr_id: u32, persisted_value_id: u32, record_count: u64, @@ -592,8 +592,8 @@ impl PersistentStore { cache: Mutex::new(LruCache::new( std::num::NonZeroUsize::new(DEFAULT_CACHE_CAPACITY).expect("cache capacity"), )), - staged_records: Mutex::new(Vec::new()), - staged_identities: Mutex::new(HashMap::new()), + staged_records: HashMap::new(), + staged_identities: HashMap::new(), persisted_attr_id, persisted_value_id, record_count, @@ -748,14 +748,8 @@ impl PersistentStore { self.cluster_count = 0; self.conflict_summary_count = 0; self.read_fault.store(false, Ordering::Release); - self.staged_records - .get_mut() - .map_err(|_| anyhow!("staged records lock poisoned"))? - .clear(); - self.staged_identities - .get_mut() - .map_err(|_| anyhow!("staged identities lock poisoned"))? - .clear(); + self.staged_records.clear(); + self.staged_identities.clear(); self.cache .get_mut() .map_err(|_| anyhow!("record cache lock poisoned"))? @@ -786,88 +780,67 @@ impl PersistentStore { Ok(()) } - /// Stage a record for later batch write. Returns (record_id, inserted). - /// The record is added to cache immediately so it's readable, but not yet persisted to DB. - pub fn stage_record_if_absent(&mut self, mut record: Record) -> Result<(RecordId, bool)> { - ::ensure_healthy(self)?; - if let Some(existing) = self.get_record_id_by_identity(&record.identity) { - self.ensure_idempotent_record(existing, &record)?; - return Ok((existing, false)); + /// Hydrate referenced durable IDs before the in-memory preparer validates them. + /// A bounded interner cache does not make a durable string ID invalid. + fn hydrate_record_interner(&mut self, record: &Record) -> Result<()> { + for descriptor in &record.descriptors { + if self.inner.interner().get_attr(descriptor.attr).is_none() { + if let Some(value) = self.lookup_interner_value(b'a', descriptor.attr.0) { + self.inner + .interner_mut() + .insert_attr_with_id(descriptor.attr, value); + } + } + if self.inner.interner().get_value(descriptor.value).is_none() { + if let Some(value) = self.lookup_interner_value(b'v', descriptor.value.0) { + self.inner + .interner_mut() + .insert_value_with_id(descriptor.value, value); + } + } } - ::ensure_healthy(self)?; + self.ensure_healthy() + } - if let Some(existing) = self - .staged_identities - .lock() - .map_err(|_| anyhow!("staged identities lock poisoned"))? - .get(&record.identity) - .copied() - { - self.ensure_idempotent_record(existing, &record)?; - return Ok((existing, false)); + fn ensure_record_id_available(&self, record_id: RecordId) -> Result<()> { + if self.get_record(record_id).is_some() { + anyhow::bail!("record ID {} already exists", record_id.0); } + self.ensure_healthy() + } - let record_id = self.inner.prepare_record(&mut record)?; - let identity = record.identity.clone(); - - // Add to cache immediately so it's readable - self.cache - .lock() - .map_err(|_| anyhow!("record cache lock poisoned"))? - .put(record_id, record.clone()); - - // Stage for later batch write - self.staged_records - .lock() - .map_err(|_| anyhow!("staged records lock poisoned"))? - .push(record); - self.staged_identities - .lock() - .map_err(|_| anyhow!("staged identities lock poisoned"))? - .insert(identity, record_id); - - Ok((record_id, true)) + /// Stage a record for later batch write. Pending records remain addressable + /// independently of the bounded cache used for already committed records. + pub fn stage_record_if_absent(&mut self, record: Record) -> Result<(RecordId, bool)> { + self.stage_record(record, false) } pub fn stage_record_with_explicit_id_if_absent( &mut self, - mut record: Record, + record: Record, ) -> Result<(RecordId, bool)> { - ::ensure_healthy(self)?; + self.stage_record(record, true) + } + + fn stage_record(&mut self, mut record: Record, preserve_id: bool) -> Result<(RecordId, bool)> { + self.ensure_healthy()?; if let Some(existing) = self.get_record_id_by_identity(&record.identity) { self.ensure_idempotent_record(existing, &record)?; return Ok((existing, false)); } - ::ensure_healthy(self)?; - if let Some(existing) = self - .staged_identities - .lock() - .map_err(|_| anyhow!("staged identities lock poisoned"))? - .get(&record.identity) - .copied() - { - self.ensure_idempotent_record(existing, &record)?; - return Ok((existing, false)); - } - if self.get_record(record.id).is_some() { - anyhow::bail!("record ID {} already exists", record.id.0); + self.ensure_healthy()?; + if preserve_id || record.id.0 != 0 { + self.ensure_record_id_available(record.id)?; } - ::ensure_healthy(self)?; - - let record_id = self.inner.prepare_record_with_explicit_id(&mut record)?; - let identity = record.identity.clone(); - self.cache - .lock() - .map_err(|_| anyhow!("record cache lock poisoned"))? - .put(record_id, record.clone()); - self.staged_records - .lock() - .map_err(|_| anyhow!("staged records lock poisoned"))? - .push(record); + self.hydrate_record_interner(&record)?; + let record_id = if preserve_id { + self.inner.prepare_record_with_explicit_id(&mut record)? + } else { + self.inner.prepare_record(&mut record)? + }; self.staged_identities - .lock() - .map_err(|_| anyhow!("staged identities lock poisoned"))? - .insert(identity, record_id); + .insert(record.identity.clone(), record_id); + self.staged_records.insert(record_id, record); Ok((record_id, true)) } @@ -886,13 +859,7 @@ impl PersistentStore { /// Flush all staged records to the database in a single batch write. pub fn flush_staged_records(&mut self) -> Result { ::ensure_healthy(self)?; - let records = { - let mut staged = self - .staged_records - .lock() - .map_err(|_| anyhow!("staged records lock poisoned"))?; - std::mem::take(&mut *staged) - }; + let records = std::mem::take(&mut self.staged_records); if records.is_empty() { return Ok(0); @@ -908,7 +875,7 @@ impl PersistentStore { .db .cf_handle(CF_RECORDS) .ok_or_else(|| anyhow!("missing records column family"))?; - for record in &records { + for record in records.values() { let key = record.id.0.to_be_bytes(); with_record_bytes(record, |bytes| { batch.put_cf(records_cf, key, bytes); @@ -923,19 +890,18 @@ impl PersistentStore { let watermark = match write_result { Ok(watermark) => watermark, Err(error) => { - *self - .staged_records - .lock() - .map_err(|_| anyhow!("staged records lock poisoned"))? = records; + self.staged_records = records; return Err(error); } }; self.commit_interner_watermark(watermark); self.record_count = next_count; - self.staged_identities - .lock() - .map_err(|_| anyhow!("staged identities lock poisoned"))? - .clear(); + self.staged_identities.clear(); + if let Ok(mut cache) = self.cache.lock() { + for (id, record) in records { + cache.put(id, record); + } + } Ok(count) } @@ -1017,6 +983,10 @@ impl RecordStore for PersistentStore { fn add_record(&mut self, record: Record) -> Result { self.ensure_healthy()?; let mut record = record; + if record.id.0 != 0 { + self.ensure_record_id_available(record.id)?; + } + self.hydrate_record_interner(&record)?; let record_id = self.inner.prepare_record(&mut record)?; let next_count = self.record_count.saturating_add(1); let mut batch = WriteBatch::default(); @@ -1047,8 +1017,8 @@ impl RecordStore for PersistentStore { return Ok(()); } - let records_cf = self - .db + let db = Arc::clone(&self.db); + let records_cf = db .cf_handle(CF_RECORDS) .ok_or_else(|| anyhow!("missing records column family"))?; @@ -1056,8 +1026,19 @@ impl RecordStore for PersistentStore { let mut prepared_records = Vec::with_capacity(records.len()); // Prepare all records (assign IDs, intern strings) without writing + let mut pending_ids = std::collections::HashSet::with_capacity(records.len()); for mut record in records { + if record.id.0 != 0 { + self.ensure_record_id_available(record.id)?; + } + self.hydrate_record_interner(&record)?; let record_id = self.inner.prepare_record(&mut record)?; + if !pending_ids.insert(record_id) { + anyhow::bail!( + "record ID {} appears more than once in a batch", + record_id.0 + ); + } let key = record_id.0.to_be_bytes(); with_record_bytes(&record, |bytes| { batch.put_cf(records_cf, key, bytes); @@ -1143,8 +1124,8 @@ impl RecordStore for PersistentStore { } // Batch insert all new records - let records_cf = self - .db + let db = Arc::clone(&self.db); + let records_cf = db .cf_handle(CF_RECORDS) .ok_or_else(|| anyhow!("missing records column family"))?; @@ -1152,8 +1133,19 @@ impl RecordStore for PersistentStore { let mut prepared_records = Vec::with_capacity(new_records.len()); let mut assigned_ids = Vec::with_capacity(new_records.len()); + let mut pending_ids = std::collections::HashSet::with_capacity(new_records.len()); for mut record in new_records { + if record.id.0 != 0 { + self.ensure_record_id_available(record.id)?; + } + self.hydrate_record_interner(&record)?; let record_id = self.inner.prepare_record(&mut record)?; + if !pending_ids.insert(record_id) { + anyhow::bail!( + "record ID {} appears more than once in a batch", + record_id.0 + ); + } let key = record_id.0.to_be_bytes(); with_record_bytes(&record, |bytes| { batch.put_cf(records_cf, key, bytes); @@ -1213,6 +1205,12 @@ impl RecordStore for PersistentStore { PersistentStore::flush_staged_records(self) } + fn discard_staged_records(&mut self) -> Result<()> { + self.staged_records.clear(); + self.staged_identities.clear(); + Ok(()) + } + fn reserve_source_records( &mut self, reservations: &[SourceRecordReservation], @@ -1335,7 +1333,14 @@ impl RecordStore for PersistentStore { Ok(()) } + fn get_record_ref(&self, id: RecordId) -> Option<&Record> { + self.staged_records.get(&id) + } + fn get_record(&self, id: RecordId) -> Option { + if let Some(record) = self.staged_records.get(&id) { + return Some(record.clone()); + } if let Ok(mut cache) = self.cache.lock() { if let Some(record) = cache.get(&id) { return Some(record.clone()); @@ -1400,6 +1405,9 @@ impl RecordStore for PersistentStore { } fn get_record_id_by_identity(&self, identity: &RecordIdentity) -> Option { + if let Some(id) = self.staged_identities.get(identity) { + return Some(*id); + } let identity_cf = match self.db.cf_handle(CF_INDEX_IDENTITY) { Some(cf) => cf, None => { @@ -1726,12 +1734,30 @@ impl RecordStore for PersistentStore { self.inner.interner_mut() } + fn lookup_attr(&self, attr: &str) -> Option { + self.inner.interner().get_attr_id(attr).or_else(|| { + self.lookup_interner_id(b'A', attr) + .map(crate::model::AttrId) + }) + } + + fn lookup_value(&self, value: &str) -> Option { + self.inner.interner().get_value_id(value).or_else(|| { + self.lookup_interner_id(b'V', value) + .map(crate::model::ValueId) + }) + } + fn intern_attr(&mut self, attr: &str) -> crate::model::AttrId { if let Some(id) = self.inner.interner().get_attr_id(attr) { return id; } if let Some(id) = self.lookup_interner_id(b'A', attr) { - return crate::model::AttrId(id); + let id = crate::model::AttrId(id); + self.inner + .interner_mut() + .insert_attr_with_id(id, attr.to_owned()); + return id; } self.inner.interner_mut().intern_attr(attr) } @@ -1741,7 +1767,11 @@ impl RecordStore for PersistentStore { return id; } if let Some(id) = self.lookup_interner_id(b'V', value) { - return crate::model::ValueId(id); + let id = crate::model::ValueId(id); + self.inner + .interner_mut() + .insert_value_with_id(id, value.to_owned()); + return id; } self.inner.interner_mut().intern_value(value) } @@ -3213,6 +3243,25 @@ impl<'a> LinkerStatePersistence<'a> { Self { db } } + /// Check the ID scheme without migrating metadata or removing old redirects. + pub fn has_stable_global_cluster_ids(&self) -> Result { + let cf = self + .db + .cf_handle(linker_cf::METADATA) + .ok_or_else(|| anyhow!("Column family {} not found", linker_cf::METADATA))?; + match self + .db + .get_cf(cf, linker_encoding::KEY_GLOBAL_CLUSTER_ID_SCHEME)? + .as_deref() + { + Some([linker_encoding::STABLE_RECORD_ANCHOR_SCHEME]) => Ok(true), + Some(value) => { + anyhow::bail!("unsupported global cluster ID scheme marker: {:?}", value) + } + None => Ok(false), + } + } + /// Prepare replay-stable global IDs, removing redirects from the old /// allocation-order scheme if this database predates the scheme marker. pub fn prepare_stable_global_cluster_ids(&self) -> Result { diff --git a/src/query.rs b/src/query.rs index 5e48485..537e438 100644 --- a/src/query.rs +++ b/src/query.rs @@ -13,6 +13,21 @@ pub struct QueryDescriptor { pub value: ValueId, } +/// A shard's contribution to a canonical entity. Descriptor intervals are +/// aligned with the request; conjunction is evaluated after global assembly. +#[derive(Debug, Clone)] +pub struct EntityFragment { + pub global_id: crate::model::GlobalClusterId, + pub descriptor_intervals: Vec>, + /// Raw, coalesced attribute observations. Trim conflicting values only + /// after every shard's contribution has been assembled by the router. + pub golden: Vec, + /// Complete temporal guards, including observations outside a query window. + pub strong_ids: Vec, + pub cluster_key: Option, + pub cluster_key_identity: Option, +} + #[derive(Debug, Clone)] pub struct QuerySelectivity { count: u64, @@ -339,7 +354,7 @@ pub fn query_master_entities_with_cache_selective( Ok(QueryOutcome::Matches(matches)) } -fn intersect_interval_sets(a: &[Interval], b: &[Interval]) -> Vec { +pub(crate) fn intersect_interval_sets(a: &[Interval], b: &[Interval]) -> Vec { let mut overlaps = Vec::new(); for interval_a in a { @@ -353,7 +368,7 @@ fn intersect_interval_sets(a: &[Interval], b: &[Interval]) -> Vec { coalesce_intervals(&overlaps) } -fn coalesce_intervals(intervals: &[Interval]) -> Vec { +pub(crate) fn coalesce_intervals(intervals: &[Interval]) -> Vec { if intervals.is_empty() { return Vec::new(); } @@ -476,7 +491,7 @@ pub fn build_record_to_cluster_map( map } -fn filter_golden_for_interval( +pub(crate) fn filter_golden_for_interval( golden: &[GoldenDescriptor], interval: Interval, ) -> Vec { diff --git a/src/sharding.rs b/src/sharding.rs index e9cdaf7..d93e122 100644 --- a/src/sharding.rs +++ b/src/sharding.rs @@ -132,9 +132,7 @@ impl BloomFilter { /// Clear all bits. pub fn clear(&mut self) { - for word in &mut self.bits { - *word = 0; - } + self.bits.fill(0); } fn hash_with_seed(&self, key: &IdentityKeySignature, seed: usize) -> u64 { @@ -476,6 +474,8 @@ pub struct ReconciliationResult { #[derive(Debug, Default)] pub(crate) struct ReconciliationCandidates { merges: Vec<(GlobalClusterId, GlobalClusterId)>, + observations: HashMap>, + merge_signatures: HashMap<(GlobalClusterId, GlobalClusterId), IdentityKeySignature>, keys_checked: usize, keys_matched: usize, merge_candidates: usize, @@ -484,8 +484,37 @@ pub(crate) struct ReconciliationCandidates { } impl ReconciliationCandidates { + pub(crate) fn cluster_ids(&self) -> Vec { + let mut clusters = self + .merges + .iter() + .flat_map(|(left, right)| [*left, *right]) + .collect::>(); + clusters.sort_by_key(GlobalClusterId::to_u64); + clusters.dedup(); + clusters + } + + pub(crate) fn add_observations( + &mut self, + cluster: GlobalClusterId, + observations: Vec, + ) { + self.observations + .entry(cluster) + .or_default() + .extend(observations); + } + pub(crate) fn extend(&mut self, mut other: Self) { self.merges.append(&mut other.merges); + for (cluster, observations) in other.observations { + self.observations + .entry(cluster) + .or_default() + .extend(observations); + } + self.merge_signatures.extend(other.merge_signatures); self.keys_checked = self.keys_checked.saturating_add(other.keys_checked); self.keys_matched = self.keys_matched.saturating_add(other.keys_matched); self.merge_candidates = self.merge_candidates.saturating_add(other.merge_candidates); @@ -496,9 +525,15 @@ impl ReconciliationCandidates { .append(&mut other.detected_conflicts); } - pub(crate) fn finish(self) -> ReconciliationResult { - let (merges, component_conflicts_blocked) = - canonicalize_merges(self.merges, &self.detected_conflicts); + pub(crate) fn finish(mut self) -> ReconciliationResult { + let (merges, component_conflicts_blocked, component_conflicts) = + canonicalize_merges_with_observations( + self.merges, + &self.detected_conflicts, + self.observations, + &self.merge_signatures, + ); + self.detected_conflicts.extend(component_conflicts); ReconciliationResult { merges_performed: merges.len(), merged_clusters: merges, @@ -555,82 +590,17 @@ impl IncrementalReconciler { usize, Vec, ) { - let mut merges = Vec::new(); - let mut merge_candidates = 0; - let mut conflicts_blocked = 0usize; - let mut detected_conflicts = Vec::new(); - - // Build a map of all signatures to their entries across all shards - let mut all_entries: HashMap> = HashMap::new(); - - for boundary in &self.shard_boundaries { - for (sig, entries) in &boundary.boundary_keys { - all_entries.entry(*sig).or_default().extend(entries.iter()); - } - } - - // Find signatures that appear in multiple shards - for (sig, entries) in all_entries { - if entries.len() < 2 { - continue; - } - - // Group by shard and find overlapping intervals - let mut shard_entries: HashMap> = HashMap::new(); - for entry in &entries { - shard_entries.entry(entry.shard_id).or_default().push(entry); - } - - // If entries exist in multiple shards, check for temporal overlap - if shard_entries.len() > 1 { - let entries_vec: Vec<_> = entries.iter().collect(); - for i in 0..entries_vec.len() { - for j in (i + 1)..entries_vec.len() { - let e1 = entries_vec[i]; - let e2 = entries_vec[j]; - - // Only merge if from different shards and intervals overlap - if e1.shard_id != e2.shard_id && is_overlapping(&e1.interval, &e2.interval) - { - merge_candidates += 1; - - // Check for same-perspective conflicts before proposing merge - if let Some(conflict_detail) = detect_cross_shard_conflict(e1, e2) { - // Record the conflict and skip this merge - conflicts_blocked += 1; - - detected_conflicts.push(CrossShardConflict { - identity_key_signature: sig, - cluster1: e1.cluster_id, - cluster2: e2.cluster_id, - interval: conflict_detail.interval, - perspective_hash: conflict_detail.perspective_hash, - strong_id_hash1: conflict_detail.strong_id_hash1, - strong_id_hash2: conflict_detail.strong_id_hash2, - }); - continue; - } - - // Merge into the lower shard_id for consistency - let (primary, secondary) = if e1.shard_id < e2.shard_id { - (e1.cluster_id, e2.cluster_id) - } else { - (e2.cluster_id, e1.cluster_id) - }; - merges.push((primary, secondary)); - } - } - } - } - } - - let (merges, component_conflicts_blocked) = - canonicalize_merges(merges, &detected_conflicts); + let keys = self + .shard_boundaries + .iter() + .flat_map(|boundary| boundary.boundary_keys.keys().copied()) + .collect(); + let result = self.reconciliation_candidates(&keys).finish(); ( - merges, - merge_candidates, - conflicts_blocked.saturating_add(component_conflicts_blocked), - detected_conflicts, + result.merged_clusters, + result.merge_candidates, + result.conflicts_blocked, + result.detected_conflicts, ) } @@ -699,6 +669,8 @@ impl IncrementalReconciler { let mut conflicts_blocked = 0usize; let mut keys_matched = 0; let mut detected_conflicts = Vec::new(); + let mut observations: HashMap> = HashMap::new(); + let mut merge_signatures = HashMap::new(); // For each dirty key, collect entries from all shards for sig in keys { @@ -753,7 +725,22 @@ impl IncrementalReconciler { } else { (e2.cluster_id, e1.cluster_id) }; + if primary == secondary { + continue; + } + for entry in [e1, e2] { + observations + .entry(entry.cluster_id) + .or_default() + .extend(entry.strong_ids.iter().cloned()); + } merges.push((primary, secondary)); + let edge = if primary.to_u64() < secondary.to_u64() { + (primary, secondary) + } else { + (secondary, primary) + }; + merge_signatures.insert(edge, *sig); } } } @@ -761,6 +748,8 @@ impl IncrementalReconciler { ReconciliationCandidates { merges, + observations, + merge_signatures, keys_checked: keys.len(), keys_matched, merge_candidates, @@ -779,10 +768,26 @@ impl IncrementalReconciler { } } +#[cfg(test)] fn canonicalize_merges( merges: Vec<(GlobalClusterId, GlobalClusterId)>, conflicts: &[CrossShardConflict], ) -> (Vec<(GlobalClusterId, GlobalClusterId)>, usize) { + let (merges, blocked, _) = + canonicalize_merges_with_observations(merges, conflicts, HashMap::new(), &HashMap::new()); + (merges, blocked) +} + +fn canonicalize_merges_with_observations( + merges: Vec<(GlobalClusterId, GlobalClusterId)>, + conflicts: &[CrossShardConflict], + mut observations: HashMap>, + merge_signatures: &HashMap<(GlobalClusterId, GlobalClusterId), IdentityKeySignature>, +) -> ( + Vec<(GlobalClusterId, GlobalClusterId)>, + usize, + Vec, +) { fn find(parent: &mut [usize], mut index: usize) -> usize { while parent[index] != index { parent[index] = parent[parent[index]]; @@ -842,6 +847,17 @@ fn canonicalize_merges( .map(|node| blocked_by_node.remove(node).unwrap_or_default()) .collect::>(); let mut component_conflicts_blocked = 0usize; + let mut component_observations = nodes + .iter() + .map(|node| { + observations + .remove(node) + .unwrap_or_default() + .into_iter() + .collect::>() + }) + .collect::>(); + let mut component_conflicts = Vec::new(); for (left, right) in edges { let mut left_root = find(&mut parent, node_index[&left]); @@ -855,6 +871,28 @@ fn canonicalize_merges( component_conflicts_blocked = component_conflicts_blocked.saturating_add(1); continue; } + // A bridge may carry no strong ID of its own. Validate accumulated + // observations before joining components, including across key chunks. + let conflict = component_observations[left_root].iter().find_map(|left| { + component_observations[right_root] + .iter() + .find_map(|right| detect_strong_id_conflict(left, right)) + }); + if let Some(conflict) = conflict { + component_conflicts_blocked = component_conflicts_blocked.saturating_add(1); + if let Some(signature) = merge_signatures.get(&(left, right)) { + component_conflicts.push(CrossShardConflict { + identity_key_signature: *signature, + cluster1: nodes[left_root], + cluster2: nodes[right_root], + interval: conflict.interval, + perspective_hash: conflict.perspective_hash, + strong_id_hash1: conflict.strong_id_hash1, + strong_id_hash2: conflict.strong_id_hash2, + }); + } + continue; + } if nodes[left_root].to_u64() > nodes[right_root].to_u64() { std::mem::swap(&mut left_root, &mut right_root); } @@ -863,6 +901,8 @@ fn canonicalize_merges( members[left_root].extend(right_members); let right_blocked = std::mem::take(&mut blocked[right_root]); blocked[left_root].extend(right_blocked); + let right_observations = std::mem::take(&mut component_observations[right_root]); + component_observations[left_root].extend(right_observations); } let mut canonical = Vec::new(); @@ -882,7 +922,7 @@ fn canonicalize_merges( } } canonical.sort_by_key(|(primary, secondary)| (primary.to_u64(), secondary.to_u64())); - (canonical, component_conflicts_blocked) + (canonical, component_conflicts_blocked, component_conflicts) } impl Default for IncrementalReconciler { @@ -913,6 +953,34 @@ fn stable_boundary_hash(domain: &[u8], values: &[&str]) -> u64 { ) } +fn detect_strong_id_conflict( + left: &BoundaryStrongId, + right: &BoundaryStrongId, +) -> Option { + if left.perspective != right.perspective + || left.attribute != right.attribute + || left.value == right.value + { + return None; + } + let interval = crate::temporal::intersect(&left.interval, &right.interval)?; + Some(ConflictDetail { + perspective_hash: stable_boundary_hash( + b"unirust.boundary-perspective.v1", + &[&left.perspective], + ), + strong_id_hash1: stable_boundary_hash( + b"unirust.boundary-strong-id.v1", + &[&left.attribute, &left.value], + ), + strong_id_hash2: stable_boundary_hash( + b"unirust.boundary-strong-id.v1", + &[&right.attribute, &right.value], + ), + interval, + }) +} + /// Check if two boundary entries have conflicting strong IDs. /// /// Two entries conflict if they share a perspective (same key in perspective_strong_ids) @@ -922,7 +990,7 @@ fn stable_boundary_hash(domain: &[u8], values: &[&str]) -> u64 { /// /// Returns `Some(ConflictDetail)` if a conflict is found, `None` otherwise. fn detect_cross_shard_conflict(e1: &BoundaryEntry, e2: &BoundaryEntry) -> Option { - if !e1.strong_ids.is_empty() && !e2.strong_ids.is_empty() { + if !e1.strong_ids.is_empty() || !e2.strong_ids.is_empty() { for strong_id_1 in &e1.strong_ids { for strong_id_2 in &e2.strong_ids { if strong_id_1.perspective != strong_id_2.perspective @@ -1725,6 +1793,7 @@ mod tests { merge_candidates: 3, conflicts_blocked: 1, detected_conflicts: vec![conflict], + ..ReconciliationCandidates::default() }); let result = combined.finish(); @@ -1735,6 +1804,41 @@ mod tests { assert_eq!(result.conflicts_blocked, 2); } + #[test] + fn exact_component_guards_survive_metadata_chunks_and_allow_adjacent_values() { + let a = GlobalClusterId::new(0, 0, 0); + let b = GlobalClusterId::new(1, 0, 0); + let c = GlobalClusterId::new(2, 0, 0); + let signature = IdentityKeySignature::from_bytes([7; 32]); + for (second_start, expected_merges) in [(25, 1), (50, 2)] { + let observation = |value: &str, start, end| BoundaryStrongId { + perspective: "hr".into(), + attribute: "ssn".into(), + value: value.into(), + interval: Interval::new(start, end).unwrap(), + }; + let mut combined = ReconciliationCandidates::default(); + combined.extend(ReconciliationCandidates { + merges: vec![(a, b)], + observations: HashMap::from([(a, HashSet::from([observation("111", 0, 50)]))]), + merge_signatures: HashMap::from([((a, b), signature)]), + ..ReconciliationCandidates::default() + }); + combined.extend(ReconciliationCandidates { + merges: vec![(b, c)], + observations: HashMap::from([( + c, + HashSet::from([observation("222", second_start, 100)]), + )]), + merge_signatures: HashMap::from([((b, c), signature)]), + ..ReconciliationCandidates::default() + }); + let result = combined.finish(); + assert_eq!(result.merges_performed, expected_merges); + assert_eq!(result.detected_conflicts.len(), 2 - expected_merges); + } + } + #[test] fn test_boundary_metadata_export_import() { use crate::model::{AttrId, ValueId}; diff --git a/src/store.rs b/src/store.rs index 8dc51fc..44d7e37 100644 --- a/src/store.rs +++ b/src/store.rs @@ -141,6 +141,16 @@ pub trait RecordStore: Send + Sync { /// Get a mutable reference to the string interner. fn interner_mut(&mut self) -> &mut StringInterner; + /// Look up a string without allocating an ID or mutating the store. + fn lookup_attr(&self, attr: &str) -> Option { + self.interner().get_attr_id(attr) + } + + /// Look up a string without allocating an ID or mutating the store. + fn lookup_value(&self, value: &str) -> Option { + self.interner().get_value_id(value) + } + /// Intern an attribute string. fn intern_attr(&mut self, attr: &str) -> AttrId { self.interner_mut().intern_attr(attr) @@ -291,6 +301,12 @@ pub trait RecordStore: Send + Sync { Ok(0) } + /// Abandon records staged by a failed ingest before they can be committed by + /// a later request. Implementations without staging have nothing to discard. + fn discard_staged_records(&mut self) -> Result<()> { + Ok(()) + } + /// Make all completed writes durable on stable storage. /// /// In-memory stores have nothing to synchronize. Persistent implementations @@ -355,6 +371,8 @@ pub struct Store { temporal_index: TemporalIndex, /// Next available record ID next_record_id: u32, + /// Inserts made through staging remain removable until the batch commits. + staged_record_ids: Vec, } pub(crate) fn records_have_same_payload(left: &Record, right: &Record) -> bool { @@ -404,6 +422,7 @@ impl Store { attribute_value_index: AttributeValueIndex::new(), temporal_index: TemporalIndex::new(), next_record_id: 0, + staged_record_ids: Vec::new(), } } @@ -418,6 +437,7 @@ impl Store { attribute_value_index: AttributeValueIndex::new(), temporal_index: TemporalIndex::new(), next_record_id, + staged_record_ids: Vec::new(), } } @@ -473,6 +493,9 @@ impl Store { /// Add a single record to the store and return its assigned ID. pub fn add_record(&mut self, mut record: Record) -> Result { + if record.id.0 != 0 && self.records.contains_key(&record.id) { + anyhow::bail!("record ID {} already exists", record.id.0); + } let record_id = self.prepare_record(&mut record)?; let identity = record.identity.clone(); @@ -487,6 +510,9 @@ impl Store { /// Insert a record with an explicit ID without assigning a new one. pub fn insert_record(&mut self, mut record: Record) -> Result { + if self.records.contains_key(&record.id) { + anyhow::bail!("record ID {} already exists", record.id.0); + } self.intern_record(&mut record); if record.id.0 == Self::EXHAUSTED_RECORD_ID { anyhow::bail!("record ID {} is reserved", Self::EXHAUSTED_RECORD_ID); @@ -884,6 +910,40 @@ impl RecordStore for Store { Store::add_record_if_absent(self, record) } + fn stage_record_if_absent(&mut self, record: Record) -> Result<(RecordId, bool)> { + let (id, inserted) = self.add_record_if_absent(record)?; + if inserted { + self.staged_record_ids.push(id); + } + Ok((id, inserted)) + } + + fn flush_staged_records(&mut self) -> Result { + let count = self.staged_record_ids.len(); + self.staged_record_ids.clear(); + Ok(count) + } + + fn discard_staged_records(&mut self) -> Result<()> { + if self.staged_record_ids.is_empty() { + return Ok(()); + } + for id in self.staged_record_ids.drain(..) { + if let Some(record) = self.records.remove(&id) { + self.identity_index.remove(&record.identity); + } + } + // Abort is an exceptional path; rebuilding avoids retaining index entries + // for records whose staging never committed. + self.attribute_value_index = AttributeValueIndex::new(); + self.temporal_index = TemporalIndex::new(); + for record in self.records.values() { + self.attribute_value_index.add_record(record); + self.temporal_index.add_record(record); + } + Ok(()) + } + fn stage_record_with_explicit_id_if_absent( &mut self, record: Record, @@ -900,6 +960,7 @@ impl RecordStore for Store { anyhow::bail!("record ID {} already exists", record.id.0); } let id = self.insert_record(record)?; + self.staged_record_ids.push(id); Ok((id, true)) } } @@ -1068,6 +1129,37 @@ mod tests { assert_eq!(store.len(), 0); } + #[test] + fn aborted_staging_removes_records_and_secondary_indexes() { + let mut store = Store::new(); + let attr = store.interner_mut().intern_attr("email"); + let value = store.interner_mut().intern_value("shared@example.com"); + let interval = Interval::new(0, 10).unwrap(); + let make_record = |source: &str| { + Record::new( + RecordId(0), + RecordIdentity::new("person".into(), "crm".into(), source.into()), + vec![Descriptor::new(attr, value, interval)], + ) + }; + let committed = store.add_record(make_record("committed")).unwrap(); + let (pending, _) = store + .stage_record_if_absent(make_record("pending")) + .unwrap(); + assert_eq!(store.get_records_in_interval(interval).len(), 2); + store.discard_staged_records().unwrap(); + assert!(store.get_record(pending).is_none()); + assert_eq!(store.get_records_in_interval(interval).len(), 1); + assert_eq!(store.get_records_with_attribute(attr)[0].id, committed); + let (retried, inserted) = store + .stage_record_if_absent(make_record("pending")) + .unwrap(); + assert!(inserted); + store.flush_staged_records().unwrap(); + store.discard_staged_records().unwrap(); + assert!(store.get_record(retried).is_some()); + } + #[test] fn record_id_allocation_fails_closed_at_u32_boundary() { let mut store = Store::new(); diff --git a/tests/audit_performance.rs b/tests/audit_performance.rs new file mode 100644 index 0000000..5aea844 --- /dev/null +++ b/tests/audit_performance.rs @@ -0,0 +1,61 @@ +//! Manual measurements for the September 2026 audit. No timing assertions. + +use std::time::Instant; + +use tempfile::tempdir; +use unirust_rs::ontology::IdentityKey; +use unirust_rs::{ + Descriptor, Interval, Ontology, PersistentStore, QueryDescriptor, QueryOutcome, Record, + RecordId, RecordIdentity, Unirust, +}; + +fn make_record(engine: &mut Unirust, id: usize) -> Record { + let attr = engine.intern_attr("email"); + let value = engine.intern_value(&format!("person-{id}@example.com")); + Record::new( + RecordId(0), + RecordIdentity::new("person".into(), "audit".into(), id.to_string()), + vec![Descriptor::new(attr, value, Interval::new(0, 100).unwrap())], + ) +} + +#[test] +#[ignore = "manual performance audit; run in release mode with --ignored --nocapture"] +fn selective_query_after_unrelated_ingest() -> anyhow::Result<()> { + for size in [5_000, 20_000, 80_000] { + let dir = tempdir()?; + let store = PersistentStore::open(dir.path())?; + let mut ontology = Ontology::new(); + ontology.add_identity_key(IdentityKey::from_names(vec!["email"], "email")); + let mut engine = Unirust::with_store(ontology, store); + for start in (0..size).step_by(1_000) { + let records = (start..start + 1_000) + .map(|id| make_record(&mut engine, id)) + .collect(); + engine.stream_records(records)?; + } + let query = [QueryDescriptor { + attr: engine.intern_attr("email"), + value: engine.intern_value("person-0@example.com"), + }]; + let mut elapsed = Vec::new(); + for step in 0..3 { + if step == 2 { + let record = make_record(&mut engine, size); + engine.stream_records(vec![record])?; + } + let start = Instant::now(); + let outcome = engine.query_master_entities(&query, Interval::new(0, 100)?)?; + elapsed.push(start.elapsed().as_secs_f64() * 1_000.0); + match outcome { + QueryOutcome::Matches(matches) => assert_eq!(matches.len(), 1), + other => panic!("unexpected query result: {other:?}"), + } + } + println!( + "records={size} cold_ms={:.3} warm_ms={:.3} after_unrelated_ingest_ms={:.3}", + elapsed[0], elapsed[1], elapsed[2] + ); + } + Ok(()) +} diff --git a/tests/distributed_apply_ontology.rs b/tests/distributed_apply_ontology.rs index 3c8939c..af80a17 100644 --- a/tests/distributed_apply_ontology.rs +++ b/tests/distributed_apply_ontology.rs @@ -101,6 +101,17 @@ async fn apply_ontology_enables_queries_distributed() -> anyhow::Result<()> { }) .await?; + let response = client.query_entities(query.clone()).await?.into_inner(); + match response.outcome { + Some(proto::query_entities_response::Outcome::Matches(matches)) => { + assert!( + matches.matches.is_empty(), + "new ontology must leave empty shards queryable" + ); + } + _ => anyhow::bail!("expected empty matches immediately after applying ontology"), + } + let record = RecordInput { index: 0, identity: Some(ProtoRecordIdentity { @@ -118,7 +129,7 @@ async fn apply_ontology_enables_queries_distributed() -> anyhow::Result<()> { client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record], }) .await?; diff --git a/tests/distributed_conflicts_presets.rs b/tests/distributed_conflicts_presets.rs index 0c60c06..b408d85 100644 --- a/tests/distributed_conflicts_presets.rs +++ b/tests/distributed_conflicts_presets.rs @@ -272,7 +272,8 @@ async fn distributed_conflict_presets_match_local() -> anyhow::Result<()> { if !inputs.is_empty() { client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: + unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: inputs.clone(), }) .await diff --git a/tests/distributed_cross_shard_reconciliation.rs b/tests/distributed_cross_shard_reconciliation.rs index a8415fa..a46c966 100644 --- a/tests/distributed_cross_shard_reconciliation.rs +++ b/tests/distributed_cross_shard_reconciliation.rs @@ -220,7 +220,7 @@ async fn cross_shard_conflict_detected_via_reconcile() -> anyhow::Result<()> { shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![shard0_rec1, shard0_rec2], }) .await?; @@ -251,7 +251,7 @@ async fn cross_shard_conflict_detected_via_reconcile() -> anyhow::Result<()> { shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![shard1_rec1, shard1_rec2], }) .await?; @@ -348,7 +348,7 @@ async fn cross_shard_merge_succeeds_without_conflict() -> anyhow::Result<()> { ); shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![shard0_rec1], }) .await?; @@ -368,7 +368,7 @@ async fn cross_shard_merge_succeeds_without_conflict() -> anyhow::Result<()> { ); shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![shard1_rec1], }) .await?; @@ -483,7 +483,7 @@ async fn cross_shard_conflicts_propagated_to_shards() -> anyhow::Result<()> { shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![shard0_rec1, shard0_rec2], }) .await?; @@ -514,7 +514,7 @@ async fn cross_shard_conflicts_propagated_to_shards() -> anyhow::Result<()> { shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![shard1_rec1, shard1_rec2], }) .await?; @@ -586,7 +586,7 @@ async fn cross_shard_merge_respects_strong_id_validity_intervals() -> anyhow::Re shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record_input( 0, "instrument", @@ -602,7 +602,7 @@ async fn cross_shard_merge_respects_strong_id_validity_intervals() -> anyhow::Re .await?; shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record_input( 1, "instrument", @@ -689,7 +689,7 @@ async fn cross_shard_reconciliation_preserves_identity_key_gaps() -> anyhow::Res // is valid only before and after the middle gap. let shard0_ingest = shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![ record_input( 0, @@ -715,7 +715,7 @@ async fn cross_shard_reconciliation_preserves_identity_key_gaps() -> anyhow::Res ); shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record_input( 2, "person", @@ -797,7 +797,7 @@ async fn boundary_metadata_includes_perspective_strong_ids() -> anyhow::Result<( router_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record1, record2], }) .await?; @@ -899,7 +899,7 @@ async fn dirty_boundary_keys_include_perspective_strong_ids() -> anyhow::Result< router_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record1, record2], }) .await?; @@ -981,7 +981,7 @@ async fn dirty_boundary_keys_are_drained_in_bounded_pages() -> anyhow::Result<() let uid2 = format!("page_{index}_b"); router_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![ record_input( index * 2, @@ -1134,7 +1134,7 @@ async fn transitive_cross_shard_conflict_detected() -> anyhow::Result<()> { ); shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![a1, a2], }) .await?; @@ -1166,7 +1166,7 @@ async fn transitive_cross_shard_conflict_detected() -> anyhow::Result<()> { ); shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![b1, b2], }) .await?; @@ -1196,7 +1196,7 @@ async fn transitive_cross_shard_conflict_detected() -> anyhow::Result<()> { ); shard2_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![c1, c2], }) .await?; @@ -1321,7 +1321,7 @@ async fn peic_many_entities_claim_same_identifier_across_shards() -> anyhow::Res ); shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![a1, a2], }) .await?; @@ -1352,7 +1352,7 @@ async fn peic_many_entities_claim_same_identifier_across_shards() -> anyhow::Res ); shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![b1, b2], }) .await?; @@ -1449,7 +1449,7 @@ async fn temporal_overlap_conflict_across_shards() -> anyhow::Result<()> { ); shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![a1, a2], }) .await?; @@ -1479,7 +1479,7 @@ async fn temporal_overlap_conflict_across_shards() -> anyhow::Result<()> { ); shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![b1, b2], }) .await?; @@ -1573,7 +1573,7 @@ async fn late_arriving_data_triggers_conflict() -> anyhow::Result<()> { ); shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![initial1, initial2], }) .await?; @@ -1618,7 +1618,7 @@ async fn late_arriving_data_triggers_conflict() -> anyhow::Result<()> { ); shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![late1, late2], }) .await?; @@ -1742,7 +1742,7 @@ async fn multi_hop_chain_conflict_across_four_shards() -> anyhow::Result<()> { ); shard0 .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![a1, a2], }) .await?; @@ -1772,7 +1772,7 @@ async fn multi_hop_chain_conflict_across_four_shards() -> anyhow::Result<()> { ); shard1 .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![b1, b2], }) .await?; @@ -1802,7 +1802,7 @@ async fn multi_hop_chain_conflict_across_four_shards() -> anyhow::Result<()> { ); shard2 .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![c1, c2], }) .await?; @@ -1827,7 +1827,7 @@ async fn multi_hop_chain_conflict_across_four_shards() -> anyhow::Result<()> { ); shard3 .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![d1, d2], }) .await?; @@ -1930,7 +1930,7 @@ async fn different_perspectives_no_false_positive_conflict() -> anyhow::Result<( ); shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![msci1, msci2], }) .await?; @@ -1960,7 +1960,7 @@ async fn different_perspectives_no_false_positive_conflict() -> anyhow::Result<( ); shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![axioma1, axioma2], }) .await?; diff --git a/tests/distributed_cross_shard_reconciliation_persistent.rs b/tests/distributed_cross_shard_reconciliation_persistent.rs index 703f25f..b56be1e 100644 --- a/tests/distributed_cross_shard_reconciliation_persistent.rs +++ b/tests/distributed_cross_shard_reconciliation_persistent.rs @@ -382,7 +382,7 @@ async fn cross_shard_merge_persistent_store() -> anyhow::Result<()> { ); shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![shard0_rec1, shard0_rec2], }) .await?; @@ -409,7 +409,7 @@ async fn cross_shard_merge_persistent_store() -> anyhow::Result<()> { ); shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![shard1_rec1, shard1_rec2], }) .await?; @@ -459,7 +459,7 @@ async fn three_shard_singleton_merge_survives_full_restart() -> anyhow::Result<( let mut shard_client = ShardServiceClient::connect(format!("http://{shard_addr}")).await?; shard_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record_input( shard_id as u32, "person", @@ -569,7 +569,7 @@ async fn source_identity_reservation_survives_routing_change_and_restart() -> an let response = router_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![original.clone()], }) .await? @@ -579,7 +579,7 @@ async fn source_identity_reservation_survives_routing_change_and_restart() -> an let error = router_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![changed.clone()], }) .await @@ -616,7 +616,7 @@ async fn source_identity_reservation_survives_routing_change_and_restart() -> an let error = router_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![changed], }) .await @@ -626,7 +626,7 @@ async fn source_identity_reservation_survives_routing_change_and_restart() -> an original.index = 1; let retry = router_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![original], }) .await? @@ -701,7 +701,7 @@ async fn router_backfills_legacy_records_before_serving_ingest() -> anyhow::Resu assert_eq!(status_before.source_reservation_backfill_version, 0); original_target_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![original], }) .await?; @@ -717,7 +717,7 @@ async fn router_backfills_legacy_records_before_serving_ingest() -> anyhow::Resu .into_inner(); assert_eq!( status.source_reservation_backfill_version, - DISTRIBUTED_PROTOCOL_VERSION + unirust_rs::distributed::SOURCE_RESERVATION_BACKFILL_VERSION ); assert_eq!(status.source_reservation_shard_count, 2); } @@ -725,7 +725,7 @@ async fn router_backfills_legacy_records_before_serving_ingest() -> anyhow::Resu let mut router_client = RouterServiceClient::connect(format!("http://{router_addr}")).await?; let error = router_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![changed], }) .await @@ -791,7 +791,7 @@ async fn reserved_ingest_retries_after_target_failure_and_full_restart() -> anyh let error = router_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record.clone()], }) .await @@ -829,7 +829,7 @@ async fn reserved_ingest_retries_after_target_failure_and_full_restart() -> anyh record.index = 1; let retry = router_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record.clone()], }) .await? @@ -842,7 +842,7 @@ async fn reserved_ingest_retries_after_target_failure_and_full_restart() -> anyh changed.descriptors[0].value = "changed-after-partial-failure@example.com".to_string(); let error = router_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![changed], }) .await @@ -875,7 +875,7 @@ async fn partial_reconciliation_blocks_traffic_and_recovers_after_full_restart( let mut client1 = ShardServiceClient::connect(format!("http://{addr1}")).await?; client0 .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record_input( 0, "person", @@ -890,7 +890,7 @@ async fn partial_reconciliation_blocks_traffic_and_recovers_after_full_restart( .await?; client1 .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record_input( 1, "person", @@ -1095,7 +1095,7 @@ async fn partial_reconciliation_can_be_retried_in_place() -> anyhow::Result<()> let mut client = ShardServiceClient::connect(format!("http://{shard_addr}")).await?; client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record_input( 0, "person", diff --git a/tests/distributed_e2e.rs b/tests/distributed_e2e.rs index ed0af5b..cf8e38c 100644 --- a/tests/distributed_e2e.rs +++ b/tests/distributed_e2e.rs @@ -182,7 +182,7 @@ async fn distributed_ingest_and_query_unbounded_intervals() -> anyhow::Result<() .collect(); let response = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records, }) .await? @@ -264,7 +264,7 @@ async fn distributed_stream_and_query() -> anyhow::Result<()> { let response = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record_a.clone(), record_b.clone()], }) .await? diff --git a/tests/distributed_entity_regressions.rs b/tests/distributed_entity_regressions.rs new file mode 100644 index 0000000..826e531 --- /dev/null +++ b/tests/distributed_entity_regressions.rs @@ -0,0 +1,702 @@ +//! Persistent distributed regressions for canonical entity queries and guards. + +use proto::router_service_server::RouterService; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; +use tempfile::TempDir; +use tokio::task::JoinHandle; +use tokio_stream::wrappers::TcpListenerStream; +use tonic::codegen::{http, Service}; +use tonic::server::NamedService; +use tonic::{transport::Server, Request}; +use unirust_rs::distributed::{ + hash_record_to_shard, proto, DistributedOntologyConfig, IdentityKeyConfig, RouterNode, + ShardNode, +}; +use unirust_rs::{StreamingTuning, TuningProfile}; + +static AUDIT_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + +#[derive(Clone)] +struct SynchronizeQueries { + inner: S, + barrier: Option>, +} + +impl Service> for SynchronizeQueries +where + S: Service< + http::Request, + Response = http::Response, + Error = std::convert::Infallible, + > + Send, + S::Future: Send + 'static, + B: Send + 'static, +{ + type Response = http::Response; + type Error = std::convert::Infallible; + type Future = Pin> + Send>>; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, request: http::Request) -> Self::Future { + let barrier = (request.uri().path() == "/unirust.ShardService/QueryEntityFragments") + .then(|| self.barrier.clone()) + .flatten(); + let future = self.inner.call(request); + Box::pin(async move { + let response = future.await?; + if let Some(barrier) = barrier { + if tokio::time::timeout(std::time::Duration::from_secs(2), barrier.wait()) + .await + .is_err() + { + return Ok(tonic::Status::deadline_exceeded( + "router did not dispatch peer query concurrently", + ) + .into_http()); + } + } + Ok(response) + }) + } +} + +impl NamedService for SynchronizeQueries { + const NAME: &'static str = S::NAME; +} + +struct Cluster { + router: Arc, + shards: Vec, + servers: Vec>, + _directory: TempDir, +} + +impl Cluster { + async fn stop(self) -> anyhow::Result<()> { + for shard in &self.shards { + shard.shutdown().await?; + } + for server in self.servers { + server.abort(); + let _ = server.await; + } + Ok(()) + } + + async fn query(&self, descriptors: &[(&str, &str)]) -> anyhow::Result> { + let response = RouterService::query_entities( + self.router.as_ref(), + Request::new(proto::QueryEntitiesRequest { + descriptors: descriptors + .iter() + .map(|(attr, value)| proto::QueryDescriptor { + attr: (*attr).into(), + value: (*value).into(), + }) + .collect(), + start: 0, + end: 100, + }), + ) + .await? + .into_inner(); + match response.outcome { + Some(proto::query_entities_response::Outcome::Matches(matches)) => Ok(matches.matches), + other => anyhow::bail!("expected query matches, got {other:?}"), + } + } +} + +fn record(index: u32, email: &str, extra: (&str, &str)) -> proto::RecordInput { + proto::RecordInput { + index, + identity: Some(proto::RecordIdentity { + entity_type: "person".into(), + perspective: format!("source-{index}"), + uid: format!("record-{index}"), + }), + descriptors: [("email", email), ("phone", "shared-phone"), extra] + .into_iter() + .map(|(attr, value)| proto::RecordDescriptor { + attr: attr.into(), + value: value.into(), + start: 0, + end: 100, + }) + .collect(), + } +} + +async fn empty_cluster(config: &DistributedOntologyConfig, count: u32) -> anyhow::Result { + empty_cluster_with_barrier(config, count, None).await +} + +async fn empty_cluster_with_barrier( + config: &DistributedOntologyConfig, + count: u32, + barrier: Option>, +) -> anyhow::Result { + let directory = tempfile::tempdir()?; + let mut urls = Vec::new(); + let mut shards = Vec::new(); + let mut servers = Vec::new(); + for shard_id in 0..count { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + urls.push(format!("http://{}", listener.local_addr()?)); + let shard = ShardNode::new_with_data_dir( + shard_id, + config.clone(), + StreamingTuning::from_profile(TuningProfile::Balanced), + Some(directory.path().join(format!("shard-{shard_id}"))), + false, + None, + )?; + let service_shard = shard.clone(); + let barrier = barrier.clone(); + servers.push(tokio::spawn(async move { + Server::builder() + .add_service(SynchronizeQueries { + inner: proto::shard_service_server::ShardServiceServer::new(service_shard), + barrier, + }) + .serve_with_incoming(TcpListenerStream::new(listener)) + .await + .expect("audit shard server"); + })); + shards.push(shard); + } + let router = RouterNode::connect(urls, config.clone()).await?; + Ok(Cluster { + router, + shards, + servers, + _directory: directory, + }) +} + +async fn reconciled_cluster() -> anyhow::Result { + let config = DistributedOntologyConfig { + identity_keys: ["email", "phone"] + .into_iter() + .map(|attr| IdentityKeyConfig { + name: format!("{attr}_key"), + attributes: vec![attr.into()], + }) + .collect(), + strong_identifiers: vec![], + constraints: vec![], + }; + let cluster = empty_cluster(&config, 2).await?; + let router = &cluster.router; + let first = record(0, "first@example.com", ("name", "Alice")); + let first_target = hash_record_to_shard(&config, &first, 2); + let second = (0..1000) + .map(|candidate| { + record( + 1, + &format!("second-{candidate}@example.com"), + ("city", "London"), + ) + }) + .find(|candidate| hash_record_to_shard(&config, candidate, 2) != first_target) + .expect("find email routed to other shard"); + let response = RouterService::ingest_records( + router.as_ref(), + Request::new(proto::IngestRecordsRequest { + records: vec![first, second], + internal_protocol_version: 0, + }), + ) + .await? + .into_inner(); + assert_ne!( + response.assignments[0].shard_id, + response.assignments[1].shard_id + ); + let reconciliation = RouterService::reconcile( + router.as_ref(), + Request::new(proto::ReconcileRequest { + shard_metadata: vec![], + }), + ) + .await? + .into_inner(); + assert_eq!(reconciliation.merges_performed, 1); + assert_eq!( + cluster.query(&[("phone", "shared-phone")]).await?.len(), + 1, + "shared identity must resolve to one entity" + ); + Ok(cluster) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn query_conjunction_across_reconciled_shards() -> anyhow::Result<()> { + let _lock = AUDIT_LOCK.lock().await; + let cluster = reconciled_cluster().await?; + let matches = cluster + .query(&[("name", "Alice"), ("city", "London")]) + .await?; + cluster.stop().await?; + assert_eq!( + matches.len(), + 1, + "the reconciled entity contains both descriptors over [0,100)" + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn query_golden_across_reconciled_shards() -> anyhow::Result<()> { + let _lock = AUDIT_LOCK.lock().await; + let cluster = reconciled_cluster().await?; + let matches = cluster.query(&[("name", "Alice")]).await?; + cluster.stop().await?; + assert_eq!(matches.len(), 1); + assert!( + matches[0] + .golden + .iter() + .any(|descriptor| descriptor.attr == "city" && descriptor.value == "London"), + "entity golden must include city from its reconciled fragment, got {:?}", + matches[0].golden + ); + Ok(()) +} + +fn config_for_keys(keys: &[&str], strong: &[&str]) -> DistributedOntologyConfig { + DistributedOntologyConfig { + identity_keys: keys + .iter() + .map(|attr| IdentityKeyConfig { + name: format!("{attr}_key"), + attributes: vec![(*attr).into()], + }) + .collect(), + strong_identifiers: strong.iter().map(|attr| (*attr).into()).collect(), + constraints: Vec::new(), + } +} + +fn record_on_shard( + config: &DistributedOntologyConfig, + target: u32, + count: usize, + extra: (&str, &str), +) -> proto::RecordInput { + (0..1000) + .map(|candidate| { + record( + target, + &format!("target-{target}-{candidate}@example.com"), + extra, + ) + }) + .find(|record| hash_record_to_shard(config, record, count) == target as usize) + .expect("find routed record") +} + +async fn ingest(cluster: &Cluster, records: Vec) -> anyhow::Result<()> { + RouterService::ingest_records( + cluster.router.as_ref(), + Request::new(proto::IngestRecordsRequest { + records, + internal_protocol_version: 0, + }), + ) + .await?; + Ok(()) +} + +async fn reconcile(cluster: &Cluster) -> anyhow::Result { + Ok(RouterService::reconcile( + cluster.router.as_ref(), + Request::new(proto::ReconcileRequest { + shard_metadata: Vec::new(), + }), + ) + .await? + .into_inner()) +} + +fn descriptor(attr: &str, value: &str, start: i64, end: i64) -> proto::RecordDescriptor { + proto::RecordDescriptor { + attr: attr.into(), + value: value.into(), + start, + end, + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn distributed_query_preserves_intersections_and_adjacent_intervals() -> anyhow::Result<()> { + let _lock = AUDIT_LOCK.lock().await; + let config = config_for_keys(&["email", "phone"], &[]); + let cluster = empty_cluster(&config, 2).await?; + let mut left = record_on_shard(&config, 0, 2, ("name", "Alice")); + left.descriptors + .iter_mut() + .find(|entry| entry.attr == "name") + .unwrap() + .end = 60; + left.descriptors.extend([ + descriptor("left", "yes", 0, 40), + descriptor("status", "active", 0, 50), + ]); + let mut right = record_on_shard(&config, 1, 2, ("city", "London")); + right + .descriptors + .iter_mut() + .find(|entry| entry.attr == "city") + .unwrap() + .start = 40; + right + .descriptors + .push(descriptor("status", "active", 50, 100)); + ingest(&cluster, vec![left, right]).await?; + assert_eq!(reconcile(&cluster).await?.merges_performed, 1); + let overlap = cluster + .query(&[("name", "Alice"), ("city", "London")]) + .await?; + assert_eq!( + overlap + .iter() + .map(|entry| (entry.start, entry.end)) + .collect::>(), + vec![(40, 60)] + ); + assert!(overlap[0] + .golden + .iter() + .all(|entry| entry.start >= 40 && entry.end <= 60)); + assert!( + cluster + .query(&[("left", "yes"), ("city", "London")]) + .await? + .is_empty(), + "adjacent intervals have no conjunction overlap" + ); + let adjacent = cluster.query(&[("status", "active")]).await?; + assert_eq!( + adjacent + .iter() + .map(|entry| (entry.start, entry.end)) + .collect::>(), + vec![(0, 100)] + ); + assert_eq!( + adjacent[0] + .golden + .iter() + .filter(|entry| entry.attr == "status") + .map(|entry| (entry.start, entry.end)) + .collect::>(), + vec![(0, 100)] + ); + assert!(cluster.query(&[("unknown", "value")]).await?.is_empty()); + cluster.stop().await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn golden_conflict_trimming_uses_all_raw_entity_fragments() -> anyhow::Result<()> { + let _lock = AUDIT_LOCK.lock().await; + let config = config_for_keys(&["email", "phone"], &[]); + let cluster = empty_cluster(&config, 2).await?; + let left = record_on_shard(&config, 0, 2, ("name", "Alice")); + let right = record_on_shard(&config, 1, 2, ("name", "Alice")); + let mut conflict = left.clone(); + conflict.index = 2; + conflict.identity.as_mut().unwrap().uid = "conflicting-name".into(); + let name = conflict + .descriptors + .iter_mut() + .find(|entry| entry.attr == "name") + .unwrap(); + name.value = "Bob".into(); + name.start = 20; + name.end = 40; + ingest(&cluster, vec![left, right, conflict]).await?; + reconcile(&cluster).await?; + let matches = cluster.query(&[("phone", "shared-phone")]).await?; + assert_eq!(matches.len(), 1); + let names = matches[0] + .golden + .iter() + .filter(|entry| entry.attr == "name") + .map(|entry| (entry.value.as_str(), entry.start, entry.end)) + .collect::>(); + assert_eq!( + names, + vec![("Alice", 0, 20), ("Alice", 40, 100)], + "a remote Alice observation must not restore the locally conflicted interval" + ); + assert!( + !matches[0].golden.iter().any(|entry| entry.attr == "email"), + "conflicting emails on separate fragments must be removed globally" + ); + cluster.stop().await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn router_dispatches_candidate_and_hydration_queries_concurrently() -> anyhow::Result<()> { + let _lock = AUDIT_LOCK.lock().await; + let config = config_for_keys(&["email", "phone"], &[]); + let cluster = + empty_cluster_with_barrier(&config, 2, Some(Arc::new(tokio::sync::Barrier::new(2)))) + .await?; + let left = record_on_shard(&config, 0, 2, ("name", "Alice")); + let right = record_on_shard(&config, 1, 2, ("city", "London")); + ingest(&cluster, vec![left, right]).await?; + reconcile(&cluster).await?; + assert_eq!( + cluster + .query(&[("name", "Alice"), ("city", "London")]) + .await? + .len(), + 1 + ); + cluster.stop().await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn protocol_upgrade_retains_v5_reservation_topology_binding() -> anyhow::Result<()> { + let _lock = AUDIT_LOCK.lock().await; + let directory = tempfile::tempdir()?; + { + let mut store = unirust_rs::PersistentStore::open(directory.path())?; + unirust_rs::store::RecordStore::mark_source_reservation_backfill(&mut store, 5, 2)?; + } + let config = DistributedOntologyConfig::empty(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let address = listener.local_addr()?; + let shard = ShardNode::new_with_data_dir( + 0, + config.clone(), + StreamingTuning::balanced(), + Some(directory.path().to_path_buf()), + false, + None, + )?; + let service_shard = shard.clone(); + let server = tokio::spawn(async move { + Server::builder() + .add_service(proto::shard_service_server::ShardServiceServer::new( + service_shard, + )) + .serve_with_incoming(TcpListenerStream::new(listener)) + .await + .expect("upgrade shard"); + }); + let error = RouterNode::connect(vec![format!("http://{address}")], config) + .await + .err() + .expect("upgrading the wire protocol must not permit changing the old two-shard topology"); + assert_eq!(error.code(), tonic::Code::FailedPrecondition); + assert!(error.message().contains("topology mismatch")); + shard.shutdown().await?; + server.abort(); + let _ = server.await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn prior_reconciliation_preserves_remote_component_strong_ids() -> anyhow::Result<()> { + let _lock = AUDIT_LOCK.lock().await; + let config = config_for_keys(&["email", "key1", "key2"], &["ssn"]); + let cluster = empty_cluster(&config, 3).await?; + let mut records = Vec::new(); + for target in 0..3 { + let mut record = record_on_shard(&config, target, 3, ("group", "audit")); + record.descriptors.retain(|entry| entry.attr != "phone"); + record.identity.as_mut().unwrap().perspective = + if target == 1 { "bridge" } else { "hr" }.into(); + match target { + 0 => record.descriptors.extend([ + descriptor("key1", "AB", 0, 100), + descriptor("ssn", "111", 0, 100), + ]), + 1 => record.descriptors.extend([ + descriptor("key1", "AB", 0, 100), + descriptor("key2", "BC", 0, 100), + ]), + _ => record.descriptors.extend([ + descriptor("key2", "BC", 0, 100), + descriptor("ssn", "222", 0, 100), + ]), + } + records.push(record); + } + let last = records.pop().unwrap(); + ingest(&cluster, records).await?; + assert_eq!(reconcile(&cluster).await?.merges_performed, 1); + ingest(&cluster, vec![last]).await?; + let result = reconcile(&cluster).await?; + assert_eq!( + result.merges_performed, 0, + "clean key1 must still contribute hr SSN111 to the canonical component" + ); + assert!(result.conflicts_blocked > 0); + let response = RouterService::query_entities( + cluster.router.as_ref(), + Request::new(proto::QueryEntitiesRequest { + descriptors: vec![proto::QueryDescriptor { + attr: "group".into(), + value: "audit".into(), + }], + start: 0, + end: 100, + }), + ) + .await? + .into_inner(); + assert!(matches!( + response.outcome, + Some(proto::query_entities_response::Outcome::Conflict(_)) + )); + cluster.stop().await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn transitive_cross_shard_merge_preserves_strong_id_guard() -> anyhow::Result<()> { + let _lock = AUDIT_LOCK.lock().await; + let config = DistributedOntologyConfig { + identity_keys: ["email", "key1", "key2"] + .into_iter() + .map(|attr| IdentityKeyConfig { + name: format!("{attr}_key"), + attributes: vec![attr.into()], + }) + .collect(), + strong_identifiers: vec!["ssn".into()], + constraints: vec![], + }; + let cluster = empty_cluster(&config, 3).await?; + let mut records = Vec::new(); + for target in 0..3 { + let extras = match target { + 0 => vec![("key1", "AB"), ("ssn", "111")], + 1 => vec![("key1", "AB"), ("key2", "BC")], + _ => vec![("key2", "BC"), ("ssn", "222")], + }; + let candidate = (0..1000) + .map(|candidate| { + let mut input = record( + target, + &format!("target-{target}-{candidate}@example.com"), + ("group", "audit"), + ); + input.identity.as_mut().unwrap().perspective = + if target == 1 { "bridge" } else { "hr" }.into(); + input + .descriptors + .retain(|descriptor| descriptor.attr != "phone"); + input.descriptors.extend(extras.iter().map(|(attr, value)| { + proto::RecordDescriptor { + attr: (*attr).into(), + value: (*value).into(), + start: 0, + end: 100, + } + })); + input + }) + .find(|input| hash_record_to_shard(&config, input, 3) == target as usize) + .expect("find target shard"); + records.push(candidate); + } + { + use unirust_rs::ontology::{IdentityKey, StrongIdentifier}; + use unirust_rs::{ + Descriptor, Interval, Ontology, PersistentStore, Record, RecordId, RecordIdentity, + Unirust, + }; + let store = PersistentStore::open(cluster._directory.path().join("local-control"))?; + let mut ontology = Ontology::new(); + for attr in ["email", "key1", "key2"] { + ontology.add_identity_key(IdentityKey::from_names(vec![attr], format!("{attr}_key"))); + } + ontology.add_strong_identifier(StrongIdentifier::from_name("ssn", "ssn")); + let mut engine = + Unirust::with_store_and_tuning(ontology, store, StreamingTuning::balanced()); + let mut local_records = Vec::new(); + for input in &records { + let identity = input.identity.as_ref().expect("audit input identity"); + let mut descriptors = Vec::new(); + for descriptor in &input.descriptors { + descriptors.push(Descriptor::new( + engine.intern_attr(&descriptor.attr), + engine.intern_value(&descriptor.value), + Interval::new(descriptor.start, descriptor.end)?, + )); + } + local_records.push(Record::new( + RecordId(0), + RecordIdentity::new( + identity.entity_type.clone(), + identity.perspective.clone(), + identity.uid.clone(), + ), + descriptors, + )); + } + engine.stream_records(local_records)?; + assert_eq!( + engine.streaming_cluster_count(), + Some(2), + "local persistent engine must preserve the strong-ID guard" + ); + let descriptor = unirust_rs::query::QueryDescriptor { + attr: engine.intern_attr("group"), + value: engine.intern_value("audit"), + }; + let local_query = engine.query_master_entities(&[descriptor], Interval::new(0, 100)?)?; + assert!( + matches!(local_query, unirust_rs::query::QueryOutcome::Conflict(_)), + "local persistent query should expose distinct overlapping entities: {local_query:?}" + ); + println!("local PersistentStore control: 2 clusters, query reports Conflict"); + } + RouterService::ingest_records( + cluster.router.as_ref(), + Request::new(proto::IngestRecordsRequest { + records, + internal_protocol_version: 0, + }), + ) + .await?; + let reconciliation = RouterService::reconcile( + cluster.router.as_ref(), + Request::new(proto::ReconcileRequest { + shard_metadata: vec![], + }), + ) + .await? + .into_inner(); + let response = RouterService::query_entities( + cluster.router.as_ref(), + Request::new(proto::QueryEntitiesRequest { + descriptors: vec![proto::QueryDescriptor { + attr: "group".into(), + value: "audit".into(), + }], + start: 0, + end: 100, + }), + ) + .await? + .into_inner(); + cluster.stop().await?; + assert_eq!(reconciliation.merges_performed, 1, + "cannot merge A(hr,ssn=111)—B(bridge,no ssn)—C(hr,ssn=222) at [0,100): reconciliation={reconciliation:?}; query={response:?}"); + assert!(reconciliation.conflicts_blocked > 0); + assert!(matches!( + response.outcome, + Some(proto::query_entities_response::Outcome::Conflict(_)) + )); + Ok(()) +} diff --git a/tests/distributed_ingest_stream.rs b/tests/distributed_ingest_stream.rs index 3d3aa9d..dd7e3f7 100644 --- a/tests/distributed_ingest_stream.rs +++ b/tests/distributed_ingest_stream.rs @@ -90,7 +90,7 @@ async fn shard_stream_ingest_accepts_chunks() -> anyhow::Result<()> { }); tx.send(IngestRecordsChunk { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record_input( 0, "person", @@ -101,7 +101,7 @@ async fn shard_stream_ingest_accepts_chunks() -> anyhow::Result<()> { }) .await?; tx.send(IngestRecordsChunk { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record_input( 1, "person", diff --git a/tests/distributed_metrics.rs b/tests/distributed_metrics.rs index 894ff8c..269b011 100644 --- a/tests/distributed_metrics.rs +++ b/tests/distributed_metrics.rs @@ -110,7 +110,7 @@ async fn metrics_report_requests() -> anyhow::Result<()> { router .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record_input( 0, "person", diff --git a/tests/distributed_rebalance.rs b/tests/distributed_rebalance.rs index dc05227..c8c28f0 100644 --- a/tests/distributed_rebalance.rs +++ b/tests/distributed_rebalance.rs @@ -114,7 +114,7 @@ async fn distributed_export_import_range() -> anyhow::Result<()> { let ingest_response = shard0 .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records, }) .await? @@ -163,7 +163,7 @@ async fn distributed_export_import_range() -> anyhow::Result<()> { shard1 .import_records(ImportRecordsRequest { records: all_records, - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, }) .await?; @@ -196,7 +196,7 @@ async fn distributed_export_import_range() -> anyhow::Result<()> { let error = shard1 .import_records(ImportRecordsRequest { records: vec![collision], - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, }) .await .expect_err("import must not overwrite an existing numeric record ID"); diff --git a/tests/distributed_rebalance_stream.rs b/tests/distributed_rebalance_stream.rs index 581ad76..48bab46 100644 --- a/tests/distributed_rebalance_stream.rs +++ b/tests/distributed_rebalance_stream.rs @@ -139,7 +139,7 @@ async fn distributed_rebalance_stream_rejects_cross_shard_copy() -> anyhow::Resu assert_eq!(records.len(), 2); router .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records, }) .await?; diff --git a/tests/distributed_reliability.rs b/tests/distributed_reliability.rs index f969964..8037f26 100644 --- a/tests/distributed_reliability.rs +++ b/tests/distributed_reliability.rs @@ -308,7 +308,7 @@ async fn shard_recovery_after_restart() -> anyhow::Result<()> { .collect(); client .ingest_records(proto::IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records, }) .await?; @@ -373,7 +373,7 @@ async fn router_handles_partial_shard_availability() -> anyhow::Result<()> { ); let response = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record], }) .await? @@ -554,7 +554,7 @@ async fn destructive_reset_is_disabled_and_preserves_records_by_default() -> any client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record_input( 0, "person", @@ -634,7 +634,7 @@ async fn incomplete_ontology_update_blocks_cluster_traffic() -> anyhow::Result<( let ingest_error = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record_input( 0, "person", @@ -782,7 +782,7 @@ async fn ingest_operations_complete_before_shutdown() -> anyhow::Result<()> { let response = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records, }) .await? @@ -833,7 +833,7 @@ async fn cluster_stats_reflect_all_shards() -> anyhow::Result<()> { client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records, }) .await?; diff --git a/tests/distributed_router_admin.rs b/tests/distributed_router_admin.rs index 88dc54b..66f9eca 100644 --- a/tests/distributed_router_admin.rs +++ b/tests/distributed_router_admin.rs @@ -136,7 +136,7 @@ async fn router_admin_rejects_non_atomic_cross_shard_copy() -> anyhow::Result<() assert_eq!(records.len(), 2); router .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records, }) .await?; diff --git a/tests/durable_ingest_regressions.rs b/tests/durable_ingest_regressions.rs new file mode 100644 index 0000000..0abf101 --- /dev/null +++ b/tests/durable_ingest_regressions.rs @@ -0,0 +1,427 @@ +use tempfile::tempdir; +use unirust_rs::model::{AttrId, ClusterId, StringInterner}; +use unirust_rs::ontology::IdentityKey; +use unirust_rs::{ + Descriptor, Interval, Ontology, PersistentStore, Record, RecordId, RecordIdentity, RecordStore, + StreamingTuning, Unirust, +}; + +fn ontology() -> Ontology { + let mut ontology = Ontology::new(); + ontology.add_identity_key(IdentityKey::from_names(vec!["email"], "email")); + ontology +} + +fn record(engine: &mut Unirust, source: &str, value: &str) -> Record { + Record::new( + RecordId(0), + RecordIdentity::new("person".into(), "crm".into(), source.into()), + vec![Descriptor::new( + engine.intern_attr("email"), + engine.intern_value(value), + Interval::new(0, 10).unwrap(), + )], + ) +} + +fn ingest_variant( + engine: &mut Unirust, + records: Vec, + variant: usize, +) -> anyhow::Result<()> { + match variant { + 0 => engine.ingest(records).map(|_| ()), + 1 => engine.stream_records(records).map(|_| ()), + 2 => engine.stream_records_with_conflicts(records).map(|_| ()), + 3 => engine.stream_records_update_graph(records).map(|_| ()), + _ => unreachable!(), + } +} + +/// Commit records and one assignment, then inject an error at the next durable +/// boundary. This exercises recovery from a partially written ingest result. +struct FailAfterRecordCommit { + inner: PersistentStore, + fail_assignment: bool, +} + +impl RecordStore for FailAfterRecordCommit { + fn add_record(&mut self, record: Record) -> anyhow::Result { + self.inner.add_record(record) + } + fn stage_record_if_absent(&mut self, record: Record) -> anyhow::Result<(RecordId, bool)> { + self.inner.stage_record_if_absent(record) + } + fn flush_staged_records(&mut self) -> anyhow::Result { + self.inner.flush_staged_records() + } + fn discard_staged_records(&mut self) -> anyhow::Result<()> { + self.inner.discard_staged_records() + } + fn get_record(&self, id: RecordId) -> Option { + self.inner.get_record(id) + } + fn get_record_ref(&self, id: RecordId) -> Option<&Record> { + self.inner.get_record_ref(id) + } + fn get_record_id_by_identity(&self, identity: &RecordIdentity) -> Option { + self.inner.get_record_id_by_identity(identity) + } + fn get_all_records(&self) -> Vec { + self.inner.get_all_records() + } + fn get_records_by_entity_type(&self, entity_type: &str) -> Vec { + self.inner.get_records_by_entity_type(entity_type) + } + fn get_records_by_perspective(&self, perspective: &str) -> Vec { + self.inner.get_records_by_perspective(perspective) + } + fn get_records_with_attribute(&self, attr: AttrId) -> Vec { + self.inner.get_records_with_attribute(attr) + } + fn get_records_in_interval(&self, interval: Interval) -> Vec { + self.inner.get_records_in_interval(interval) + } + fn interner(&self) -> &StringInterner { + self.inner.interner() + } + fn interner_mut(&mut self) -> &mut StringInterner { + self.inner.interner_mut() + } + fn len(&self) -> usize { + self.inner.len() + } + fn is_empty(&self) -> bool { + self.inner.is_empty() + } + fn records_in_id_range( + &self, + start: RecordId, + end: RecordId, + max_results: usize, + ) -> Vec { + self.inner.records_in_id_range(start, end, max_results) + } + fn record_id_bounds(&self) -> Option<(RecordId, RecordId)> { + self.inner.record_id_bounds() + } + fn shared_db(&self) -> Option> { + self.inner.shared_db() + } + fn sync(&self) -> anyhow::Result<()> { + self.inner.sync() + } + fn set_cluster_count(&mut self, count: usize) -> anyhow::Result<()> { + self.inner.set_cluster_count(count) + } + fn set_cluster_assignments_batch( + &mut self, + assignments: &[(RecordId, ClusterId)], + ) -> anyhow::Result<()> { + if self.fail_assignment { + self.fail_assignment = false; + let (record_id, cluster_id) = assignments[0]; + self.inner.set_cluster_assignment(record_id, cluster_id)?; + self.inner.sync()?; + anyhow::bail!("injected error after durable record commit"); + } + self.inner.set_cluster_assignments_batch(assignments) + } +} + +#[test] +fn partial_assignment_write_recovers_all_committed_records() -> anyhow::Result<()> { + for retry_before_reopen in [false, true] { + let dir = tempdir()?; + let tuning = StreamingTuning { + use_persistent_dsu: true, + use_tiered_index: true, + ..StreamingTuning::default() + }; + let store = FailAfterRecordCommit { + inner: PersistentStore::open(dir.path())?, + fail_assignment: true, + }; + let mut engine = Unirust::with_store_and_tuning(ontology(), store, tuning.clone()); + let first = record(&mut engine, "first", "shared@example.com"); + let second = record(&mut engine, "second", "shared@example.com"); + let error = engine.ingest(vec![first.clone(), second]).unwrap_err(); + assert!(error + .to_string() + .contains("injected error after durable record commit")); + assert_eq!(engine.record_count(), 2); + if retry_before_reopen { + let result = engine.ingest(vec![first])?; + assert_eq!(result.cluster_count, 1); + assert_eq!(engine.record_count(), 2); + } + drop(engine); + let store = PersistentStore::open(dir.path())?; + let mut recovered = Unirust::with_store_and_tuning(ontology(), store, tuning); + recovered.initialize_streaming()?; + assert_eq!(recovered.record_count(), 2); + let clusters = recovered.clusters()?; + assert_eq!(clusters.len(), 1); + for stored in recovered.store().get_all_records() { + assert_eq!(clusters.get_clusters_for_record(stored.id).len(), 1); + } + } + Ok(()) +} + +#[test] +fn failed_batches_discard_pending_records_and_partial_resolution() -> anyhow::Result<()> { + for persistent_dsu in [false, true] { + for variant in 0..4 { + let dir = tempdir()?; + let tuning = StreamingTuning { + use_persistent_dsu: persistent_dsu, + use_tiered_index: persistent_dsu, + ..StreamingTuning::default() + }; + let store = PersistentStore::open(dir.path())?; + let mut engine = Unirust::with_store_and_tuning(ontology(), store, tuning.clone()); + let original = record(&mut engine, "original", "original@example.com"); + ingest_variant(&mut engine, vec![original], variant)?; + let pending = record(&mut engine, "pending", "shared@example.com"); + let pending_identity = pending.identity.clone(); + let changed = record(&mut engine, "original", "changed@example.com"); + let error = ingest_variant(&mut engine, vec![pending.clone(), changed], variant) + .expect_err("changed immutable source identity must fail"); + assert!(error.to_string().contains("different payload")); + assert!(engine.get_record_by_identity(&pending_identity)?.is_none()); + assert_eq!(engine.record_count(), 1); + + let next = record(&mut engine, "next", "shared@example.com"); + ingest_variant(&mut engine, vec![next], variant)?; + assert!(engine.get_record_by_identity(&pending_identity)?.is_none()); + ingest_variant(&mut engine, vec![pending], variant)?; + assert_eq!(engine.record_count(), 3); + let clusters = engine.clusters()?; + assert_eq!( + clusters.len(), + 2, + "retried record must resolve with the shared email" + ); + for stored in engine.store().get_all_records() { + assert_eq!(clusters.get_clusters_for_record(stored.id).len(), 1); + } + drop(engine); + + let store = PersistentStore::open(dir.path())?; + let mut recovered = Unirust::with_store_and_tuning(ontology(), store, tuning); + recovered.initialize_streaming()?; + assert_eq!(recovered.record_count(), 3); + assert_eq!(recovered.clusters()?.len(), 2); + } + } + Ok(()) +} + +#[test] +fn failed_first_batch_discards_already_flushed_dsu_nodes() -> anyhow::Result<()> { + let dir = tempdir()?; + let tuning = StreamingTuning { + use_persistent_dsu: true, + use_tiered_index: true, + dsu_config: Some(unirust_rs::advanced::PersistentDSUConfig { + dirty_buffer_size: 1, + ..Default::default() + }), + ..StreamingTuning::default() + }; + let store = PersistentStore::open(dir.path())?; + let mut engine = Unirust::with_store_and_tuning(ontology(), store, tuning); + let first = record(&mut engine, "first", "first@example.com"); + let changed = record(&mut engine, "first", "changed@example.com"); + assert!(engine + .stream_records_update_graph(vec![first, changed]) + .is_err()); + assert_eq!(engine.record_count(), 0); + let next = record(&mut engine, "next", "next@example.com"); + let result = engine.ingest(vec![next])?; + assert_eq!(result.cluster_count, 1); + assert_eq!(engine.clusters()?.len(), 1); + Ok(()) +} + +#[test] +fn duplicate_explicit_id_does_not_overwrite_durable_record() -> anyhow::Result<()> { + let dir = tempdir()?; + let store = PersistentStore::open(dir.path())?; + let mut engine = Unirust::with_store(ontology(), store); + let mut original = record(&mut engine, "original", "original@example.com"); + original.id = RecordId(42); + let original_identity = original.identity.clone(); + engine.ingest(vec![original])?; + let mut replacement = record(&mut engine, "replacement", "replacement@example.com"); + replacement.id = RecordId(42); + let replacement_identity = replacement.identity.clone(); + for variant in 0..4 { + let error = ingest_variant(&mut engine, vec![replacement.clone()], variant) + .expect_err("an occupied record ID must be rejected"); + assert!(error.to_string().contains("already exists")); + } + drop(engine); + let engine = Unirust::with_store(ontology(), PersistentStore::open(dir.path())?); + let recovered = engine.get_record_by_identity(&original_identity)?.unwrap(); + assert_eq!(recovered.identity, original_identity); + assert!(engine + .get_record_by_identity(&replacement_identity)? + .is_none()); + assert_eq!(engine.record_count(), 1); + Ok(()) +} + +#[test] +fn persistent_write_apis_reject_occupied_and_repeated_explicit_ids() -> anyhow::Result<()> { + for variant in 0..3 { + let dir = tempdir()?; + let mut store = PersistentStore::open(dir.path())?; + let original = Record::new( + RecordId(42), + RecordIdentity::new("person".into(), "crm".into(), "original".into()), + vec![], + ); + store.add_record(original.clone())?; + let replacement = Record::new( + RecordId(42), + RecordIdentity::new("person".into(), "crm".into(), "replacement".into()), + vec![], + ); + let result = match variant { + 0 => store.add_record(replacement).map(|_| ()), + 1 => store.add_records(vec![replacement]), + 2 => store.add_records_if_absent(vec![replacement]).map(|_| ()), + _ => unreachable!(), + }; + assert!(result.is_err()); + assert_eq!( + store.get_record(RecordId(42)).unwrap().identity, + original.identity + ); + assert_eq!(store.len(), 1); + + let first = Record::new( + RecordId(43), + RecordIdentity::new("person".into(), "crm".into(), "first".into()), + vec![], + ); + let second = Record::new( + RecordId(43), + RecordIdentity::new("person".into(), "crm".into(), "second".into()), + vec![], + ); + let result = match variant { + 0 | 1 => store.add_records(vec![first, second]), + 2 => store.add_records_if_absent(vec![first, second]).map(|_| ()), + _ => unreachable!(), + }; + assert!(result.is_err()); + assert!(store.get_record(RecordId(43)).is_none()); + assert_eq!(store.len(), 1); + } + Ok(()) +} + +#[test] +fn bounded_interner_caches_preserve_durable_descriptors_and_resolution() -> anyhow::Result<()> { + const CHILD: &str = "UNIRUST_DURABLE_INTERNER_TEST_CHILD"; + if std::env::var_os(CHILD).is_none() { + // Give cache configuration to an isolated child process: no test changes + // environment variables while another thread can be reading them. + let status = std::process::Command::new(std::env::current_exe()?) + .args([ + "--exact", + "bounded_interner_caches_preserve_durable_descriptors_and_resolution", + "--nocapture", + ]) + .env(CHILD, "1") + .env("UNIRUST_INTERNER_CACHE_ATTRS", "0") + .env("UNIRUST_INTERNER_CACHE_VALUES", "0") + .status()?; + assert!(status.success(), "bounded-cache child process failed"); + return Ok(()); + } + let dir = tempdir()?; + let descriptors; + { + let store = PersistentStore::open(dir.path())?; + let mut engine = Unirust::with_store(ontology(), store); + let original = record(&mut engine, "original", "known@example.com"); + descriptors = original.descriptors.clone(); + engine.ingest(vec![original])?; + } + let store = PersistentStore::open(dir.path())?; + assert!(store.interner().get_attr_id("email").is_none()); + assert!(store.interner().get_value_id("known@example.com").is_none()); + assert_eq!(store.lookup_attr("email"), Some(descriptors[0].attr)); + assert_eq!( + store.lookup_value("known@example.com"), + Some(descriptors[0].value) + ); + assert!(store.interner().get_attr_id("email").is_none()); + assert!(store.interner().get_value_id("known@example.com").is_none()); + let mut engine = Unirust::with_store(ontology(), store); + // The record deliberately reuses durable IDs without going through string + // interning, exercising record preparation as well as the lookup APIs. + let reused = Record::new( + RecordId(0), + RecordIdentity::new("person".into(), "crm".into(), "reused".into()), + descriptors, + ); + let assignments = engine.ingest(vec![reused])?.assignments; + let named = record(&mut engine, "named", "known@example.com"); + engine.ingest(vec![named])?; + let persisted = engine.get_record(assignments[0].record_id).unwrap(); + assert_eq!( + engine + .resolve_attr(persisted.descriptors[0].attr) + .as_deref(), + Some("email") + ); + assert_eq!( + engine + .resolve_value(persisted.descriptors[0].value) + .as_deref(), + Some("known@example.com") + ); + assert_eq!(engine.clusters()?.len(), 1); + drop(engine); + let store = PersistentStore::open(dir.path())?; + let mut recovered = Unirust::with_store(ontology(), store); + recovered.initialize_streaming()?; + assert_eq!(recovered.record_count(), 3); + assert_eq!(recovered.clusters()?.len(), 1); + Ok(()) +} + +#[test] +fn ingest_batch_larger_than_record_cache_resolves_every_record() -> anyhow::Result<()> { + const COUNT: usize = 100_001; + let dir = tempdir()?; + let store = PersistentStore::open(dir.path())?; + let mut engine = Unirust::with_store(Ontology::new(), store); + let records = (0..COUNT) + .map(|index| { + Record::new( + RecordId(0), + RecordIdentity::new("person".into(), "crm".into(), format!("record-{index}")), + vec![], + ) + }) + .collect::>(); + let first = records[0].clone(); + let assignments = engine.stream_records(records)?; + assert_eq!(assignments.len(), COUNT); + assert_eq!(engine.record_count(), COUNT); + assert_eq!(engine.streaming_cluster_count(), Some(COUNT)); + let retry = engine.stream_records(vec![first])?; + assert_eq!(retry[0].record_id, assignments[0].record_id); + assert_eq!(retry[0].cluster_id, assignments[0].cluster_id); + drop(engine); + let store = PersistentStore::open(dir.path())?; + assert_eq!(store.len(), COUNT); + assert!(store.get_record(assignments[0].record_id).is_some()); + Ok(()) +} diff --git a/tests/external_backup_restore.rs b/tests/external_backup_restore.rs index ee89d92..cf9a139 100644 --- a/tests/external_backup_restore.rs +++ b/tests/external_backup_restore.rs @@ -216,13 +216,13 @@ async fn verified_export_survives_loss_of_checkpoint_volumes() -> anyhow::Result client0 .ingest_records(IngestRecordsRequest { records: vec![shard0_record], - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, }) .await?; client1 .ingest_records(IngestRecordsRequest { records: vec![shard1_record], - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, }) .await?; @@ -349,7 +349,7 @@ async fn external_checkpoint_restores_a_lost_shard_volume() -> anyhow::Result<() client .ingest_records(IngestRecordsRequest { records: vec![record(0, "backup@example.com")], - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, }) .await?; let before = client.get_stats(StatsRequest {}).await?.into_inner(); @@ -391,7 +391,7 @@ async fn external_checkpoint_restores_a_lost_shard_volume() -> anyhow::Result<() client .ingest_records(IngestRecordsRequest { records: vec![record(1, "backup@example.com")], - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, }) .await?; let retry_stats = client.get_stats(StatsRequest {}).await?.into_inner(); @@ -400,7 +400,7 @@ async fn external_checkpoint_restores_a_lost_shard_volume() -> anyhow::Result<() let error = client .ingest_records(IngestRecordsRequest { records: vec![record(2, "changed@example.com")], - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, }) .await .expect_err("restored immutable source identity must reject a changed payload"); diff --git a/tests/process_crash_recovery.rs b/tests/process_crash_recovery.rs index 6273517..249a8ab 100644 --- a/tests/process_crash_recovery.rs +++ b/tests/process_crash_recovery.rs @@ -153,7 +153,7 @@ async fn acknowledged_ingest_survives_process_kill_and_restart() -> anyhow::Resu .collect(); let response = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records, }) .await? @@ -201,7 +201,7 @@ async fn acknowledged_ingest_survives_process_kill_and_restart() -> anyhow::Resu let linked_after_restart = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record( 128, "process-kill-after-restart".to_string(), @@ -235,7 +235,7 @@ async fn acknowledged_ingest_survives_process_kill_and_restart() -> anyhow::Resu let linked_after_shutdown = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record( 129, "graceful-shutdown-after-restart".to_string(), diff --git a/tests/resolution_capacity_regressions.rs b/tests/resolution_capacity_regressions.rs new file mode 100644 index 0000000..65f4a75 --- /dev/null +++ b/tests/resolution_capacity_regressions.rs @@ -0,0 +1,365 @@ +use tempfile::tempdir; +use unirust_rs::dsu::DsuBackend; +use unirust_rs::index::{TierConfig, TieredIdentityKeyIndex}; +use unirust_rs::linker::StreamingLinker; +use unirust_rs::model::KeyValue; +use unirust_rs::ontology::IdentityKey; +use unirust_rs::store::RecordStore; +use unirust_rs::{ + Descriptor, Interval, Ontology, PersistentStore, Record, RecordId, RecordIdentity, + StreamingTuning, Unirust, +}; + +#[test] +fn repeated_single_entity_never_becomes_hot_or_stops_linking() -> anyhow::Result<()> { + let dir = tempdir()?; + let store = PersistentStore::open(dir.path())?; + let mut ontology = Ontology::new(); + ontology.add_identity_key(IdentityKey::from_names(vec!["email"], "email")); + let mut engine = Unirust::with_store_and_tuning(ontology, store, StreamingTuning::balanced()); + let attr = engine.intern_attr("email"); + let value = engine.intern_value("same@example.com"); + let mut latest = None; + // A batch exercises production parallel linking. Each record has the exact same + // key and interval, and there are no conflicting strong IDs. + let records = (0..300) + .map(|i| { + Record::new( + RecordId(0), + RecordIdentity::new("person".into(), "crm".into(), format!("r{i}")), + vec![Descriptor::new(attr, value, Interval::new(0, 100).unwrap())], + ) + }) + .collect(); + engine.stream_records(records)?; + eprintln!( + "after first batch: clusters={:?}, metrics={:?}", + engine.streaming_cluster_count(), + engine.linker_metrics_snapshot() + ); + for i in 300..302 { + latest = Some(engine.stream_records(vec![Record::new( + RecordId(0), + RecordIdentity::new("person".into(), "crm".into(), format!("r{i}")), + vec![Descriptor::new(attr, value, Interval::new(0, 100)?)], + )])?); + } + eprintln!( + "after later batches: clusters={:?}, metrics={:?}, latest={latest:?}", + engine.streaming_cluster_count(), + engine.linker_metrics_snapshot() + ); + assert_eq!(engine.streaming_cluster_count(), Some(1)); + assert_eq!(engine.linker_metrics_snapshot().hot_key_exits, 0); + Ok(()) +} + +#[test] +fn memory_saver_merges_two_matching_perspectives() -> anyhow::Result<()> { + let dir = tempdir()?; + let store = PersistentStore::open(dir.path())?; + let mut ontology = Ontology::new(); + ontology.add_identity_key(IdentityKey::from_names(vec!["email"], "email")); + let mut engine = + Unirust::with_store_and_tuning(ontology, store, StreamingTuning::memory_saver()); + let attr = engine.intern_attr("email"); + let value = engine.intern_value("same@example.com"); + for perspective in ["crm", "billing"] { + engine.stream_records(vec![Record::new( + RecordId(0), + RecordIdentity::new("person".into(), perspective.into(), "first".into()), + vec![Descriptor::new(attr, value, Interval::new(0, 100)?)], + )])?; + } + eprintln!( + "memory_saver clusters={:?}, metrics={:?}", + engine.streaming_cluster_count(), + engine.linker_metrics_snapshot() + ); + assert_eq!(engine.streaming_cluster_count(), Some(1)); + Ok(()) +} + +#[test] +fn cached_key_insertion_enforces_tier_capacity() -> anyhow::Result<()> { + let dir = tempdir()?; + let mut store = PersistentStore::open(dir.path())?; + let attr = store.intern_attr("email"); + let key = IdentityKey::new(vec![attr], "email".into()); + let config = TierConfig { + hot_tier_capacity: 2, + warm_tier_capacity: 20, + hot_threshold: 1.1, + tier_management_interval_secs: 0, + ..TierConfig::default() + }; + let mut cached_index = + TieredIdentityKeyIndex::with_config(config.clone(), Some(store.db_shared())); + let mut direct_index = TieredIdentityKeyIndex::with_config(config, Some(store.db_shared())); + let mut ontology = Ontology::new(); + ontology.add_identity_key(key.clone()); + for i in 0..10 { + let value = store.intern_value(&format!("person{i}@example.com")); + let interval = Interval::new(0, 100)?; + let id = store.add_record(Record::new( + RecordId(0), + RecordIdentity::new("person".into(), "crm".into(), format!("r{i}")), + vec![Descriptor::new(attr, value, interval)], + ))?; + cached_index.add_record_with_cached_keys( + id, + id, + "person", + vec![(&key, vec![(vec![KeyValue::new(attr, value)], interval)])], + )?; + direct_index.add_record(&store.get_record(id).unwrap(), &ontology)?; + } + eprintln!( + "cached insertion: {:?}, ordinary insertion: {:?}", + cached_index.tier_stats(), + direct_index.tier_stats() + ); + assert_eq!(direct_index.tier_stats().hot_keys, 2); + assert_eq!(cached_index.tier_stats().hot_keys, 2); + Ok(()) +} + +#[test] +fn tier_eviction_and_repeated_updates_preserve_temporal_candidates_after_restart( +) -> anyhow::Result<()> { + let dir = tempdir()?; + let mut store = PersistentStore::open(dir.path())?; + let attr = store.intern_attr("email"); + let key = IdentityKey::new(vec![attr], "email".into()); + let config = TierConfig { + hot_tier_capacity: 1, + warm_tier_capacity: 1, + tier_management_interval_secs: 0, + ..TierConfig::default() + }; + let mut index = TieredIdentityKeyIndex::with_config(config.clone(), Some(store.db_shared())); + let mut expected = Vec::new(); + // Revisit old keys after both hot and warm eviction. Each value has disjoint + // temporal observations, so promotion must preserve the entire old bucket. + for round in 0..3 { + for n in 0..12 { + let value = store.intern_value(&format!("person{n}@example.com")); + let interval = Interval::new(round * 20, round * 20 + 10)?; + let id = store.add_record(Record::new( + RecordId(0), + RecordIdentity::new("person".into(), "crm".into(), format!("{round}-{n}")), + vec![Descriptor::new(attr, value, interval)], + ))?; + index.add_record_with_cached_keys( + id, + id, + "person", + vec![(&key, vec![(vec![KeyValue::new(attr, value)], interval)])], + )?; + expected.push((value, interval, id)); + assert!(index.tier_stats().hot_keys <= 1); + assert!(index.tier_stats().warm_keys <= 1); + } + } + index.flush_warm_to_cold()?; + drop(index); + drop(store); + let store = PersistentStore::open(dir.path())?; + let mut index = TieredIdentityKeyIndex::with_config(config, Some(store.db_shared())); + let mut dsu = DsuBackend::in_memory(); + for (value, interval, id) in expected { + let candidates = index.find_matching_clusters_overlapping( + &mut dsu, + "person", + &[KeyValue::new(attr, value)], + interval, + )?; + assert_eq!(candidates, &[(id, interval)]); + let records = index.find_matching_records("person", &[KeyValue::new(attr, value)])?; + assert_eq!(records.len(), 3); + } + Ok(()) +} + +#[test] +fn long_multi_source_history_retains_membership_and_guards_through_tiers_and_recovery( +) -> anyhow::Result<()> { + let dir = tempdir()?; + let store = PersistentStore::open(dir.path())?; + let tuning = StreamingTuning { + use_tiered_index: true, + tier_config: Some(TierConfig { + hot_tier_capacity: 2, + warm_tier_capacity: 2, + tier_management_interval_secs: 0, + ..TierConfig::default() + }), + ..StreamingTuning::billion_scale() + }; + let mut ontology = Ontology::new(); + ontology.add_identity_key(IdentityKey::from_names(vec!["email"], "email")); + ontology.add_strong_identifier(unirust_rs::ontology::StrongIdentifier::from_name( + "ssn", "ssn", + )); + let mut engine = Unirust::with_store_and_tuning(ontology.clone(), store, tuning.clone()); + let email = engine.intern_attr("email"); + let ssn = engine.intern_attr("ssn"); + let id_a = engine.intern_value("111"); + let id_b = engine.intern_value("222"); + for batch in 0..6 { + let mut records = Vec::new(); + for offset in 0..500 { + let n = batch * 500 + offset; + let value = engine.intern_value(&format!("person{}@example.com", n % 10)); + records.push(Record::new( + RecordId(0), + RecordIdentity::new( + "person".into(), + if n % 20 < 10 { "crm" } else { "billing" }.into(), + format!("r{n}"), + ), + vec![ + Descriptor::new(email, value, Interval::new(0, 100)?), + Descriptor::new(ssn, id_a, Interval::new(0, 100)?), + ], + )); + } + engine.stream_records(records)?; + assert_eq!(engine.streaming_cluster_count(), Some(10)); + } + let value = engine.intern_value("person0@example.com"); + engine.stream_records(vec![Record::new( + RecordId(0), + RecordIdentity::new("person".into(), "crm".into(), "conflicting".into()), + vec![ + Descriptor::new(email, value, Interval::new(0, 100)?), + Descriptor::new(ssn, id_b, Interval::new(0, 100)?), + ], + )])?; + assert_eq!(engine.streaming_cluster_count(), Some(11)); + drop(engine); + let store = PersistentStore::open(dir.path())?; + let mut engine = Unirust::with_store_and_tuning(ontology, store, tuning); + let email = engine.intern_attr("email"); + let ssn = engine.intern_attr("ssn"); + let value = engine.intern_value("person0@example.com"); + let id_a = engine.intern_value("111"); + engine.stream_records(vec![Record::new( + RecordId(0), + RecordIdentity::new("person".into(), "crm".into(), "after-restart".into()), + vec![ + Descriptor::new(email, value, Interval::new(0, 100)?), + Descriptor::new(ssn, id_a, Interval::new(0, 100)?), + ], + )])?; + assert_eq!(engine.streaming_cluster_count(), Some(11)); + let clusters = engine.clusters()?; + assert_eq!(clusters.clusters.len(), 11); + assert_eq!( + clusters + .clusters + .iter() + .map(|cluster| cluster.records.len()) + .sum::(), + 3002 + ); + Ok(()) +} + +#[test] +fn deferred_merge_updates_readonly_membership_and_global_ids() -> anyhow::Result<()> { + let dir = tempdir()?; + let mut store = PersistentStore::open(dir.path())?; + let attr = store.intern_attr("email"); + let value = store.intern_value("same@example.com"); + let mut ontology = Ontology::new(); + ontology.add_identity_key(IdentityKey::new(vec![attr], "email".into())); + let tuning = StreamingTuning { + candidate_cap: 1, + adaptive_candidate_cap: false, + stochastic_sampling: false, + enable_boundary_tracking: true, + ..StreamingTuning::balanced() + }; + let mut linker = StreamingLinker::new(&store, &ontology, &tuning)?; + let mut ids = Vec::new(); + for (i, (start, end)) in [(0, 10), (20, 30), (0, 30)].into_iter().enumerate() { + let id = store.add_record(Record::new( + RecordId(0), + RecordIdentity::new("person".into(), "crm".into(), format!("r{i}")), + vec![Descriptor::new(attr, value, Interval::new(start, end)?)], + ))?; + linker.link_record(&store, &ontology, id)?; + linker.global_cluster_id_for(id); + ids.push(id); + } + assert_eq!(linker.cluster_count(), 3); + linker.reconcile_pending(&store, &ontology)?; + assert_eq!(linker.cluster_count(), 1); + for id in &ids { + let cluster = linker.cluster_for_record(*id).unwrap(); + let mut actual = cluster.records; + actual.sort_by_key(|id| id.0); + assert_eq!(actual, ids); + assert_eq!(Some(cluster.id), Some(linker.cluster_id_for(*id))); + assert_eq!( + linker.global_cluster_id_for_readonly(*id), + linker.global_cluster_id_for_readonly(ids[0]) + ); + } + Ok(()) +} + +#[test] +fn corrupt_cold_bucket_returns_error_instead_of_missing_candidates() -> anyhow::Result<()> { + let dir = tempdir()?; + let mut store = PersistentStore::open(dir.path())?; + let attr = store.intern_attr("email"); + let value = store.intern_value("first@example.com"); + let key = IdentityKey::new(vec![attr], "email".into()); + let config = TierConfig { + hot_tier_capacity: 0, + warm_tier_capacity: 1, + ..TierConfig::default() + }; + let mut index = TieredIdentityKeyIndex::with_config(config, Some(store.db_shared())); + for (i, value) in [value, store.intern_value("second@example.com")] + .into_iter() + .enumerate() + { + let interval = Interval::new(0, 100)?; + let id = store.add_record(Record::new( + RecordId(0), + RecordIdentity::new("person".into(), "crm".into(), format!("r{i}")), + vec![Descriptor::new(attr, value, interval)], + ))?; + index.add_record_with_cached_keys( + id, + id, + "person", + vec![(&key, vec![(vec![KeyValue::new(attr, value)], interval)])], + )?; + } + let encoded = unirust_rs::persistence::index_encoding::encode_identity_key( + "person", + &[KeyValue::new(attr, value)], + ); + let cf = store + .db() + .cf_handle(unirust_rs::persistence::index_cf::IDENTITY_KEYS) + .unwrap(); + store.db().put_cf(cf, encoded, [0xff])?; + let mut dsu = DsuBackend::in_memory(); + assert!(index + .find_matching_clusters_overlapping( + &mut dsu, + "person", + &[KeyValue::new(attr, value)], + Interval::new(0, 100)? + ) + .is_err()); + assert!(index + .find_matching_records("person", &[KeyValue::new(attr, value)]) + .is_err()); + Ok(()) +} diff --git a/tests/selective_query_regressions.rs b/tests/selective_query_regressions.rs new file mode 100644 index 0000000..32a2801 --- /dev/null +++ b/tests/selective_query_regressions.rs @@ -0,0 +1,318 @@ +use tempfile::tempdir; +use unirust_rs::advanced::GlobalClusterId; +use unirust_rs::ontology::IdentityKey; +use unirust_rs::persistence::{dsu_cf, index_cf, linker_cf, linker_encoding}; +use unirust_rs::{ + Descriptor, Interval, Ontology, PersistentStore, QueryOutcome, Record, RecordId, + RecordIdentity, RecordStore, StreamingTuning, Unirust, +}; + +fn ontology() -> Ontology { + let mut ontology = Ontology::new(); + ontology.add_identity_key(IdentityKey::from_names(vec!["email"], "email")); + ontology.add_identity_key(IdentityKey::from_names(vec!["phone"], "phone")); + ontology +} + +fn record(engine: &mut Unirust, uid: &str, values: &[(&str, &str)], interval: Interval) -> Record { + Record::new( + RecordId(0), + RecordIdentity::new("person".into(), "crm".into(), uid.into()), + values + .iter() + .map(|(attr, value)| { + Descriptor::new( + engine.intern_attr(attr), + engine.intern_value(value), + interval, + ) + }) + .collect(), + ) +} + +#[test] +fn selective_queries_follow_bridge_merges_and_match_recovery() -> anyhow::Result<()> { + for persistent_dsu in [false, true] { + let dir = tempdir()?; + let tuning = StreamingTuning { + use_persistent_dsu: persistent_dsu, + use_tiered_index: persistent_dsu, + ..StreamingTuning::default() + }; + let interval = Interval::new(0, 100)?; + let mut engine = Unirust::with_store_and_tuning( + ontology(), + PersistentStore::open(dir.path())?, + tuning.clone(), + ); + let records = vec![ + record( + &mut engine, + "email", + &[("email", "alice@example.com"), ("name", "Alice")], + interval, + ), + record( + &mut engine, + "phone", + &[("phone", "123"), ("city", "London")], + interval, + ), + ]; + engine.stream_records(records)?; + let query = [ + engine.lookup_query_descriptor("name", "Alice").unwrap(), + engine.lookup_query_descriptor("city", "London").unwrap(), + ]; + assert_eq!( + engine.query(&query, interval)?, + QueryOutcome::Matches(Vec::new()) + ); + let bridge = record( + &mut engine, + "bridge", + &[("email", "alice@example.com"), ("phone", "123")], + interval, + ); + engine.stream_records(vec![bridge])?; + let result = engine.query(&query, interval)?; + let QueryOutcome::Matches(matches) = &result else { + panic!("unexpected conflict: {result:?}") + }; + assert_eq!(matches.len(), 1); + assert!(matches[0] + .golden + .iter() + .any(|d| d.attr == "city" && d.value == "London")); + let unrelated = record( + &mut engine, + "unrelated", + &[("email", "bob@example.com")], + interval, + ); + engine.stream_records(vec![unrelated])?; + assert_eq!(engine.query(&query, interval)?, result); + drop(engine); + + let mut recovered = + Unirust::with_store_and_tuning(ontology(), PersistentStore::open(dir.path())?, tuning); + assert_eq!(recovered.query(&query, interval)?, result); + recovered.initialize_streaming()?; + assert_eq!(recovered.query(&query, interval)?, result); + } + Ok(()) +} + +#[test] +fn selective_query_labels_agree_with_full_store_labels() -> anyhow::Result<()> { + let dir = tempdir()?; + let interval = Interval::new(0, 100)?; + let mut config = ontology(); + config.add_identity_key(IdentityKey::from_names(vec!["name", "email"], "name_email")); + let mut engine = Unirust::with_store(config.clone(), PersistentStore::open(dir.path())?); + // These different identity values generate the same short human-readable token. + let records = vec![ + record( + &mut engine, + "one", + &[("name", "Alice"), ("email", "one@example.com")], + interval, + ), + record( + &mut engine, + "two", + &[("name", "Alice"), ("email", "two@example.com")], + interval, + ), + ]; + engine.stream_records(records)?; + let query = [engine + .lookup_query_descriptor("email", "one@example.com") + .unwrap()]; + let result = engine.query(&query, interval)?; + drop(engine); + let recovered = Unirust::with_store(config, PersistentStore::open(dir.path())?); + assert_eq!(recovered.query(&query, interval)?, result); + Ok(()) +} + +#[test] +fn restored_cluster_ids_invalidate_cached_composite_labels() -> anyhow::Result<()> { + for persistent_dsu in [false, true] { + let dir = tempdir()?; + let interval = Interval::new(0, 100)?; + let mut config = Ontology::new(); + config.add_identity_key(IdentityKey::from_names(vec!["name", "email"], "name_email")); + let tuning = StreamingTuning { + use_persistent_dsu: persistent_dsu, + use_tiered_index: persistent_dsu, + ..StreamingTuning::default() + }; + let mut engine = Unirust::with_store_and_tuning( + config.clone(), + PersistentStore::open(dir.path())?, + tuning.clone(), + ); + let mut first = record( + &mut engine, + "first", + &[("name", "Alice"), ("email", "alice@example.com")], + interval, + ); + first.id = RecordId(100); + let mut second = record( + &mut engine, + "second", + &[("name", "Bob"), ("email", "bob@example.com")], + interval, + ); + second.id = RecordId(1); + engine.stream_records(vec![first, second])?; + let query = [engine + .lookup_query_descriptor("email", "alice@example.com") + .unwrap()]; + let expected = engine.query(&query, interval)?; + engine.checkpoint_linker_state()?; + drop(engine); + + let mut recovered = + Unirust::with_store_and_tuning(config, PersistentStore::open(dir.path())?, tuning); + recovered.initialize_streaming()?; + // Replay visits ID 1 before ID 100, reversing allocation of local cluster + // IDs. Warm the composite-label cache under those temporary assignments. + let before_restore = recovered.query(&query, interval)?; + let (QueryOutcome::Matches(before), QueryOutcome::Matches(original)) = + (&before_restore, &expected) + else { + panic!("expected one matching entity before and after replay"); + }; + assert_eq!(before.len(), 1); + assert_eq!(original.len(), 1); + assert_ne!(before[0].cluster_id, original[0].cluster_id); + assert_eq!(before[0].cluster_key, original[0].cluster_key); + assert_eq!(recovered.restore_linker_state()?, 2); + assert_eq!(recovered.query(&query, interval)?, expected); + } + Ok(()) +} + +#[test] +fn fragment_recovery_preserves_durable_derived_state_and_redirects() -> anyhow::Result<()> { + let dir = tempdir()?; + let interval = Interval::new(0, 100)?; + let tuning = StreamingTuning { + shard_id: 1, + use_persistent_dsu: true, + use_tiered_index: true, + ..StreamingTuning::default() + }; + let canonical = GlobalClusterId::new(0, 20, 0); + { + let mut engine = Unirust::with_store_and_tuning( + ontology(), + PersistentStore::open(dir.path())?, + tuning.clone(), + ); + let source = record( + &mut engine, + "alice", + &[("email", "alice@example.com")], + interval, + ); + engine.ingest(vec![source])?; + engine.apply_cross_shard_merge(canonical, GlobalClusterId::new(1, 0, 0))?; + } + let store = PersistentStore::open(dir.path())?; + let db = store.shared_db().unwrap(); + let parent_cf = db.cf_handle(dsu_cf::PARENT).unwrap(); + let index_cf = db.cf_handle(index_cf::IDENTITY_KEYS).unwrap(); + let sentinel_id = 777u32.to_be_bytes(); + db.put_cf(parent_cf, sentinel_id, sentinel_id)?; + db.put_cf(index_cf, b"query-recovery-sentinel", b"preserve")?; + let sequence = db.latest_sequence_number(); + let engine = Unirust::with_store_and_tuning(ontology(), store, tuning); + let query = [engine + .lookup_query_descriptor("email", "alice@example.com") + .unwrap()]; + let candidates = engine.query_entity_fragments(&query, interval, &[], false)?; + assert_eq!(candidates.len(), 1); + assert_eq!(candidates[0].global_id, canonical); + let hydrated = engine.query_entity_fragments(&query, interval, &[canonical], false)?; + assert_eq!(hydrated.len(), 1); + assert_eq!(hydrated[0].global_id, canonical); + assert!(hydrated[0] + .golden + .iter() + .any(|descriptor| descriptor.value == "alice@example.com")); + assert_eq!(db.get_cf(parent_cf, sentinel_id)?.unwrap(), sentinel_id); + assert_eq!( + db.get_cf(index_cf, b"query-recovery-sentinel")?.unwrap(), + b"preserve" + ); + assert_eq!( + db.latest_sequence_number(), + sequence, + "fragment queries must not write RocksDB" + ); + Ok(()) +} + +#[test] +fn fragment_queries_allow_fresh_and_reset_stores_without_writes() -> anyhow::Result<()> { + let dir = tempdir()?; + let store = PersistentStore::open(dir.path())?; + let db = store.shared_db().unwrap(); + let mut engine = Unirust::with_store(ontology(), store); + let interval = Interval::new(0, 100)?; + let sequence = db.latest_sequence_number(); + assert!(engine + .query_entity_fragments(&[], interval, &[], false)? + .is_empty()); + assert_eq!(db.latest_sequence_number(), sequence); + + engine.initialize_streaming()?; + engine.reset_with_ontology(ontology())?; + let sequence = db.latest_sequence_number(); + assert!(engine + .query_entity_fragments(&[], interval, &[], false)? + .is_empty()); + assert_eq!(db.latest_sequence_number(), sequence); + Ok(()) +} + +#[test] +fn fragment_queries_reject_legacy_or_invalid_scheme_without_migration() -> anyhow::Result<()> { + for corrupt_marker in [false, true] { + let dir = tempdir()?; + let store = PersistentStore::open(dir.path())?; + let db = store.shared_db().unwrap(); + let metadata = db.cf_handle(linker_cf::METADATA).unwrap(); + let marker = linker_encoding::KEY_GLOBAL_CLUSTER_ID_SCHEME; + if corrupt_marker { + db.put_cf(metadata, marker, b"invalid")?; + } + let legacy_key = + linker_encoding::encode_cross_shard_merge_key(GlobalClusterId::new(1, 3, 0)); + let legacy_value = linker_encoding::encode_global_cluster_id(GlobalClusterId::new(0, 2, 0)); + db.put_cf(metadata, &legacy_key, legacy_value)?; + let sequence = db.latest_sequence_number(); + let engine = Unirust::with_store(ontology(), store); + let error = engine + .query_entity_fragments(&[], Interval::new(0, 100)?, &[], false) + .unwrap_err(); + let expected = if corrupt_marker { + "unsupported global cluster ID scheme" + } else { + "initialize_streaming is required" + }; + assert!(error.to_string().contains(expected), "{error}"); + assert_eq!(db.get_cf(metadata, legacy_key)?.unwrap(), legacy_value); + assert_eq!( + db.get_cf(metadata, marker)?, + corrupt_marker.then(|| b"invalid".to_vec()) + ); + assert_eq!(db.latest_sequence_number(), sequence); + } + Ok(()) +} diff --git a/tests/temporal_evolution.rs b/tests/temporal_evolution.rs index b1399c3..db2174b 100644 --- a/tests/temporal_evolution.rs +++ b/tests/temporal_evolution.rs @@ -229,19 +229,19 @@ async fn incremental_ingestion_equals_batch_ingestion() -> anyhow::Result<()> { // Ingest incrementally client_inc .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record_c1.clone()], }) .await?; client_inc .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record_c2.clone()], }) .await?; client_inc .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record_c3.clone()], }) .await?; @@ -249,7 +249,7 @@ async fn incremental_ingestion_equals_batch_ingestion() -> anyhow::Result<()> { // Batch: Ingest all 3 at once client_batch .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record_c1, record_c2, record_c3], }) .await?; @@ -370,7 +370,8 @@ async fn ingestion_order_independence() -> anyhow::Result<()> { for &idx in order { client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: + unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![records[idx].clone()], }) .await?; @@ -447,19 +448,19 @@ async fn temporal_descriptor_evolution() -> anyhow::Result<()> { // Ingest over time client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![jan], }) .await?; client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![feb], }) .await?; client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![mar], }) .await?; @@ -572,7 +573,7 @@ async fn attribute_value_changes_over_time() -> anyhow::Result<()> { client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![initial, with_secondary], }) .await?; @@ -652,7 +653,7 @@ async fn multi_perspective_incremental_merge() -> anyhow::Result<()> { ); client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![crm_record], }) .await?; @@ -683,7 +684,7 @@ async fn multi_perspective_incremental_merge() -> anyhow::Result<()> { ); client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![erp_record], }) .await?; @@ -720,7 +721,7 @@ async fn multi_perspective_incremental_merge() -> anyhow::Result<()> { ); client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![hr_record], }) .await?; @@ -768,7 +769,7 @@ async fn idempotent_re_ingestion() -> anyhow::Result<()> { // First ingestion let resp1 = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record.clone()], }) .await? @@ -778,7 +779,7 @@ async fn idempotent_re_ingestion() -> anyhow::Result<()> { // Second ingestion (identical) let resp2 = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record.clone()], }) .await? @@ -794,7 +795,7 @@ async fn idempotent_re_ingestion() -> anyhow::Result<()> { // Third ingestion let resp3 = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record], }) .await? @@ -847,13 +848,13 @@ async fn overlapping_temporal_ranges_extend_validity() -> anyhow::Result<()> { client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![jan_mar], }) .await?; client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![feb_apr], }) .await?; @@ -932,7 +933,7 @@ async fn temporal_gap_between_records() -> anyhow::Result<()> { client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![q1, q3], }) .await?; @@ -1084,7 +1085,8 @@ async fn same_cluster_regardless_of_insertion_order() -> anyhow::Result<()> { for &idx in order { let resp = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: + unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![records[idx].clone()], }) .await? @@ -1159,7 +1161,7 @@ async fn batch_ingestion_consistent_cluster() -> anyhow::Result<()> { let batch_resp = batch_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: records.clone(), }) .await? @@ -1186,7 +1188,7 @@ async fn batch_ingestion_consistent_cluster() -> anyhow::Result<()> { for record in &records { let resp = seq_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![record.clone()], }) .await? @@ -1271,7 +1273,8 @@ async fn independent_entities_stay_separate_any_order() -> anyhow::Result<()> { for &idx in order { client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: + unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![records[idx].clone()], }) .await?; @@ -1358,7 +1361,7 @@ async fn late_arrival_merges_into_existing_cluster() -> anyhow::Result<()> { let resp1 = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![rec1, rec2], }) .await? @@ -1385,7 +1388,7 @@ async fn late_arrival_merges_into_existing_cluster() -> anyhow::Result<()> { let resp2 = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 5, + internal_protocol_version: unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION, records: vec![late_rec], }) .await?