diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b01358d..68fd406 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -62,3 +62,17 @@ jobs: run: | bash -n scripts/*.sh docker compose -f compose.yaml config --quiet + + container-image: + runs-on: ubuntu-latest + steps: + - name: Free disk space + run: | + sudo rm -rf /usr/share/dotnet + sudo rm -rf /usr/local/lib/android + sudo rm -rf /opt/ghc + sudo rm -rf /opt/hostedtoolcache/CodeQL + - name: Checkout + uses: actions/checkout@v4 + - name: Build production image + run: docker build --file Containerfile --tag unirust-ci . diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 7caa983..a045a51 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -26,6 +26,8 @@ jobs: run: cargo clippy --locked --all-targets --all-features -- -D warnings - 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/Cargo.toml b/Cargo.toml index 6831ade..4e60773 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,8 +7,8 @@ description = "High-performance temporal entity resolution engine with conflict authors = ["unirust contributors"] license = "MIT" readme = "README.md" -repository = "https://github.com/unirust/unirust" -documentation = "https://docs.rs/unirust" +repository = "https://github.com/script3r/unirust" +documentation = "https://docs.rs/unirust-rs" keywords = ["entity-resolution", "temporal", "data-deduplication", "conflict-resolution"] categories = ["data-structures", "algorithms"] diff --git a/Containerfile b/Containerfile index 1ddd933..17944fc 100644 --- a/Containerfile +++ b/Containerfile @@ -4,7 +4,7 @@ # # Usage: # podman build -t unirust -f Containerfile . -# podman run --rm -p 50061:50061 -v unirust-data:/data unirust shard +# podman run --rm -p 50061:50061 -v unirust-data:/data -v unirust-backup:/backup unirust shard # podman run --rm -p 50060:50060 unirust router --shards shard-0:50061 # # Or use with compose.yaml for full cluster deployment @@ -17,6 +17,7 @@ WORKDIR /app RUN apt-get update && apt-get install -y --no-install-recommends \ build-essential \ cmake \ + libclang-dev \ protobuf-compiler \ && rm -rf /var/lib/apt/lists/* @@ -31,6 +32,8 @@ RUN cargo build --release --locked --features test-support \ --bin unirust_shard \ --bin unirust_router \ --bin unirust_healthcheck \ + --bin unirust_backup \ + --bin unirust_rebalance \ --bin unirust_client \ --bin unirust_loadtest @@ -49,6 +52,8 @@ RUN useradd -m -u 1000 unirust COPY --from=builder /app/target/release/unirust_shard /usr/local/bin/ COPY --from=builder /app/target/release/unirust_router /usr/local/bin/ COPY --from=builder /app/target/release/unirust_healthcheck /usr/local/bin/ +COPY --from=builder /app/target/release/unirust_backup /usr/local/bin/ +COPY --from=builder /app/target/release/unirust_rebalance /usr/local/bin/ COPY --from=builder /app/target/release/unirust_client /usr/local/bin/ COPY --from=builder /app/target/release/unirust_loadtest /usr/local/bin/ @@ -82,6 +87,14 @@ case "$1" in shift exec unirust_loadtest "$@" ;; + backup) + shift + exec unirust_backup "$@" + ;; + rebalance) + shift + exec unirust_rebalance "$@" + ;; *) exec "$@" ;; @@ -91,4 +104,4 @@ EOF RUN chmod +x /usr/local/bin/entrypoint.sh ENTRYPOINT ["/usr/local/bin/entrypoint.sh"] -CMD ["shard", "--shard-id", "0", "--data-dir", "/data"] +CMD ["shard", "--shard-id", "0", "--data-dir", "/data", "--backup-dir", "/backup"] diff --git a/README.md b/README.md index 948f9ea..952d1f0 100644 --- a/README.md +++ b/README.md @@ -33,7 +33,7 @@ Unirust will: ### Installation ```bash -git clone https://github.com/unirust/unirust.git +git clone https://github.com/script3r/unirust.git cd unirust cargo build --release ``` @@ -42,13 +42,13 @@ cargo build --release ```bash # Start a single shard -./target/release/unirust_shard --listen 127.0.0.1:50061 --shard-id 0 +./target/release/unirust_shard --listen 127.0.0.1:50061 --shard-id 0 --ephemeral # In another terminal, start the router ./target/release/unirust_router --listen 127.0.0.1:50060 --shards 127.0.0.1:50061 ``` -### Multi-Shard Cluster (Production) +### Local Multi-Shard Cluster ```bash # Use the cluster script @@ -61,11 +61,11 @@ SHARDS=5 ONTOLOGY=/etc/unirust/ontology.json ./scripts/cluster.sh restart UNIRUST_CONFIRM_RESET=1 ./scripts/cluster.sh reset # Or start manually: -./target/release/unirust_shard --listen 127.0.0.1:50061 --shard-id 0 --data-dir /data/shard0 -./target/release/unirust_shard --listen 127.0.0.1:50062 --shard-id 1 --data-dir /data/shard1 -./target/release/unirust_shard --listen 127.0.0.1:50063 --shard-id 2 --data-dir /data/shard2 -./target/release/unirust_shard --listen 127.0.0.1:50064 --shard-id 3 --data-dir /data/shard3 -./target/release/unirust_shard --listen 127.0.0.1:50065 --shard-id 4 --data-dir /data/shard4 +./target/release/unirust_shard --listen 127.0.0.1:50061 --shard-id 0 --data-dir /data/shard0 --backup-dir /backup/shard0 +./target/release/unirust_shard --listen 127.0.0.1:50062 --shard-id 1 --data-dir /data/shard1 --backup-dir /backup/shard1 +./target/release/unirust_shard --listen 127.0.0.1:50063 --shard-id 2 --data-dir /data/shard2 --backup-dir /backup/shard2 +./target/release/unirust_shard --listen 127.0.0.1:50064 --shard-id 3 --data-dir /data/shard3 --backup-dir /backup/shard3 +./target/release/unirust_shard --listen 127.0.0.1:50065 --shard-id 4 --data-dir /data/shard4 --backup-dir /backup/shard4 ./target/release/unirust_router --listen 127.0.0.1:50060 \ --shards 127.0.0.1:50061,127.0.0.1:50062,127.0.0.1:50063,127.0.0.1:50064,127.0.0.1:50065 @@ -120,10 +120,20 @@ profile = "high-throughput" listen = "0.0.0.0:50061" id = 0 data_dir = "/var/lib/unirust/shard-0" +backup_dir = "/var/backups/unirust/shard-0" +tls_cert = "/etc/unirust/tls/shard-0.crt" +tls_key = "/etc/unirust/tls/shard-0.key" +tls_client_ca = "/etc/unirust/tls/clients-ca.crt" [router] listen = "0.0.0.0:50060" -shards = ["shard-0:50061", "shard-1:50061", "shard-2:50061"] +shards = ["https://shard-0:50061", "https://shard-1:50061", "https://shard-2:50061"] +tls_cert = "/etc/unirust/tls/router.crt" +tls_key = "/etc/unirust/tls/router.key" +tls_client_ca = "/etc/unirust/tls/clients-ca.crt" +shard_tls_ca = "/etc/unirust/tls/shards-ca.crt" +shard_tls_cert = "/etc/unirust/tls/router-client.crt" +shard_tls_key = "/etc/unirust/tls/router-client.key" [storage] block_cache_mb = 1024 @@ -140,10 +150,29 @@ write_buffer_mb = 256 | `UNIRUST_SHARD_ID` | Shard ID | | `UNIRUST_SHARD_DATA_DIR` | Persistent shard data directory | | `UNIRUST_SHARD_BACKUP_DIR` | External checkpoint root | +| `UNIRUST_SHARD_TLS_CERT` | Shard server certificate | +| `UNIRUST_SHARD_TLS_KEY` | Shard server private key | +| `UNIRUST_SHARD_TLS_CLIENT_CA` | CA for required shard client certificates | +| `UNIRUST_SHARD_REPLICA` | Passive replica endpoint for this primary | +| `UNIRUST_SHARD_REPLICA_MODE` | Run as a passive replica | +| `UNIRUST_SHARD_REPLICATION_TOKEN_FILE` | Shared secret file for one replica pair | +| `UNIRUST_SHARD_ALLOW_INSECURE_REPLICATION` | Permit plaintext replication for isolated development | +| `UNIRUST_SHARD_REPLICA_CONNECT_TIMEOUT_SECS` | Replica connection timeout | +| `UNIRUST_SHARD_REPLICA_REQUEST_TIMEOUT_SECS` | Per-RPC replica timeout | +| `UNIRUST_SHARD_REPLICA_TCP_KEEPALIVE_SECS` | Replica TCP keepalive interval | +| `UNIRUST_SHARD_REPLICA_TLS_CA` | CA used to verify the replica | +| `UNIRUST_SHARD_REPLICA_TLS_CERT` | Primary certificate presented to the replica | +| `UNIRUST_SHARD_REPLICA_TLS_KEY` | Primary client private key | | `UNIRUST_ROUTER_SHARDS` | Comma-separated shard addresses | | `UNIRUST_ROUTER_CHECKPOINT_INTERVAL_SECS` | Coordinated checkpoint interval (`0` disables) | | `UNIRUST_ROUTER_SHARD_CONNECT_TIMEOUT_SECS` | Shard connection timeout | | `UNIRUST_ROUTER_SHARD_REQUEST_TIMEOUT_SECS` | Per-RPC shard timeout | +| `UNIRUST_ROUTER_TLS_CERT` | Router server certificate | +| `UNIRUST_ROUTER_TLS_KEY` | Router server private key | +| `UNIRUST_ROUTER_TLS_CLIENT_CA` | CA for required router client certificates | +| `UNIRUST_ROUTER_SHARD_TLS_CA` | CA used to verify shard certificates | +| `UNIRUST_ROUTER_SHARD_TLS_CERT` | Router certificate presented to shards | +| `UNIRUST_ROUTER_SHARD_TLS_KEY` | Router client private key | ### Tuning Profiles @@ -259,18 +288,33 @@ Cross-shard copy imports therefore fail closed; online scale-out and scale-in require a future transactional relocation protocol or an offline rebuild. Cross-shard redirects are durably applied to every shard and are idempotent. If -any shard fails while a reconciliation result is being applied, the router -latches the cluster closed for ingest, query, and administrative traffic rather -than serving a partially updated global view. Retrying `Reconcile` repairs the -retained dirty keys in place. After a router or full-cluster restart, router -startup performs that repair before returning a serviceable node and clears the -dirty generation only after every shard converges. +any shard fails, or the initiating request is cancelled, while a reconciliation +result is being applied, the router latches the cluster closed for ingest, query, +and administrative traffic rather than serving a partially updated global view. +Retrying `Reconcile` repairs the retained dirty keys in place. After a router or +full-cluster restart, router startup performs that repair before returning a +serviceable node and clears the dirty generation only after every shard +converges. + +Cluster-wide ontology replacement has the same fail-closed cancellation +semantics. Readiness stays failed after an ambiguous partial update until +`SetOntology` is retried with the intended configuration or every shard is +recovered offline to one configuration. Router startup also rejects mismatched +shard ontologies. The shard reconstructs all derived linker state before opening its gRPC -listener. Recovery currently scans all persisted records, so recovery time is -O(record count). This is a correctness-first crash-recovery path, not a bounded -recovery-time guarantee. Measure restart time at the intended dataset size and -set orchestration startup probes accordingly. +listener. Recovery scans persisted records in ordered, bounded batches through +the normal entity-resolution path, so recovery time remains O(record count). +This is a correctness-first crash-recovery path, not a bounded recovery-time +guarantee. Measure restart time at the intended dataset size and set +orchestration startup probes accordingly. + +Global cluster IDs are anchored to durable record IDs so replay order cannot +change cross-shard identity. On the first startup of a database created before +this scheme marker existed, the shard atomically removes allocation-order +redirects that cannot be trusted after replay; router startup reconstructs them +from authoritative records before becoming ready. Upgrade every shard and the +router together for this transition rather than mixing versions. `LinkerStateConfig` cache capacities are not enforced because the current LRU backend has no durable spill/read-through path. Evicting cluster IDs, strong-ID @@ -308,10 +352,74 @@ deployment and cluster script use these semantic probes. Router-to-shard calls are bounded by configurable transport settings under `[router]`: `shard_connect_timeout_secs` (10 seconds by default), `shard_request_timeout_secs` (120 seconds), and -`shard_tcp_keepalive_secs` (30 seconds). A deadline error is not proof that a -mutation was rolled back; the shard may have committed just before the response -was lost. Retry ingest with the same immutable source-record identity and -payload so the operation is resolved idempotently. +`shard_tcp_keepalive_secs` (30 seconds). A deadline error or client cancellation +is not proof that a mutation was rolled back; an admitted ingest continues to a +durable outcome after its caller disconnects. Retry ingest with the same +immutable source-record identity and payload so the operation is resolved +idempotently. + +### Synchronous Replication + +A logical shard can run as one primary and one passive replica on distinct +persistent volumes and failure domains. Both processes use the same shard ID, +ontology, config version, and replication token. Bootstrap both volumes from +the same committed checkpoint, or start with two empty volumes. Primary startup +computes a SHA-256 digest over every logical RocksDB key/value pair on both +nodes and refuses traffic unless their complete durable states match. Pairing +startup is O(database size), so measure it against the production dataset. +Set `UNIRUST_SHARD_REPLICA_REQUEST_TIMEOUT_SECS` above the measured worst-case +pairing digest time. + +Generate a separate secret for each pair, store at least 32 random bytes in a +file readable only by the service account, and mount the same content on both +nodes. Start the passive replica first: + +```bash +unirust_shard \ + --shard-id 0 \ + --data-dir /var/lib/unirust/shard-0-replica \ + --backup-dir /var/backups/unirust/shard-0-replica \ + --replica-mode \ + --replication-token-file /etc/unirust/replication/shard-0.token \ + --tls-cert /etc/unirust/tls/shard-0-replica.crt \ + --tls-key /etc/unirust/tls/shard-0-replica.key \ + --tls-client-ca /etc/unirust/tls/primaries-ca.crt +``` + +Then start the primary with its normal shard server credentials plus: + +```bash +unirust_shard \ + --shard-id 0 \ + --data-dir /var/lib/unirust/shard-0-primary \ + --backup-dir /var/backups/unirust/shard-0-primary \ + --replica https://shard-0-replica:50061 \ + --replication-token-file /etc/unirust/replication/shard-0.token \ + --replica-tls-ca /etc/unirust/tls/replicas-ca.crt \ + --replica-tls-cert /etc/unirust/tls/shard-0-primary.crt \ + --replica-tls-key /etc/unirust/tls/shard-0-primary.key +``` + +Every durable mutation is applied to the replica first, then locally, under one +per-pair serialization gate. The primary acknowledges only after both results +match. An unavailable or ambiguous replica result latches the primary +unhealthy and blocks reads and writes until operators reconcile the pair. +Replication therefore protects acknowledged writes from one volume loss, but +adds replica latency and requires both nodes to be available for writes. + +Failover is manual because Unirust does not implement leader election or +quorum fencing: + +1. Prove the old primary is stopped or isolated from clients and the replica. +2. Stop the passive process and restart its volume without `--replica-mode`. +3. Point the router at the promoted endpoint and restart the router. +4. Rebootstrap the old primary from a checkpoint of the promoted node before + attaching it as a new passive replica. + +Never serve the old primary and promoted replica simultaneously. Doing so can +create split brain. Online reset is disabled while a primary has a replica; +stop and rebootstrap both members together. Keep verified off-host backups for +correlated failures, operator mistakes, and storage corruption. ### External Backups @@ -381,27 +489,78 @@ generation and shard count, then verifies the topology, ontology, and protocol versions before accepting traffic. Mixed restored/unrestored volumes and different committed generations fail closed. -The built-in scheduler only creates local coordinated checkpoints. Unirust does -not replicate, encrypt, transfer, retain, or verify off-host backups. Without -storage-layer replication, the recovery point is the last completed coordinated -checkpoint, so acknowledged writes after it can be lost in a volume failure. -Production operators must provide off-host replication, retention, monitoring, -and periodic restore drills according to their RPO and RTO. Process crashes and -ordinary restarts remain covered by the synced ingest WAL independently of this -backup path. +Export a complete committed generation to storage mounted from another failure +domain. All shard checkpoint paths must be accessible to the export process: -## Deployment Security +```bash +unirust_backup export \ + --destination /mnt/off-host/unirust/backup-2026-07-24T1300Z \ + --checkpoint /var/backups/unirust/shard-0/backup-2026-07-24T1300Z \ + --checkpoint /var/backups/unirust/shard-1/backup-2026-07-24T1300Z \ + --checkpoint /var/backups/unirust/shard-2/backup-2026-07-24T1300Z + +unirust_backup verify \ + --backup /mnt/off-host/unirust/backup-2026-07-24T1300Z +``` + +Export validates that every shard is present exactly once and belongs to the +same generation and topology. It copies into a sibling staging directory, +records every file length and SHA-256 digest in a binary manifest, opens every +copied RocksDB checkpoint read-only with paranoid checks, syncs the tree, and +publishes it with one rename. Verification rejects modified, missing, extra, or +symlinked content. Restore from the exported `shard-0`, `shard-1`, and +`shard-2` directories, not from the deleted local checkpoint roots. -The gRPC services do not currently terminate TLS or authenticate callers. -Production deployments must place the router and every shard on private -networks and enforce authenticated TLS with a service mesh or reverse proxy. -Never expose a shard port directly to an untrusted network. Destructive RPCs -remain disabled by default, but ingest, query, export, reconciliation, and other -administrative surfaces are otherwise unauthenticated. +Retention only removes generations after every entry in its root verifies: + +```bash +unirust_backup prune --root /mnt/off-host/unirust --retain 14 +``` + +The built-in scheduler creates coordinated source checkpoints; it does not +automatically run the export command. Schedule export and verification after +checkpoint completion, monitor both, and run periodic restore drills. The +destination filesystem must provide independent storage and encryption at rest. +Without an enabled synchronous replica, the recovery point for a lost volume +remains the last successfully exported generation and acknowledged writes +after it can be lost. Process crashes and ordinary restarts remain covered +independently by the synced ingest and RocksDB WALs. + +## Deployment Security -The shard binary requires a persistent `--data-dir`. An in-memory shard can only -be started with the explicit `--ephemeral` flag and loses all records when the -process exits. +Router and shard binaries support mutually authenticated TLS. Each server +requires its certificate, private key, and client CA as an all-or-none group. +The router-to-shard client likewise requires a CA, client certificate, and +private key together, and secured shard addresses must use `https://`. Startup +fails closed for partial certificate groups, unreadable or empty PEM files, +plaintext endpoints paired with TLS configuration, or HTTPS endpoints without +explicit trust configuration. Certificate SANs must match the endpoint host. + +Client tools use `--tls-ca`, `--tls-cert`, and `--tls-key`; the semantic health +probe uses `--ca-cert`, `--client-cert`, and `--client-key`. All three options +are required together. Certificate rotation currently requires a process +restart. Production deployments must enable native mTLS on the router and every +shard, or enforce equivalent authenticated TLS through a service mesh. Keep +shard ports private and never expose plaintext gRPC to an untrusted network. +The supplied Compose file is a local plaintext example and binds its router +port to loopback only. + +Replication additionally uses a per-pair shared token. The shard binary rejects +plaintext replication by default; `--allow-insecure-replication` exists only +for isolated development. Protect token files as credentials, rotate them by +stopping and restarting both members, and use different tokens for every pair. + +The shard binary requires a persistent `--data-dir` and a non-overlapping +`--backup-dir`. Mount them from independent storage: separate directory names +alone do not protect against volume loss. An in-memory shard can only be started +with the explicit `--ephemeral` flag and loses all records when the process +exits. `--allow-colocated-checkpoints` is a development-only escape hatch and +does not provide volume-loss recovery. + +Router and shard servers cap each encoded or decoded gRPC message at 4 MiB, +limit each connection to 128 concurrent requests, and shed excess load. Use the +streaming ingest, import, and export RPCs for larger transfers. Bound connection +counts and request rates at the load balancer as well. ## Performance diff --git a/compose.yaml b/compose.yaml index a7db127..745f7a7 100644 --- a/compose.yaml +++ b/compose.yaml @@ -101,7 +101,7 @@ services: volumes: - ./examples/loadtest-ontology.json:/etc/unirust/ontology.json:ro ports: - - "50060:50060" + - "127.0.0.1:50060:50060" depends_on: shard-0: condition: service_healthy diff --git a/examples/cluster.rs b/examples/cluster.rs index e05f59f..a1fd593 100644 --- a/examples/cluster.rs +++ b/examples/cluster.rs @@ -27,9 +27,9 @@ //! SHARDS=3 ./scripts/cluster.sh start //! //! # Option 2: Start manually -//! ./target/release/unirust_shard --listen 127.0.0.1:50061 --shard-id 0 --data-dir /tmp/shard0 -//! ./target/release/unirust_shard --listen 127.0.0.1:50062 --shard-id 1 --data-dir /tmp/shard1 -//! ./target/release/unirust_shard --listen 127.0.0.1:50063 --shard-id 2 --data-dir /tmp/shard2 +//! ./target/release/unirust_shard --listen 127.0.0.1:50061 --shard-id 0 --data-dir /tmp/shard0 --backup-dir /tmp/shard0-backup +//! ./target/release/unirust_shard --listen 127.0.0.1:50062 --shard-id 1 --data-dir /tmp/shard1 --backup-dir /tmp/shard1-backup +//! ./target/release/unirust_shard --listen 127.0.0.1:50063 --shard-id 2 --data-dir /tmp/shard2 --backup-dir /tmp/shard2-backup //! ./target/release/unirust_router --listen 127.0.0.1:50060 \ //! --shards 127.0.0.1:50061,127.0.0.1:50062,127.0.0.1:50063 //! ``` @@ -73,13 +73,13 @@ async fn main() -> anyhow::Result<()> { eprintln!(" SHARDS=3 ./scripts/cluster.sh start\n"); eprintln!("Or manually:"); eprintln!( - " ./target/release/unirust_shard --listen 127.0.0.1:50061 --shard-id 0 --data-dir /tmp/shard0" + " ./target/release/unirust_shard --listen 127.0.0.1:50061 --shard-id 0 --data-dir /tmp/shard0 --backup-dir /tmp/shard0-backup" ); eprintln!( - " ./target/release/unirust_shard --listen 127.0.0.1:50062 --shard-id 1 --data-dir /tmp/shard1" + " ./target/release/unirust_shard --listen 127.0.0.1:50062 --shard-id 1 --data-dir /tmp/shard1 --backup-dir /tmp/shard1-backup" ); eprintln!( - " ./target/release/unirust_shard --listen 127.0.0.1:50063 --shard-id 2 --data-dir /tmp/shard2" + " ./target/release/unirust_shard --listen 127.0.0.1:50063 --shard-id 2 --data-dir /tmp/shard2 --backup-dir /tmp/shard2-backup" ); eprintln!(" ./target/release/unirust_router --listen 127.0.0.1:50060 \\"); eprintln!(" --shards 127.0.0.1:50061,127.0.0.1:50062,127.0.0.1:50063"); diff --git a/examples/unirust.toml b/examples/unirust.toml index 1a0be98..9450707 100644 --- a/examples/unirust.toml +++ b/examples/unirust.toml @@ -15,12 +15,28 @@ profile = "high-throughput" listen = "0.0.0.0:50061" # Shard ID (0-indexed, unique per shard) id = 0 -# Data directory for persistence (comment out for in-memory only) +# Data directory for persistence (`--ephemeral` is required when omitted) # data_dir = "/var/lib/unirust/shard-0" # Checkpoint root on an independent volume (use a unique path per shard) # backup_dir = "/var/backups/unirust/shard-0" # Path to ontology configuration (JSON) # ontology = "/etc/unirust/ontology.json" +# Mutually authenticated TLS server credentials (all three are required together) +# tls_cert = "/etc/unirust/tls/shard-0.crt" +# tls_key = "/etc/unirust/tls/shard-0.key" +# tls_client_ca = "/etc/unirust/tls/clients-ca.crt" +# Primary-only synchronous replica settings. Bootstrap both volumes from the +# same committed checkpoint before enabling these options. +# replica = "https://shard-0-replica:50061" +# replication_token_file = "/etc/unirust/replication/shard-0.token" +# replica_tls_ca = "/etc/unirust/tls/replicas-ca.crt" +# replica_tls_cert = "/etc/unirust/tls/shard-0-primary.crt" +# replica_tls_key = "/etc/unirust/tls/shard-0-primary.key" +# On the passive process, omit `replica` and set: +# replica_mode = true +# Plaintext replication is rejected unless explicitly enabled for isolated +# development: +# allow_insecure_replication = false # Run repair on startup (recovers from unclean shutdown) repair = false @@ -44,6 +60,14 @@ shard_request_timeout_secs = 120 shard_tcp_keepalive_secs = 30 # Automatic coordinated checkpoint interval (0 disables) # checkpoint_interval_secs = 3600 +# Mutually authenticated TLS for clients connecting to the router +# tls_cert = "/etc/unirust/tls/router.crt" +# tls_key = "/etc/unirust/tls/router.key" +# tls_client_ca = "/etc/unirust/tls/clients-ca.crt" +# Router identity and trust for https:// shard addresses +# shard_tls_ca = "/etc/unirust/tls/shards-ca.crt" +# shard_tls_cert = "/etc/unirust/tls/router-client.crt" +# shard_tls_key = "/etc/unirust/tls/router-client.key" # Advanced storage tuning (RocksDB) # These settings control the underlying storage engine performance. diff --git a/proto/unirust.proto b/proto/unirust.proto index 1c94ad9..485a988 100644 --- a/proto/unirust.proto +++ b/proto/unirust.proto @@ -255,7 +255,16 @@ message IngestRecordsFromUrlRequest { // --- Config Version --- -message ConfigVersionRequest {} +message ConfigVersionRequest { + bool include_durable_state_digest = 1; +} + +enum ShardRole { + SHARD_ROLE_UNSPECIFIED = 0; + SHARD_ROLE_STANDALONE = 1; + SHARD_ROLE_PRIMARY = 2; + SHARD_ROLE_REPLICA = 3; +} message ConfigVersionResponse { string version = 1; @@ -267,6 +276,9 @@ message ConfigVersionResponse { // Empty/zero for a volume that has never been restored from a checkpoint. string restore_generation = 7; uint32 restore_shard_count = 8; + ShardRole shard_role = 9; + // SHA-256 over every logical RocksDB column-family key/value pair. + bytes durable_state_digest = 10; } // --- Durable Source Identity Reservation --- @@ -456,13 +468,20 @@ message BoundaryKeyEntries { message GetBoundaryMetadataRequest { uint64 since_version = 1; + // Return metadata only for these signatures. Empty requests return shard + // identity/version without materializing the full boundary index. + repeated IdentityKeySignature signatures = 2; } message GetBoundaryMetadataResponse { BoundaryMetadata metadata = 1; } -message GetDirtyBoundaryKeysRequest {} +message GetDirtyBoundaryKeysRequest { + // Exclusive lexicographic cursor over the 32-byte signature. + bytes after_signature = 1; + uint32 limit = 2; +} message DirtyBoundaryKey { IdentityKeySignature signature = 1; @@ -472,6 +491,8 @@ message DirtyBoundaryKey { message GetDirtyBoundaryKeysResponse { repeated DirtyBoundaryKey dirty_keys = 1; uint32 shard_id = 2; + bytes next_after_signature = 3; + bool has_more = 4; } message ClearDirtyKeysRequest { @@ -511,6 +532,16 @@ message ApplyMergeResponse { string error = 3; } +message ApplyMergesRequest { + repeated ClusterMerge merges = 1; +} + +message ApplyMergesResponse { + bool success = 1; + uint64 records_updated = 2; + string error = 3; +} + // --- Cross-Shard Conflict Propagation --- // A cross-shard conflict detected during reconciliation. @@ -580,6 +611,7 @@ service ShardService { // Cross-shard reconciliation rpc GetBoundaryMetadata(GetBoundaryMetadataRequest) returns (GetBoundaryMetadataResponse); rpc ApplyMerge(ApplyMergeRequest) returns (ApplyMergeResponse); + rpc ApplyMerges(ApplyMergesRequest) returns (ApplyMergesResponse); rpc GetDirtyBoundaryKeys(GetDirtyBoundaryKeysRequest) returns (GetDirtyBoundaryKeysResponse); rpc ClearDirtyKeys(ClearDirtyKeysRequest) returns (ClearDirtyKeysResponse); diff --git a/scripts/podman_cluster.sh b/scripts/podman_cluster.sh index 71e6983..0d8e522 100755 --- a/scripts/podman_cluster.sh +++ b/scripts/podman_cluster.sh @@ -9,8 +9,10 @@ ROUTER_PORT="${ROUTER_PORT:-50060}" ONTOLOGY="${ONTOLOGY:-$ROOT_DIR/examples/loadtest-ontology.json}" CONFIG_VERSION="${CONFIG_VERSION:-}" PROFILE="${PROFILE:-high-throughput}" +CHECKPOINT_INTERVAL_SECS="${CHECKPOINT_INTERVAL_SECS:-3600}" CONTAINER_PREFIX="${CONTAINER_PREFIX:-unirust}" VOLUME_PREFIX="${VOLUME_PREFIX:-unirust-data}" +BACKUP_VOLUME_PREFIX="${BACKUP_VOLUME_PREFIX:-unirust-backup}" STOP_WAIT_SECS="${STOP_WAIT_SECS:-10}" UNIRUST_CONFIRM_RESET="${UNIRUST_CONFIRM_RESET:-0}" @@ -61,11 +63,13 @@ start_cluster() { --network "$NETWORK" \ --restart=unless-stopped \ -v "${VOLUME_PREFIX}-shard-$i:/data" \ + -v "${BACKUP_VOLUME_PREFIX}-shard-$i:/backup" \ "${ontology_mount[@]}" \ "$IMAGE" \ shard \ --shard-id "$i" \ --data-dir /data \ + --backup-dir /backup \ --profile "$PROFILE" \ "${ontology_args[@]}" \ "${version_args[@]}" @@ -90,12 +94,14 @@ start_cluster() { "$IMAGE" \ router \ --shards "$shard_list" \ + --checkpoint-interval-secs "$CHECKPOINT_INTERVAL_SECS" \ "${ontology_args[@]}" \ "${version_args[@]}" echo "Cluster started." echo "Router listening on localhost:${ROUTER_PORT}" echo "Persistent volumes use prefix: ${VOLUME_PREFIX}-shard-" + echo "Checkpoint volumes use prefix: ${BACKUP_VOLUME_PREFIX}-shard-" } stop_cluster() { @@ -125,11 +131,12 @@ reset_cluster() { fi stop_cluster while IFS= read -r volume; do - if [[ "$volume" == "${VOLUME_PREFIX}-shard-"* ]]; then + if [[ "$volume" == "${VOLUME_PREFIX}-shard-"* || "$volume" == "${BACKUP_VOLUME_PREFIX}-shard-"* ]]; then podman volume rm "$volume" fi done < <(podman volume ls --format "{{.Name}}") echo "Deleted persistent volumes with prefix: ${VOLUME_PREFIX}-shard-" + echo "Deleted checkpoint volumes with prefix: ${BACKUP_VOLUME_PREFIX}-shard-" } status_cluster() { diff --git a/src/backup.rs b/src/backup.rs new file mode 100644 index 0000000..01ae8bf --- /dev/null +++ b/src/backup.rs @@ -0,0 +1,512 @@ +use crate::persistence::{ + verify_cluster_checkpoint, ClusterCheckpointManifest, CHECKPOINT_PROTOCOL_VERSION, +}; +use anyhow::{anyhow, Context, Result}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use std::collections::BTreeSet; +use std::fs; +use std::io::{Read, Write}; +use std::path::{Component, Path, PathBuf}; +use std::time::{SystemTime, UNIX_EPOCH}; + +const BACKUP_FORMAT_VERSION: u32 = 1; +const BACKUP_MANIFEST_FILE: &str = "UNIRUST_BACKUP_MANIFEST"; +const BACKUP_COMMITTED_FILE: &str = "UNIRUST_BACKUP_COMMITTED"; +const COPY_BUFFER_SIZE: usize = 1024 * 1024; +const MAX_BACKUP_MANIFEST_BYTES: u64 = 256 * 1024 * 1024; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +struct BackupFile { + path: String, + length: u64, + sha256: [u8; 32], +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +struct BackupShard { + shard_id: u32, + files: Vec, +} + +struct ValidatedCheckpointSet { + generation: String, + shard_count: u32, + checkpoints: Vec<(PathBuf, ClusterCheckpointManifest)>, +} + +/// Integrity and topology metadata for one exported cluster checkpoint. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct ClusterBackupManifest { + format_version: u32, + checkpoint_protocol_version: u32, + generation: String, + shard_count: u32, + created_unix_nanos: u128, + shards: Vec, +} + +impl ClusterBackupManifest { + pub fn generation(&self) -> &str { + &self.generation + } + + pub fn shard_count(&self) -> u32 { + self.shard_count + } + + pub fn created_unix_nanos(&self) -> u128 { + self.created_unix_nanos + } + + fn validate(&self) -> Result<()> { + if self.format_version != BACKUP_FORMAT_VERSION { + anyhow::bail!( + "unsupported cluster backup format version {}", + self.format_version + ); + } + if self.checkpoint_protocol_version != CHECKPOINT_PROTOCOL_VERSION { + anyhow::bail!( + "unsupported checkpoint protocol version {}", + self.checkpoint_protocol_version + ); + } + if self.generation.is_empty() { + anyhow::bail!("cluster backup generation is empty"); + } + if self.shard_count == 0 || self.shards.len() != self.shard_count as usize { + anyhow::bail!( + "cluster backup contains {} shards but declares {}", + self.shards.len(), + self.shard_count + ); + } + for (expected, shard) in self.shards.iter().enumerate() { + if shard.shard_id != expected as u32 { + anyhow::bail!( + "cluster backup shard sequence contains {} at position {}", + shard.shard_id, + expected + ); + } + if shard.files.is_empty() { + anyhow::bail!("cluster backup shard {} has no files", shard.shard_id); + } + let mut previous = None; + for file in &shard.files { + validate_relative_file_path(&file.path)?; + if previous.is_some_and(|value: &str| value >= file.path.as_str()) { + anyhow::bail!( + "cluster backup shard {} file list is not strictly ordered", + shard.shard_id + ); + } + previous = Some(file.path.as_str()); + } + } + Ok(()) + } +} + +#[cfg(unix)] +fn sync_directory(path: &Path) -> Result<()> { + fs::File::open(path)?.sync_all()?; + Ok(()) +} + +#[cfg(not(unix))] +fn sync_directory(_path: &Path) -> Result<()> { + Ok(()) +} + +fn validate_relative_file_path(path: &str) -> Result<()> { + let path = Path::new(path); + if path.as_os_str().is_empty() + || path + .components() + .any(|component| !matches!(component, Component::Normal(_))) + { + anyhow::bail!("backup file path must be a safe relative path"); + } + Ok(()) +} + +fn sorted_entries(path: &Path) -> Result> { + let mut entries = fs::read_dir(path)?.collect::>>()?; + entries.sort_by_key(fs::DirEntry::file_name); + Ok(entries) +} + +fn digest_file(path: &Path) -> Result<(u64, [u8; 32])> { + let metadata = fs::symlink_metadata(path)?; + if metadata.file_type().is_symlink() || !metadata.is_file() { + anyhow::bail!("backup entry {} must be a real file", path.display()); + } + let mut source = fs::File::open(path)?; + let mut digest = Sha256::new(); + let mut length = 0u64; + let mut buffer = vec![0u8; COPY_BUFFER_SIZE]; + loop { + let read = source.read(&mut buffer)?; + if read == 0 { + break; + } + digest.update(&buffer[..read]); + length = length + .checked_add(read as u64) + .ok_or_else(|| anyhow!("backup file length overflow"))?; + } + Ok((length, digest.finalize().into())) +} + +fn copy_file(source: &Path, destination: &Path) -> Result<(u64, [u8; 32])> { + let metadata = fs::symlink_metadata(source)?; + if metadata.file_type().is_symlink() || !metadata.is_file() { + anyhow::bail!("checkpoint entry {} must be a real file", source.display()); + } + let mut source_file = fs::File::open(source)?; + let mut destination_file = fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(destination)?; + let mut digest = Sha256::new(); + let mut length = 0u64; + let mut buffer = vec![0u8; COPY_BUFFER_SIZE]; + loop { + let read = source_file.read(&mut buffer)?; + if read == 0 { + break; + } + destination_file.write_all(&buffer[..read])?; + digest.update(&buffer[..read]); + length = length + .checked_add(read as u64) + .ok_or_else(|| anyhow!("backup file length overflow"))?; + } + destination_file.sync_all()?; + fs::set_permissions(destination, metadata.permissions())?; + Ok((length, digest.finalize().into())) +} + +fn copy_tree( + source: &Path, + destination: &Path, + relative: &Path, + files: &mut Vec, +) -> Result<()> { + fs::create_dir(destination)?; + for entry in sorted_entries(source)? { + let file_type = entry.file_type()?; + let source_path = entry.path(); + let destination_path = destination.join(entry.file_name()); + let relative_path = relative.join(entry.file_name()); + if file_type.is_dir() { + copy_tree(&source_path, &destination_path, &relative_path, files)?; + } else if file_type.is_file() { + let path = relative_path + .to_str() + .ok_or_else(|| anyhow!("checkpoint contains a non-UTF-8 file name"))? + .to_string(); + validate_relative_file_path(&path)?; + let (length, sha256) = copy_file(&source_path, &destination_path)?; + files.push(BackupFile { + path, + length, + sha256, + }); + } else { + anyhow::bail!( + "checkpoint contains unsupported entry {}", + source_path.display() + ); + } + } + sync_directory(destination)?; + Ok(()) +} + +fn hash_tree( + root: &Path, + current: &Path, + relative: &Path, + files: &mut Vec, +) -> Result<()> { + let metadata = fs::symlink_metadata(current)?; + if metadata.file_type().is_symlink() || !metadata.is_dir() { + anyhow::bail!( + "backup directory {} must be a real directory", + current.display() + ); + } + for entry in sorted_entries(current)? { + let file_type = entry.file_type()?; + let path = entry.path(); + let relative_path = relative.join(entry.file_name()); + if file_type.is_dir() { + hash_tree(root, &path, &relative_path, files)?; + } else if file_type.is_file() { + let relative_string = relative_path + .to_str() + .ok_or_else(|| anyhow!("backup contains a non-UTF-8 file name"))? + .to_string(); + validate_relative_file_path(&relative_string)?; + let (length, sha256) = digest_file(&path)?; + files.push(BackupFile { + path: relative_string, + length, + sha256, + }); + } else { + anyhow::bail!("backup contains unsupported entry {}", path.display()); + } + } + if !current.starts_with(root) { + anyhow::bail!("backup traversal escaped its shard directory"); + } + Ok(()) +} + +fn write_marker(path: &Path, bytes: &[u8]) -> Result<()> { + let mut file = fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(path)?; + file.write_all(bytes)?; + file.sync_all()?; + Ok(()) +} + +fn read_marker(path: &Path) -> Result> { + let metadata = fs::symlink_metadata(path)?; + if metadata.file_type().is_symlink() || !metadata.is_file() { + anyhow::bail!("backup marker {} must be a real file", path.display()); + } + if metadata.len() > MAX_BACKUP_MANIFEST_BYTES { + anyhow::bail!("backup marker {} is unreasonably large", path.display()); + } + Ok(fs::read(path)?) +} + +fn validate_checkpoint_set(checkpoints: &[PathBuf]) -> Result { + if checkpoints.is_empty() { + anyhow::bail!("at least one committed shard checkpoint is required"); + } + let mut validated = Vec::with_capacity(checkpoints.len()); + for checkpoint in checkpoints { + let manifest = verify_cluster_checkpoint(checkpoint) + .with_context(|| format!("invalid checkpoint {}", checkpoint.display()))?; + validated.push((checkpoint.canonicalize()?, manifest)); + } + validated.sort_by_key(|(_, manifest)| manifest.shard_id()); + let generation = validated[0].1.generation().to_string(); + let shard_count = validated[0].1.shard_count(); + if validated.len() != shard_count as usize { + anyhow::bail!( + "checkpoint set contains {} shards but generation {} requires {}", + validated.len(), + generation, + shard_count + ); + } + for (expected, (_, manifest)) in validated.iter().enumerate() { + if manifest.generation() != generation { + anyhow::bail!( + "checkpoint set mixes generations {} and {}", + generation, + manifest.generation() + ); + } + if manifest.shard_count() != shard_count { + anyhow::bail!("checkpoint set mixes shard counts"); + } + if manifest.shard_id() != expected as u32 { + anyhow::bail!( + "checkpoint set is missing shard {} or contains a duplicate", + expected + ); + } + } + Ok(ValidatedCheckpointSet { + generation, + shard_count, + checkpoints: validated, + }) +} + +/// Copy a complete coordinated checkpoint to another filesystem with an +/// integrity manifest and atomic publication. +pub fn export_cluster_backup( + checkpoints: &[PathBuf], + destination: &Path, +) -> Result { + let validated = validate_checkpoint_set(checkpoints)?; + if destination.exists() { + anyhow::bail!( + "backup destination {} already exists", + destination.display() + ); + } + let destination_name = destination + .file_name() + .and_then(|name| name.to_str()) + .ok_or_else(|| anyhow!("backup destination must have a UTF-8 directory name"))?; + let parent = destination + .parent() + .filter(|path| !path.as_os_str().is_empty()) + .unwrap_or_else(|| Path::new(".")); + fs::create_dir_all(parent)?; + let parent = parent.canonicalize()?; + let destination = parent.join(destination_name); + let staging = parent.join(format!(".{destination_name}.tmp")); + if staging.exists() { + anyhow::bail!( + "backup staging directory {} already exists; inspect and remove it before retrying", + staging.display() + ); + } + for (checkpoint, _) in &validated.checkpoints { + if destination.starts_with(checkpoint) || checkpoint.starts_with(&destination) { + anyhow::bail!("backup destination and checkpoint sources must be disjoint"); + } + } + + fs::create_dir(&staging)?; + let export_result = (|| -> Result { + let mut shards = Vec::with_capacity(validated.checkpoints.len()); + for (checkpoint, checkpoint_manifest) in validated.checkpoints { + let shard_id = checkpoint_manifest.shard_id(); + let shard_dir = staging.join(format!("shard-{shard_id}")); + let mut files = Vec::new(); + copy_tree(&checkpoint, &shard_dir, Path::new(""), &mut files)?; + files.sort_by(|left, right| left.path.cmp(&right.path)); + shards.push(BackupShard { shard_id, files }); + } + let manifest = ClusterBackupManifest { + format_version: BACKUP_FORMAT_VERSION, + checkpoint_protocol_version: CHECKPOINT_PROTOCOL_VERSION, + generation: validated.generation, + shard_count: validated.shard_count, + created_unix_nanos: SystemTime::now().duration_since(UNIX_EPOCH)?.as_nanos(), + shards, + }; + manifest.validate()?; + let bytes = bincode::serialize(&manifest)?; + write_marker(&staging.join(BACKUP_MANIFEST_FILE), &bytes)?; + write_marker(&staging.join(BACKUP_COMMITTED_FILE), &bytes)?; + sync_directory(&staging)?; + let verified = verify_cluster_backup(&staging)?; + if verified != manifest { + anyhow::bail!("staged backup verification returned different metadata"); + } + Ok(manifest) + })(); + + let manifest = export_result?; + fs::rename(&staging, &destination)?; + sync_directory(&parent)?; + let published = verify_cluster_backup(&destination)?; + if published != manifest { + anyhow::bail!("published backup verification returned different metadata"); + } + Ok(manifest) +} + +/// Verify all bytes, shard identities, topology, and RocksDB contents in an +/// exported cluster backup. +pub fn verify_cluster_backup(backup: &Path) -> Result { + let metadata = fs::symlink_metadata(backup)?; + if metadata.file_type().is_symlink() || !metadata.is_dir() { + anyhow::bail!("cluster backup must be a real directory"); + } + let backup = backup.canonicalize()?; + let manifest_bytes = read_marker(&backup.join(BACKUP_MANIFEST_FILE))?; + let committed_bytes = read_marker(&backup.join(BACKUP_COMMITTED_FILE))?; + if committed_bytes != manifest_bytes { + anyhow::bail!("cluster backup commit marker does not match its manifest"); + } + let manifest: ClusterBackupManifest = bincode::deserialize(&manifest_bytes)?; + manifest.validate()?; + + let expected_root_entries = manifest + .shards + .iter() + .map(|shard| format!("shard-{}", shard.shard_id)) + .chain([ + BACKUP_MANIFEST_FILE.to_string(), + BACKUP_COMMITTED_FILE.to_string(), + ]) + .collect::>(); + let actual_root_entries = sorted_entries(&backup)? + .into_iter() + .map(|entry| { + entry + .file_name() + .into_string() + .map_err(|_| anyhow!("backup contains a non-UTF-8 root entry")) + }) + .collect::>>()?; + if actual_root_entries != expected_root_entries { + anyhow::bail!("cluster backup contains missing or unexpected root entries"); + } + + for shard in &manifest.shards { + let shard_dir = backup.join(format!("shard-{}", shard.shard_id)); + let mut actual_files = Vec::new(); + hash_tree(&shard_dir, &shard_dir, Path::new(""), &mut actual_files)?; + actual_files.sort_by(|left, right| left.path.cmp(&right.path)); + if actual_files != shard.files { + anyhow::bail!("backup shard {} failed file verification", shard.shard_id); + } + let checkpoint = verify_cluster_checkpoint(&shard_dir)?; + if checkpoint.generation() != manifest.generation + || checkpoint.shard_count() != manifest.shard_count + || checkpoint.shard_id() != shard.shard_id + { + anyhow::bail!( + "backup shard {} checkpoint provenance does not match the cluster manifest", + shard.shard_id + ); + } + } + Ok(manifest) +} + +/// Retain the newest verified backup sets and remove older verified sets. +/// +/// The operation fails before deleting anything if any entry under `root` is +/// incomplete, unexpected, or fails verification. +pub fn prune_verified_cluster_backups(root: &Path, retain: usize) -> Result> { + if retain == 0 { + anyhow::bail!("backup retention must keep at least one generation"); + } + let metadata = fs::symlink_metadata(root)?; + if metadata.file_type().is_symlink() || !metadata.is_dir() { + anyhow::bail!("backup retention root must be a real directory"); + } + let root = root.canonicalize()?; + let mut backups = Vec::new(); + for entry in sorted_entries(&root)? { + let path = entry.path(); + let file_type = entry.file_type()?; + if !file_type.is_dir() || file_type.is_symlink() { + anyhow::bail!( + "backup retention root contains unexpected entry {}", + path.display() + ); + } + let manifest = verify_cluster_backup(&path).with_context(|| { + format!("refusing retention with invalid backup {}", path.display()) + })?; + backups.push((manifest.created_unix_nanos(), path)); + } + backups.sort_by(|left, right| left.0.cmp(&right.0).then_with(|| left.1.cmp(&right.1))); + let remove_count = backups.len().saturating_sub(retain); + let mut removed = Vec::with_capacity(remove_count); + for (_, path) in backups.into_iter().take(remove_count) { + fs::remove_dir_all(&path)?; + sync_directory(&root)?; + removed.push(path); + } + Ok(removed) +} diff --git a/src/bin/unirust_backup.rs b/src/bin/unirust_backup.rs new file mode 100644 index 0000000..c628156 --- /dev/null +++ b/src/bin/unirust_backup.rs @@ -0,0 +1,143 @@ +use anyhow::{anyhow, Context, Result}; +use std::path::PathBuf; +use unirust_rs::{export_cluster_backup, prune_verified_cluster_backups, verify_cluster_backup}; + +fn usage() -> &'static str { + "Usage: + unirust_backup export --destination --checkpoint [--checkpoint ...] + unirust_backup verify --backup + unirust_backup prune --root --retain " +} + +fn required_value(args: &[String], index: &mut usize, flag: &str) -> Result { + *index += 1; + args.get(*index) + .filter(|value| !value.starts_with("--")) + .cloned() + .ok_or_else(|| anyhow!("{flag} requires a value")) +} + +fn export(args: &[String]) -> Result<()> { + let mut destination = None; + let mut checkpoints = Vec::new(); + let mut index = 0; + while index < args.len() { + match args[index].as_str() { + "--destination" => { + destination = Some(PathBuf::from(required_value( + args, + &mut index, + "--destination", + )?)); + } + "--checkpoint" => checkpoints.push(PathBuf::from(required_value( + args, + &mut index, + "--checkpoint", + )?)), + "--help" | "-h" => { + println!("{}", usage()); + return Ok(()); + } + flag => anyhow::bail!("unknown export argument {flag}\n{}", usage()), + } + index += 1; + } + let destination = + destination.ok_or_else(|| anyhow!("export requires --destination\n{}", usage()))?; + let manifest = export_cluster_backup(&checkpoints, &destination)?; + println!( + "exported generation {} with {} shards to {}", + manifest.generation(), + manifest.shard_count(), + destination.display() + ); + Ok(()) +} + +fn verify(args: &[String]) -> Result<()> { + let mut backup = None; + let mut index = 0; + while index < args.len() { + match args[index].as_str() { + "--backup" => { + backup = Some(PathBuf::from(required_value(args, &mut index, "--backup")?)); + } + "--help" | "-h" => { + println!("{}", usage()); + return Ok(()); + } + flag => anyhow::bail!("unknown verify argument {flag}\n{}", usage()), + } + index += 1; + } + let backup = backup.ok_or_else(|| anyhow!("verify requires --backup\n{}", usage()))?; + let manifest = verify_cluster_backup(&backup)?; + println!( + "verified generation {} with {} shards at {}", + manifest.generation(), + manifest.shard_count(), + backup.display() + ); + Ok(()) +} + +fn prune(args: &[String]) -> Result<()> { + let mut root = None; + let mut retain = None; + let mut index = 0; + while index < args.len() { + match args[index].as_str() { + "--root" => { + root = Some(PathBuf::from(required_value(args, &mut index, "--root")?)); + } + "--retain" => { + let value = required_value(args, &mut index, "--retain")?; + retain = Some( + value + .parse::() + .with_context(|| format!("invalid --retain value {value}"))?, + ); + } + "--help" | "-h" => { + println!("{}", usage()); + return Ok(()); + } + flag => anyhow::bail!("unknown prune argument {flag}\n{}", usage()), + } + index += 1; + } + let root = root.ok_or_else(|| anyhow!("prune requires --root\n{}", usage()))?; + let retain = retain.ok_or_else(|| anyhow!("prune requires --retain\n{}", usage()))?; + let removed = prune_verified_cluster_backups(&root, retain)?; + println!( + "removed {} verified backup generation(s) from {}", + removed.len(), + root.display() + ); + Ok(()) +} + +fn run() -> Result<()> { + let args = std::env::args().skip(1).collect::>(); + let Some(command) = args.first() else { + anyhow::bail!("{}", usage()); + }; + match command.as_str() { + "export" => export(&args[1..]), + "verify" => verify(&args[1..]), + "prune" => prune(&args[1..]), + "--help" | "-h" => { + println!("{}", usage()); + Ok(()) + } + command => anyhow::bail!("unknown command {command}\n{}", usage()), + } +} + +fn main() { + if let Err(error) = run() { + eprintln!("backup operation failed: {error:#}"); + std::process::exit(1); + } +} diff --git a/src/bin/unirust_rebalance.rs b/src/bin/unirust_rebalance.rs index 99cc9c0..4daedb4 100644 --- a/src/bin/unirust_rebalance.rs +++ b/src/bin/unirust_rebalance.rs @@ -1,11 +1,14 @@ +use std::path::PathBuf; use tokio_stream::wrappers::ReceiverStream; use tokio_stream::StreamExt; +use tonic::transport::{Channel, ClientTlsConfig, Endpoint}; use unirust_rs::distributed::proto::router_service_client::RouterServiceClient; use unirust_rs::distributed::proto::shard_service_client::ShardServiceClient; use unirust_rs::distributed::proto::{ ExportRecordsRequest, ImportRecordsChunk, ImportRecordsRequest, RecordIdRangeRequest, RouterExportRecordsRequest, RouterImportRecordsRequest, RouterRecordIdRangeRequest, }; +use unirust_rs::transport_security::load_client_mtls; fn parse_arg(flag: &str) -> Option { let mut args = std::env::args(); @@ -21,17 +24,39 @@ fn has_flag(flag: &str) -> bool { std::env::args().any(|arg| arg == flag) } -fn normalize_addr(addr: impl AsRef) -> String { +fn normalize_addr(addr: impl AsRef, mtls: bool) -> anyhow::Result { let addr = addr.as_ref(); if addr.starts_with("http://") || addr.starts_with("https://") { - addr.to_string() + if mtls && addr.starts_with("http://") { + anyhow::bail!("mTLS rebalance connections require https:// endpoints"); + } + if !mtls && addr.starts_with("https://") { + anyhow::bail!("https:// endpoints require --tls-ca, --tls-cert, and --tls-key"); + } + Ok(addr.to_string()) + } else if mtls { + Ok(format!("https://{addr}")) } else { - format!("http://{}", addr) + Ok(format!("http://{addr}")) } } +async fn connect(addr: &str, mtls: Option<&ClientTlsConfig>) -> anyhow::Result { + let endpoint = Endpoint::from_shared(normalize_addr(addr, mtls.is_some())?)?; + let endpoint = if let Some(mtls) = mtls { + endpoint.tls_config(mtls.clone())? + } else { + endpoint + }; + Ok(endpoint.connect().await?) +} + #[tokio::main] async fn main() -> anyhow::Result<()> { + let tls_ca = parse_arg("--tls-ca").map(PathBuf::from); + let tls_cert = parse_arg("--tls-cert").map(PathBuf::from); + let tls_key = parse_arg("--tls-key").map(PathBuf::from); + let mtls = load_client_mtls(tls_ca.as_deref(), tls_cert.as_deref(), tls_key.as_deref())?; let router = parse_arg("--router"); if has_flag("--range") { let shard_arg = parse_arg("--shard"); @@ -41,7 +66,8 @@ async fn main() -> anyhow::Result<()> { .as_deref() .ok_or_else(|| anyhow::anyhow!("--shard-id is required with --router"))? .parse()?; - let mut client = RouterServiceClient::connect(normalize_addr(router)).await?; + let channel = connect(&router, mtls.as_ref()).await?; + let mut client = RouterServiceClient::new(channel); client .get_record_id_range(RouterRecordIdRangeRequest { shard_id }) .await? @@ -50,7 +76,8 @@ async fn main() -> anyhow::Result<()> { let shard = shard_arg .as_deref() .ok_or_else(|| anyhow::anyhow!("--shard is required for --range"))?; - let mut client = ShardServiceClient::connect(normalize_addr(shard)).await?; + let channel = connect(shard, mtls.as_ref()).await?; + let mut client = ShardServiceClient::new(channel); client .get_record_id_range(RecordIdRangeRequest {}) .await? @@ -105,19 +132,25 @@ async fn main() -> anyhow::Result<()> { let use_router = router.is_some(); let mut router_client = if let Some(router) = router { - Some(RouterServiceClient::connect(normalize_addr(router)).await?) + Some(RouterServiceClient::new( + connect(&router, mtls.as_ref()).await?, + )) } else { None }; let mut source_client = if use_router { None } else { - Some(ShardServiceClient::connect(normalize_addr(&source)).await?) + Some(ShardServiceClient::new( + connect(&source, mtls.as_ref()).await?, + )) }; let mut target_client = if use_router { None } else { - Some(ShardServiceClient::connect(normalize_addr(&target)).await?) + Some(ShardServiceClient::new( + connect(&target, mtls.as_ref()).await?, + )) }; let source_shard_id = if use_router { Some(source.parse::()?) diff --git a/src/bin/unirust_router.rs b/src/bin/unirust_router.rs index 1925326..95ca00a 100644 --- a/src/bin/unirust_router.rs +++ b/src/bin/unirust_router.rs @@ -7,6 +7,9 @@ use unirust_rs::distributed::{ }; use unirust_rs::transport_security::{load_client_mtls, load_server_mtls}; +const MAX_GRPC_MESSAGE_BYTES: usize = 4 * 1024 * 1024; +const MAX_CONCURRENT_REQUESTS_PER_CONNECTION: usize = 128; + fn parse_arg(flag: &str) -> Option { let mut args = std::env::args(); while let Some(arg) = args.next() { @@ -55,11 +58,27 @@ ENVIRONMENT: UNIRUST_ROUTER_SHARDS Comma-separated shard addresses UNIRUST_ROUTER_CHECKPOINT_INTERVAL_SECS Automatic coordinated checkpoint interval + UNIRUST_ROUTER_TLS_CERT PEM server certificate + UNIRUST_ROUTER_TLS_KEY PEM server private key + UNIRUST_ROUTER_TLS_CLIENT_CA + PEM CA for required client certificates + UNIRUST_ROUTER_SHARD_TLS_CA + PEM CA used to verify shards + UNIRUST_ROUTER_SHARD_TLS_CERT + PEM client certificate presented to shards + UNIRUST_ROUTER_SHARD_TLS_KEY + PEM router client private key CONFIG FILE (unirust.toml): [router] listen = "0.0.0.0:50060" - shards = ["shard-0:50061", "shard-1:50061", "shard-2:50061", "shard-3:50061", "shard-4:50061"] + shards = ["https://shard-0:50061", "https://shard-1:50061"] + tls_cert = "/etc/unirust/tls/router.crt" + tls_key = "/etc/unirust/tls/router.key" + tls_client_ca = "/etc/unirust/tls/clients-ca.crt" + shard_tls_ca = "/etc/unirust/tls/shards-ca.crt" + shard_tls_cert = "/etc/unirust/tls/router-client.crt" + shard_tls_key = "/etc/unirust/tls/router-client.key" "# ); } @@ -239,14 +258,18 @@ async fn main() -> anyhow::Result<()> { }; println!("Unirust router listening on {}", config.router.listen); - let mut server = Server::builder(); + let mut server = Server::builder() + .concurrency_limit_per_connection(MAX_CONCURRENT_REQUESTS_PER_CONNECTION) + .load_shed(true); if let Some(server_mtls) = server_mtls { server = server.tls_config(server_mtls)?; } server - .add_service(proto::router_service_server::RouterServiceServer::new( - router, - )) + .add_service( + proto::router_service_server::RouterServiceServer::new(router) + .max_decoding_message_size(MAX_GRPC_MESSAGE_BYTES) + .max_encoding_message_size(MAX_GRPC_MESSAGE_BYTES), + ) .serve_with_shutdown(config.router.listen, shutdown_signal()) .await?; if let Some(checkpoint_task) = checkpoint_task { diff --git a/src/bin/unirust_shard.rs b/src/bin/unirust_shard.rs index 4d4f11e..1c2cb1a 100644 --- a/src/bin/unirust_shard.rs +++ b/src/bin/unirust_shard.rs @@ -1,11 +1,16 @@ use std::fs; -use tonic::transport::Server; +use std::time::Duration; + +use tonic::transport::{Endpoint, Server}; use unirust_rs::config::{ConfigOverrides, Profile, ShardOverrides, UniConfig}; use unirust_rs::distributed::{proto, DistributedOntologyConfig, ShardNode}; -use unirust_rs::transport_security::load_server_mtls; +use unirust_rs::transport_security::{load_client_mtls, load_replication_token, load_server_mtls}; use unirust_rs::{restore_checkpoint_for_shard, StreamingTuning}; +const MAX_GRPC_MESSAGE_BYTES: usize = 4 * 1024 * 1024; +const MAX_CONCURRENT_REQUESTS_PER_CONNECTION: usize = 128; + fn parse_arg(flag: &str) -> Option { let mut args = std::env::args(); while let Some(arg) = args.next() { @@ -41,12 +46,26 @@ OPTIONS: billion-scale-high-performance --repair Run repair on startup --ephemeral Allow an in-memory shard; all data is lost on process exit + --allow-colocated-checkpoints + Permit checkpoints without an independent path (development only) --allow-destructive-admin Enable destructive admin RPCs such as Reset --tls-cert PEM shard server certificate --tls-key PEM shard server private key --tls-client-ca PEM CA used to require router client certificates + --replica Passive replica endpoint for synchronous replication + --replica-mode Run as a passive replica; routers reject this endpoint + --allow-insecure-replication + Permit plaintext replication for isolated development + --replication-token-file + Shared secret file (at least 32 bytes) for this replica pair + --replica-tls-ca + PEM CA used to verify the replica + --replica-tls-cert + PEM primary client certificate presented to the replica + --replica-tls-key + PEM primary client private key --config-version Config version for compatibility checking -h, --help Print help @@ -58,6 +77,23 @@ ENVIRONMENT: UNIRUST_SHARD_DATA_DIR Data directory UNIRUST_SHARD_BACKUP_DIR Checkpoint root + UNIRUST_SHARD_TLS_CERT PEM server certificate + UNIRUST_SHARD_TLS_KEY PEM server private key + UNIRUST_SHARD_TLS_CLIENT_CA + PEM CA for required client certificates + UNIRUST_SHARD_REPLICA Passive replica endpoint + UNIRUST_SHARD_REPLICA_MODE + Run as a passive replica + UNIRUST_SHARD_ALLOW_INSECURE_REPLICATION + Permit plaintext replication for isolated development + UNIRUST_SHARD_REPLICATION_TOKEN_FILE + Shared secret file for the primary/replica pair + UNIRUST_SHARD_REPLICA_TLS_CA + PEM CA used to verify the replica + UNIRUST_SHARD_REPLICA_TLS_CERT + PEM primary certificate presented to the replica + UNIRUST_SHARD_REPLICA_TLS_KEY + PEM primary client private key CONFIG FILE (unirust.toml): profile = "billion-scale-high-performance" @@ -67,6 +103,14 @@ CONFIG FILE (unirust.toml): id = 0 data_dir = "/var/lib/unirust" backup_dir = "/var/backups/unirust/shard-0" + tls_cert = "/etc/unirust/tls/shard-0.crt" + tls_key = "/etc/unirust/tls/shard-0.key" + tls_client_ca = "/etc/unirust/tls/clients-ca.crt" + replica = "https://shard-0-replica:50061" + replication_token_file = "/etc/unirust/replication/shard-0.token" + replica_tls_ca = "/etc/unirust/tls/replicas-ca.crt" + replica_tls_cert = "/etc/unirust/tls/shard-0-primary.crt" + replica_tls_key = "/etc/unirust/tls/shard-0-primary.key" "# ); } @@ -166,6 +210,27 @@ async fn main() -> anyhow::Result<()> { if let Some(path) = parse_arg("--tls-client-ca") { shard_overrides.tls_client_ca = Some(path.into()); } + if let Some(replica) = parse_arg("--replica") { + shard_overrides.replica = Some(replica); + } + if has_flag("--replica-mode") { + shard_overrides.replica_mode = Some(true); + } + if has_flag("--allow-insecure-replication") { + shard_overrides.allow_insecure_replication = Some(true); + } + if let Some(path) = parse_arg("--replication-token-file") { + shard_overrides.replication_token_file = Some(path.into()); + } + if let Some(path) = parse_arg("--replica-tls-ca") { + shard_overrides.replica_tls_ca = Some(path.into()); + } + if let Some(path) = parse_arg("--replica-tls-cert") { + shard_overrides.replica_tls_cert = Some(path.into()); + } + if let Some(path) = parse_arg("--replica-tls-key") { + shard_overrides.replica_tls_key = Some(path.into()); + } if shard_overrides.listen.is_some() || shard_overrides.id.is_some() @@ -176,6 +241,13 @@ async fn main() -> anyhow::Result<()> { || shard_overrides.tls_cert.is_some() || shard_overrides.tls_key.is_some() || shard_overrides.tls_client_ca.is_some() + || shard_overrides.replica.is_some() + || shard_overrides.replica_mode.is_some() + || shard_overrides.allow_insecure_replication.is_some() + || shard_overrides.replication_token_file.is_some() + || shard_overrides.replica_tls_ca.is_some() + || shard_overrides.replica_tls_cert.is_some() + || shard_overrides.replica_tls_key.is_some() { overrides.shard = Some(shard_overrides); } @@ -190,12 +262,45 @@ async fn main() -> anyhow::Result<()> { config.shard.tls_key.as_deref(), config.shard.tls_client_ca.as_deref(), )?; + let replica_mtls = load_client_mtls( + config.shard.replica_tls_ca.as_deref(), + config.shard.replica_tls_cert.as_deref(), + config.shard.replica_tls_key.as_deref(), + )?; + let replication_token = config + .shard + .replication_token_file + .as_deref() + .map(load_replication_token) + .transpose()?; if config.shard.data_dir.is_none() && !has_flag("--ephemeral") { anyhow::bail!( "persistent shard storage is required; configure --data-dir (or use --ephemeral \ explicitly for disposable development data)" ); } + if let Some(data_dir) = config.shard.data_dir.as_deref() { + if !has_flag("--allow-colocated-checkpoints") { + let backup_dir = config.shard.backup_dir.as_deref().ok_or_else(|| { + anyhow::anyhow!( + "persistent shards require --backup-dir on independent storage (or use \ + --allow-colocated-checkpoints explicitly for development)" + ) + })?; + let data_dir = std::path::absolute(data_dir)?; + let backup_dir = std::path::absolute(backup_dir)?; + if data_dir == backup_dir + || data_dir.starts_with(&backup_dir) + || backup_dir.starts_with(&data_dir) + { + anyhow::bail!( + "shard data and checkpoint paths must not overlap: data={}, backup={}", + data_dir.display(), + backup_dir.display() + ); + } + } + } if let Some(source) = parse_arg("--restore-from") { let data_dir = config .shard @@ -222,7 +327,7 @@ async fn main() -> anyhow::Result<()> { let config_version = parse_arg("--config-version").or(config.shard.config_version.clone()); // Create shard node - let shard = ShardNode::new_with_storage_paths( + let mut shard = ShardNode::new_with_storage_paths( config.shard.id as u32, ontology, tuning, @@ -230,21 +335,74 @@ async fn main() -> anyhow::Result<()> { config.shard.backup_dir.clone(), config.shard.repair, config_version, - )? - .with_destructive_admin(has_flag("--allow-destructive-admin")); + )?; + if config.shard.replica_mode { + shard = shard.into_replica( + replication_token + .clone() + .ok_or_else(|| anyhow::anyhow!("replica mode requires a replication token"))?, + )?; + } + if let Some(replica) = &config.shard.replica { + let replica = if replica.starts_with("http://") || replica.starts_with("https://") { + replica.clone() + } else if replica_mtls.is_some() { + format!("https://{replica}") + } else { + format!("http://{replica}") + }; + let uses_https = replica.starts_with("https://"); + let endpoint = Endpoint::from_shared(replica)? + .connect_timeout(Duration::from_secs( + config.shard.replica_connect_timeout_secs, + )) + .timeout(Duration::from_secs( + config.shard.replica_request_timeout_secs, + )) + .tcp_keepalive(Some(Duration::from_secs( + config.shard.replica_tcp_keepalive_secs, + ))); + let endpoint = match (replica_mtls, uses_https) { + (Some(tls), true) => endpoint.tls_config(tls)?, + (Some(_), false) => { + anyhow::bail!("shard-to-replica mTLS requires an https:// replica endpoint"); + } + (None, true) => { + anyhow::bail!( + "https:// replica endpoints require explicit replica TLS certificate \ + configuration" + ); + } + (None, false) => endpoint, + }; + let channel = endpoint.connect().await?; + shard = shard + .with_replica( + proto::shard_service_client::ShardServiceClient::new(channel), + replication_token + .clone() + .ok_or_else(|| anyhow::anyhow!("replication requires a shared token"))?, + ) + .await?; + } + let shard = shard.with_destructive_admin(has_flag("--allow-destructive-admin")); println!( "Unirust shard {} listening on {}", config.shard.id, config.shard.listen ); - let mut server = Server::builder(); + let mut server = Server::builder() + .concurrency_limit_per_connection(MAX_CONCURRENT_REQUESTS_PER_CONNECTION) + .load_shed(true); if let Some(server_mtls) = server_mtls { server = server.tls_config(server_mtls)?; } server - .add_service(proto::shard_service_server::ShardServiceServer::new( - shard.clone(), - )) + .add_service( + proto::shard_service_server::ShardServiceServer::new(shard.clone()) + .max_decoding_message_size(MAX_GRPC_MESSAGE_BYTES) + .max_encoding_message_size(MAX_GRPC_MESSAGE_BYTES), + ) .serve_with_shutdown(config.shard.listen, shutdown_signal()) .await?; shard.shutdown().await?; diff --git a/src/config/mod.rs b/src/config/mod.rs index d527917..cf25426 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -44,6 +44,16 @@ fn config_env_path(key: &str) -> Option<&'static str> { "SHARD_TLS_CERT" => Some("shard.tls_cert"), "SHARD_TLS_KEY" => Some("shard.tls_key"), "SHARD_TLS_CLIENT_CA" => Some("shard.tls_client_ca"), + "SHARD_REPLICA" => Some("shard.replica"), + "SHARD_REPLICA_MODE" => Some("shard.replica_mode"), + "SHARD_ALLOW_INSECURE_REPLICATION" => Some("shard.allow_insecure_replication"), + "SHARD_REPLICATION_TOKEN_FILE" => Some("shard.replication_token_file"), + "SHARD_REPLICA_CONNECT_TIMEOUT_SECS" => Some("shard.replica_connect_timeout_secs"), + "SHARD_REPLICA_REQUEST_TIMEOUT_SECS" => Some("shard.replica_request_timeout_secs"), + "SHARD_REPLICA_TCP_KEEPALIVE_SECS" => Some("shard.replica_tcp_keepalive_secs"), + "SHARD_REPLICA_TLS_CA" => Some("shard.replica_tls_ca"), + "SHARD_REPLICA_TLS_CERT" => Some("shard.replica_tls_cert"), + "SHARD_REPLICA_TLS_KEY" => Some("shard.replica_tls_key"), "ROUTER_LISTEN" => Some("router.listen"), "ROUTER_SHARDS_FILE" => Some("router.shards_file"), "ROUTER_ONTOLOGY" => Some("router.ontology"), @@ -192,6 +202,68 @@ impl UniConfig { config.shard.tls_client_ca.as_ref(), ], )?; + validate_mtls_group( + "shard-to-replica mTLS", + [ + config.shard.replica_tls_ca.as_ref(), + config.shard.replica_tls_cert.as_ref(), + config.shard.replica_tls_key.as_ref(), + ], + )?; + if config.shard.replica_mode && config.shard.replica.is_some() { + return Err(ConfigError { + message: "a passive replica cannot configure another replica".to_string(), + }); + } + if (config.shard.replica_mode || config.shard.replica.is_some()) + && config.shard.data_dir.is_none() + { + return Err(ConfigError { + message: "shard replication requires persistent data_dir storage".to_string(), + }); + } + if (config.shard.replica_mode || config.shard.replica.is_some()) + && config.shard.replication_token_file.is_none() + { + return Err(ConfigError { + message: "shard replication requires replication_token_file".to_string(), + }); + } + if config.shard.replica.is_some() + && config.shard.replica_tls_ca.is_none() + && !config.shard.allow_insecure_replication + { + return Err(ConfigError { + message: "primary-to-replica transport requires mTLS; use \ + allow_insecure_replication only for isolated development" + .to_string(), + }); + } + if config.shard.replica_mode + && config.shard.tls_cert.is_none() + && !config.shard.allow_insecure_replication + { + return Err(ConfigError { + message: "replica mode requires shard server mTLS; use \ + allow_insecure_replication only for isolated development" + .to_string(), + }); + } + if config.shard.replica.is_some() + && (config.shard.replica_connect_timeout_secs == 0 + || config.shard.replica_request_timeout_secs == 0 + || config.shard.replica_tcp_keepalive_secs == 0) + { + return Err(ConfigError { + message: "shard replica transport timeouts must be greater than zero".to_string(), + }); + } + let replica_tls_configured = config.shard.replica_tls_ca.is_some(); + if config.shard.replica.is_none() && replica_tls_configured { + return Err(ConfigError { + message: "shard replica TLS credentials require a replica endpoint".to_string(), + }); + } validate_mtls_group( "router mTLS", [ @@ -277,6 +349,26 @@ pub struct ShardConfig { pub tls_key: Option, /// PEM CA used to require and verify client certificates pub tls_client_ca: Option, + /// Passive replica endpoint for synchronous mutation replication + pub replica: Option, + /// Run as a passive replica that routers must not target + pub replica_mode: bool, + /// Permit a plaintext transport in isolated development only + pub allow_insecure_replication: bool, + /// File containing at least 32 bytes shared by one primary/replica pair + pub replication_token_file: Option, + /// Maximum time to establish the replica connection + pub replica_connect_timeout_secs: u64, + /// Maximum time for one replica mutation + pub replica_request_timeout_secs: u64, + /// TCP keepalive interval for the replica connection + pub replica_tcp_keepalive_secs: u64, + /// PEM CA used to verify the replica server + pub replica_tls_ca: Option, + /// PEM primary client certificate presented to the replica + pub replica_tls_cert: Option, + /// PEM primary client private key + pub replica_tls_key: Option, } impl Default for ShardConfig { @@ -292,6 +384,16 @@ impl Default for ShardConfig { tls_cert: None, tls_key: None, tls_client_ca: None, + replica: None, + replica_mode: false, + allow_insecure_replication: false, + replication_token_file: None, + replica_connect_timeout_secs: DEFAULT_SHARD_CONNECT_TIMEOUT_SECS, + replica_request_timeout_secs: DEFAULT_SHARD_REQUEST_TIMEOUT_SECS, + replica_tcp_keepalive_secs: DEFAULT_SHARD_TCP_KEEPALIVE_SECS, + replica_tls_ca: None, + replica_tls_cert: None, + replica_tls_key: None, } } } @@ -437,6 +539,20 @@ pub struct ShardOverrides { pub tls_key: Option, #[serde(skip_serializing_if = "Option::is_none")] pub tls_client_ca: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub replica: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub replica_mode: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub allow_insecure_replication: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub replication_token_file: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub replica_tls_ca: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub replica_tls_cert: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub replica_tls_key: Option, } #[derive(Debug, Clone, Default, Serialize, Deserialize)] diff --git a/src/distributed.rs b/src/distributed.rs index a362987..d4ee4b4 100644 --- a/src/distributed.rs +++ b/src/distributed.rs @@ -10,7 +10,7 @@ use crate::persistence::{ CHECKPOINT_PROTOCOL_VERSION, }; use crate::query::{QueryDescriptor, QueryOutcome}; -use crate::sharding::{BloomFilter, IdentityKeySignature}; +use crate::sharding::{BloomFilter, IdentityKeySignature, ReconciliationCandidates}; use crate::store::{SourceRecordReservation, SourceReservationError, StoreMetrics}; use crate::temporal::Interval; use crate::{PersistentStore, StreamingTuning, Unirust}; @@ -31,9 +31,10 @@ use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::{Arc, RwLock as StdRwLock}; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use tokio::sync::{mpsc, oneshot}; -use tokio::sync::{Mutex, RwLock}; +use tokio::sync::{Mutex, RwLock, Semaphore}; use tokio_stream::wrappers::ReceiverStream; use tokio_stream::{Stream, StreamExt}; +use tonic::transport::Channel; use tonic::{Request, Response, Status}; /// WAL record format for binary serialization. @@ -67,7 +68,8 @@ 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 = 3; +pub const DISTRIBUTED_PROTOCOL_VERSION: u32 = 5; +const REPLICATION_TOKEN_HEADER: &str = "x-unirust-replication-token"; #[allow(clippy::result_large_err)] fn validate_distributed_protocol(protocol_version: u32) -> Result<(), Status> { @@ -926,6 +928,20 @@ fn global_cluster_id_to_proto(id: GlobalClusterId) -> proto::GlobalClusterId { } } +#[allow(clippy::result_large_err)] +fn global_cluster_id_from_proto( + id: &proto::GlobalClusterId, + field: &str, +) -> Result { + Ok(GlobalClusterId::new( + u16::try_from(id.shard_id) + .map_err(|_| Status::invalid_argument(format!("{field} shard_id exceeds u16")))?, + id.local_id, + u16::try_from(id.version) + .map_err(|_| Status::invalid_argument(format!("{field} version exceeds u16")))?, + )) +} + fn cross_shard_conflict_to_proto( conflict: &crate::sharding::CrossShardConflict, ) -> proto::CrossShardConflict { @@ -1045,11 +1061,18 @@ pub struct ShardNode { checkpoint_root: Option, ingest_wal: Option>, ingest_wal_lock: Arc>, + ingest_commit_slots: Arc, mutation_gate: Arc>, ingest_txs: Vec>, ingest_worker_handles: Arc>>>, config_version: String, metrics: Arc, + shard_role: proto::ShardRole, + replica_client: Option>, + replication_token: Option, + replication_consistent: Arc, + mutation_consistent: Arc, + replication_gate: Arc>, /// Destructive RPCs are disabled unless explicitly enabled by the embedding process. allow_destructive_admin: bool, /// Cross-shard conflicts detected during reconciliation. @@ -1061,6 +1084,10 @@ pub struct ShardNode { const INGEST_QUEUE_CAPACITY: usize = 1024; // Increased from 128 for high-throughput const DEFAULT_INGEST_WORKERS: usize = 32; // Increased from 16 for high-throughput const EXPORT_DEFAULT_LIMIT: usize = 1000; +const DIRTY_KEY_PAGE_LIMIT: usize = 50_000; +const RECONCILIATION_KEY_CHUNK: usize = 10_000; +const MERGE_APPLICATION_CHUNK: usize = 50_000; +const CONFLICT_APPLICATION_CHUNK: usize = 10_000; /// Check if partitioned processing is enabled (default: true) /// Set UNIRUST_PARTITIONED=0 to disable @@ -1100,6 +1127,159 @@ struct IngestJob { respond_to: oneshot::Sender, Status>>, } +struct ReplicationAttempt { + consistency: Arc, + armed: bool, +} + +struct DurableMutationAttempt { + consistency: Arc, + armed: bool, +} + +struct ConsistencyAttempt { + consistency: Arc, + armed: bool, +} + +impl DurableMutationAttempt { + fn new(consistency: Arc, armed: bool) -> Self { + Self { consistency, armed } + } + + fn finish(mut self) { + self.armed = false; + } +} + +impl Drop for DurableMutationAttempt { + fn drop(&mut self) { + if self.armed { + self.consistency.store(false, Ordering::Release); + } + } +} + +impl ConsistencyAttempt { + fn new(consistency: Arc) -> Self { + Self { + consistency, + armed: true, + } + } + + fn finish(mut self) { + self.consistency.store(true, Ordering::Release); + self.armed = false; + } +} + +impl Drop for ConsistencyAttempt { + fn drop(&mut self) { + if self.armed { + self.consistency.store(false, Ordering::Release); + } + } +} + +impl ReplicationAttempt { + fn new(consistency: Arc, armed: bool) -> Self { + Self { consistency, armed } + } + + fn finish(mut self, responses_match: bool, operation: &str) -> Result<(), Status> { + if !responses_match { + return Err(Status::data_loss(format!( + "primary and replica returned different results for {operation}; traffic is \ + blocked" + ))); + } + self.armed = false; + Ok(()) + } +} + +impl Drop for ReplicationAttempt { + fn drop(&mut self) { + if self.armed { + self.consistency.store(false, Ordering::Release); + } + } +} + +fn replication_error_is_ambiguous(code: tonic::Code) -> bool { + matches!( + code, + tonic::Code::Cancelled + | tonic::Code::Unknown + | tonic::Code::DeadlineExceeded + | tonic::Code::Aborted + | tonic::Code::Internal + | tonic::Code::Unavailable + | tonic::Code::DataLoss + ) +} + +fn authenticated_replica_request( + payload: T, + replication_token: &str, +) -> Result, Status> { + let token = tonic::metadata::MetadataValue::try_from(replication_token) + .map_err(|_| Status::internal("replication token metadata is invalid"))?; + let mut request = Request::new(payload); + request + .metadata_mut() + .insert(REPLICATION_TOKEN_HEADER, token); + Ok(request) +} + +macro_rules! replicate_to_replica { + ($self:expr, $method:ident, $payload:expr) => {{ + $self.ensure_replication_consistent()?; + let replication_serial = if $self.replica_client.is_some() { + Some($self.replication_gate.lock().await) + } else { + None + }; + $self.ensure_replication_consistent()?; + let replication_attempt = ReplicationAttempt::new( + $self.replication_consistent.clone(), + $self.replica_client.is_some(), + ); + let replica_response = if let Some(mut replica) = $self.replica_client.clone() { + let token = $self.replication_token.as_deref().ok_or_else(|| { + Status::failed_precondition("primary replication token is missing") + })?; + let replica_request = authenticated_replica_request(($payload).clone(), token)?; + match replica.$method(replica_request).await { + Ok(response) => Some(response.into_inner()), + Err(error) => { + if replication_error_is_ambiguous(error.code()) { + $self.replication_consistent.store(false, Ordering::Release); + return Err(Status::aborted(format!( + "replica {operation} failed and may have committed; traffic is \ + blocked: {error}", + operation = stringify!($method) + ))); + } + $self.replication_consistent.store(false, Ordering::Release); + return Err(Status::new( + error.code(), + format!( + "replica rejected {operation}: {}", + error.message(), + operation = stringify!($method) + ), + )); + } + } + } else { + None + }; + (replication_serial, replica_response, replication_attempt) + }}; +} + impl ShardNode { pub fn new( shard_id: u32, @@ -1257,11 +1437,18 @@ impl ShardNode { checkpoint_root: Some(checkpoint_root), ingest_wal, ingest_wal_lock: Arc::new(Mutex::new(())), + ingest_commit_slots: Arc::new(Semaphore::new(1)), mutation_gate: Arc::new(tokio::sync::RwLock::new(())), ingest_txs, ingest_worker_handles: Arc::new(Mutex::new(ingest_worker_handles)), config_version, metrics: Arc::new(PerfMetrics::new()), + shard_role: proto::ShardRole::Standalone, + replica_client: None, + replication_token: None, + replication_consistent: Arc::new(AtomicBool::new(true)), + mutation_consistent: Arc::new(AtomicBool::new(true)), + replication_gate: Arc::new(Mutex::new(())), allow_destructive_admin: false, cross_shard_conflicts: Arc::new(parking_lot::RwLock::new( recovered_cross_shard_conflicts, @@ -1317,11 +1504,18 @@ impl ShardNode { checkpoint_root: None, ingest_wal, ingest_wal_lock: Arc::new(Mutex::new(())), + ingest_commit_slots: Arc::new(Semaphore::new(INGEST_QUEUE_CAPACITY)), mutation_gate: Arc::new(tokio::sync::RwLock::new(())), ingest_txs, ingest_worker_handles: Arc::new(Mutex::new(ingest_worker_handles)), config_version, metrics: Arc::new(PerfMetrics::new()), + shard_role: proto::ShardRole::Standalone, + replica_client: None, + replication_token: None, + replication_consistent: Arc::new(AtomicBool::new(true)), + mutation_consistent: Arc::new(AtomicBool::new(true)), + replication_gate: Arc::new(Mutex::new(())), allow_destructive_admin: false, cross_shard_conflicts: Arc::new(parking_lot::RwLock::new(Vec::new())), }) @@ -1336,6 +1530,287 @@ impl ShardNode { self } + /// Configure this persistent node as a passive replica. A replica can only + /// be targeted by a primary and is rejected as a router shard endpoint. + pub fn into_replica(mut self, replication_token: String) -> AnyResult { + if self.data_dir.is_none() { + anyhow::bail!("replica mode requires persistent shard storage"); + } + if self.replica_client.is_some() { + anyhow::bail!("a replica cannot have another replica"); + } + if replication_token.is_empty() { + anyhow::bail!("replica mode requires a nonempty replication token"); + } + self.shard_role = proto::ShardRole::Replica; + self.replication_token = Some(replication_token); + Ok(self) + } + + /// Attach a passive replica after proving that its topology, configuration, + /// restore provenance, and complete durable key/value state match. + pub async fn with_replica( + mut self, + mut replica: proto::shard_service_client::ShardServiceClient, + replication_token: String, + ) -> Result { + if self.data_dir.is_none() { + return Err(Status::failed_precondition( + "synchronous replication requires persistent primary storage", + )); + } + if self.shard_role != proto::ShardRole::Standalone || self.replica_client.is_some() { + return Err(Status::failed_precondition( + "replication can only be attached to a standalone shard", + )); + } + if replication_token.is_empty() { + return Err(Status::invalid_argument( + "replication requires a nonempty shared token", + )); + } + let health = replica + .health_check(Request::new(proto::HealthCheckRequest {})) + .await + .map_err(|error| Status::unavailable(format!("replica health check failed: {error}")))? + .into_inner(); + if health.status != "ok" { + return Err(Status::unavailable(format!( + "replica reported unhealthy status {}", + health.status + ))); + } + let request = authenticated_replica_request( + proto::ConfigVersionRequest { + include_durable_state_digest: true, + }, + &replication_token, + )?; + let response = replica + .get_config_version(request) + .await + .map_err(|error| { + Status::new( + error.code(), + format!("replica configuration check failed: {}", error.message()), + ) + })? + .into_inner(); + validate_distributed_protocol(response.protocol_version)?; + validate_checkpoint_protocol(response.checkpoint_protocol_version)?; + if proto::ShardRole::try_from(response.shard_role) != Ok(proto::ShardRole::Replica) { + return Err(Status::failed_precondition( + "configured replication target is not running in replica mode", + )); + } + if response.version != self.config_version { + return Err(Status::failed_precondition(format!( + "replica config version mismatch: primary {}, replica {}", + self.config_version, response.version + ))); + } + let replica_ontology = response.ontology_config.ok_or_else(|| { + Status::failed_precondition("replica did not report its ontology configuration") + })?; + if map_proto_config(&replica_ontology) != *self.ontology_config.lock().await { + return Err(Status::failed_precondition( + "replica ontology does not match the primary", + )); + } + let local_restore = self + .restored_checkpoint + .as_ref() + .map(|manifest| (manifest.generation(), manifest.shard_count())); + let replica_restore = match ( + response.restore_generation.is_empty(), + response.restore_shard_count, + ) { + (true, 0) => None, + (false, shard_count) if shard_count > 0 => { + Some((response.restore_generation.as_str(), shard_count)) + } + _ => { + return Err(Status::data_loss( + "replica reported incomplete restore checkpoint provenance", + )); + } + }; + if local_restore != replica_restore { + return Err(Status::failed_precondition( + "primary and replica restore provenance do not match", + )); + } + let local_digest = self + .durable_state_digest()? + .ok_or_else(|| Status::failed_precondition("primary does not expose durable state"))?; + if response.durable_state_digest.as_slice() != local_digest { + return Err(Status::failed_precondition( + "replica durable state does not match the primary; bootstrap both from the same \ + checkpoint before enabling replication", + )); + } + let metadata = replica + .get_boundary_metadata(Request::new(proto::GetBoundaryMetadataRequest { + since_version: u64::MAX, + signatures: Vec::new(), + })) + .await + .map_err(|error| { + Status::unavailable(format!("replica shard identity check failed: {error}")) + })? + .into_inner() + .metadata + .ok_or_else(|| Status::data_loss("replica returned no boundary metadata"))?; + if metadata.shard_id != self.shard_id { + return Err(Status::failed_precondition(format!( + "replica reports shard {}, primary is shard {}", + metadata.shard_id, self.shard_id + ))); + } + self.replica_client = Some(replica); + self.replication_token = Some(replication_token); + self.shard_role = proto::ShardRole::Primary; + Ok(self) + } + + fn durable_state_digest(&self) -> Result, Status> { + let db = read_unirust!(self).store().shared_db(); + db.map(|db| { + crate::persistence::durable_state_digest(&db) + .map_err(|error| Status::internal(error.to_string())) + }) + .transpose() + } + + fn ensure_replication_consistent(&self) -> Result<(), Status> { + if !self.mutation_consistent.load(Ordering::Acquire) { + return Err(Status::failed_precondition( + "a durable mutation failed after it may have changed local state; traffic is \ + blocked until the shard restarts and recovers", + )); + } + if !self.replication_consistent.load(Ordering::Acquire) { + return Err(Status::failed_precondition( + "primary and replica may have diverged; traffic is blocked until their complete \ + durable states are reconciled", + )); + } + Ok(()) + } + + fn begin_durable_mutation(&self) -> DurableMutationAttempt { + DurableMutationAttempt::new(self.mutation_consistent.clone(), self.data_dir.is_some()) + } + + async fn commit_ingest_batch( + &self, + records: Vec, + operation: &'static str, + ) -> Result, Status> { + let _mutation_guard = self.mutation_gate.read().await; + let start = Instant::now(); + let record_count = records.len(); + let _wal_guard = if self.ingest_wal.is_some() { + Some(self.ingest_wal_lock.lock().await) + } else { + None + }; + self.ensure_store_healthy()?; + self.ensure_ingest_payloads_idempotent(&records)?; + + let replication_request = proto::IngestRecordsRequest { + records: records.clone(), + internal_protocol_version: DISTRIBUTED_PROTOCOL_VERSION, + }; + let (_replication_serial, replica_response, replication_attempt) = + replicate_to_replica!(self, ingest_records, replication_request); + let local_attempt = self.begin_durable_mutation(); + let partitioned_arc = self.partitioned.read().clone(); + let assignments = if let Some(partitioned) = + partitioned_arc.as_ref().filter(|_| records.len() >= 100) + { + dispatch_ingest_partitioned( + partitioned, + &self.concurrent_interner, + self.ingest_wal.as_deref(), + self.shard_id, + records, + ) + .await? + } else { + dispatch_ingest_records(&self.ingest_txs, self.ingest_wal.as_deref(), records).await? + }; + + local_attempt.finish(); + replication_attempt.finish( + replica_response + .as_ref() + .is_none_or(|replica| replica.assignments == assignments), + operation, + )?; + self.metrics + .record_ingest(record_count, start.elapsed().as_micros() as u64); + Ok(assignments) + } + + async fn await_ingest_commit( + &self, + records: Vec, + operation: &'static str, + ) -> Result, Status> { + let commit_slot = self + .ingest_commit_slots + .clone() + .acquire_owned() + .await + .map_err(|_| Status::unavailable("ingest commit queue is closed"))?; + let shard = self.clone(); + tokio::spawn(async move { + let _commit_slot = commit_slot; + shard.commit_ingest_batch(records, operation).await + }) + .await + .map_err(|error| { + if self.data_dir.is_some() { + self.mutation_consistent.store(false, Ordering::Release); + } + if self.replica_client.is_some() { + self.replication_consistent.store(false, Ordering::Release); + } + Status::internal(format!("ingest commit task failed: {error}")) + })? + } + + fn authorize_mutation(&self, request: &Request) -> Result<(), Status> { + let supplied = request + .metadata() + .get(REPLICATION_TOKEN_HEADER) + .and_then(|value| value.to_str().ok()); + match self.shard_role { + proto::ShardRole::Replica => { + if supplied == self.replication_token.as_deref() { + Ok(()) + } else { + Err(Status::permission_denied( + "passive replicas only accept mutations forwarded by their primary", + )) + } + } + proto::ShardRole::Standalone | proto::ShardRole::Primary => { + if supplied.is_some() { + Err(Status::permission_denied( + "replicated mutation metadata is only valid on a passive replica", + )) + } else { + Ok(()) + } + } + proto::ShardRole::Unspecified => { + Err(Status::failed_precondition("shard role is not configured")) + } + } + } + /// Drain mutations and durably flush derived state before process exit. pub async fn shutdown(&self) -> AnyResult<()> { let _mutation_guard = self.mutation_gate.write().await; @@ -1367,6 +1842,7 @@ impl ShardNode { } fn ensure_no_pending_ingest(&self) -> Result<(), Status> { + self.ensure_replication_consistent()?; if self .ingest_wal .as_ref() @@ -1389,6 +1865,7 @@ impl ShardNode { } fn ensure_recovery_healthy(&self) -> Result<(), Status> { + self.ensure_replication_consistent()?; if self .ingest_wal .as_ref() @@ -1408,6 +1885,7 @@ impl ShardNode { } fn ensure_store_healthy(&self) -> Result<(), Status> { + self.ensure_replication_consistent()?; let unirust = read_unirust!(self); unirust.ensure_store_healthy().map_err(|err| { Status::failed_precondition(format!( @@ -1570,19 +2048,31 @@ impl ShardNode { } #[allow(clippy::result_large_err)] - fn build_import_records( - unirust: &mut Unirust, + fn validate_import_records( + unirust: &Unirust, snapshots: &[proto::RecordSnapshot], - ) -> Result, Status> { - let mut records = Vec::with_capacity(snapshots.len()); + ) -> Result, Status> { + let mut records_to_build = Vec::with_capacity(snapshots.len()); let mut batch_ids: HashMap = HashMap::new(); let mut batch_identities: HashMap<(String, String, String), u32> = HashMap::new(); - for snapshot in snapshots { + for (index, snapshot) in snapshots.iter().enumerate() { let identity = snapshot .identity .as_ref() .ok_or_else(|| Status::invalid_argument("record identity is required"))?; + if identity.entity_type.is_empty() + || identity.perspective.is_empty() + || identity.uid.is_empty() + { + return Err(Status::invalid_argument( + "record identity fields must not be empty", + )); + } + for descriptor in &snapshot.descriptors { + Interval::new(descriptor.start, descriptor.end) + .map_err(|err| Status::invalid_argument(err.to_string()))?; + } if let Some(previous) = batch_ids.get(&snapshot.record_id) { if **previous != *snapshot { return Err(Status::already_exists(format!( @@ -1628,6 +2118,27 @@ impl ShardNode { existing_id.0, snapshot.record_id ))); } + records_to_build.push(index); + } + unirust + .ensure_store_healthy() + .map_err(|err| Status::failed_precondition(err.to_string()))?; + Ok(records_to_build) + } + + #[allow(clippy::result_large_err)] + fn build_import_records( + unirust: &mut Unirust, + snapshots: &[proto::RecordSnapshot], + records_to_build: &[usize], + ) -> Result, Status> { + let mut records = Vec::with_capacity(records_to_build.len()); + for &index in records_to_build { + let snapshot = &snapshots[index]; + let identity = snapshot + .identity + .as_ref() + .ok_or_else(|| Status::internal("validated record identity is missing"))?; records.push(Self::build_record_with_id( unirust, snapshot.record_id, @@ -1635,9 +2146,6 @@ impl ShardNode { &snapshot.descriptors, )?); } - unirust - .ensure_store_healthy() - .map_err(|err| Status::failed_precondition(err.to_string()))?; Ok(records) } @@ -1755,36 +2263,18 @@ async fn dispatch_ingest_records( } let worker_count = ingest_txs.len().max(1); - let mut batches = vec![Vec::new(); worker_count]; - for record in records { - let idx = ingest_worker_index(&record, worker_count); - batches[idx].push(record); - } - - let mut receivers = Vec::new(); - for (idx, batch) in batches.into_iter().enumerate() { - if batch.is_empty() { - continue; - } - let (tx, rx) = oneshot::channel(); - let job = IngestJob { - records: batch, + let worker_index = ingest_worker_index(&records[0], worker_count); + let (tx, rx) = oneshot::channel(); + ingest_txs[worker_index] + .send(IngestJob { + records, respond_to: tx, - }; - ingest_txs[idx] - .send(job) - .await - .map_err(|_| Status::unavailable("ingest queue unavailable"))?; - receivers.push(rx); - } - - let mut assignments = Vec::new(); - for rx in receivers { - let batch_assignments = rx - .await - .map_err(|_| Status::internal("ingest worker dropped"))??; - assignments.extend(batch_assignments); - } + }) + .await + .map_err(|_| Status::unavailable("ingest queue unavailable"))?; + let mut assignments = rx + .await + .map_err(|_| Status::internal("ingest worker dropped"))??; if let Some(wal) = ingest_wal { wal.clear()?; @@ -2037,9 +2527,11 @@ impl proto::shard_service_server::ShardService for ShardNode { &self, request: Request, ) -> Result, Status> { + self.authorize_mutation(&request)?; let _mutation_guard = self.mutation_gate.read().await; self.ensure_store_healthy()?; let request = request.into_inner(); + let replication_request = request.clone(); let mut reservations = Vec::with_capacity(request.reservations.len()); let mut indices = Vec::with_capacity(request.reservations.len()); @@ -2072,25 +2564,30 @@ impl proto::shard_service_server::ShardService for ShardNode { }); } + let (_replication_serial, replica_response, replication_attempt) = + replicate_to_replica!(self, reserve_source_records, replication_request); + let local_attempt = self.begin_durable_mutation(); let targets = { let mut unirust = write_unirust!(self); - unirust - .store_mut() - .reserve_source_records(&reservations) - .map_err(|err| { - if let Some(conflict) = err.downcast_ref::() { - match conflict { - SourceReservationError::PayloadConflict => { - Status::already_exists(conflict.to_string()) - } - SourceReservationError::TargetConflict { .. } => { - Status::failed_precondition(conflict.to_string()) - } + unirust.store_mut().reserve_source_records(&reservations) + }; + let targets = match targets { + Ok(targets) => targets, + Err(err) => { + if let Some(conflict) = err.downcast_ref::() { + // The store validates the complete batch before its atomic write. + local_attempt.finish(); + return Err(match conflict { + SourceReservationError::PayloadConflict => { + Status::already_exists(conflict.to_string()) } - } else { - Status::internal(err.to_string()) - } - })? + SourceReservationError::TargetConflict { .. } => { + Status::failed_precondition(conflict.to_string()) + } + }); + } + return Err(Status::internal(err.to_string())); + } }; let reservations = indices .into_iter() @@ -2102,15 +2599,22 @@ impl proto::shard_service_server::ShardService for ShardNode { }, ) .collect(); - Ok(Response::new(proto::ReserveSourceRecordsResponse { - reservations, - })) + let response = proto::ReserveSourceRecordsResponse { reservations }; + local_attempt.finish(); + replication_attempt.finish( + replica_response + .as_ref() + .is_none_or(|replica| replica == &response), + "source reservation", + )?; + Ok(Response::new(response)) } async fn mark_source_reservations_backfilled( &self, request: Request, ) -> Result, Status> { + self.authorize_mutation(&request)?; let _mutation_guard = self.mutation_gate.read().await; self.ensure_store_healthy()?; let request = request.into_inner(); @@ -2126,24 +2630,35 @@ impl proto::shard_service_server::ShardService for ShardNode { )); } + let (_replication_serial, replica_response, replication_attempt) = + replicate_to_replica!(self, mark_source_reservations_backfilled, request); + let local_attempt = self.begin_durable_mutation(); let mut unirust = write_unirust!(self); unirust .store_mut() .mark_source_reservation_backfill(request.protocol_version, request.shard_count) .map_err(|err| Status::internal(err.to_string()))?; - Ok(Response::new( - proto::MarkSourceReservationsBackfilledResponse {}, - )) + let response = proto::MarkSourceReservationsBackfilledResponse {}; + local_attempt.finish(); + replication_attempt.finish( + replica_response + .as_ref() + .is_none_or(|replica| replica == &response), + "source reservation migration marker", + )?; + Ok(Response::new(response)) } async fn set_ontology( &self, request: Request, ) -> Result, Status> { + self.authorize_mutation(&request)?; let _mutation_guard = self.mutation_gate.write().await; self.ensure_no_pending_ingest()?; + let request = request.into_inner(); + let replication_request = request.clone(); let config = request - .into_inner() .config .ok_or_else(|| Status::invalid_argument("ontology config is required"))?; @@ -2164,6 +2679,9 @@ impl proto::shard_service_server::ShardService for ShardNode { ))); } + let (_replication_serial, replica_response, replication_attempt) = + replicate_to_replica!(self, set_ontology, replication_request); + let local_attempt = self.begin_durable_mutation(); if let Some(path) = &self.data_dir { // Must acquire write lock and drop old Unirust BEFORE opening new store // to release the RocksDB file lock @@ -2227,47 +2745,27 @@ impl proto::shard_service_server::ShardService for ShardNode { } *config_guard = config; - Ok(Response::new(proto::ApplyOntologyResponse {})) + local_attempt.finish(); + let response = proto::ApplyOntologyResponse {}; + replication_attempt.finish( + replica_response + .as_ref() + .is_none_or(|replica| replica == &response), + "ontology update", + )?; + Ok(Response::new(response)) } async fn ingest_records( &self, request: Request, ) -> Result, Status> { - let _mutation_guard = self.mutation_gate.read().await; - let start = Instant::now(); + self.authorize_mutation(&request)?; let request = request.into_inner(); validate_distributed_protocol(request.internal_protocol_version)?; - let records = request.records; - let record_count = records.len(); - let _wal_guard = if self.ingest_wal.is_some() { - Some(self.ingest_wal_lock.lock().await) - } else { - None - }; - self.ensure_store_healthy()?; - self.ensure_ingest_payloads_idempotent(&records)?; - - // Use partitioned processing for large batches (high-throughput mode) - // Small batches use sequential path for correctness (query support) - // Clone the Arc and release the lock before awaiting (guards aren't Send) - let partitioned_arc = self.partitioned.read().clone(); - let use_partitioned = partitioned_arc.is_some() && records.len() >= 100; - let assignments = if use_partitioned { - dispatch_ingest_partitioned( - partitioned_arc.as_ref().unwrap(), - &self.concurrent_interner, - self.ingest_wal.as_deref(), - self.shard_id, - records, - ) - .await? - } else { - dispatch_ingest_records(&self.ingest_txs, self.ingest_wal.as_deref(), records).await? - }; - - self.metrics - .record_ingest(record_count, start.elapsed().as_micros() as u64); + let assignments = self + .await_ingest_commit(request.records, "record ingest") + .await?; Ok(Response::new(proto::IngestRecordsResponse { assignments })) } @@ -2275,11 +2773,9 @@ impl proto::shard_service_server::ShardService for ShardNode { &self, request: Request>, ) -> Result, Status> { - let _mutation_guard = self.mutation_gate.read().await; - let start = Instant::now(); + self.authorize_mutation(&request)?; let mut stream = request.into_inner(); let mut assignments = Vec::new(); - let mut record_count = 0usize; while let Some(chunk) = stream .message() @@ -2290,38 +2786,12 @@ impl proto::shard_service_server::ShardService for ShardNode { continue; } validate_distributed_protocol(chunk.internal_protocol_version)?; - record_count += chunk.records.len(); - let _wal_guard = if self.ingest_wal.is_some() { - Some(self.ingest_wal_lock.lock().await) - } else { - None - }; - self.ensure_store_healthy()?; - self.ensure_ingest_payloads_idempotent(&chunk.records)?; - - // Use partitioned processing for large batches (high-throughput mode) - // Small batches use sequential path for correctness - // Clone the Arc and release the lock before awaiting (guards aren't Send) - let partitioned_arc = self.partitioned.read().clone(); - let use_partitioned = partitioned_arc.is_some() && chunk.records.len() >= 100; - let batch_assignments = if use_partitioned { - dispatch_ingest_partitioned( - partitioned_arc.as_ref().unwrap(), - &self.concurrent_interner, - self.ingest_wal.as_deref(), - self.shard_id, - chunk.records, - ) - .await? - } else { - dispatch_ingest_records(&self.ingest_txs, self.ingest_wal.as_deref(), chunk.records) - .await? - }; + let batch_assignments = self + .await_ingest_commit(chunk.records, "streamed record ingest") + .await?; assignments.extend(batch_assignments); } - self.metrics - .record_ingest(record_count, start.elapsed().as_micros() as u64); Ok(Response::new(proto::IngestRecordsResponse { assignments })) } @@ -2499,6 +2969,21 @@ impl proto::shard_service_server::ShardService for ShardNode { _request: Request, ) -> Result, Status> { self.ensure_recovery_healthy()?; + if let Some(mut replica) = self.replica_client.clone() { + let response = replica + .health_check(Request::new(proto::HealthCheckRequest {})) + .await + .map_err(|error| { + Status::unavailable(format!("replica health check failed: {error}")) + })? + .into_inner(); + if response.status != "ok" { + return Err(Status::unavailable(format!( + "replica reported unhealthy status {}", + response.status + ))); + } + } Ok(Response::new(proto::HealthCheckResponse { status: "ok".to_string(), })) @@ -2506,8 +2991,17 @@ impl proto::shard_service_server::ShardService for ShardNode { async fn get_config_version( &self, - _request: Request, + request: Request, ) -> Result, Status> { + let include_durable_state_digest = request.get_ref().include_durable_state_digest; + if include_durable_state_digest && self.shard_role == proto::ShardRole::Replica { + self.authorize_mutation(&request)?; + } + let _mutation_guard = if include_durable_state_digest { + Some(self.mutation_gate.write().await) + } else { + None + }; let ontology_config = self.ontology_config.lock().await; let source_reservation_backfill = read_unirust!(self) .store() @@ -2534,6 +3028,14 @@ impl proto::shard_service_server::ShardService for ShardNode { .as_ref() .map(ClusterCheckpointManifest::shard_count) .unwrap_or_default(), + shard_role: self.shard_role as i32, + durable_state_digest: if include_durable_state_digest { + self.durable_state_digest()? + .map(|digest| digest.to_vec()) + .unwrap_or_default() + } else { + Vec::new() + }, })) } @@ -2562,6 +3064,7 @@ impl proto::shard_service_server::ShardService for ShardNode { &self, request: Request, ) -> Result, Status> { + self.authorize_mutation(&request)?; let _mutation_guard = self.mutation_gate.write().await; self.ensure_no_pending_ingest()?; let request = request.into_inner(); @@ -2598,6 +3101,8 @@ impl proto::shard_service_server::ShardService for ShardNode { )); } + let (_replication_serial, replica_response, replication_attempt) = + replicate_to_replica!(self, checkpoint, request); if request.finalize { commit_cluster_checkpoint(&target, &request.path, self.shard_id, request.shard_count) .map_err(|err| Status::failed_precondition(err.to_string()))?; @@ -2620,11 +3125,20 @@ impl proto::shard_service_server::ShardService for ShardNode { prepare_cluster_checkpoint(&target, &request.path, self.shard_id, request.shard_count) .map_err(|err| Status::internal(err.to_string()))?; } - Ok(Response::new(proto::CheckpointResponse { + let response = proto::CheckpointResponse { paths: vec![target.to_string_lossy().to_string()], generation: request.path, committed: request.finalize, - })) + }; + replication_attempt.finish( + replica_response.as_ref().is_none_or(|replica| { + replica.generation == response.generation + && replica.committed == response.committed + && replica.paths.len() == 1 + }), + "checkpoint", + )?; + Ok(Response::new(response)) } async fn get_record_id_range( @@ -2792,6 +3306,7 @@ impl proto::shard_service_server::ShardService for ShardNode { &self, request: Request, ) -> Result, Status> { + self.authorize_mutation(&request)?; let _mutation_guard = self.mutation_gate.read().await; self.ensure_no_pending_ingest()?; let request = request.into_inner(); @@ -2799,20 +3314,34 @@ impl proto::shard_service_server::ShardService for ShardNode { if request.records.is_empty() { return Ok(Response::new(proto::ImportRecordsResponse { imported: 0 })); } + let (_replication_serial, replica_response, replication_attempt) = + replicate_to_replica!(self, import_records, request); let mut unirust = write_unirust!(self); - let records = Self::build_import_records(&mut unirust, &request.records)?; + let records_to_build = Self::validate_import_records(&unirust, &request.records)?; + let local_attempt = self.begin_durable_mutation(); + let records = + Self::build_import_records(&mut unirust, &request.records, &records_to_build)?; unirust .ingest_with_explicit_ids(records) .map_err(|err| Status::internal(err.to_string()))?; - Ok(Response::new(proto::ImportRecordsResponse { + let response = proto::ImportRecordsResponse { imported: request.records.len() as u64, - })) + }; + local_attempt.finish(); + replication_attempt.finish( + replica_response + .as_ref() + .is_none_or(|replica| replica == &response), + "record import", + )?; + Ok(Response::new(response)) } async fn import_records_stream( &self, request: Request>, ) -> Result, Status> { + self.authorize_mutation(&request)?; let _mutation_guard = self.mutation_gate.read().await; self.ensure_no_pending_ingest()?; let mut stream = request.into_inner(); @@ -2826,11 +3355,30 @@ impl proto::shard_service_server::ShardService for ShardNode { continue; } validate_distributed_protocol(chunk.internal_protocol_version)?; + let replication_request = proto::ImportRecordsRequest { + records: chunk.records.clone(), + internal_protocol_version: DISTRIBUTED_PROTOCOL_VERSION, + }; + let (_replication_serial, replica_response, replication_attempt) = + replicate_to_replica!(self, import_records, replication_request); let mut unirust = write_unirust!(self); - let records = Self::build_import_records(&mut unirust, &chunk.records)?; + let records_to_build = Self::validate_import_records(&unirust, &chunk.records)?; + let local_attempt = self.begin_durable_mutation(); + let records = + Self::build_import_records(&mut unirust, &chunk.records, &records_to_build)?; unirust .ingest_with_explicit_ids(records) .map_err(|err| Status::internal(err.to_string()))?; + let response = proto::ImportRecordsResponse { + imported: chunk.records.len() as u64, + }; + local_attempt.finish(); + replication_attempt.finish( + replica_response + .as_ref() + .is_none_or(|replica| replica == &response), + "streamed record import", + )?; imported += chunk.records.len() as u64; } Ok(Response::new(proto::ImportRecordsResponse { imported })) @@ -2840,9 +3388,14 @@ impl proto::shard_service_server::ShardService for ShardNode { &self, request: Request, ) -> Result, Status> { + self.authorize_mutation(&request)?; let _mutation_guard = self.mutation_gate.read().await; self.ensure_no_pending_ingest()?; let request = request.into_inner(); + let replication_request = request.clone(); + let (_replication_serial, replica_response, replication_attempt) = + replicate_to_replica!(self, list_conflicts, replication_request); + let local_attempt = self.begin_durable_mutation(); let mut summaries = if let Some(cache) = { let guard = read_unirust!(self); guard.cached_conflict_summaries() @@ -2917,13 +3470,21 @@ impl proto::shard_service_server::ShardService for ShardNode { .map(to_proto_conflict_summary) .collect(), }; + local_attempt.finish(); + replication_attempt.finish( + replica_response + .as_ref() + .is_none_or(|replica| replica == &response), + "conflict materialization", + )?; Ok(Response::new(response)) } async fn reset( &self, - _request: Request, + request: Request, ) -> Result, Status> { + self.authorize_mutation(&request)?; if !self.allow_destructive_admin { return Err(Status::permission_denied( "destructive reset RPC is disabled; stop the cluster and use the confirmed \ @@ -2932,6 +3493,13 @@ impl proto::shard_service_server::ShardService for ShardNode { } let _mutation_guard = self.mutation_gate.write().await; self.ensure_no_pending_ingest()?; + if self.replica_client.is_some() { + return Err(Status::failed_precondition( + "online reset is disabled while replication is configured; stop both nodes and \ + reset or rebootstrap them together", + )); + } + let local_attempt = self.begin_durable_mutation(); let config = self.ontology_config.lock().await.clone(); { let mut guard = write_unirust!(self); @@ -2955,6 +3523,7 @@ impl proto::shard_service_server::ShardService for ShardNode { *self.partitioned.write() = None; } self.cross_shard_conflicts.write().clear(); + local_attempt.finish(); Ok(Response::new(proto::Empty {})) } @@ -2964,17 +3533,25 @@ impl proto::shard_service_server::ShardService for ShardNode { ) -> Result, Status> { let _mutation_guard = self.mutation_gate.read().await; self.ensure_no_pending_ingest()?; - let since_version = request.into_inner().since_version; + let request = request.into_inner(); + let mut signatures = std::collections::HashSet::with_capacity(request.signatures.len()); + for signature in request.signatures { + let bytes = <[u8; 32]>::try_from(signature.signature.as_slice()) + .map_err(|_| Status::invalid_argument("boundary signature must be 32 bytes"))?; + signatures.insert(IdentityKeySignature::from_bytes(bytes)); + } - // Export boundary metadata from the streaming linker let unirust = read_unirust!(self); - let boundary_index = unirust.export_boundary_index(); + let boundary_index = if signatures.is_empty() { + None + } else { + unirust.export_boundary_index_for_signatures(&signatures) + }; let metadata = match boundary_index { Some(index) => { let exported = index.export_metadata(); - // Only return if version is newer than requested - if exported.version <= since_version { + if exported.version <= request.since_version { proto::BoundaryMetadata { shard_id: self.shard_id, version: exported.version, @@ -3031,9 +3608,11 @@ impl proto::shard_service_server::ShardService for ShardNode { &self, request: Request, ) -> Result, Status> { + self.authorize_mutation(&request)?; let _mutation_guard = self.mutation_gate.read().await; self.ensure_no_pending_ingest()?; let req = request.into_inner(); + let replication_request = req; let primary = req .primary @@ -3042,110 +3621,204 @@ impl proto::shard_service_server::ShardService for ShardNode { .secondary .ok_or_else(|| Status::invalid_argument("secondary cluster ID is required"))?; - let primary_id = GlobalClusterId::new( - u16::try_from(primary.shard_id) - .map_err(|_| Status::invalid_argument("primary shard_id exceeds u16"))?, - primary.local_id, - u16::try_from(primary.version) - .map_err(|_| Status::invalid_argument("primary version exceeds u16"))?, - ); - let secondary_id = GlobalClusterId::new( - u16::try_from(secondary.shard_id) - .map_err(|_| Status::invalid_argument("secondary shard_id exceeds u16"))?, - secondary.local_id, - u16::try_from(secondary.version) - .map_err(|_| Status::invalid_argument("secondary version exceeds u16"))?, - ); + let primary_id = global_cluster_id_from_proto(&primary, "primary")?; + let secondary_id = global_cluster_id_from_proto(&secondary, "secondary")?; if primary_id == secondary_id { return Err(Status::invalid_argument( "primary and secondary cluster IDs must differ", )); } + if primary_id.to_u64() >= secondary_id.to_u64() { + return Err(Status::invalid_argument( + "cross-shard merge must redirect the higher cluster ID to the lower canonical ID", + )); + } + let (_replication_serial, replica_response, replication_attempt) = + replicate_to_replica!(self, apply_merge, replication_request); + let local_attempt = self.begin_durable_mutation(); // Get write lock on unirust let mut unirust = write_unirust!(self); // Apply the merge via the streaming linker's DSU - let records_updated = match unirust.apply_cross_shard_merge(primary_id, secondary_id) { - Ok(count) => count, - Err(err) => { - return Ok(Response::new(proto::ApplyMergeResponse { - success: false, - records_updated: 0, - error: err.to_string(), - })); - } + let response = match unirust.apply_cross_shard_merge(primary_id, secondary_id) { + Ok(records_updated) => proto::ApplyMergeResponse { + success: true, + records_updated: records_updated as u32, + error: String::new(), + }, + Err(err) => proto::ApplyMergeResponse { + success: false, + records_updated: 0, + error: err.to_string(), + }, }; + if response.success { + local_attempt.finish(); + } + replication_attempt.finish( + replica_response + .as_ref() + .is_none_or(|replica| replica == &response), + "cross-shard merge", + )?; + Ok(Response::new(response)) + } - Ok(Response::new(proto::ApplyMergeResponse { - success: true, - records_updated: records_updated as u32, - error: String::new(), - })) + async fn apply_merges( + &self, + request: Request, + ) -> Result, Status> { + self.authorize_mutation(&request)?; + let _mutation_guard = self.mutation_gate.read().await; + self.ensure_no_pending_ingest()?; + let req = request.into_inner(); + let replication_request = req.clone(); + let mut merges = Vec::with_capacity(req.merges.len()); + let mut primaries_by_secondary = std::collections::HashMap::with_capacity(req.merges.len()); + for merge in &req.merges { + let primary = merge + .primary + .as_ref() + .ok_or_else(|| Status::invalid_argument("merge primary is required"))?; + let secondary = merge + .secondary + .as_ref() + .ok_or_else(|| Status::invalid_argument("merge secondary is required"))?; + let primary = global_cluster_id_from_proto(primary, "primary")?; + let secondary = global_cluster_id_from_proto(secondary, "secondary")?; + if primary == secondary { + return Err(Status::invalid_argument( + "primary and secondary cluster IDs must differ", + )); + } + if primary.to_u64() >= secondary.to_u64() { + return Err(Status::invalid_argument( + "cross-shard merge must redirect the higher cluster ID to the lower canonical ID", + )); + } + if let Some(existing) = primaries_by_secondary.insert(secondary, primary) { + if existing != primary { + return Err(Status::invalid_argument( + "cross-shard merge batch redirects one secondary to multiple primaries", + )); + } + continue; + } + merges.push((primary, secondary)); + } + + let (_replication_serial, replica_response, replication_attempt) = + replicate_to_replica!(self, apply_merges, replication_request); + let local_attempt = self.begin_durable_mutation(); + let mut unirust = write_unirust!(self); + let response = match unirust.apply_cross_shard_merges(&merges) { + Ok(records_updated) => proto::ApplyMergesResponse { + success: true, + records_updated: records_updated as u64, + error: String::new(), + }, + Err(err) => proto::ApplyMergesResponse { + success: false, + records_updated: 0, + error: err.to_string(), + }, + }; + if response.success { + local_attempt.finish(); + } + replication_attempt.finish( + replica_response + .as_ref() + .is_none_or(|replica| replica == &response), + "cross-shard merge batch", + )?; + Ok(Response::new(response)) } async fn get_dirty_boundary_keys( &self, - _request: Request, + request: Request, ) -> Result, Status> { let _mutation_guard = self.mutation_gate.read().await; - self.ensure_no_pending_ingest()?; + let _wal_guard = if self.ingest_wal.is_some() { + Some(self.ingest_wal_lock.lock().await) + } else { + None + }; + if self + .ingest_wal + .as_ref() + .is_some_and(|wal| wal.has_pending()) + { + return Err(Status::failed_precondition( + "ingest recovery is pending; restart the shard", + )); + } + self.ensure_store_healthy()?; + let request = request.into_inner(); + let after_signature = if request.after_signature.is_empty() { + None + } else { + Some( + <[u8; 32]>::try_from(request.after_signature.as_slice()) + .map(IdentityKeySignature::from_bytes) + .map_err(|_| Status::invalid_argument("dirty-key cursor must be 32 bytes"))?, + ) + }; + let limit = if request.limit == 0 { + DIRTY_KEY_PAGE_LIMIT + } else { + usize::try_from(request.limit) + .map_err(|_| Status::invalid_argument("dirty-key limit exceeds usize"))? + }; + if limit > DIRTY_KEY_PAGE_LIMIT { + return Err(Status::invalid_argument(format!( + "dirty-key limit must not exceed {DIRTY_KEY_PAGE_LIMIT}" + ))); + } let unirust = read_unirust!(self); let partitioned_guard = self.partitioned.read(); - // Get dirty keys from both partitioned and non-partitioned paths - let mut all_dirty_keys = unirust.get_dirty_boundary_keys().unwrap_or_default(); - let mut combined_index = unirust.export_boundary_index(); - - // Add dirty keys from partitioned processor + let candidate_limit = limit.saturating_add(1); + let mut all_dirty_keys = unirust + .dirty_boundary_key_candidates(after_signature, candidate_limit) + .into_iter() + .collect::>(); if let Some(partitioned) = partitioned_guard.as_ref() { - all_dirty_keys.extend(partitioned.get_all_dirty_boundary_keys()); - let partitioned_index = partitioned.export_boundary_index(); - if let Some(ref mut index) = combined_index { - index.merge_from(&partitioned_index); - } else { - combined_index = Some(partitioned_index); - } + all_dirty_keys.extend( + partitioned.dirty_boundary_key_candidates(after_signature, candidate_limit), + ); } - - let mut dirty_key_entries = Vec::new(); - if let Some(index) = combined_index { - for sig in &all_dirty_keys { - if let Some(entries) = index.get_boundaries(sig) { - let proto_entries: Vec = entries - .iter() - .map(|e| proto::ClusterBoundaryEntry { - cluster_id: Some(proto::GlobalClusterId { - shard_id: e.cluster_id.shard_id as u32, - local_id: e.cluster_id.local_id, - version: e.cluster_id.version as u32, - }), - interval_start: e.interval.start, - interval_end: e.interval.end, - shard_id: e.shard_id as u32, - perspective_strong_ids: e.perspective_strong_ids.clone(), - strong_ids: e - .strong_ids - .iter() - .cloned() - .map(boundary_strong_id_to_proto) - .collect(), - }) - .collect(); - - dirty_key_entries.push(proto::DirtyBoundaryKey { - signature: Some(proto::IdentityKeySignature { - signature: sig.to_bytes().to_vec(), - }), - entries: proto_entries, - }); - } - } + let mut page = all_dirty_keys + .into_iter() + .take(candidate_limit) + .collect::>(); + let has_more = page.len() > limit; + if has_more { + page.pop(); } + let next_after_signature = if has_more { + page.last() + .map_or_else(Vec::new, |signature| signature.to_bytes().to_vec()) + } else { + Vec::new() + }; + let dirty_keys = page + .into_iter() + .map(|signature| proto::DirtyBoundaryKey { + signature: Some(proto::IdentityKeySignature { + signature: signature.to_bytes().to_vec(), + }), + entries: Vec::new(), + }) + .collect(); Ok(Response::new(proto::GetDirtyBoundaryKeysResponse { - dirty_keys: dirty_key_entries, + dirty_keys, shard_id: self.shard_id, + next_after_signature, + has_more, })) } @@ -3153,11 +3826,11 @@ impl proto::shard_service_server::ShardService for ShardNode { &self, request: Request, ) -> Result, Status> { + self.authorize_mutation(&request)?; let _mutation_guard = self.mutation_gate.read().await; self.ensure_no_pending_ingest()?; let req = request.into_inner(); - let mut unirust = write_unirust!(self); - let partitioned_guard = self.partitioned.read(); + let replication_request = req.clone(); let keys: Vec = req .keys @@ -3173,6 +3846,10 @@ impl proto::shard_service_server::ShardService for ShardNode { }) .collect(); + let (_replication_serial, replica_response, replication_attempt) = + replicate_to_replica!(self, clear_dirty_keys, replication_request); + let mut unirust = write_unirust!(self); + let partitioned_guard = self.partitioned.read(); let keys_cleared = keys.len() as u32; unirust.clear_dirty_boundary_keys(&keys); @@ -3181,18 +3858,25 @@ impl proto::shard_service_server::ShardService for ShardNode { partitioned.clear_dirty_boundary_keys(&keys); } - Ok(Response::new(proto::ClearDirtyKeysResponse { - keys_cleared, - })) + let response = proto::ClearDirtyKeysResponse { keys_cleared }; + replication_attempt.finish( + replica_response + .as_ref() + .is_none_or(|replica| replica == &response), + "dirty boundary key clear", + )?; + Ok(Response::new(response)) } async fn store_cross_shard_conflicts( &self, request: Request, ) -> Result, Status> { + self.authorize_mutation(&request)?; let _mutation_guard = self.mutation_gate.write().await; self.ensure_no_pending_ingest()?; let req = request.into_inner(); + let replication_request = req.clone(); let conflicts: Vec = req .conflicts .into_iter() @@ -3247,6 +3931,9 @@ impl proto::shard_service_server::ShardService for ShardNode { }) .collect::>()?; + let (_replication_serial, replica_response, replication_attempt) = + replicate_to_replica!(self, store_cross_shard_conflicts, replication_request); + let local_attempt = self.begin_durable_mutation(); // Store conflicts - keep only those relevant to this shard let shard_id = self.shard_id as u16; let mut conflict_storage = self.cross_shard_conflicts.write(); @@ -3268,9 +3955,15 @@ impl proto::shard_service_server::ShardService for ShardNode { *conflict_storage = updated; } - Ok(Response::new(proto::StoreCrossShardConflictsResponse { - stored_count, - })) + let response = proto::StoreCrossShardConflictsResponse { stored_count }; + local_attempt.finish(); + replication_attempt.finish( + replica_response + .as_ref() + .is_none_or(|replica| replica == &response), + "cross-shard conflict storage", + )?; + Ok(Response::new(response)) } } @@ -3391,10 +4084,13 @@ impl ReconciliationCoordinator { /// Take and clear dirty keys for reconciliation. fn take_dirty_keys(&mut self) -> Vec { self.oldest_dirty = None; - self.last_reconcile = Instant::now(); self.dirty_keys.drain().map(|(k, _)| k).collect() } + fn mark_reconciled(&mut self) { + self.last_reconcile = Instant::now(); + } + /// Check if we should reconcile based on adaptive conditions. fn should_reconcile(&self, current_ingest_rate: f64) -> bool { let dirty_count = self.dirty_keys.len(); @@ -3568,12 +4264,36 @@ impl RouterNode { for (expected_shard_id, client) in shard_clients.iter().enumerate() { let mut client = client.clone(); let response = client - .get_config_version(Request::new(proto::ConfigVersionRequest {})) + .get_config_version(Request::new(proto::ConfigVersionRequest { + include_durable_state_digest: false, + })) .await .map_err(|err| Status::unavailable(err.to_string()))? .into_inner(); validate_distributed_protocol(response.protocol_version)?; validate_checkpoint_protocol(response.checkpoint_protocol_version)?; + match proto::ShardRole::try_from(response.shard_role) { + Ok(proto::ShardRole::Standalone | proto::ShardRole::Primary) => {} + Ok(proto::ShardRole::Replica) => { + return Err(Status::failed_precondition(format!( + "shard endpoint at index {expected_shard_id} is a passive replica; \ + promote it before routing traffic" + ))); + } + _ => { + return Err(Status::data_loss(format!( + "shard endpoint at index {expected_shard_id} reported an invalid role" + ))); + } + } + if !response.durable_state_digest.is_empty() + && response.durable_state_digest.len() != 32 + { + return Err(Status::data_loss(format!( + "shard endpoint at index {expected_shard_id} reported an invalid durable \ + state digest" + ))); + } let shard_restore_state = match ( response.restore_generation.is_empty(), response.restore_shard_count, @@ -3644,7 +4364,8 @@ impl RouterNode { ); let metadata = client .get_boundary_metadata(Request::new(proto::GetBoundaryMetadataRequest { - since_version: 0, + since_version: u64::MAX, + signatures: Vec::new(), })) .await .map_err(|err| Status::unavailable(err.to_string()))? @@ -3803,60 +4524,17 @@ impl RouterNode { } async fn recover_dirty_reconciliation_on_startup(&self) -> Result<(), Status> { - let mut dirty_keys = std::collections::HashSet::new(); - for (expected_shard_id, client) in self.shard_clients.iter().enumerate() { - let mut client = client.clone(); - let response = client - .get_dirty_boundary_keys(Request::new(proto::GetDirtyBoundaryKeysRequest {})) - .await - .map_err(|err| Status::unavailable(err.to_string()))? - .into_inner(); - if response.shard_id as usize != expected_shard_id { - return Err(Status::data_loss(format!( - "dirty-key response from endpoint {expected_shard_id} reports shard {}", - response.shard_id - ))); - } - for dirty_key in response.dirty_keys { - let signature = dirty_key.signature.ok_or_else(|| { - Status::data_loss("dirty-key response is missing a signature") - })?; - let signature = <[u8; 32]>::try_from(signature.signature.as_slice()) - .map_err(|_| Status::data_loss("dirty-key signature must be 32 bytes"))?; - dirty_keys.insert(IdentityKeySignature::from_bytes(signature)); - } - } + let dirty_keys = self.fetch_dirty_boundary_keys(true).await?; if dirty_keys.is_empty() { return Ok(()); } - let mut dirty_keys = dirty_keys.into_iter().collect::>(); - dirty_keys.sort_by_key(|signature| *signature.to_bytes()); tracing::warn!( dirty_keys = dirty_keys.len(), "router startup is repairing retained cross-shard reconciliation work" ); self.reconcile_dirty_keys(&dirty_keys).await?; - - let keys = dirty_keys - .iter() - .map(|signature| proto::IdentityKeySignature { - signature: signature.to_bytes().to_vec(), - }) - .collect::>(); - for (shard_id, client) in self.shard_clients.iter().enumerate() { - let mut client = client.clone(); - client - .clear_dirty_keys(Request::new(proto::ClearDirtyKeysRequest { - keys: keys.clone(), - })) - .await - .map_err(|err| { - Status::unavailable(format!( - "failed to clear recovered dirty keys on shard {shard_id}: {err}" - )) - })?; - } + self.clear_dirty_boundary_keys(&dirty_keys).await?; Ok(()) } @@ -4160,17 +4838,36 @@ impl RouterNode { Ok(summaries) } - async fn fetch_boundary_metadata(&self) -> Result, Status> { + async fn fetch_boundary_metadata( + &self, + keys: &[IdentityKeySignature], + ) -> Result, Status> { + let signatures = keys + .iter() + .map(|signature| proto::IdentityKeySignature { + signature: signature.to_bytes().to_vec(), + }) + .collect::>(); + let requests = self.shard_clients.iter().cloned().enumerate().map( + |(expected_shard_id, mut client)| { + let signatures = signatures.clone(); + async move { + let response = client + .get_boundary_metadata(Request::new(proto::GetBoundaryMetadataRequest { + since_version: 0, + signatures, + })) + .await + .map_err(|err| Status::unavailable(err.to_string()))? + .into_inner(); + Ok::<_, Status>((expected_shard_id, response)) + } + }, + ); + let responses = futures::future::join_all(requests).await; let mut metadata = Vec::with_capacity(self.shard_clients.len()); - for (expected_shard_id, client) in self.shard_clients.iter().enumerate() { - let mut client = client.clone(); - let response = client - .get_boundary_metadata(Request::new(proto::GetBoundaryMetadataRequest { - since_version: 0, - })) - .await - .map_err(|err| Status::unavailable(err.to_string()))? - .into_inner(); + for response in responses { + let (expected_shard_id, response) = response?; let shard_metadata = response .metadata .ok_or_else(|| Status::data_loss("shard returned no boundary metadata"))?; @@ -4187,24 +4884,116 @@ impl RouterNode { Ok(metadata) } - async fn apply_merge_to_shard( + async fn fetch_dirty_boundary_keys( + &self, + fetch_all_pages: bool, + ) -> Result, Status> { + let requests = self.shard_clients.iter().cloned().enumerate().map( + |(expected_shard_id, mut client)| async move { + let mut shard_keys = Vec::new(); + let mut after_signature = Vec::new(); + loop { + let response = client + .get_dirty_boundary_keys(Request::new(proto::GetDirtyBoundaryKeysRequest { + after_signature: after_signature.clone(), + limit: DIRTY_KEY_PAGE_LIMIT as u32, + })) + .await + .map_err(|err| Status::unavailable(err.to_string()))? + .into_inner(); + if response.shard_id as usize != expected_shard_id { + return Err(Status::data_loss(format!( + "dirty-key response from endpoint {expected_shard_id} reports shard {}", + response.shard_id + ))); + } + for dirty_key in response.dirty_keys { + if !dirty_key.entries.is_empty() { + return Err(Status::data_loss( + "dirty-key response unexpectedly included boundary metadata", + )); + } + let signature = dirty_key.signature.ok_or_else(|| { + Status::data_loss("dirty-key response is missing a signature") + })?; + let signature = <[u8; 32]>::try_from(signature.signature.as_slice()) + .map_err(|_| { + Status::data_loss("dirty-key signature must be 32 bytes") + })?; + shard_keys.push(IdentityKeySignature::from_bytes(signature)); + } + if !fetch_all_pages || !response.has_more { + break; + } + if response.next_after_signature.len() != 32 + || response.next_after_signature <= after_signature + { + return Err(Status::data_loss( + "shard returned a non-progressing dirty-key cursor", + )); + } + after_signature = response.next_after_signature; + } + Ok::<_, Status>(shard_keys) + }, + ); + let responses = futures::future::join_all(requests).await; + let mut dirty_keys = std::collections::HashSet::new(); + for shard_keys in responses { + dirty_keys.extend(shard_keys?); + } + let mut dirty_keys = dirty_keys.into_iter().collect::>(); + dirty_keys.sort_by_key(|signature| *signature.to_bytes()); + Ok(dirty_keys) + } + + async fn clear_dirty_boundary_keys(&self, keys: &[IdentityKeySignature]) -> Result<(), Status> { + for chunk in keys.chunks(RECONCILIATION_KEY_CHUNK) { + let keys = chunk + .iter() + .map(|signature| proto::IdentityKeySignature { + signature: signature.to_bytes().to_vec(), + }) + .collect::>(); + for (shard_id, client) in self.shard_clients.iter().enumerate() { + let mut client = client.clone(); + client + .clear_dirty_keys(Request::new(proto::ClearDirtyKeysRequest { + keys: keys.clone(), + })) + .await + .map_err(|err| { + Status::unavailable(format!( + "failed to clear reconciled dirty keys on shard {shard_id}: {err}" + )) + })?; + } + } + Ok(()) + } + + async fn apply_merges_to_shard( &self, shard_id: u16, - primary: GlobalClusterId, - secondary: GlobalClusterId, - ) -> Result { + merges: &[(GlobalClusterId, GlobalClusterId)], + ) -> Result { let mut client = self.shard_client(u32::from(shard_id))?; let response = client - .apply_merge(Request::new(proto::ApplyMergeRequest { - primary: Some(global_cluster_id_to_proto(primary)), - secondary: Some(global_cluster_id_to_proto(secondary)), + .apply_merges(Request::new(proto::ApplyMergesRequest { + merges: merges + .iter() + .map(|(primary, secondary)| proto::ClusterMerge { + primary: Some(global_cluster_id_to_proto(*primary)), + secondary: Some(global_cluster_id_to_proto(*secondary)), + }) + .collect(), })) .await .map_err(|err| Status::unavailable(err.to_string()))? .into_inner(); if !response.success { return Err(Status::aborted(format!( - "shard {shard_id} rejected cross-shard merge: {}", + "shard {shard_id} rejected cross-shard merge batch: {}", response.error ))); } @@ -4230,23 +5019,26 @@ impl RouterNode { async fn apply_reconciliation_result( &self, result: &crate::sharding::ReconciliationResult, - ) -> Result<(), Status> { + ) -> Result { + let attempt = ConsistencyAttempt::new(self.reconciliation_consistent.clone()); let outcome = self.apply_reconciliation_result_inner(result).await; - self.reconciliation_consistent - .store(outcome.is_ok(), Ordering::Release); + if outcome.is_ok() { + attempt.finish(); + } outcome } async fn apply_reconciliation_result_inner( &self, result: &crate::sharding::ReconciliationResult, - ) -> Result<(), Status> { - for (primary, secondary) in &result.merged_clusters { + ) -> Result { + let mut records_updated = 0u64; + for merges in result.merged_clusters.chunks(MERGE_APPLICATION_CHUNK) { for shard_id in 0..self.shard_clients.len() { let shard_id = u16::try_from(shard_id) .map_err(|_| Status::internal("shard index exceeds u16"))?; - self.apply_merge_to_shard(shard_id, *primary, *secondary) - .await?; + records_updated = records_updated + .saturating_add(self.apply_merges_to_shard(shard_id, merges).await?); } } @@ -4266,9 +5058,12 @@ impl RouterNode { } } for (shard_id, conflicts) in conflicts_by_shard { - self.store_conflicts_on_shard(shard_id, conflicts).await?; + for chunk in conflicts.chunks(CONFLICT_APPLICATION_CHUNK) { + self.store_conflicts_on_shard(shard_id, chunk.to_vec()) + .await?; + } } - Ok(()) + Ok(records_updated) } /// Get read access to the cluster locality index. @@ -4392,74 +5187,24 @@ impl RouterNode { } async fn run_adaptive_reconciliation_once(&self) { - let poll_complete = { + let dirty_keys = { let _mutation_guard = self.mutation_gate.read().await; if self.ensure_ontology_consistent().is_err() { return; } - - // 1. Poll all shards for dirty boundary keys without pausing ingest. - let mut poll_complete = true; - for (expected_shard_id, client) in self.shard_clients.iter().enumerate() { - let mut client = client.clone(); - match client - .get_dirty_boundary_keys(Request::new(proto::GetDirtyBoundaryKeysRequest {})) - .await - { - Ok(response) => { - let resp = response.into_inner(); - let shard_id = match u16::try_from(resp.shard_id) { - Ok(shard_id) => shard_id, - Err(_) => { - tracing::error!( - shard_id = resp.shard_id, - "dirty-key response shard_id exceeds u16" - ); - poll_complete = false; - continue; - } - }; - if usize::from(shard_id) != expected_shard_id { - tracing::error!( - expected_shard_id, - reported_shard_id = shard_id, - "dirty-key response came from the wrong shard" - ); - poll_complete = false; - continue; - } - - for dirty_key in resp.dirty_keys { - let Some(sig_proto) = &dirty_key.signature else { - tracing::error!("dirty-key response is missing a signature"); - poll_complete = false; - continue; - }; - let Ok(sig_bytes) = - <[u8; 32]>::try_from(sig_proto.signature.as_slice()) - else { - tracing::error!("dirty-key signature must be 32 bytes"); - poll_complete = false; - continue; - }; - let sig = IdentityKeySignature::from_bytes(sig_bytes); - let mut coordinator = self.reconciliation_coordinator.lock().await; - coordinator.add_dirty_keys_from_shard(shard_id, vec![sig]); - } - } - Err(err) => { - tracing::warn!(error = %err, "failed to poll shard dirty keys"); - poll_complete = false; - } + match self.fetch_dirty_boundary_keys(false).await { + Ok(dirty_keys) => dirty_keys, + Err(err) => { + tracing::warn!(error = %err, "failed to poll shard dirty keys"); + return; } } - poll_complete }; - if !poll_complete { - return; + if !dirty_keys.is_empty() { + let mut coordinator = self.reconciliation_coordinator.lock().await; + coordinator.add_dirty_keys_from_shard(0, dirty_keys); } - // 2. Check if we should reconcile let should_run = { let coordinator = self.reconciliation_coordinator.lock().await; let ingest_rate = self.current_ingest_rate(); @@ -4471,56 +5216,56 @@ impl RouterNode { // Pause router-mediated ingest while refetching metadata, applying merges, // and clearing exactly the reconciled dirty-key generation. let _mutation_guard = self.mutation_gate.write().await; + let reconciliation_started = Instant::now(); if self.ensure_ontology_consistent().is_err() { return; } - // 3. Take dirty keys and perform targeted reconciliation - let dirty_keys = { + { let mut coordinator = self.reconciliation_coordinator.lock().await; - coordinator.take_dirty_keys() - }; - - if !dirty_keys.is_empty() { - let result = match self.reconcile_dirty_keys(&dirty_keys).await { - Ok(result) => result, - Err(err) => { - tracing::error!( - error = %err, - "adaptive reconciliation failed; shard dirty keys retained" - ); - return; - } - }; - - // Update metrics - if result.conflicts_blocked > 0 { - self.metrics - .cross_shard_conflicts - .fetch_add(result.conflicts_blocked as u64, Ordering::Relaxed); + coordinator.take_dirty_keys(); + } + let dirty_keys = match self.fetch_dirty_boundary_keys(true).await { + Ok(dirty_keys) => dirty_keys, + Err(err) => { + tracing::error!( + error = %err, + "adaptive reconciliation failed to refetch authoritative dirty keys" + ); + return; } - - // 4. Clear dirty keys on shards - let proto_keys: Vec = dirty_keys - .iter() - .map(|sig| proto::IdentityKeySignature { - signature: sig.to_bytes().to_vec(), - }) - .collect(); - - for client in &self.shard_clients { - let mut client = client.clone(); - if let Err(err) = client - .clear_dirty_keys(Request::new(proto::ClearDirtyKeysRequest { - keys: proto_keys.clone(), - })) - .await - { - tracing::warn!( - error = %err, - "failed to clear reconciled dirty keys; retry is safe" - ); - } + }; + let result = match self.reconcile_dirty_keys(&dirty_keys).await { + Ok(result) => result, + Err(err) => { + tracing::error!( + error = %err, + "adaptive reconciliation failed; shard dirty keys retained" + ); + return; } + }; + if let Err(err) = self.clear_dirty_boundary_keys(&dirty_keys).await { + tracing::error!( + error = %err, + "adaptive reconciliation failed to clear dirty keys; retry is safe" + ); + return; + } + tracing::info!( + dirty_keys = dirty_keys.len(), + merges = result.merges_performed, + conflicts_blocked = result.conflicts_blocked, + elapsed_ms = reconciliation_started.elapsed().as_millis(), + "adaptive cross-shard reconciliation completed" + ); + self.reconciliation_coordinator + .lock() + .await + .mark_reconciled(); + if result.conflicts_blocked > 0 { + self.metrics + .cross_shard_conflicts + .fetch_add(result.conflicts_blocked as u64, Ordering::Relaxed); } } } @@ -4532,14 +5277,36 @@ impl RouterNode { ) -> Result { use crate::sharding::IncrementalReconciler; - let mut reconciler = IncrementalReconciler::new(); - for metadata in self.fetch_boundary_metadata().await? { - reconciler.add_shard_boundary(boundary_index_from_metadata(&metadata)?); - } - - let key_set: std::collections::HashSet<_> = keys.iter().copied().collect(); - let result = reconciler.reconcile_keys(&key_set); - self.apply_reconciliation_result(&result).await?; + let mut candidates = ReconciliationCandidates::default(); + let mut metadata_elapsed = Duration::ZERO; + let mut reconcile_elapsed = Duration::ZERO; + for chunk in keys.chunks(RECONCILIATION_KEY_CHUNK) { + let metadata_started = Instant::now(); + let mut reconciler = IncrementalReconciler::new(); + for metadata in self.fetch_boundary_metadata(chunk).await? { + reconciler.add_shard_boundary(boundary_index_from_metadata(&metadata)?); + } + metadata_elapsed += metadata_started.elapsed(); + + let reconcile_started = Instant::now(); + let key_set = chunk.iter().copied().collect(); + candidates.extend(reconciler.reconciliation_candidates(&key_set)); + reconcile_elapsed += reconcile_started.elapsed(); + } + let reconcile_started = Instant::now(); + let result = candidates.finish(); + reconcile_elapsed += reconcile_started.elapsed(); + let apply_started = Instant::now(); + let records_updated = self.apply_reconciliation_result(&result).await?; + tracing::info!( + dirty_keys = keys.len(), + merges = result.merges_performed, + records_updated, + metadata_ms = metadata_elapsed.as_millis(), + reconcile_ms = reconcile_elapsed.as_millis(), + apply_ms = apply_started.elapsed().as_millis(), + "cross-shard reconciliation phase timings" + ); Ok(result) } @@ -4699,6 +5466,7 @@ impl proto::router_service_server::RouterService for RouterNode { } } self.invalidate_global_conflict_cache(); + let consistency_attempt = ConsistencyAttempt::new(self.cluster_consistent.clone()); for client in &self.shard_clients { let mut client = client.clone(); @@ -4718,17 +5486,17 @@ impl proto::router_service_server::RouterService for RouterNode { } } if !rollback_complete { - self.cluster_consistent.store(false, Ordering::Release); return Err(Status::aborted(format!( "ontology update failed and rollback was incomplete; cluster is blocked: \ {apply_error}" ))); } + consistency_attempt.finish(); return Err(apply_error); } } *self.ontology_config.write().await = mapped; - self.cluster_consistent.store(true, Ordering::Release); + consistency_attempt.finish(); Ok(Response::new(proto::ApplyOntologyResponse {})) } @@ -4913,6 +5681,8 @@ impl proto::router_service_server::RouterService for RouterNode { .as_ref() .map(|_| self.shard_clients.len() as u32) .unwrap_or_default(), + shard_role: proto::ShardRole::Unspecified as i32, + durable_state_digest: Vec::new(), })) } @@ -5282,15 +6052,9 @@ impl proto::router_service_server::RouterService for RouterNode { metadata from its configured shards", )); } - let shard_metadata = self.fetch_boundary_metadata().await?; - - let mut reconciler = crate::sharding::IncrementalReconciler::new(); - for metadata in &shard_metadata { - reconciler.add_shard_boundary(boundary_index_from_metadata(metadata)?); - } - - let result = reconciler.reconcile(); - self.apply_reconciliation_result(&result).await?; + let dirty_keys = self.fetch_dirty_boundary_keys(true).await?; + let result = self.reconcile_dirty_keys(&dirty_keys).await?; + self.clear_dirty_boundary_keys(&dirty_keys).await?; if result.conflicts_blocked > 0 { self.metrics.cross_shard_conflicts.fetch_add( @@ -5683,7 +6447,7 @@ mod wal_tests { ShardService::ingest_records( &shard, Request::new(proto::IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![original.clone()], }), ) @@ -5695,7 +6459,7 @@ mod wal_tests { let error = ShardService::ingest_records( &shard, Request::new(proto::IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![conflicting], }), ) @@ -5715,6 +6479,57 @@ mod wal_tests { shard.shutdown().await.expect("shutdown"); } + #[tokio::test] + async fn invalid_persistent_import_does_not_latch_shard_closed() { + use proto::shard_service_server::ShardService; + + let temp_dir = tempfile::tempdir().expect("temporary shard directory"); + let shard = ShardNode::new_with_data_dir( + 0, + DistributedOntologyConfig::empty(), + StreamingTuning::balanced(), + Some(temp_dir.path().to_path_buf()), + false, + None, + ) + .expect("persistent shard"); + let error = ShardService::import_records( + &shard, + Request::new(proto::ImportRecordsRequest { + records: vec![proto::RecordSnapshot { + record_id: 7, + identity: Some(proto::RecordIdentity { + entity_type: "person".to_string(), + perspective: "crm".to_string(), + uid: "invalid-import".to_string(), + }), + descriptors: vec![proto::RecordDescriptor { + attr: "email".to_string(), + value: "invalid@example.com".to_string(), + start: 10, + end: 10, + }], + }], + internal_protocol_version: DISTRIBUTED_PROTOCOL_VERSION, + }), + ) + .await + .expect_err("invalid import must fail"); + assert_eq!(error.code(), tonic::Code::InvalidArgument); + + ShardService::health_check(&shard, Request::new(proto::HealthCheckRequest {})) + .await + .expect("fully validated client errors must not poison readiness"); + shard.shutdown().await.expect("shutdown"); + } + + #[test] + fn incomplete_durable_mutation_latches_consistency() { + let consistency = Arc::new(AtomicBool::new(true)); + drop(DurableMutationAttempt::new(consistency.clone(), true)); + assert!(!consistency.load(Ordering::Acquire)); + } + #[tokio::test] async fn shard_rejects_ingest_from_an_older_router_protocol() { use proto::shard_service_server::ShardService; diff --git a/src/lib.rs b/src/lib.rs index 61f4036..483eb14 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -32,6 +32,7 @@ //! - `checkpoint()` - Persist state to disk //! - `stats()` - Get metrics and statistics +pub mod backup; pub mod coalescer; pub mod config; pub mod conflicts; @@ -77,9 +78,14 @@ pub use graph::KnowledgeGraph; pub use query::{QueryDescriptor, QueryMatch, QueryOutcome}; // Storage (for custom backends) +pub use backup::{ + export_cluster_backup, prune_verified_cluster_backups, verify_cluster_backup, + ClusterBackupManifest, +}; pub use persistence::{ read_cluster_checkpoint_manifest, restore_checkpoint, restore_checkpoint_for_shard, - ClusterCheckpointManifest, PersistentStore, CHECKPOINT_PROTOCOL_VERSION, + verify_cluster_checkpoint, ClusterCheckpointManifest, PersistentStore, + CHECKPOINT_PROTOCOL_VERSION, }; pub use store::{RecordStore, Store}; @@ -310,6 +316,12 @@ impl Unirust { if let Some(db) = db.as_ref() { let persistence = LinkerStatePersistence::new(db); + if persistence.prepare_stable_global_cluster_ids()? { + tracing::warn!( + "removed legacy allocation-order cross-shard redirects; \ + router reconciliation will rebuild them from durable records" + ); + } linker.restore_cross_shard_merges(&persistence)?; } Ok(linker) @@ -1211,14 +1223,38 @@ impl Unirust { self.streaming.as_ref().map(|s| s.export_boundary_index()) } + /// Export boundary metadata only for the requested signatures. + #[doc(hidden)] + pub fn export_boundary_index_for_signatures( + &self, + signatures: &std::collections::HashSet, + ) -> Option { + self.streaming + .as_ref() + .map(|streaming| streaming.export_boundary_index_for_signatures(signatures)) + } + /// Get dirty boundary keys from the streaming linker. /// Returns keys that have been modified since last reconciliation. pub fn get_dirty_boundary_keys( &self, - ) -> Option> { + ) -> Option> { self.streaming.as_ref().map(|s| s.get_dirty_boundary_keys()) } + /// Return bounded ordered dirty-key candidates without cloning the full set. + #[doc(hidden)] + pub fn dirty_boundary_key_candidates( + &self, + after: Option, + limit: usize, + ) -> Vec { + self.streaming + .as_ref() + .map(|streaming| streaming.dirty_boundary_key_candidates(after, limit)) + .unwrap_or_default() + } + /// Get dirty boundary key count. pub fn dirty_boundary_key_count(&self) -> usize { self.streaming @@ -1260,41 +1296,70 @@ impl Unirust { &mut self, primary: GlobalClusterId, secondary: GlobalClusterId, + ) -> anyhow::Result { + self.apply_cross_shard_merges(&[(primary, secondary)]) + } + + /// Apply a batch of canonical cross-shard redirects durably. + #[doc(hidden)] + pub fn apply_cross_shard_merges( + &mut self, + merges: &[(GlobalClusterId, GlobalClusterId)], ) -> anyhow::Result { // Initialize streaming linker if not already done if self.streaming.is_none() { self.streaming = Some(self.create_streaming_linker()?); } - let (primary, secondary) = { + let normalized = { let streaming = self.streaming.as_ref().ok_or_else(|| { anyhow::anyhow!("Streaming not initialized - call enable_streaming() first") })?; - ( - streaming.resolve_global_cluster_id(primary), - streaming.resolve_global_cluster_id(secondary), - ) + let mut by_secondary = std::collections::HashMap::with_capacity(merges.len()); + for (primary, secondary) in merges { + let primary = streaming.resolve_global_cluster_id(*primary); + let secondary = streaming.resolve_global_cluster_id(*secondary); + if primary == secondary { + continue; + } + if primary.to_u64() >= secondary.to_u64() { + anyhow::bail!( + "cross-shard merge must redirect cluster {} to lower canonical cluster {}", + secondary.to_u64(), + primary.to_u64() + ); + } + if let Some(existing) = by_secondary.insert(secondary, primary) { + if existing != primary { + anyhow::bail!( + "cross-shard merge batch redirects cluster {} to multiple primaries", + secondary.to_u64() + ); + } + } + } + by_secondary + .into_iter() + .map(|(secondary, primary)| (primary, secondary)) + .collect::>() }; - if primary == secondary { + if normalized.is_empty() { return Ok(0); } - if primary.to_u64() >= secondary.to_u64() { - anyhow::bail!( - "cross-shard merge must redirect cluster {} to lower canonical cluster {}", - secondary.to_u64(), - primary.to_u64() - ); - } if let Some(db) = self.store.shared_db() { let persistence = LinkerStatePersistence::new(&db); - persistence.save_cross_shard_merge(secondary, primary)?; + let durable_merges = normalized + .iter() + .map(|(primary, secondary)| (*secondary, *primary)) + .collect::>(); + persistence.save_cross_shard_merges(&durable_merges)?; self.store.sync()?; } let streaming = self.streaming.as_mut().ok_or_else(|| { anyhow::anyhow!("Streaming not initialized - call enable_streaming() first") })?; - Ok(streaming.apply_cross_shard_merge(primary, secondary)) + Ok(streaming.apply_cross_shard_merges(&normalized)) } /// Resolve a global cluster ID through durable cross-shard redirects. diff --git a/src/linker.rs b/src/linker.rs index 274bb08..05afcc1 100644 --- a/src/linker.rs +++ b/src/linker.rs @@ -26,7 +26,7 @@ use rayon::prelude::*; use rocksdb::DB; use rustc_hash::{FxHashMap, FxHashSet}; use smallvec::SmallVec; -use std::collections::{HashMap, HashSet}; +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use std::num::NonZeroUsize; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; @@ -106,6 +106,23 @@ impl BoundaryData { } } +fn register_boundary_data( + index: &mut crate::sharding::ClusterBoundaryIndex, + signature: ShardingKeySignature, + global_id: GlobalClusterId, + data: &BoundaryData, +) { + for interval in &data.intervals { + index.register_boundary_key_with_conflict_data( + signature, + global_id, + *interval, + data.perspective_strong_ids.clone(), + data.strong_ids.clone(), + ); + } +} + /// Metrics for observability of the streaming linker. /// All counters are atomic for thread-safe access without locking. #[derive(Debug, Default)] @@ -398,6 +415,8 @@ pub struct StreamingLinker { next_cluster_id: u32, /// Global cluster IDs for cross-shard tracking (LRU-bounded when config provided) global_cluster_ids: LinkerState, + /// Minimum durable record ID in each local cluster, independent of replay order. + global_cluster_anchors: LinkerState, /// Shard ID for this linker (used in GlobalClusterId generation). shard_id: u16, /// Strong ID summaries for conflict detection (LRU-bounded when config provided) @@ -413,9 +432,9 @@ pub struct StreamingLinker { key_stats: FxHashMap, /// Boundary keys for cross-shard reconciliation (deduplicated by key+cluster). /// Maps (signature, cluster) -> (interval, perspective_strong_ids) for that combination. - boundary_signatures: HashMap<(ShardingKeySignature, GlobalClusterId), BoundaryData>, + boundary_signatures: BTreeMap<(ShardingKeySignature, GlobalClusterId), BoundaryData>, /// Dirty boundary keys - signatures modified since last reconciliation. - dirty_boundary_keys: HashSet, + dirty_boundary_keys: BTreeSet, /// Cross-shard merge mappings: secondary -> primary. /// Used to redirect cluster ID lookups after cross-shard reconciliation. cross_shard_merges: HashMap, @@ -531,6 +550,7 @@ impl StreamingLinker { cluster_ids, next_cluster_id: 0, global_cluster_ids, + global_cluster_anchors: LinkerState::unbounded(), shard_id, strong_id_summaries, tainted_identity_keys: FxHashSet::default(), @@ -538,22 +558,35 @@ impl StreamingLinker { pending_keys: FxHashSet::default(), tuning: tuning.clone(), key_stats: FxHashMap::default(), - boundary_signatures: HashMap::new(), - dirty_boundary_keys: HashSet::new(), + boundary_signatures: BTreeMap::new(), + dirty_boundary_keys: BTreeSet::new(), cross_shard_merges: HashMap::new(), metrics: LinkerMetrics::new(), partition_opts: None, }; if !store.is_empty() { - let mut record_ids: Vec = Vec::new(); - store.for_each_record(&mut |record| { - record_ids.push(record.id); - }); - record_ids.sort_by_key(|record_id| record_id.0); - - for record_id in record_ids { - streamer.link_record(store, ontology, record_id)?; + const RECOVERY_BATCH_SIZE: usize = 4_096; + let mut batch = Vec::with_capacity(RECOVERY_BATCH_SIZE); + let mut link_batch = |batch: &mut Vec| -> Result<()> { + let records = batch.iter().collect::>(); + streamer.link_records_batch_parallel_with_interner( + &records, + ontology, + store.interner(), + )?; + batch.clear(); + Ok(()) + }; + store.try_for_each_record_ordered(&mut |record| { + batch.push(record); + if batch.len() == RECOVERY_BATCH_SIZE { + link_batch(&mut batch)?; + } + Ok(()) + })?; + if !batch.is_empty() { + link_batch(&mut batch)?; } } @@ -615,6 +648,7 @@ impl StreamingLinker { let _guard = crate::profile::profile_scope("link_record"); if !self.dsu.has_record(record_id)? { self.dsu.add_record(record_id)?; + self.global_cluster_anchors.insert(record_id, record_id.0); } // Try to get a reference first (avoids cloning), fall back to cloning if not available. @@ -825,6 +859,7 @@ impl StreamingLinker { cluster_ids, next_cluster_id, global_cluster_ids, + global_cluster_anchors, strong_id_summaries, tainted_identity_keys, record_perspectives, @@ -917,6 +952,7 @@ impl StreamingLinker { ); reconcile_global_cluster_ids( global_cluster_ids, + global_cluster_anchors, cross_shard_merges, boundary_signatures, dirty_boundary_keys, @@ -1095,6 +1131,7 @@ impl StreamingLinker { // Initialize record in DSU if needed if !self.dsu.has_record(record_id)? { self.dsu.add_record(record_id)?; + self.global_cluster_anchors.insert(record_id, record_id.0); } // Store perspective and summary @@ -1141,6 +1178,7 @@ impl StreamingLinker { // Merge with candidates let mut root_a = self.dsu.find(record_id).unwrap_or(record_id); + let key_is_tainted = self.tainted_identity_keys.contains(&key_signature); for (candidate_id, candidate_interval) in candidates { if candidate_id == record_id { continue; @@ -1154,6 +1192,33 @@ impl StreamingLinker { continue; } + let candidate_perspective = self.record_perspectives.get(&candidate_id); + let same_perspective_conflict = candidate_perspective + .map(|perspective| { + perspective == &extraction.perspective + && same_perspective_conflict_for_clusters( + &self.strong_id_summaries, + root_a, + root_b, + perspective, + ) + }) + .unwrap_or(false); + if same_perspective_conflict { + self.metrics + .conflicts_detected + .fetch_add(1, Ordering::Relaxed); + self.tainted_identity_keys.insert(key_signature.clone()); + continue; + } + if key_is_tainted + && candidate_perspective + .map(|perspective| perspective != &extraction.perspective) + .unwrap_or(false) + { + continue; + } + // Check for conflicts if would_create_conflict_in_clusters( &self.strong_id_summaries, @@ -1187,6 +1252,7 @@ impl StreamingLinker { ); 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, @@ -1299,8 +1365,11 @@ impl StreamingLinker { if let Some(global_id) = self.global_cluster_ids.get(&root) { return *global_id; } - let local_id = self.get_or_assign_cluster_id(root); - let global_id = GlobalClusterId::from_local(self.shard_id, local_id); + let local_id = self + .global_cluster_anchors + .get_copy(&root) + .unwrap_or(root.0); + let global_id = GlobalClusterId::new(self.shard_id, local_id, 0); self.global_cluster_ids.insert(root, global_id); global_id } @@ -1327,10 +1396,35 @@ impl StreamingLinker { } /// Get the set of dirty boundary keys (modified since last reconciliation). - pub fn get_dirty_boundary_keys(&self) -> HashSet { + pub fn get_dirty_boundary_keys(&self) -> BTreeSet { self.dirty_boundary_keys.clone() } + /// Return the first ordered dirty signatures after an exclusive cursor. + pub fn dirty_boundary_key_candidates( + &self, + after: Option, + limit: usize, + ) -> Vec { + match after { + Some(cursor) => self + .dirty_boundary_keys + .range(( + std::ops::Bound::Excluded(cursor), + std::ops::Bound::Unbounded, + )) + .take(limit) + .copied() + .collect(), + None => self + .dirty_boundary_keys + .iter() + .take(limit) + .copied() + .collect(), + } + } + /// Get the count of dirty boundary keys. pub fn dirty_boundary_key_count(&self) -> usize { self.dirty_boundary_keys.len() @@ -1346,16 +1440,36 @@ impl StreamingLinker { /// Export boundary signatures to a ClusterBoundaryIndex. /// This can be used for cross-shard reconciliation. pub fn export_boundary_index(&self) -> crate::sharding::ClusterBoundaryIndex { + self.export_boundary_index_for_keys(None) + } + + /// Export boundary metadata only for the requested signatures. + pub fn export_boundary_index_for_signatures( + &self, + signatures: &HashSet, + ) -> crate::sharding::ClusterBoundaryIndex { + self.export_boundary_index_for_keys(Some(signatures)) + } + + fn export_boundary_index_for_keys( + &self, + signatures: Option<&HashSet>, + ) -> crate::sharding::ClusterBoundaryIndex { let mut index = crate::sharding::ClusterBoundaryIndex::new_small(self.shard_id); - for ((sig, global_id), data) in &self.boundary_signatures { - for interval in &data.intervals { - index.register_boundary_key_with_conflict_data( - *sig, - *global_id, - *interval, - data.perspective_strong_ids.clone(), - data.strong_ids.clone(), - ); + if let Some(signatures) = signatures { + let first_cluster = GlobalClusterId::new(0, 0, 0); + let last_cluster = GlobalClusterId::new(u16::MAX, u32::MAX, u16::MAX); + for signature in signatures { + for ((signature, global_id), data) in self.boundary_signatures.range(( + std::ops::Bound::Included((*signature, first_cluster)), + std::ops::Bound::Included((*signature, last_cluster)), + )) { + register_boundary_data(&mut index, *signature, *global_id, data); + } + } + } else { + for ((signature, global_id), data) in &self.boundary_signatures { + register_boundary_data(&mut index, *signature, *global_id, data); } } index @@ -1366,10 +1480,10 @@ impl StreamingLinker { pub fn drain_boundaries(&mut self) -> Vec<(ShardingKeySignature, GlobalClusterId, Interval)> { std::mem::take(&mut self.boundary_signatures) .into_iter() - .flat_map(|((sig, global_id), data)| { + .flat_map(|((signature, global_id), data)| { data.intervals .into_iter() - .map(move |interval| (sig, global_id, interval)) + .map(move |interval| (signature, global_id, interval)) }) .collect() } @@ -1396,8 +1510,6 @@ impl StreamingLinker { Some(sig) => sig, None => return, }; - let key = (sharding_sig, global_id); - // Compute perspective_strong_ids from the cluster's summary // Uses actual string values (via interner) for cross-shard consistency let perspective_strong_ids = self @@ -1413,7 +1525,7 @@ impl StreamingLinker { // Merge intervals for the same (signature, cluster) combination self.boundary_signatures - .entry(key) + .entry((sharding_sig, global_id)) .and_modify(|existing| { existing.intervals.push(interval); coalesce_boundary_intervals(&mut existing.intervals); @@ -1438,36 +1550,48 @@ impl StreamingLinker { primary: GlobalClusterId, secondary: GlobalClusterId, ) -> usize { - // Record the merge mapping - self.cross_shard_merges.insert(secondary, primary); + self.apply_cross_shard_merges(&[(primary, secondary)]) + } - // Update any existing boundary signatures that reference the secondary cluster - // Need to collect keys first to avoid borrow issues - let keys_to_update: Vec<_> = self + /// Apply canonical cross-shard redirects with one pass over local linker state. + pub fn apply_cross_shard_merges( + &mut self, + merges: &[(GlobalClusterId, GlobalClusterId)], + ) -> usize { + let redirects = merges + .iter() + .map(|(primary, secondary)| (*secondary, *primary)) + .collect::>(); + self.cross_shard_merges.extend( + redirects + .iter() + .map(|(secondary, primary)| (*secondary, *primary)), + ); + + let keys_to_update = self .boundary_signatures .keys() - .filter(|(_, gid)| *gid == secondary) - .cloned() - .collect(); - - for (sig, _) in keys_to_update { - if let Some(data) = self.boundary_signatures.remove(&(sig, secondary)) { - let new_key = (sig, primary); - // Merge with existing BoundaryData for primary if present + .filter_map(|(signature, cluster_id)| { + redirects + .get(cluster_id) + .map(|primary| (*signature, *cluster_id, *primary)) + }) + .collect::>(); + for (signature, secondary, primary) in keys_to_update { + if let Some(data) = self.boundary_signatures.remove(&(signature, secondary)) { self.boundary_signatures - .entry(new_key) + .entry((signature, primary)) .and_modify(|existing| existing.merge_from(&data)) .or_insert(data); - // Mark this key as dirty for adaptive reconciliation - self.dirty_boundary_keys.insert(sig); + self.dirty_boundary_keys.insert(signature); } } // Update global_cluster_ids map: any root pointing to secondary should now point to primary let mut updated_count = 0; for (_root, global_id) in self.global_cluster_ids.iter_mut() { - if *global_id == secondary { - *global_id = primary; + if let Some(primary) = redirects.get(global_id) { + *global_id = *primary; updated_count += 1; } } @@ -1702,8 +1826,12 @@ impl StreamingLinker { ) -> Result { let mappings = persistence.load_cross_shard_merges()?; let count = mappings.len(); - for (secondary, primary) in mappings { - self.apply_cross_shard_merge(primary, secondary); + if !mappings.is_empty() { + let merges = mappings + .into_iter() + .map(|(secondary, primary)| (primary, secondary)) + .collect::>(); + self.apply_cross_shard_merges(&merges); } Ok(count) } @@ -2109,9 +2237,10 @@ fn reconcile_cluster_ids( #[allow(clippy::too_many_arguments)] fn reconcile_global_cluster_ids( global_cluster_ids: &mut LinkerState, + global_cluster_anchors: &mut LinkerState, cross_shard_merges: &mut HashMap, - boundary_signatures: &mut HashMap<(ShardingKeySignature, GlobalClusterId), BoundaryData>, - dirty_boundary_keys: &mut HashSet, + boundary_signatures: &mut BTreeMap<(ShardingKeySignature, GlobalClusterId), BoundaryData>, + dirty_boundary_keys: &mut BTreeSet, shard_id: u16, root_a: RecordId, root_b: RecordId, @@ -2132,6 +2261,17 @@ fn reconcile_global_cluster_ids( current } + let anchor_a = global_cluster_anchors.get_copy(&root_a).unwrap_or(root_a.0); + let anchor_b = global_cluster_anchors.get_copy(&root_b).unwrap_or(root_b.0); + if root_a != new_root { + global_cluster_anchors.remove(&root_a); + } + if root_b != new_root { + global_cluster_anchors.remove(&root_b); + } + let stable_local_id = anchor_a.min(anchor_b); + global_cluster_anchors.insert(new_root, stable_local_id); + let id_a = global_cluster_ids.get_copy(&root_a); let id_b = global_cluster_ids.get_copy(&root_b); if id_a.is_none() && id_b.is_none() { @@ -2149,6 +2289,7 @@ fn reconcile_global_cluster_ids( .flatten() .flat_map(|id| [id, resolve(cross_shard_merges, id)]) .collect::>(); + ids.push(GlobalClusterId::new(shard_id, stable_local_id, 0)); ids.sort_by_key(GlobalClusterId::to_u64); ids.dedup(); let canonical = ids @@ -2173,7 +2314,7 @@ fn reconcile_global_cluster_ids( let keys_to_update = boundary_signatures .keys() .filter(|(_, global_id)| ids.contains(global_id) && *global_id != canonical) - .cloned() + .copied() .collect::>(); for (signature, old_id) in keys_to_update { if let Some(data) = boundary_signatures.remove(&(signature, old_id)) { diff --git a/src/model.rs b/src/model.rs index 71c3b73..b3fbfee 100644 --- a/src/model.rs +++ b/src/model.rs @@ -38,7 +38,7 @@ impl fmt::Display for ClusterId { /// - Tracking which shard owns a cluster /// - Detecting stale references after cross-shard merges /// - Efficient comparison and hashing -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] pub struct GlobalClusterId { /// The shard that owns this cluster pub shard_id: u16, diff --git a/src/partitioned.rs b/src/partitioned.rs index 0a205aa..a96b748 100644 --- a/src/partitioned.rs +++ b/src/partitioned.rs @@ -529,10 +529,18 @@ impl Partition { /// Get the dirty boundary keys for adaptive reconciliation pub fn get_dirty_boundary_keys( &self, - ) -> std::collections::HashSet { + ) -> std::collections::BTreeSet { self.linker.get_dirty_boundary_keys() } + pub fn dirty_boundary_key_candidates( + &self, + after: Option, + limit: usize, + ) -> Vec { + self.linker.dirty_boundary_key_candidates(after, limit) + } + /// Clear specific dirty boundary keys after reconciliation pub fn clear_dirty_boundary_keys(&mut self, keys: &[crate::sharding::IdentityKeySignature]) { self.linker.clear_dirty_boundary_keys(keys); @@ -1120,14 +1128,26 @@ impl ParallelPartitionedUnirust { /// Get all dirty boundary keys across all partitions pub fn get_all_dirty_boundary_keys( &self, - ) -> std::collections::HashSet { - let mut all_keys = std::collections::HashSet::new(); + ) -> std::collections::BTreeSet { + let mut all_keys = std::collections::BTreeSet::new(); for partition in &self.partitions { all_keys.extend(partition.lock().get_dirty_boundary_keys()); } all_keys } + pub fn dirty_boundary_key_candidates( + &self, + after: Option, + limit: usize, + ) -> Vec { + let mut candidates = std::collections::BTreeSet::new(); + for partition in &self.partitions { + candidates.extend(partition.lock().dirty_boundary_key_candidates(after, limit)); + } + candidates.into_iter().take(limit).collect() + } + /// Clear dirty boundary keys on all partitions pub fn clear_dirty_boundary_keys(&self, keys: &[crate::sharding::IdentityKeySignature]) { for partition in &self.partitions { diff --git a/src/persistence.rs b/src/persistence.rs index 375d068..6e89b2f 100644 --- a/src/persistence.rs +++ b/src/persistence.rs @@ -10,6 +10,7 @@ use rocksdb::{ Direction, IteratorMode, Options, SliceTransform, WriteBatch, WriteOptions, DB, }; use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; use std::cell::RefCell; use std::collections::HashMap; use std::fs; @@ -296,6 +297,15 @@ pub fn read_cluster_checkpoint_manifest(checkpoint: &Path) -> Result Result { + let manifest = read_cluster_checkpoint_manifest(checkpoint)?; + let checkpoint = checkpoint.canonicalize()?; + validate_rocksdb_checkpoint(&checkpoint)?; + Ok(manifest) +} + /// Return the committed checkpoint provenance copied into a restored RocksDB /// data directory. A lone or corrupt marker is a recovery integrity failure, /// not an unrestored directory. @@ -335,7 +345,7 @@ pub fn restore_checkpoint_for_shard( destination: &Path, expected_shard_id: Option, ) -> Result<()> { - let manifest = read_cluster_checkpoint_manifest(source)?; + let manifest = verify_cluster_checkpoint(source)?; if let Some(expected_shard_id) = expected_shard_id.filter(|expected| *expected != manifest.shard_id) { @@ -347,7 +357,6 @@ pub fn restore_checkpoint_for_shard( } let source = source.canonicalize()?; - validate_rocksdb_checkpoint(&source)?; if destination.exists() { let metadata = fs::symlink_metadata(destination)?; @@ -409,6 +418,29 @@ const CF_LINKER_CLUSTER_IDS: &str = "linker_cluster_ids"; const CF_LINKER_GLOBAL_IDS: &str = "linker_global_ids"; const CF_LINKER_METADATA: &str = "linker_metadata"; +const DURABLE_STATE_COLUMN_FAMILIES: &[&str] = &[ + CF_RECORDS, + CF_METADATA, + CF_INTERNER, + CF_INDEX_ATTR_VALUE, + CF_INDEX_ENTITY_TYPE, + CF_INDEX_PERSPECTIVE, + CF_INDEX_TEMPORAL_BUCKET, + CF_INDEX_IDENTITY, + CF_CONFLICT_SUMMARIES, + CF_CLUSTER_ASSIGNMENTS, + CF_SOURCE_RESERVATIONS, + CF_DSU_PARENT, + CF_DSU_RANK, + CF_DSU_GUARDS, + CF_DSU_METADATA, + CF_INDEX_IDENTITY_KEYS, + CF_INDEX_KEY_STATS, + CF_LINKER_CLUSTER_IDS, + CF_LINKER_GLOBAL_IDS, + CF_LINKER_METADATA, +]; + const RESET_DATA_CFS: &[&str] = &[ CF_RECORDS, CF_INTERNER, @@ -1421,6 +1453,24 @@ impl RecordStore for PersistentStore { } } + fn try_for_each_record_ordered(&self, f: &mut dyn FnMut(Record) -> Result<()>) -> Result<()> { + self.ensure_healthy()?; + let records_cf = self + .db + .cf_handle(CF_RECORDS) + .ok_or_else(|| anyhow!("missing records column family"))?; + for entry in self.db.iterator_cf(records_cf, IteratorMode::Start) { + let (_key, value) = entry.inspect_err(|_| { + self.mark_read_fault(); + })?; + let record = bincode::deserialize(&value).inspect_err(|_| { + self.mark_read_fault(); + })?; + f(record)?; + } + self.ensure_healthy() + } + fn get_records_by_entity_type(&self, entity_type: &str) -> Vec { let cf = match self.db.cf_handle(CF_INDEX_ENTITY_TYPE) { Some(cf) => cf, @@ -2295,6 +2345,29 @@ fn open_db(path: impl AsRef) -> Result { Ok(DB::open_cf_descriptors(&base, path, cfs)?) } +pub(crate) fn durable_state_digest(db: &DB) -> Result<[u8; 32]> { + fn update_field(digest: &mut Sha256, bytes: &[u8]) { + digest.update((bytes.len() as u64).to_be_bytes()); + digest.update(bytes); + } + + db.flush_wal(true)?; + let mut digest = Sha256::new(); + digest.update(b"unirust-durable-state-v1"); + for &name in DURABLE_STATE_COLUMN_FAMILIES { + update_field(&mut digest, name.as_bytes()); + let cf = db + .cf_handle(name) + .ok_or_else(|| anyhow!("missing column family {name}"))?; + for entry in db.iterator_cf(cf, IteratorMode::Start) { + let (key, value) = entry?; + update_field(&mut digest, &key); + update_field(&mut digest, &value); + } + } + Ok(digest.finalize().into()) +} + fn encode_attr_value_index(attr: u32, value: u32, start: i64, end: i64, record_id: u32) -> Vec { let mut key = Vec::with_capacity(4 + 4 + 8 + 8 + 4); key.extend_from_slice(&attr.to_be_bytes()); @@ -3076,6 +3149,8 @@ pub mod linker_encoding { /// Prefix for durable cross-shard merge mappings. pub const CROSS_SHARD_MERGE_PREFIX: &[u8] = b"cross_shard_merge/"; + pub const KEY_GLOBAL_CLUSTER_ID_SCHEME: &[u8] = b"global_cluster_id_scheme"; + pub const STABLE_RECORD_ANCHOR_SCHEME: u8 = 1; pub fn encode_cross_shard_merge_key(secondary: GlobalClusterId) -> Vec { let mut key = Vec::with_capacity(CROSS_SHARD_MERGE_PREFIX.len() + 8); @@ -3104,6 +3179,58 @@ impl<'a> LinkerStatePersistence<'a> { Self { db } } + /// 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 { + let cf = self + .db + .cf_handle(linker_cf::METADATA) + .ok_or_else(|| anyhow::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]) => return Ok(false), + Some(value) => { + anyhow::bail!("unsupported global cluster ID scheme marker: {:?}", value); + } + None => {} + } + + let had_legacy_redirects = self + .db + .iterator_cf( + &cf, + IteratorMode::From( + linker_encoding::CROSS_SHARD_MERGE_PREFIX, + Direction::Forward, + ), + ) + .next() + .transpose()? + .is_some_and(|(key, _)| key.starts_with(linker_encoding::CROSS_SHARD_MERGE_PREFIX)); + + let mut prefix_end = linker_encoding::CROSS_SHARD_MERGE_PREFIX.to_vec(); + let last = prefix_end + .last_mut() + .ok_or_else(|| anyhow!("cross-shard merge prefix must not be empty"))?; + *last = last + .checked_add(1) + .ok_or_else(|| anyhow!("cross-shard merge prefix has no upper bound"))?; + + let mut batch = WriteBatch::default(); + batch.delete_range_cf(&cf, linker_encoding::CROSS_SHARD_MERGE_PREFIX, &prefix_end); + batch.put_cf( + &cf, + linker_encoding::KEY_GLOBAL_CLUSTER_ID_SCHEME, + [linker_encoding::STABLE_RECORD_ANCHOR_SCHEME], + ); + self.db.write(batch)?; + self.db.flush_wal(true)?; + Ok(had_legacy_redirects) + } + /// Flush cluster ID mappings to the database. pub fn flush_cluster_ids(&self, mappings: I) -> Result<()> where @@ -3165,16 +3292,28 @@ impl<'a> LinkerStatePersistence<'a> { &self, secondary: GlobalClusterId, primary: GlobalClusterId, + ) -> Result<()> { + self.save_cross_shard_merges(&[(secondary, primary)]) + } + + /// Persist cross-shard redirects in one RocksDB write batch. + pub fn save_cross_shard_merges( + &self, + merges: &[(GlobalClusterId, GlobalClusterId)], ) -> Result<()> { let cf = self .db .cf_handle(linker_cf::METADATA) .ok_or_else(|| anyhow::anyhow!("Column family {} not found", linker_cf::METADATA))?; - self.db.put_cf( - &cf, - linker_encoding::encode_cross_shard_merge_key(secondary), - linker_encoding::encode_global_cluster_id(primary), - )?; + let mut batch = WriteBatch::default(); + for (secondary, primary) in merges { + batch.put_cf( + &cf, + linker_encoding::encode_cross_shard_merge_key(*secondary), + linker_encoding::encode_global_cluster_id(*primary), + ); + } + self.db.write(batch)?; Ok(()) } diff --git a/src/sharding.rs b/src/sharding.rs index 0cf3c28..e9cdaf7 100644 --- a/src/sharding.rs +++ b/src/sharding.rs @@ -22,7 +22,7 @@ fn hash_length_prefixed(hasher: &mut Sha256, value: &[u8]) { /// A signature of identity key values for fast lookup. /// This is a 32-byte hash of the identity key's attribute-value pairs. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] pub struct IdentityKeySignature(pub [u8; 32]); impl IdentityKeySignature { @@ -473,6 +473,46 @@ pub struct ReconciliationResult { pub detected_conflicts: Vec, } +#[derive(Debug, Default)] +pub(crate) struct ReconciliationCandidates { + merges: Vec<(GlobalClusterId, GlobalClusterId)>, + keys_checked: usize, + keys_matched: usize, + merge_candidates: usize, + conflicts_blocked: usize, + detected_conflicts: Vec, +} + +impl ReconciliationCandidates { + pub(crate) fn extend(&mut self, mut other: Self) { + self.merges.append(&mut other.merges); + 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); + self.conflicts_blocked = self + .conflicts_blocked + .saturating_add(other.conflicts_blocked); + self.detected_conflicts + .append(&mut other.detected_conflicts); + } + + pub(crate) fn finish(self) -> ReconciliationResult { + let (merges, component_conflicts_blocked) = + canonicalize_merges(self.merges, &self.detected_conflicts); + ReconciliationResult { + merges_performed: merges.len(), + merged_clusters: merges, + keys_checked: self.keys_checked, + keys_matched: self.keys_matched, + merge_candidates: self.merge_candidates, + conflicts_blocked: self + .conflicts_blocked + .saturating_add(component_conflicts_blocked), + detected_conflicts: self.detected_conflicts, + } + } +} + /// Incremental reconciler for cross-shard cluster merges. /// /// Instead of loading all records from all shards (O(n)), @@ -645,10 +685,15 @@ impl IncrementalReconciler { /// This is more efficient than full reconciliation when only a subset /// of keys have changed. Returns the reconciliation result. pub fn reconcile_keys(&mut self, keys: &HashSet) -> ReconciliationResult { - if keys.is_empty() { - return ReconciliationResult::default(); - } + let result = self.reconciliation_candidates(keys).finish(); + self.pending_merges = result.merged_clusters.clone(); + result + } + pub(crate) fn reconciliation_candidates( + &self, + keys: &HashSet, + ) -> ReconciliationCandidates { let mut merges = Vec::new(); let mut merge_candidates = 0; let mut conflicts_blocked = 0usize; @@ -714,14 +759,8 @@ impl IncrementalReconciler { } } - let (merges, component_conflicts_blocked) = - canonicalize_merges(merges, &detected_conflicts); - conflicts_blocked = conflicts_blocked.saturating_add(component_conflicts_blocked); - self.pending_merges = merges.clone(); - - ReconciliationResult { - merges_performed: merges.len(), - merged_clusters: merges, + ReconciliationCandidates { + merges, keys_checked: keys.len(), keys_matched, merge_candidates, @@ -1657,6 +1696,45 @@ mod tests { assert_eq!(blocked, 1); } + #[test] + fn reconciliation_candidates_preserve_global_component_conflicts_across_batches() { + let a = GlobalClusterId::new(0, 10, 0); + let b = GlobalClusterId::new(1, 20, 0); + let c = GlobalClusterId::new(2, 30, 0); + let conflict = CrossShardConflict { + identity_key_signature: IdentityKeySignature::from_bytes([9; 32]), + cluster1: a, + cluster2: c, + interval: Interval::new(0, 100).unwrap(), + perspective_hash: 1, + strong_id_hash1: 2, + strong_id_hash2: 3, + }; + let mut combined = ReconciliationCandidates::default(); + combined.extend(ReconciliationCandidates { + merges: vec![(a, b)], + keys_checked: 1, + keys_matched: 1, + merge_candidates: 1, + ..ReconciliationCandidates::default() + }); + combined.extend(ReconciliationCandidates { + merges: vec![(b, c)], + keys_checked: 2, + keys_matched: 2, + merge_candidates: 3, + conflicts_blocked: 1, + detected_conflicts: vec![conflict], + }); + + let result = combined.finish(); + assert_eq!(result.merged_clusters, vec![(a, b)]); + assert_eq!(result.keys_checked, 3); + assert_eq!(result.keys_matched, 3); + assert_eq!(result.merge_candidates, 4); + assert_eq!(result.conflicts_blocked, 2); + } + #[test] fn test_boundary_metadata_export_import() { use crate::model::{AttrId, ValueId}; diff --git a/src/store.rs b/src/store.rs index 7db5a57..8dc51fc 100644 --- a/src/store.rs +++ b/src/store.rs @@ -86,6 +86,19 @@ pub trait RecordStore: Send + Sync { } } + /// Apply a fallible function to every record in ascending record-ID order. + /// + /// Persistent implementations should stream their ordered keyspace instead + /// of materializing all records so linker recovery stays sequential. + fn try_for_each_record_ordered(&self, f: &mut dyn FnMut(Record) -> Result<()>) -> Result<()> { + let mut records = self.get_all_records(); + records.sort_by_key(|record| record.id.0); + for record in records { + f(record)?; + } + Ok(()) + } + /// Get records for a specific entity type. fn get_records_by_entity_type(&self, entity_type: &str) -> Vec; diff --git a/src/transport_security.rs b/src/transport_security.rs index 2e5749a..97a01b4 100644 --- a/src/transport_security.rs +++ b/src/transport_security.rs @@ -12,6 +12,24 @@ fn read_nonempty(path: &Path, description: &str) -> Result> { Ok(bytes) } +/// Load a replication secret without retaining the raw file contents in +/// process configuration or logs. +pub fn load_replication_token(path: &Path) -> Result { + use sha2::{Digest, Sha256}; + + let token = read_nonempty(path, "replication token")?; + if token.len() < 32 { + anyhow::bail!("replication token must contain at least 32 bytes"); + } + let digest = Sha256::digest(token); + let mut encoded = String::with_capacity(digest.len() * 2); + for byte in digest { + use std::fmt::Write as _; + write!(&mut encoded, "{byte:02x}")?; + } + Ok(encoded) +} + /// Load a server identity and require client certificates signed by `client_ca`. pub fn load_server_mtls( cert: Option<&Path>, diff --git a/tests/backup_cli_safety.rs b/tests/backup_cli_safety.rs new file mode 100644 index 0000000..8260dbf --- /dev/null +++ b/tests/backup_cli_safety.rs @@ -0,0 +1,40 @@ +use std::process::Command; + +#[test] +fn backup_cli_requires_an_explicit_operation() { + let output = Command::new(env!("CARGO_BIN_EXE_unirust_backup")) + .output() + .expect("run backup binary"); + assert!(!output.status.success()); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!(stderr.contains("Usage:")); +} + +#[test] +fn backup_cli_rejects_incomplete_export_arguments() { + let output = Command::new(env!("CARGO_BIN_EXE_unirust_backup")) + .args(["export", "--destination", "/tmp/not-created"]) + .output() + .expect("run backup binary"); + assert!(!output.status.success()); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!(stderr.contains("at least one committed shard checkpoint")); +} + +#[test] +fn backup_cli_rejects_zero_retention() { + let root = tempfile::tempdir().expect("temporary retention root"); + let output = Command::new(env!("CARGO_BIN_EXE_unirust_backup")) + .args([ + "prune", + "--root", + root.path().to_str().expect("UTF-8 path"), + "--retain", + "0", + ]) + .output() + .expect("run backup binary"); + assert!(!output.status.success()); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!(stderr.contains("keep at least one generation")); +} diff --git a/tests/distributed_apply_ontology.rs b/tests/distributed_apply_ontology.rs index ae9775e..3c8939c 100644 --- a/tests/distributed_apply_ontology.rs +++ b/tests/distributed_apply_ontology.rs @@ -118,7 +118,7 @@ async fn apply_ontology_enables_queries_distributed() -> anyhow::Result<()> { client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![record], }) .await?; diff --git a/tests/distributed_conflicts_presets.rs b/tests/distributed_conflicts_presets.rs index 6aa89d9..0c60c06 100644 --- a/tests/distributed_conflicts_presets.rs +++ b/tests/distributed_conflicts_presets.rs @@ -272,7 +272,7 @@ async fn distributed_conflict_presets_match_local() -> anyhow::Result<()> { if !inputs.is_empty() { client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: inputs.clone(), }) .await diff --git a/tests/distributed_cross_shard_reconciliation.rs b/tests/distributed_cross_shard_reconciliation.rs index 2ac957f..a8415fa 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: 3, + internal_protocol_version: 5, 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: 3, + internal_protocol_version: 5, 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: 3, + internal_protocol_version: 5, 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: 3, + internal_protocol_version: 5, 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: 3, + internal_protocol_version: 5, 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: 3, + internal_protocol_version: 5, 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: 3, + internal_protocol_version: 5, 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: 3, + internal_protocol_version: 5, 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: 3, + internal_protocol_version: 5, 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: 3, + internal_protocol_version: 5, records: vec![record_input( 2, "person", @@ -797,14 +797,31 @@ async fn boundary_metadata_includes_perspective_strong_ids() -> anyhow::Result<( router_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![record1, record2], }) .await?; - // Get boundary metadata from shard + let dirty_keys = shard_client + .get_dirty_boundary_keys(proto::GetDirtyBoundaryKeysRequest { + after_signature: Vec::new(), + limit: 100, + }) + .await? + .into_inner() + .dirty_keys; + let signatures = dirty_keys + .into_iter() + .filter_map(|dirty_key| dirty_key.signature) + .collect::>(); + assert!(!signatures.is_empty(), "expected dirty boundary keys"); + + // Boundary metadata is fetched only for the requested dirty signatures. let metadata_response = shard_client - .get_boundary_metadata(proto::GetBoundaryMetadataRequest { since_version: 0 }) + .get_boundary_metadata(proto::GetBoundaryMetadataRequest { + since_version: 0, + signatures, + }) .await? .into_inner(); @@ -832,7 +849,7 @@ async fn boundary_metadata_includes_perspective_strong_ids() -> anyhow::Result<( Ok(()) } -/// Test dirty boundary keys also include exact temporal strong-ID observations. +/// Test dirty boundary keys require a bounded metadata lookup for strong-ID observations. #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn dirty_boundary_keys_include_perspective_strong_ids() -> anyhow::Result<()> { let _test_guard = PERSISTENT_TEST_LOCK.lock().await; @@ -882,22 +899,47 @@ async fn dirty_boundary_keys_include_perspective_strong_ids() -> anyhow::Result< router_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![record1, record2], }) .await?; // Get dirty boundary keys from shard let dirty_keys_response = shard_client - .get_dirty_boundary_keys(proto::GetDirtyBoundaryKeysRequest {}) + .get_dirty_boundary_keys(proto::GetDirtyBoundaryKeysRequest { + after_signature: Vec::new(), + limit: 100, + }) .await? .into_inner(); - - let exact_strong_ids = dirty_keys_response + assert!( + dirty_keys_response + .dirty_keys + .iter() + .all(|dirty_key| dirty_key.entries.is_empty()), + "dirty-key pages must not duplicate boundary metadata" + ); + let signatures = dirty_keys_response .dirty_keys .iter() - .flat_map(|dirty_key| &dirty_key.entries) - .flat_map(|entry| &entry.strong_ids) + .filter_map(|dirty_key| dirty_key.signature.clone()) + .collect::>(); + assert!(!signatures.is_empty(), "expected dirty boundary keys"); + + let metadata = shard_client + .get_boundary_metadata(proto::GetBoundaryMetadataRequest { + since_version: 0, + signatures, + }) + .await? + .into_inner() + .metadata + .expect("metadata should exist"); + let exact_strong_ids = metadata + .entries + .iter() + .flat_map(|key_entries| &key_entries.entries) + .flat_map(|entry| entry.strong_ids.iter()) .collect::>(); assert!( @@ -908,13 +950,91 @@ async fn dirty_boundary_keys_include_perspective_strong_ids() -> anyhow::Result< && strong_id.interval_start == 0 && strong_id.interval_end == 100 }), - "expected exact temporal strong-ID dirty-key metadata; dirty keys: {:?}", - dirty_keys_response.dirty_keys + "expected exact temporal strong-ID boundary metadata; entries: {:?}", + metadata.entries ); Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn dirty_boundary_keys_are_drained_in_bounded_pages() -> anyhow::Result<()> { + let _test_guard = PERSISTENT_TEST_LOCK.lock().await; + let config = build_instrument_config(); + let empty_config = DistributedOntologyConfig::empty(); + + let (shard0_addr, _shard0_handle) = spawn_shard(0, empty_config.clone()).await?; + let (router_addr, _router_handle) = + spawn_router(vec![shard0_addr], empty_config.clone()).await?; + let mut router_client = RouterServiceClient::connect(format!("http://{}", router_addr)).await?; + let mut shard_client = ShardServiceClient::connect(format!("http://{}", shard0_addr)).await?; + + router_client + .set_ontology(ApplyOntologyRequest { + config: Some(to_proto_config(&config)), + }) + .await?; + + for index in 0..5 { + let isin = format!("PAGE{index:03}"); + let uid1 = format!("page_{index}_a"); + let uid2 = format!("page_{index}_b"); + router_client + .ingest_records(IngestRecordsRequest { + internal_protocol_version: 5, + records: vec![ + record_input( + index * 2, + "instrument", + "page-test", + &uid1, + vec![("isin", &isin, 0, 100), ("country", "US", 0, 100)], + ), + record_input( + index * 2 + 1, + "instrument", + "page-test", + &uid2, + vec![("isin", &isin, 0, 100), ("country", "US", 0, 100)], + ), + ], + }) + .await?; + } + + let mut seen = std::collections::HashSet::new(); + let mut after_signature = Vec::new(); + loop { + let page = shard_client + .get_dirty_boundary_keys(proto::GetDirtyBoundaryKeysRequest { + after_signature: after_signature.clone(), + limit: 2, + }) + .await? + .into_inner(); + assert!(page.dirty_keys.len() <= 2); + if page.dirty_keys.is_empty() { + assert!(!page.has_more); + break; + } + for key in page.dirty_keys.into_iter().map(|dirty_key| { + assert!(dirty_key.entries.is_empty()); + dirty_key.signature.expect("signature should exist") + }) { + assert!(seen.insert(key.signature), "duplicate dirty key"); + } + if !page.has_more { + break; + } + assert_eq!(page.next_after_signature.len(), 32); + assert!(page.next_after_signature > after_signature); + after_signature = page.next_after_signature; + } + assert_eq!(seen.len(), 5); + + Ok(()) +} + // ============================================================================= // ADVANCED DISTRIBUTED CONFLICT SCENARIOS // ============================================================================= @@ -1014,7 +1134,7 @@ async fn transitive_cross_shard_conflict_detected() -> anyhow::Result<()> { ); shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![a1, a2], }) .await?; @@ -1046,7 +1166,7 @@ async fn transitive_cross_shard_conflict_detected() -> anyhow::Result<()> { ); shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![b1, b2], }) .await?; @@ -1076,7 +1196,7 @@ async fn transitive_cross_shard_conflict_detected() -> anyhow::Result<()> { ); shard2_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![c1, c2], }) .await?; @@ -1201,7 +1321,7 @@ async fn peic_many_entities_claim_same_identifier_across_shards() -> anyhow::Res ); shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![a1, a2], }) .await?; @@ -1232,7 +1352,7 @@ async fn peic_many_entities_claim_same_identifier_across_shards() -> anyhow::Res ); shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![b1, b2], }) .await?; @@ -1329,7 +1449,7 @@ async fn temporal_overlap_conflict_across_shards() -> anyhow::Result<()> { ); shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![a1, a2], }) .await?; @@ -1359,7 +1479,7 @@ async fn temporal_overlap_conflict_across_shards() -> anyhow::Result<()> { ); shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![b1, b2], }) .await?; @@ -1453,7 +1573,7 @@ async fn late_arriving_data_triggers_conflict() -> anyhow::Result<()> { ); shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![initial1, initial2], }) .await?; @@ -1498,7 +1618,7 @@ async fn late_arriving_data_triggers_conflict() -> anyhow::Result<()> { ); shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![late1, late2], }) .await?; @@ -1622,7 +1742,7 @@ async fn multi_hop_chain_conflict_across_four_shards() -> anyhow::Result<()> { ); shard0 .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![a1, a2], }) .await?; @@ -1652,7 +1772,7 @@ async fn multi_hop_chain_conflict_across_four_shards() -> anyhow::Result<()> { ); shard1 .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![b1, b2], }) .await?; @@ -1682,7 +1802,7 @@ async fn multi_hop_chain_conflict_across_four_shards() -> anyhow::Result<()> { ); shard2 .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![c1, c2], }) .await?; @@ -1707,7 +1827,7 @@ async fn multi_hop_chain_conflict_across_four_shards() -> anyhow::Result<()> { ); shard3 .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![d1, d2], }) .await?; @@ -1810,7 +1930,7 @@ async fn different_perspectives_no_false_positive_conflict() -> anyhow::Result<( ); shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![msci1, msci2], }) .await?; @@ -1840,7 +1960,7 @@ async fn different_perspectives_no_false_positive_conflict() -> anyhow::Result<( ); shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![axioma1, axioma2], }) .await?; diff --git a/tests/distributed_cross_shard_reconciliation_persistent.rs b/tests/distributed_cross_shard_reconciliation_persistent.rs index ea3d9c6..703f25f 100644 --- a/tests/distributed_cross_shard_reconciliation_persistent.rs +++ b/tests/distributed_cross_shard_reconciliation_persistent.rs @@ -103,11 +103,11 @@ where } fn call(&mut self, request: http::Request) -> Self::Future { - if request.uri().path() == "/unirust.ShardService/ApplyMerge" + if request.uri().path() == "/unirust.ShardService/ApplyMerges" && self.fail_next.swap(false, Ordering::AcqRel) { return Box::pin(async { - Ok(tonic::Status::unavailable("injected ApplyMerge failure").into_http()) + Ok(tonic::Status::unavailable("injected ApplyMerges failure").into_http()) }); } Box::pin(self.inner.call(request)) @@ -211,6 +211,41 @@ async fn spawn_apply_merge_failing_shard_at( Ok((addr, shard, handle)) } +async fn spawn_apply_merge_stalling_shard_at( + shard_id: u32, + config: DistributedOntologyConfig, + data_path: PathBuf, +) -> anyhow::Result<( + SocketAddr, + ShardNode, + JoinHandle<()>, + support::StallControls, +)> { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let addr = listener.local_addr()?; + let shard = ShardNode::new_with_data_dir( + shard_id, + config, + StreamingTuning::from_profile(TuningProfile::Balanced), + Some(data_path), + false, + None, + )?; + let (service, controls) = support::stall_response( + proto::shard_service_server::ShardServiceServer::new(shard.clone()), + "/unirust.ShardService/ApplyMerges", + true, + ); + let handle = tokio::spawn(async move { + Server::builder() + .add_service(service) + .serve_with_incoming(TcpListenerStream::new(listener)) + .await + .expect("shard server"); + }); + Ok((addr, shard, handle, controls)) +} + async fn spawn_ingest_failing_shard_at( shard_id: u32, config: DistributedOntologyConfig, @@ -347,7 +382,7 @@ async fn cross_shard_merge_persistent_store() -> anyhow::Result<()> { ); shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![shard0_rec1, shard0_rec2], }) .await?; @@ -374,7 +409,7 @@ async fn cross_shard_merge_persistent_store() -> anyhow::Result<()> { ); shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![shard1_rec1, shard1_rec2], }) .await?; @@ -424,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: 3, + internal_protocol_version: 5, records: vec![record_input( shard_id as u32, "person", @@ -534,7 +569,7 @@ async fn source_identity_reservation_survives_routing_change_and_restart() -> an let response = router_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![original.clone()], }) .await? @@ -544,7 +579,7 @@ async fn source_identity_reservation_survives_routing_change_and_restart() -> an let error = router_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![changed.clone()], }) .await @@ -581,7 +616,7 @@ async fn source_identity_reservation_survives_routing_change_and_restart() -> an let error = router_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![changed], }) .await @@ -591,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: 3, + internal_protocol_version: 5, records: vec![original], }) .await? @@ -658,13 +693,15 @@ async fn router_backfills_legacy_records_before_serving_ingest() -> anyhow::Resu let mut original_target_client = ShardServiceClient::connect(format!("http://{}", shard_addrs[original_target])).await?; let status_before = original_target_client - .get_config_version(proto::ConfigVersionRequest {}) + .get_config_version(proto::ConfigVersionRequest { + include_durable_state_digest: false, + }) .await? .into_inner(); assert_eq!(status_before.source_reservation_backfill_version, 0); original_target_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![original], }) .await?; @@ -673,7 +710,9 @@ async fn router_backfills_legacy_records_before_serving_ingest() -> anyhow::Resu for shard_addr in &shard_addrs { let mut shard_client = ShardServiceClient::connect(format!("http://{shard_addr}")).await?; let status = shard_client - .get_config_version(proto::ConfigVersionRequest {}) + .get_config_version(proto::ConfigVersionRequest { + include_durable_state_digest: false, + }) .await? .into_inner(); assert_eq!( @@ -686,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: 3, + internal_protocol_version: 5, records: vec![changed], }) .await @@ -752,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: 3, + internal_protocol_version: 5, records: vec![record.clone()], }) .await @@ -790,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: 3, + internal_protocol_version: 5, records: vec![record.clone()], }) .await? @@ -803,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: 3, + internal_protocol_version: 5, records: vec![changed], }) .await @@ -836,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: 3, + internal_protocol_version: 5, records: vec![record_input( 0, "person", @@ -851,7 +890,7 @@ async fn partial_reconciliation_blocks_traffic_and_recovers_after_full_restart( .await?; client1 .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![record_input( 1, "person", @@ -947,6 +986,94 @@ async fn partial_reconciliation_blocks_traffic_and_recovers_after_full_restart( Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn cancelled_partial_reconciliation_latches_router_closed() -> anyhow::Result<()> { + use proto::router_service_server::RouterService; + + let _test_guard = PERSISTENT_TEST_LOCK.lock().await; + let root = tempfile::tempdir()?; + let config = build_email_config(); + let (addr0, shard0, handle0) = + spawn_shard_at(0, config.clone(), root.path().join("cancel-shard-0")).await?; + let (addr1, shard1, handle1, controls) = + spawn_apply_merge_stalling_shard_at(1, config.clone(), root.path().join("cancel-shard-1")) + .await?; + let router = RouterNode::connect( + vec![format!("http://{addr0}"), format!("http://{addr1}")], + config, + ) + .await?; + + for (shard_addr, uid) in [(addr0, "cancel-merge-0"), (addr1, "cancel-merge-1")] { + let mut client = ShardServiceClient::connect(format!("http://{shard_addr}")).await?; + client + .ingest_records(IngestRecordsRequest { + internal_protocol_version: DISTRIBUTED_PROTOCOL_VERSION, + records: vec![record_input( + 0, + "person", + "hr", + uid, + vec![ + ("email", "cancel-merge@example.com", 0, 100), + ("ssn", "1234", 0, 100), + ], + )], + }) + .await?; + } + + let reconcile_router = router.clone(); + let reconcile_task = tokio::spawn(async move { + RouterService::reconcile( + reconcile_router.as_ref(), + tonic::Request::new(proto::ReconcileRequest { + shard_metadata: Vec::new(), + }), + ) + .await + }); + controls + .wait_until_committed(Duration::from_secs(10)) + .await?; + reconcile_task.abort(); + assert!(reconcile_task + .await + .expect_err("reconciliation task must be cancelled") + .is_cancelled()); + controls.release(); + + let blocked = RouterService::health_check( + router.as_ref(), + tonic::Request::new(proto::HealthCheckRequest {}), + ) + .await + .expect_err("router must fail closed after cancellation during reconciliation apply"); + assert_eq!(blocked.code(), tonic::Code::FailedPrecondition); + + RouterService::reconcile( + router.as_ref(), + tonic::Request::new(proto::ReconcileRequest { + shard_metadata: Vec::new(), + }), + ) + .await?; + RouterService::health_check( + router.as_ref(), + tonic::Request::new(proto::HealthCheckRequest {}), + ) + .await?; + + drop(router); + shard0.shutdown().await?; + shard1.shutdown().await?; + handle0.abort(); + handle1.abort(); + let _ = handle0.await; + let _ = handle1.await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn partial_reconciliation_can_be_retried_in_place() -> anyhow::Result<()> { use proto::router_service_server::RouterService; @@ -968,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: 3, + internal_protocol_version: 5, records: vec![record_input( 0, "person", diff --git a/tests/distributed_e2e.rs b/tests/distributed_e2e.rs index e6e89e4..d84883c 100644 --- a/tests/distributed_e2e.rs +++ b/tests/distributed_e2e.rs @@ -208,7 +208,7 @@ async fn distributed_stream_and_query() -> anyhow::Result<()> { let response = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![record_a.clone(), record_b.clone()], }) .await? diff --git a/tests/distributed_ingest_cancellation.rs b/tests/distributed_ingest_cancellation.rs new file mode 100644 index 0000000..6f6cd78 --- /dev/null +++ b/tests/distributed_ingest_cancellation.rs @@ -0,0 +1,115 @@ +use std::time::Duration; + +use tempfile::tempdir; +use tonic::Request; +use unirust_rs::distributed::proto::shard_service_server::ShardService; +use unirust_rs::distributed::proto::{ + HealthCheckRequest, IngestRecordsRequest, RecordDescriptor, RecordIdentity, RecordInput, + StatsRequest, +}; +use unirust_rs::distributed::{DistributedOntologyConfig, ShardNode, DISTRIBUTED_PROTOCOL_VERSION}; +use unirust_rs::StreamingTuning; + +fn records(count: u32) -> Vec { + (0..count) + .map(|index| RecordInput { + index, + identity: Some(RecordIdentity { + entity_type: "person".to_string(), + perspective: "crm".to_string(), + uid: format!("cancel-{index}"), + }), + descriptors: vec![RecordDescriptor { + attr: "email".to_string(), + value: format!("cancel-{index}@example.com"), + start: 0, + end: 10, + }], + }) + .collect() +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn cancelled_ingest_finishes_commit_and_survives_restart() -> anyhow::Result<()> { + const RECORD_COUNT: u32 = 5_000; + + let temp_dir = tempdir()?; + let shard = ShardNode::new_with_data_dir( + 0, + DistributedOntologyConfig::empty(), + StreamingTuning::balanced(), + Some(temp_dir.path().to_path_buf()), + false, + None, + )?; + let request_records = records(RECORD_COUNT); + let service = shard.clone(); + let request_task = tokio::spawn({ + let request_records = request_records.clone(); + async move { + ShardService::ingest_records( + &service, + Request::new(IngestRecordsRequest { + internal_protocol_version: DISTRIBUTED_PROTOCOL_VERSION, + records: request_records, + }), + ) + .await + } + }); + + let wal_path = temp_dir.path().join("ingest_wal.bin"); + tokio::time::timeout(Duration::from_secs(10), async { + while !wal_path.exists() { + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + .await?; + request_task.abort(); + assert!(request_task + .await + .expect_err("request task must be cancelled") + .is_cancelled()); + + tokio::time::timeout(Duration::from_secs(30), async { + while wal_path.exists() { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await?; + + let retry = ShardService::ingest_records( + &shard, + Request::new(IngestRecordsRequest { + internal_protocol_version: DISTRIBUTED_PROTOCOL_VERSION, + records: request_records, + }), + ) + .await? + .into_inner(); + assert_eq!(retry.assignments.len(), RECORD_COUNT as usize); + let stats = ShardService::get_stats(&shard, Request::new(StatsRequest {})) + .await? + .into_inner(); + assert_eq!(stats.record_count, u64::from(RECORD_COUNT)); + ShardService::health_check(&shard, Request::new(HealthCheckRequest {})).await?; + + shard.shutdown().await?; + drop(shard); + + let restarted = ShardNode::new_with_data_dir( + 0, + DistributedOntologyConfig::empty(), + StreamingTuning::balanced(), + Some(temp_dir.path().to_path_buf()), + false, + None, + )?; + let stats = ShardService::get_stats(&restarted, Request::new(StatsRequest {})) + .await? + .into_inner(); + assert_eq!(stats.record_count, u64::from(RECORD_COUNT)); + ShardService::health_check(&restarted, Request::new(HealthCheckRequest {})).await?; + restarted.shutdown().await?; + Ok(()) +} diff --git a/tests/distributed_ingest_stream.rs b/tests/distributed_ingest_stream.rs index 2024906..3d3aa9d 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: 3, + internal_protocol_version: 5, 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: 3, + internal_protocol_version: 5, records: vec![record_input( 1, "person", diff --git a/tests/distributed_metrics.rs b/tests/distributed_metrics.rs index dc8eca9..894ff8c 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: 3, + internal_protocol_version: 5, records: vec![record_input( 0, "person", diff --git a/tests/distributed_rebalance.rs b/tests/distributed_rebalance.rs index 22141ef..dc05227 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: 3, + internal_protocol_version: 5, records, }) .await? @@ -163,7 +163,7 @@ async fn distributed_export_import_range() -> anyhow::Result<()> { shard1 .import_records(ImportRecordsRequest { records: all_records, - internal_protocol_version: 3, + internal_protocol_version: 5, }) .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: 3, + internal_protocol_version: 5, }) .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 003ce27..581ad76 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: 3, + internal_protocol_version: 5, records, }) .await?; diff --git a/tests/distributed_reliability.rs b/tests/distributed_reliability.rs index e074005..f969964 100644 --- a/tests/distributed_reliability.rs +++ b/tests/distributed_reliability.rs @@ -134,6 +134,37 @@ async fn spawn_set_ontology_failing_shard( Ok((addr, handle)) } +async fn spawn_set_ontology_stalling_shard( + shard_id: u32, + config: DistributedOntologyConfig, +) -> anyhow::Result<(SocketAddr, JoinHandle<()>, support::StallControls)> { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let addr = listener.local_addr()?; + let data_dir = tempdir()?; + let shard = ShardNode::new_with_data_dir( + shard_id, + config, + StreamingTuning::from_profile(TuningProfile::Balanced), + Some(data_dir.path().to_path_buf()), + false, + None, + )?; + let (service, controls) = support::stall_response( + proto::shard_service_server::ShardServiceServer::new(shard), + "/unirust.ShardService/SetOntology", + true, + ); + let handle = tokio::spawn(async move { + let _data_dir = data_dir; + Server::builder() + .add_service(service) + .serve_with_incoming(TcpListenerStream::new(listener)) + .await + .expect("shard server"); + }); + Ok((addr, handle, controls)) +} + async fn spawn_stoppable_shard( shard_id: u32, config: DistributedOntologyConfig, @@ -277,7 +308,7 @@ async fn shard_recovery_after_restart() -> anyhow::Result<()> { .collect(); client .ingest_records(proto::IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records, }) .await?; @@ -342,7 +373,7 @@ async fn router_handles_partial_shard_availability() -> anyhow::Result<()> { ); let response = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![record], }) .await? @@ -523,7 +554,7 @@ async fn destructive_reset_is_disabled_and_preserves_records_by_default() -> any client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![record_input( 0, "person", @@ -603,7 +634,7 @@ async fn incomplete_ontology_update_blocks_cluster_traffic() -> anyhow::Result<( let ingest_error = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![record_input( 0, "person", @@ -620,6 +651,100 @@ async fn incomplete_ontology_update_blocks_cluster_traffic() -> anyhow::Result<( Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn cancelled_ontology_update_blocks_traffic_until_retried() -> anyhow::Result<()> { + use proto::router_service_server::RouterService; + + let _test_guard = PERSISTENT_TEST_LOCK.lock().await; + let empty = DistributedOntologyConfig::empty(); + let intended = support::build_iam_config(); + let (shard0_addr, shard0_handle, controls) = + spawn_set_ontology_stalling_shard(0, empty.clone()).await?; + let (shard1_addr, shard1_handle) = spawn_shard(1, empty.clone()).await?; + let router = RouterNode::connect( + vec![ + format!("http://{shard0_addr}"), + format!("http://{shard1_addr}"), + ], + empty.clone(), + ) + .await?; + + let update_router = router.clone(); + let update_task = tokio::spawn({ + let intended = intended.clone(); + async move { + RouterService::set_ontology( + update_router.as_ref(), + tonic::Request::new(ApplyOntologyRequest { + config: Some(support::to_proto_config(&intended)), + }), + ) + .await + } + }); + controls + .wait_until_committed(Duration::from_secs(10)) + .await?; + update_task.abort(); + assert!(update_task + .await + .expect_err("ontology update task must be cancelled") + .is_cancelled()); + controls.release(); + + let health_error = RouterService::health_check( + router.as_ref(), + tonic::Request::new(proto::HealthCheckRequest {}), + ) + .await + .expect_err("router must fail closed after a cancelled partial ontology update"); + assert_eq!(health_error.code(), tonic::Code::FailedPrecondition); + + let mut shard0 = ShardServiceClient::connect(format!("http://{shard0_addr}")).await?; + let mut shard1 = ShardServiceClient::connect(format!("http://{shard1_addr}")).await?; + let shard0_config = shard0 + .get_config_version(proto::ConfigVersionRequest { + include_durable_state_digest: false, + }) + .await? + .into_inner() + .ontology_config + .expect("shard 0 ontology"); + let shard1_config = shard1 + .get_config_version(proto::ConfigVersionRequest { + include_durable_state_digest: false, + }) + .await? + .into_inner() + .ontology_config + .expect("shard 1 ontology"); + assert_eq!(shard0_config, support::to_proto_config(&intended)); + assert_eq!(shard1_config, support::to_proto_config(&empty)); + + RouterService::set_ontology( + router.as_ref(), + tonic::Request::new(ApplyOntologyRequest { + config: Some(support::to_proto_config(&intended)), + }), + ) + .await?; + RouterService::health_check( + router.as_ref(), + tonic::Request::new(proto::HealthCheckRequest {}), + ) + .await?; + + drop(shard0); + drop(shard1); + drop(router); + shard0_handle.abort(); + shard1_handle.abort(); + let _ = shard0_handle.await; + let _ = shard1_handle.await; + Ok(()) +} + // Note: entity_resolution_consistent_after_shard_restart test is complex due to // RocksDB lock management. The shard_recovery_after_restart test covers the // core functionality of shard restart and data persistence. @@ -657,7 +782,7 @@ async fn ingest_operations_complete_before_shutdown() -> anyhow::Result<()> { let response = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records, }) .await? @@ -708,7 +833,7 @@ async fn cluster_stats_reflect_all_shards() -> anyhow::Result<()> { client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records, }) .await?; diff --git a/tests/distributed_router_admin.rs b/tests/distributed_router_admin.rs index 15f498b..88dc54b 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: 3, + internal_protocol_version: 5, records, }) .await?; diff --git a/tests/external_backup_restore.rs b/tests/external_backup_restore.rs index 4fd8fc5..ee89d92 100644 --- a/tests/external_backup_restore.rs +++ b/tests/external_backup_restore.rs @@ -22,8 +22,9 @@ use unirust_rs::distributed::{ IdentityKeyConfig, RouterNode, ShardNode, }; use unirust_rs::{ - read_cluster_checkpoint_manifest, restore_checkpoint, restore_checkpoint_for_shard, - StreamingTuning, TuningProfile, + export_cluster_backup, prune_verified_cluster_backups, read_cluster_checkpoint_manifest, + restore_checkpoint, restore_checkpoint_for_shard, verify_cluster_backup, StreamingTuning, + TuningProfile, }; static PERSISTENT_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); @@ -171,6 +172,166 @@ async fn stop_shard(shard: ShardNode, handle: JoinHandle<()>) -> anyhow::Result< Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn verified_export_survives_loss_of_checkpoint_volumes() -> anyhow::Result<()> { + use proto::router_service_server::RouterService; + use std::io::Write as _; + + let _test_guard = PERSISTENT_TEST_LOCK.lock().await; + let data_volume = TempDir::new()?; + let checkpoint_volume = TempDir::new()?; + let off_host_volume = TempDir::new()?; + let checkpoint_roots = (0..2) + .map(|shard_id| checkpoint_volume.path().join(format!("shard-{shard_id}"))) + .collect::>(); + + let (addr0, shard0, handle0) = spawn_shard( + 0, + data_volume.path().join("original-0"), + checkpoint_roots[0].clone(), + config(), + ) + .await?; + let (addr1, shard1, handle1) = spawn_shard( + 1, + data_volume.path().join("original-1"), + checkpoint_roots[1].clone(), + config(), + ) + .await?; + let mut client0 = ShardServiceClient::connect(format!("http://{addr0}")).await?; + let mut client1 = ShardServiceClient::connect(format!("http://{addr1}")).await?; + let mut shard0_record = record(0, "shard-0@example.com"); + shard0_record + .identity + .as_mut() + .expect("record identity") + .uid = "backup-source-0".to_string(); + let mut shard1_record = record(1, "shard-1@example.com"); + shard1_record + .identity + .as_mut() + .expect("record identity") + .uid = "backup-source-1".to_string(); + client0 + .ingest_records(IngestRecordsRequest { + records: vec![shard0_record], + internal_protocol_version: 5, + }) + .await?; + client1 + .ingest_records(IngestRecordsRequest { + records: vec![shard1_record], + internal_protocol_version: 5, + }) + .await?; + + let router = RouterNode::connect( + vec![format!("http://{addr0}"), format!("http://{addr1}")], + config(), + ) + .await?; + let checkpoint = RouterService::checkpoint( + router.as_ref(), + Request::new(CheckpointRequest { + path: "off-host-generation".to_string(), + checkpoint_protocol_version: 0, + shard_count: 0, + finalize: false, + }), + ) + .await? + .into_inner(); + let checkpoint_paths = checkpoint + .paths + .into_iter() + .map(PathBuf::from) + .collect::>(); + let export_root = off_host_volume.path().join("exports"); + std::fs::create_dir(&export_root)?; + let first_export = export_root.join("backup-a"); + let manifest = export_cluster_backup(&checkpoint_paths, &first_export)?; + assert_eq!(manifest.generation(), "off-host-generation"); + assert_eq!(manifest.shard_count(), 2); + assert_eq!(verify_cluster_backup(&first_export)?, manifest); + + drop(client0); + drop(client1); + drop(router); + stop_shard(shard0, handle0).await?; + stop_shard(shard1, handle1).await?; + for root in &checkpoint_roots { + std::fs::remove_dir_all(root)?; + } + + let restored0 = data_volume.path().join("restored-0"); + let restored1 = data_volume.path().join("restored-1"); + restore_checkpoint_for_shard(&first_export.join("shard-0"), &restored0, Some(0))?; + restore_checkpoint_for_shard(&first_export.join("shard-1"), &restored1, Some(1))?; + let recovery_backups = data_volume.path().join("recovery-backups"); + let (restored_addr0, restored_shard0, restored_handle0) = + spawn_shard(0, restored0, recovery_backups.join("shard-0"), config()).await?; + let (restored_addr1, restored_shard1, restored_handle1) = + spawn_shard(1, restored1, recovery_backups.join("shard-1"), config()).await?; + let restored_router = RouterNode::connect( + vec![ + format!("http://{restored_addr0}"), + format!("http://{restored_addr1}"), + ], + config(), + ) + .await?; + let mut restored_client0 = + ShardServiceClient::connect(format!("http://{restored_addr0}")).await?; + let mut restored_client1 = + ShardServiceClient::connect(format!("http://{restored_addr1}")).await?; + assert_eq!( + restored_client0 + .get_stats(StatsRequest {}) + .await? + .into_inner() + .record_count, + 1 + ); + assert_eq!( + restored_client1 + .get_stats(StatsRequest {}) + .await? + .into_inner() + .record_count, + 1 + ); + drop(restored_client0); + drop(restored_client1); + drop(restored_router); + stop_shard(restored_shard0, restored_handle0).await?; + stop_shard(restored_shard1, restored_handle1).await?; + + std::thread::sleep(Duration::from_millis(1)); + let second_export = export_root.join("backup-b"); + export_cluster_backup(&checkpoint_paths, &second_export) + .expect_err("deleted checkpoint volumes must not be exportable"); + export_cluster_backup( + &[first_export.join("shard-0"), first_export.join("shard-1")], + &second_export, + )?; + let removed = prune_verified_cluster_backups(&export_root, 1)?; + assert_eq!(removed.len(), 1); + assert_eq!(removed[0].file_name(), first_export.file_name()); + assert!(second_export.is_dir()); + + let mut current = std::fs::OpenOptions::new() + .append(true) + .open(second_export.join("shard-0/CURRENT"))?; + current.write_all(b"tampered")?; + current.sync_all()?; + verify_cluster_backup(&second_export) + .expect_err("modified backup bytes must fail verification"); + prune_verified_cluster_backups(&export_root, 1) + .expect_err("retention must fail closed when any backup is corrupt"); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn external_checkpoint_restores_a_lost_shard_volume() -> anyhow::Result<()> { use proto::router_service_server::RouterService; @@ -188,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: 3, + internal_protocol_version: 5, }) .await?; let before = client.get_stats(StatsRequest {}).await?.into_inner(); @@ -230,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: 3, + internal_protocol_version: 5, }) .await?; let retry_stats = client.get_stats(StatsRequest {}).await?.into_inner(); @@ -239,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: 3, + internal_protocol_version: 5, }) .await .expect_err("restored immutable source identity must reject a changed payload"); diff --git a/tests/linker_state_recovery.rs b/tests/linker_state_recovery.rs index 00c86f6..8662e8f 100644 --- a/tests/linker_state_recovery.rs +++ b/tests/linker_state_recovery.rs @@ -8,7 +8,7 @@ use tempfile::tempdir; use unirust_rs::advanced::{ClusterId, GlobalClusterId, LinkerStateConfig}; use unirust_rs::model::{Descriptor, Record, RecordIdentity}; use unirust_rs::ontology::{IdentityKey, StrongIdentifier}; -use unirust_rs::persistence::linker_encoding; +use unirust_rs::persistence::{linker_encoding, LinkerStatePersistence}; use unirust_rs::temporal::Interval; use unirust_rs::test_support::{default_ontology, generate_dataset}; use unirust_rs::{Ontology, PersistentStore, RecordId, StreamingTuning, Unirust}; @@ -346,6 +346,81 @@ fn cross_shard_merge_chain_is_persisted_as_direct_canonical_redirects() -> anyho Ok(()) } +#[test] +fn global_cluster_ids_are_stable_across_ingest_order() -> anyhow::Result<()> { + fn ingest_in_order(path: &std::path::Path, ids: &[u32]) -> anyhow::Result { + let store = PersistentStore::open(path)?; + let mut ontology = Ontology::new(); + ontology.add_identity_key(IdentityKey::from_names(vec!["email"], "email")); + let mut tuning = StreamingTuning::balanced(); + tuning.shard_id = 3; + let mut unirust = Unirust::with_store_and_tuning(ontology, store, tuning); + let email_attr = unirust.intern_attr("email"); + let email_value = unirust.intern_value("shared@example.com"); + let records = ids + .iter() + .map(|id| { + Record::new( + RecordId(*id), + RecordIdentity::new("person".to_string(), "crm".to_string(), id.to_string()), + vec![Descriptor::new( + email_attr, + email_value, + Interval::new(0, 10).expect("valid interval"), + )], + ) + }) + .collect(); + unirust.ingest_with_explicit_ids(records)?; + unirust.global_cluster_id_for_record(RecordId(ids[0])) + } + + let first = tempdir()?; + let second = tempdir()?; + let forward = ingest_in_order(first.path(), &[10, 20])?; + let reverse = ingest_in_order(second.path(), &[20, 10])?; + + assert_eq!(forward, GlobalClusterId::new(3, 10, 0)); + assert_eq!(reverse, forward); + + let store = PersistentStore::open(first.path())?; + let mut ontology = Ontology::new(); + ontology.add_identity_key(IdentityKey::from_names(vec!["email"], "email")); + let mut tuning = StreamingTuning::balanced(); + tuning.shard_id = 3; + let mut restarted = Unirust::with_store_and_tuning(ontology, store, tuning); + restarted.initialize_streaming()?; + assert_eq!( + restarted.global_cluster_id_for_record(RecordId(20))?, + forward + ); + Ok(()) +} + +#[test] +fn legacy_allocation_order_redirects_are_removed_before_recovery() -> anyhow::Result<()> { + let dir = tempdir()?; + let legacy_primary = GlobalClusterId::new(0, 10, 0); + let legacy_secondary = GlobalClusterId::new(1, 20, 0); + { + let store = PersistentStore::open(dir.path())?; + LinkerStatePersistence::new(store.db()) + .save_cross_shard_merge(legacy_secondary, legacy_primary)?; + } + + let store = PersistentStore::open(dir.path())?; + let mut unirust = Unirust::with_store(default_ontology_for_empty_store(), store); + unirust.initialize_streaming()?; + assert_eq!(unirust.cross_shard_merge_count(), 0); + + drop(unirust); + let store = PersistentStore::open(dir.path())?; + assert!(LinkerStatePersistence::new(store.db()) + .load_cross_shard_merges()? + .is_empty()); + Ok(()) +} + fn default_ontology_for_empty_store() -> unirust_rs::Ontology { unirust_rs::Ontology::new() } diff --git a/tests/process_crash_recovery.rs b/tests/process_crash_recovery.rs index c02e9d7..6273517 100644 --- a/tests/process_crash_recovery.rs +++ b/tests/process_crash_recovery.rs @@ -17,7 +17,12 @@ struct ShardProcess { } impl ShardProcess { - fn spawn(listen: SocketAddr, data_dir: &Path, ontology_path: &Path) -> anyhow::Result { + fn spawn( + listen: SocketAddr, + data_dir: &Path, + backup_dir: &Path, + ontology_path: &Path, + ) -> anyhow::Result { let child = Command::new(env!("CARGO_BIN_EXE_unirust_shard")) .args([ "--listen", @@ -28,6 +33,10 @@ impl ShardProcess { data_dir .to_str() .ok_or_else(|| anyhow::anyhow!("data directory is not valid UTF-8"))?, + "--backup-dir", + backup_dir + .to_str() + .ok_or_else(|| anyhow::anyhow!("backup directory is not valid UTF-8"))?, "--ontology", ontology_path .to_str() @@ -122,6 +131,7 @@ fn record(index: u32, uid: String, email: String) -> RecordInput { async fn acknowledged_ingest_survives_process_kill_and_restart() -> anyhow::Result<()> { let temp_dir = tempdir()?; let data_dir = temp_dir.path().join("shard-data"); + let backup_dir = temp_dir.path().join("shard-backup"); let ontology_path = temp_dir.path().join("ontology.json"); std::fs::write( &ontology_path, @@ -129,7 +139,7 @@ async fn acknowledged_ingest_survives_process_kill_and_restart() -> anyhow::Resu )?; let first_addr = available_addr()?; - let mut process = ShardProcess::spawn(first_addr, &data_dir, &ontology_path)?; + let mut process = ShardProcess::spawn(first_addr, &data_dir, &backup_dir, &ontology_path)?; let mut client = wait_for_shard(first_addr, &mut process).await?; let records = (0..128) @@ -143,7 +153,7 @@ async fn acknowledged_ingest_survives_process_kill_and_restart() -> anyhow::Resu .collect(); let response = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records, }) .await? @@ -164,7 +174,7 @@ async fn acknowledged_ingest_survives_process_kill_and_restart() -> anyhow::Resu process.kill_and_wait()?; let second_addr = available_addr()?; - let mut restarted = ShardProcess::spawn(second_addr, &data_dir, &ontology_path)?; + let mut restarted = ShardProcess::spawn(second_addr, &data_dir, &backup_dir, &ontology_path)?; let mut client = wait_for_shard(second_addr, &mut restarted).await?; let stats = client.get_stats(StatsRequest {}).await?.into_inner(); @@ -191,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: 3, + internal_protocol_version: 5, records: vec![record( 128, "process-kill-after-restart".to_string(), @@ -217,7 +227,7 @@ async fn acknowledged_ingest_survives_process_kill_and_restart() -> anyhow::Resu let third_addr = available_addr()?; let mut restarted_after_shutdown = - ShardProcess::spawn(third_addr, &data_dir, &ontology_path)?; + ShardProcess::spawn(third_addr, &data_dir, &backup_dir, &ontology_path)?; let mut client = wait_for_shard(third_addr, &mut restarted_after_shutdown).await?; let stats = client.get_stats(StatsRequest {}).await?.into_inner(); assert_eq!(stats.record_count, 129); @@ -225,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: 3, + internal_protocol_version: 5, records: vec![record( 129, "graceful-shutdown-after-restart".to_string(), diff --git a/tests/process_replication_recovery.rs b/tests/process_replication_recovery.rs new file mode 100644 index 0000000..7986110 --- /dev/null +++ b/tests/process_replication_recovery.rs @@ -0,0 +1,301 @@ +use std::net::{SocketAddr, TcpListener}; +use std::path::Path; +use std::process::{Child, Command, Stdio}; +use std::time::Duration; + +use tempfile::tempdir; +use tokio::time::sleep; +use unirust_rs::distributed::proto::{ + shard_service_client::ShardServiceClient, IngestRecordsRequest, QueryDescriptor, + QueryEntitiesRequest, RecordDescriptor, RecordIdentity, RecordInput, StatsRequest, +}; +use unirust_rs::distributed::DISTRIBUTED_PROTOCOL_VERSION; + +mod support; + +const REPLICATION_SECRET: &[u8; 32] = b"replication-test-secret-32-bytes"; + +struct ShardProcess { + child: Child, +} + +struct ShardPaths<'a> { + data: &'a Path, + backup: &'a Path, + ontology: &'a Path, +} + +impl ShardProcess { + fn spawn_replica( + listen: SocketAddr, + paths: ShardPaths<'_>, + token: &Path, + ) -> anyhow::Result { + Self::spawn( + listen, + paths, + &[ + "--replica-mode".to_string(), + "--allow-insecure-replication".to_string(), + "--replication-token-file".to_string(), + path_arg(token)?, + ], + ) + } + + fn spawn_primary( + listen: SocketAddr, + paths: ShardPaths<'_>, + token: &Path, + replica: SocketAddr, + ) -> anyhow::Result { + Self::spawn( + listen, + paths, + &[ + "--replica".to_string(), + format!("http://{replica}"), + "--allow-insecure-replication".to_string(), + "--replication-token-file".to_string(), + path_arg(token)?, + ], + ) + } + + fn spawn_standalone(listen: SocketAddr, paths: ShardPaths<'_>) -> anyhow::Result { + Self::spawn(listen, paths, &[]) + } + + fn spawn( + listen: SocketAddr, + paths: ShardPaths<'_>, + extra_args: &[String], + ) -> anyhow::Result { + let mut command = Command::new(env!("CARGO_BIN_EXE_unirust_shard")); + command.args([ + "--listen", + &listen.to_string(), + "--shard-id", + "0", + "--data-dir", + &path_arg(paths.data)?, + "--backup-dir", + &path_arg(paths.backup)?, + "--ontology", + &path_arg(paths.ontology)?, + "--profile", + "balanced", + ]); + let child = command + .args(extra_args) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn()?; + Ok(Self { child }) + } + + fn kill_and_wait(&mut self) -> anyhow::Result<()> { + if self.child.try_wait()?.is_none() { + self.child.kill()?; + } + let _ = self.child.wait()?; + Ok(()) + } +} + +impl Drop for ShardProcess { + fn drop(&mut self) { + let _ = self.kill_and_wait(); + } +} + +fn path_arg(path: &Path) -> anyhow::Result { + path.to_str() + .map(str::to_string) + .ok_or_else(|| anyhow::anyhow!("path is not valid UTF-8: {}", path.display())) +} + +fn available_addr() -> anyhow::Result { + let listener = TcpListener::bind("127.0.0.1:0")?; + let addr = listener.local_addr()?; + drop(listener); + Ok(addr) +} + +async fn wait_for_shard( + addr: SocketAddr, + process: &mut ShardProcess, +) -> anyhow::Result> { + let endpoint = format!("http://{addr}"); + for _ in 0..200 { + if let Some(status) = process.child.try_wait()? { + anyhow::bail!("shard exited before becoming ready: {status}"); + } + if let Ok(mut client) = ShardServiceClient::connect(endpoint.clone()).await { + if client + .health_check(unirust_rs::distributed::proto::HealthCheckRequest {}) + .await + .is_ok() + { + return Ok(client); + } + } + sleep(Duration::from_millis(50)).await; + } + anyhow::bail!("shard did not become ready at {endpoint}") +} + +fn record(index: u32, source: u32) -> RecordInput { + RecordInput { + index, + identity: Some(RecordIdentity { + entity_type: "person".to_string(), + perspective: "crm".to_string(), + uid: format!("replicated-process-{source}"), + }), + descriptors: vec![RecordDescriptor { + attr: "email".to_string(), + value: format!("entity-{}@example.com", source / 2), + start: 0, + end: 10, + }], + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn acknowledged_replica_batch_survives_sigkill_and_primary_volume_loss() -> anyhow::Result<()> +{ + let temp = tempdir()?; + let primary_data = temp.path().join("primary-data"); + let replica_data = temp.path().join("replica-data"); + let primary_backup = temp.path().join("primary-backup"); + let replica_backup = temp.path().join("replica-backup"); + let ontology = temp.path().join("ontology.json"); + let token = temp.path().join("replication.token"); + std::fs::write(&ontology, serde_json::to_vec(&support::build_iam_config())?)?; + std::fs::write(&token, REPLICATION_SECRET)?; + + let replica_addr = available_addr()?; + let mut replica = ShardProcess::spawn_replica( + replica_addr, + ShardPaths { + data: &replica_data, + backup: &replica_backup, + ontology: &ontology, + }, + &token, + )?; + let _ = wait_for_shard(replica_addr, &mut replica).await?; + + let primary_addr = available_addr()?; + let mut primary = ShardProcess::spawn_primary( + primary_addr, + ShardPaths { + data: &primary_data, + backup: &primary_backup, + ontology: &ontology, + }, + &token, + replica_addr, + )?; + let mut client = wait_for_shard(primary_addr, &mut primary).await?; + let response = client + .ingest_records(IngestRecordsRequest { + records: (0..128).map(|source| record(source, source)).collect(), + internal_protocol_version: DISTRIBUTED_PROTOCOL_VERSION, + }) + .await? + .into_inner(); + assert_eq!(response.assignments.len(), 128); + assert_eq!( + client + .get_stats(StatsRequest {}) + .await? + .into_inner() + .record_count, + 128 + ); + + drop(client); + primary.kill_and_wait()?; + replica.kill_and_wait()?; + + let restarted_replica_addr = available_addr()?; + let mut restarted_replica = ShardProcess::spawn_replica( + restarted_replica_addr, + ShardPaths { + data: &replica_data, + backup: &replica_backup, + ontology: &ontology, + }, + &token, + )?; + let _ = wait_for_shard(restarted_replica_addr, &mut restarted_replica).await?; + let restarted_primary_addr = available_addr()?; + let mut restarted_primary = ShardProcess::spawn_primary( + restarted_primary_addr, + ShardPaths { + data: &primary_data, + backup: &primary_backup, + ontology: &ontology, + }, + &token, + restarted_replica_addr, + )?; + let mut restarted_client = + wait_for_shard(restarted_primary_addr, &mut restarted_primary).await?; + let stats = restarted_client + .get_stats(StatsRequest {}) + .await? + .into_inner(); + assert_eq!(stats.record_count, 128); + assert_eq!(stats.cluster_count, 64); + + drop(restarted_client); + restarted_primary.kill_and_wait()?; + restarted_replica.kill_and_wait()?; + std::fs::remove_dir_all(&primary_data)?; + + let promoted_addr = available_addr()?; + let mut promoted = ShardProcess::spawn_standalone( + promoted_addr, + ShardPaths { + data: &replica_data, + backup: &replica_backup, + ontology: &ontology, + }, + )?; + let mut promoted_client = wait_for_shard(promoted_addr, &mut promoted).await?; + let query = promoted_client + .query_entities(QueryEntitiesRequest { + descriptors: vec![QueryDescriptor { + attr: "email".to_string(), + value: "entity-0@example.com".to_string(), + }], + start: 0, + end: 10, + }) + .await? + .into_inner(); + assert!(query.outcome.is_some()); + + let retry = promoted_client + .ingest_records(IngestRecordsRequest { + records: vec![record(999, 0)], + internal_protocol_version: DISTRIBUTED_PROTOCOL_VERSION, + }) + .await? + .into_inner(); + assert_eq!(retry.assignments.len(), 1); + assert_eq!( + promoted_client + .get_stats(StatsRequest {}) + .await? + .into_inner() + .record_count, + 128 + ); + assert!(!primary_data.exists()); + Ok(()) +} diff --git a/tests/shard_cli_safety.rs b/tests/shard_cli_safety.rs index ddbada3..f4cb64a 100644 --- a/tests/shard_cli_safety.rs +++ b/tests/shard_cli_safety.rs @@ -21,9 +21,11 @@ fn shard_binary_refuses_implicit_ephemeral_storage() { fn shard_binary_reads_persistent_path_from_environment() { let source = tempdir().expect("temporary checkpoint source"); let destination = source.path().join("replacement"); + let backup = tempdir().expect("temporary backup directory"); let output = Command::new(env!("CARGO_BIN_EXE_unirust_shard")) .env_clear() .env("UNIRUST_SHARD_DATA_DIR", &destination) + .env("UNIRUST_SHARD_BACKUP_DIR", backup.path()) .args([ "--restore-from", source.path().to_str().expect("UTF-8 path"), @@ -43,6 +45,47 @@ fn shard_binary_reads_persistent_path_from_environment() { ); } +#[test] +fn shard_binary_requires_independent_checkpoint_path() { + let data = tempdir().expect("temporary data directory"); + let output = Command::new(env!("CARGO_BIN_EXE_unirust_shard")) + .env_clear() + .args(["--data-dir", data.path().to_str().expect("UTF-8 path")]) + .output() + .expect("run shard binary"); + + assert!(!output.status.success()); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!( + stderr.contains("persistent shards require --backup-dir"), + "unexpected stderr: {stderr}" + ); +} + +#[test] +fn shard_binary_rejects_overlapping_checkpoint_path() { + let root = tempdir().expect("temporary root"); + let data = root.path().join("data"); + let backup = data.join("checkpoints"); + let output = Command::new(env!("CARGO_BIN_EXE_unirust_shard")) + .env_clear() + .args([ + "--data-dir", + data.to_str().expect("UTF-8 path"), + "--backup-dir", + backup.to_str().expect("UTF-8 path"), + ]) + .output() + .expect("run shard binary"); + + assert!(!output.status.success()); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!( + stderr.contains("data and checkpoint paths must not overlap"), + "unexpected stderr: {stderr}" + ); +} + #[test] fn shard_binary_rejects_unknown_shard_environment_variable() { let output = Command::new(env!("CARGO_BIN_EXE_unirust_shard")) @@ -74,3 +117,80 @@ fn shard_binary_rejects_partial_mtls_configuration() { "unexpected stderr: {stderr}" ); } + +#[test] +fn shard_binary_rejects_replica_mode_without_a_token() { + let data = tempdir().expect("temporary data directory"); + let output = Command::new(env!("CARGO_BIN_EXE_unirust_shard")) + .env_clear() + .args([ + "--data-dir", + data.path().to_str().expect("UTF-8 path"), + "--allow-colocated-checkpoints", + "--replica-mode", + ]) + .output() + .expect("run shard binary"); + + assert!(!output.status.success()); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!( + stderr.contains("replication_token_file"), + "unexpected stderr: {stderr}" + ); +} + +#[test] +fn shard_binary_rejects_short_replication_token() { + let data = tempdir().expect("temporary data directory"); + let token_dir = tempdir().expect("temporary token directory"); + let token = token_dir.path().join("replication.token"); + std::fs::write(&token, b"too short").expect("write token"); + let output = Command::new(env!("CARGO_BIN_EXE_unirust_shard")) + .env_clear() + .args([ + "--data-dir", + data.path().to_str().expect("UTF-8 path"), + "--allow-colocated-checkpoints", + "--replica-mode", + "--allow-insecure-replication", + "--replication-token-file", + token.to_str().expect("UTF-8 path"), + ]) + .output() + .expect("run shard binary"); + + assert!(!output.status.success()); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!( + stderr.contains("replication token must contain at least 32 bytes"), + "unexpected stderr: {stderr}" + ); +} + +#[test] +fn shard_binary_requires_mtls_for_replication_by_default() { + let data = tempdir().expect("temporary data directory"); + let token_dir = tempdir().expect("temporary token directory"); + let token = token_dir.path().join("replication.token"); + std::fs::write(&token, [b'x'; 32]).expect("write token"); + let output = Command::new(env!("CARGO_BIN_EXE_unirust_shard")) + .env_clear() + .args([ + "--data-dir", + data.path().to_str().expect("UTF-8 path"), + "--allow-colocated-checkpoints", + "--replica-mode", + "--replication-token-file", + token.to_str().expect("UTF-8 path"), + ]) + .output() + .expect("run shard binary"); + + assert!(!output.status.success()); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!( + stderr.contains("replica mode requires shard server mTLS"), + "unexpected stderr: {stderr}" + ); +} diff --git a/tests/support/mod.rs b/tests/support/mod.rs index ac3cae1..691a749 100644 --- a/tests/support/mod.rs +++ b/tests/support/mod.rs @@ -1,3 +1,12 @@ +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::task::{Context, Poll}; +use std::time::Duration; + +use tokio::sync::Semaphore; +use tonic::codegen::{http, Service}; +use tonic::server::NamedService; use unirust_rs::distributed::proto::{ ConstraintConfig as ProtoConstraintConfig, ConstraintKind as ProtoConstraintKind, IdentityKeyConfig as ProtoIdentityKeyConfig, OntologyConfig, @@ -6,6 +15,102 @@ use unirust_rs::distributed::{ ConstraintConfig, ConstraintKind, DistributedOntologyConfig, IdentityKeyConfig, }; +#[derive(Clone)] +#[allow(dead_code)] +pub struct StallResponse { + inner: S, + path: &'static str, + stall_next: Arc, + committed: Arc, + release: Arc, +} + +impl Service> for StallResponse +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 + 'static>>; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, request: http::Request) -> Self::Future { + let should_stall = + request.uri().path() == self.path && self.stall_next.swap(false, Ordering::AcqRel); + let future = self.inner.call(request); + let committed = self.committed.clone(); + let release = self.release.clone(); + Box::pin(async move { + let response = future.await?; + if should_stall { + committed.add_permits(1); + if let Ok(permit) = release.acquire_owned().await { + permit.forget(); + } + } + Ok(response) + }) + } +} + +impl NamedService for StallResponse +where + S: NamedService, +{ + const NAME: &'static str = S::NAME; +} + +#[allow(dead_code)] +pub struct StallControls { + committed: Arc, + release: Arc, +} + +#[allow(dead_code)] +impl StallControls { + pub async fn wait_until_committed(&self, timeout: Duration) -> anyhow::Result<()> { + let permit = + tokio::time::timeout(timeout, self.committed.clone().acquire_owned()).await??; + permit.forget(); + Ok(()) + } + + pub fn release(&self) { + self.release.add_permits(1); + } +} + +#[allow(dead_code)] +pub fn stall_response( + inner: S, + path: &'static str, + initially_armed: bool, +) -> (StallResponse, StallControls) { + let stall_next = Arc::new(AtomicBool::new(initially_armed)); + let committed = Arc::new(Semaphore::new(0)); + let release = Arc::new(Semaphore::new(0)); + ( + StallResponse { + inner, + path, + stall_next: stall_next.clone(), + committed: committed.clone(), + release: release.clone(), + }, + StallControls { committed, release }, + ) +} + #[allow(dead_code)] pub fn build_iam_config() -> DistributedOntologyConfig { DistributedOntologyConfig { diff --git a/tests/synchronous_replication.rs b/tests/synchronous_replication.rs new file mode 100644 index 0000000..9a7c814 --- /dev/null +++ b/tests/synchronous_replication.rs @@ -0,0 +1,541 @@ +use std::net::SocketAddr; +use std::path::PathBuf; +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use tempfile::TempDir; +use tokio::task::JoinHandle; +use tokio::time::{sleep, Duration}; +use tokio_stream::wrappers::TcpListenerStream; +use tonic::codegen::{http, Service}; +use tonic::server::NamedService; +use tonic::transport::Server; +use tonic::Request; +use unirust_rs::distributed::proto::{ + self, shard_service_client::ShardServiceClient, ApplyOntologyRequest, ConfigVersionRequest, + IngestRecordsRequest, RecordDescriptor, RecordIdentity, RecordInput, StatsRequest, +}; +use unirust_rs::distributed::{ + DistributedOntologyConfig, IdentityKeyConfig, RouterNode, ShardNode, + DISTRIBUTED_PROTOCOL_VERSION, +}; +use unirust_rs::{StreamingTuning, TuningProfile}; + +mod support; + +static PERSISTENT_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); +const REPLICATION_TOKEN: &str = "42fd59df99d782a1e9d614e5f2f9d3a425da036468fceb91c0dd3bcc6bf3d729"; +const WRONG_REPLICATION_TOKEN: &str = + "f12307f3555950adee23c22dc0ee37c597451f44418595d6f0774258313a649e"; + +#[derive(Clone)] +struct FailIngest { + inner: S, + fail: Arc, +} + +impl Service> for FailIngest +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 + 'static>>; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, request: http::Request) -> Self::Future { + if self.fail.load(Ordering::Acquire) + && request.uri().path() == "/unirust.ShardService/IngestRecords" + { + return Box::pin(async { + Ok(tonic::Status::unavailable("injected replica outage").into_http()) + }); + } + Box::pin(self.inner.call(request)) + } +} + +impl NamedService for FailIngest +where + S: NamedService, +{ + const NAME: &'static str = S::NAME; +} + +async fn spawn_node(node: ShardNode) -> anyhow::Result<(SocketAddr, JoinHandle<()>)> { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let addr = listener.local_addr()?; + let service_node = node.clone(); + let handle = tokio::spawn(async move { + Server::builder() + .add_service(proto::shard_service_server::ShardServiceServer::new( + service_node, + )) + .serve_with_incoming(TcpListenerStream::new(listener)) + .await + .expect("shard server"); + }); + Ok((addr, handle)) +} + +async fn spawn_failable_node( + node: ShardNode, +) -> anyhow::Result<(SocketAddr, JoinHandle<()>, Arc)> { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let addr = listener.local_addr()?; + let fail = Arc::new(AtomicBool::new(false)); + let service = FailIngest { + inner: proto::shard_service_server::ShardServiceServer::new(node), + fail: fail.clone(), + }; + let handle = tokio::spawn(async move { + Server::builder() + .add_service(service) + .serve_with_incoming(TcpListenerStream::new(listener)) + .await + .expect("shard server"); + }); + Ok((addr, handle, fail)) +} + +async fn spawn_stalling_node( + node: ShardNode, + path: &'static str, +) -> anyhow::Result<(SocketAddr, JoinHandle<()>, support::StallControls)> { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let addr = listener.local_addr()?; + let (service, controls) = support::stall_response( + proto::shard_service_server::ShardServiceServer::new(node), + path, + true, + ); + let handle = tokio::spawn(async move { + Server::builder() + .add_service(service) + .serve_with_incoming(TcpListenerStream::new(listener)) + .await + .expect("shard server"); + }); + Ok((addr, handle, controls)) +} + +async fn stop_node(node: ShardNode, handle: JoinHandle<()>) -> anyhow::Result<()> { + node.shutdown().await?; + handle.abort(); + let _ = handle.await; + drop(node); + sleep(Duration::from_millis(50)).await; + Ok(()) +} + +fn config() -> DistributedOntologyConfig { + DistributedOntologyConfig { + identity_keys: vec![IdentityKeyConfig { + name: "email".to_string(), + attributes: vec!["email".to_string()], + }], + strong_identifiers: Vec::new(), + constraints: Vec::new(), + } +} + +fn node(data_dir: PathBuf, backup_dir: PathBuf) -> anyhow::Result { + ShardNode::new_with_storage_paths( + 0, + config(), + StreamingTuning::from_profile(TuningProfile::Balanced), + Some(data_dir), + Some(backup_dir), + false, + None, + ) +} + +fn authenticated_request(payload: T, token: &str) -> Request { + let mut request = Request::new(payload); + request + .metadata_mut() + .insert("x-unirust-replication-token", token.parse().unwrap()); + request +} + +fn record(index: u32) -> RecordInput { + RecordInput { + index, + identity: Some(RecordIdentity { + entity_type: "person".to_string(), + perspective: "crm".to_string(), + uid: "replicated-source".to_string(), + }), + descriptors: vec![RecordDescriptor { + attr: "email".to_string(), + value: "replicated@example.com".to_string(), + start: 0, + end: 100, + }], + } +} + +fn distinct_record(index: u32, source: u32) -> RecordInput { + RecordInput { + index, + identity: Some(RecordIdentity { + entity_type: "person".to_string(), + perspective: "crm".to_string(), + uid: format!("replicated-source-{source}"), + }), + descriptors: vec![RecordDescriptor { + attr: "email".to_string(), + value: format!("replicated-{source}@example.com"), + start: 0, + end: 100, + }], + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn acknowledged_ingest_survives_primary_volume_loss_and_manual_failover() -> anyhow::Result<()> +{ + use proto::router_service_server::RouterService; + + let _test_guard = PERSISTENT_TEST_LOCK.lock().await; + let volumes = TempDir::new()?; + let primary_data = volumes.path().join("primary-data"); + let replica_data = volumes.path().join("replica-data"); + let primary_backup = volumes.path().join("primary-backup"); + let replica_backup = volumes.path().join("replica-backup"); + + let replica_node = node(replica_data.clone(), replica_backup.clone())? + .into_replica(REPLICATION_TOKEN.into())?; + let (replica_addr, replica_handle) = spawn_node(replica_node.clone()).await?; + let mut replica_client = ShardServiceClient::connect(format!("http://{replica_addr}")).await?; + + let direct_error = replica_client + .ingest_records(IngestRecordsRequest { + records: vec![record(0)], + internal_protocol_version: DISTRIBUTED_PROTOCOL_VERSION, + }) + .await + .expect_err("a passive replica must reject direct mutations"); + assert_eq!(direct_error.code(), tonic::Code::PermissionDenied); + + let primary_node = node(primary_data.clone(), primary_backup)? + .with_replica(replica_client.clone(), REPLICATION_TOKEN.into()) + .await?; + let (primary_addr, primary_handle) = spawn_node(primary_node.clone()).await?; + let router = RouterNode::connect(vec![format!("http://{primary_addr}")], config()).await?; + let response = RouterService::ingest_records( + router.as_ref(), + Request::new(IngestRecordsRequest { + records: (0..128) + .map(|source| distinct_record(source, source)) + .collect(), + internal_protocol_version: 0, + }), + ) + .await? + .into_inner(); + assert_eq!(response.assignments.len(), 128); + + let mut primary_client = ShardServiceClient::connect(format!("http://{primary_addr}")).await?; + let primary_stats = primary_client + .get_stats(StatsRequest {}) + .await? + .into_inner(); + let replica_stats = replica_client + .get_stats(StatsRequest {}) + .await? + .into_inner(); + assert_eq!(primary_stats.record_count, 128); + assert_eq!(primary_stats, replica_stats); + + let primary_config = primary_client + .get_config_version(ConfigVersionRequest { + include_durable_state_digest: true, + }) + .await? + .into_inner(); + let replica_config = replica_client + .get_config_version(authenticated_request( + ConfigVersionRequest { + include_durable_state_digest: true, + }, + REPLICATION_TOKEN, + )) + .await? + .into_inner(); + assert_eq!( + proto::ShardRole::try_from(primary_config.shard_role)?, + proto::ShardRole::Primary + ); + assert_eq!( + proto::ShardRole::try_from(replica_config.shard_role)?, + proto::ShardRole::Replica + ); + assert_eq!( + primary_config.durable_state_digest, + replica_config.durable_state_digest + ); + + let replica_router_error = + RouterNode::connect(vec![format!("http://{replica_addr}")], config()) + .await + .err() + .expect("routers must reject passive replicas"); + assert_eq!(replica_router_error.code(), tonic::Code::FailedPrecondition); + + drop(router); + drop(primary_client); + drop(replica_client); + stop_node(primary_node, primary_handle).await?; + stop_node(replica_node, replica_handle).await?; + + let restarted_replica = node(replica_data.clone(), replica_backup.clone())? + .into_replica(REPLICATION_TOKEN.into())?; + let (restarted_replica_addr, restarted_replica_handle) = + spawn_node(restarted_replica.clone()).await?; + let restarted_replica_client = + ShardServiceClient::connect(format!("http://{restarted_replica_addr}")).await?; + let restarted_primary = node(primary_data.clone(), volumes.path().join("primary-backup"))? + .with_replica(restarted_replica_client, REPLICATION_TOKEN.into()) + .await?; + let (restarted_primary_addr, restarted_primary_handle) = + spawn_node(restarted_primary.clone()).await?; + let mut restarted_primary_client = + ShardServiceClient::connect(format!("http://{restarted_primary_addr}")).await?; + assert_eq!( + restarted_primary_client + .health_check(proto::HealthCheckRequest {}) + .await? + .into_inner() + .status, + "ok" + ); + drop(restarted_primary_client); + stop_node(restarted_primary, restarted_primary_handle).await?; + stop_node(restarted_replica, restarted_replica_handle).await?; + + std::fs::remove_dir_all(&primary_data)?; + + let promoted_node = node(replica_data, replica_backup)?; + let (promoted_addr, promoted_handle) = spawn_node(promoted_node.clone()).await?; + let promoted_router = + RouterNode::connect(vec![format!("http://{promoted_addr}")], config()).await?; + let retry = RouterService::ingest_records( + promoted_router.as_ref(), + Request::new(IngestRecordsRequest { + records: vec![distinct_record(999, 0)], + internal_protocol_version: 0, + }), + ) + .await? + .into_inner(); + assert_eq!(retry.assignments.len(), 1); + let mut promoted_client = + ShardServiceClient::connect(format!("http://{promoted_addr}")).await?; + let promoted_stats = promoted_client + .get_stats(StatsRequest {}) + .await? + .into_inner(); + assert_eq!(promoted_stats.record_count, 128); + + drop(promoted_router); + drop(promoted_client); + stop_node(promoted_node, promoted_handle).await?; + assert!(!primary_data.exists()); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn cancelled_replica_response_latches_primary_closed() -> anyhow::Result<()> { + use proto::shard_service_server::ShardService; + + let _test_guard = PERSISTENT_TEST_LOCK.lock().await; + let volumes = TempDir::new()?; + let replica_node = node( + volumes.path().join("replica-data"), + volumes.path().join("replica-backup"), + )? + .into_replica(REPLICATION_TOKEN.into())?; + let (replica_addr, replica_handle, controls) = + spawn_stalling_node(replica_node.clone(), "/unirust.ShardService/SetOntology").await?; + let mut replica_client = ShardServiceClient::connect(format!("http://{replica_addr}")).await?; + let primary_node = node( + volumes.path().join("primary-data"), + volumes.path().join("primary-backup"), + )? + .with_replica(replica_client.clone(), REPLICATION_TOKEN.into()) + .await?; + + let mut updated_config = config(); + updated_config + .strong_identifiers + .push("employee_id".to_string()); + let primary_service = primary_node.clone(); + let update_task = tokio::spawn({ + let updated_config = updated_config.clone(); + async move { + ShardService::set_ontology( + &primary_service, + Request::new(ApplyOntologyRequest { + config: Some(support::to_proto_config(&updated_config)), + }), + ) + .await + } + }); + + controls + .wait_until_committed(Duration::from_secs(10)) + .await?; + update_task.abort(); + assert!(update_task + .await + .expect_err("primary update task must be cancelled") + .is_cancelled()); + controls.release(); + + let health_error = + ShardService::health_check(&primary_node, Request::new(proto::HealthCheckRequest {})) + .await + .expect_err("an ambiguous replica response must latch the primary closed"); + assert_eq!(health_error.code(), tonic::Code::FailedPrecondition); + + let primary_config = ShardService::get_config_version( + &primary_node, + Request::new(ConfigVersionRequest { + include_durable_state_digest: false, + }), + ) + .await? + .into_inner() + .ontology_config + .expect("primary ontology"); + let replica_config = replica_client + .get_config_version(ConfigVersionRequest { + include_durable_state_digest: false, + }) + .await? + .into_inner() + .ontology_config + .expect("replica ontology"); + assert_eq!(primary_config, support::to_proto_config(&config())); + assert_eq!(replica_config, support::to_proto_config(&updated_config)); + + drop(replica_client); + primary_node.shutdown().await?; + drop(primary_node); + stop_node(replica_node, replica_handle).await?; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn replica_outage_prevents_local_commit_and_latches_primary_closed() -> anyhow::Result<()> { + let _test_guard = PERSISTENT_TEST_LOCK.lock().await; + let volumes = TempDir::new()?; + let primary_data = volumes.path().join("primary-data"); + let replica_data = volumes.path().join("replica-data"); + let replica_node = node(replica_data, volumes.path().join("replica-backup"))? + .into_replica(REPLICATION_TOKEN.into())?; + let (replica_addr, replica_handle, fail_replica) = + spawn_failable_node(replica_node.clone()).await?; + let replica_client = ShardServiceClient::connect(format!("http://{replica_addr}")).await?; + let primary_node = node(primary_data, volumes.path().join("primary-backup"))? + .with_replica(replica_client, REPLICATION_TOKEN.into()) + .await?; + let (primary_addr, primary_handle) = spawn_node(primary_node.clone()).await?; + let mut primary_client = ShardServiceClient::connect(format!("http://{primary_addr}")).await?; + + fail_replica.store(true, Ordering::Release); + let error = primary_client + .ingest_records(IngestRecordsRequest { + records: vec![record(0)], + internal_protocol_version: DISTRIBUTED_PROTOCOL_VERSION, + }) + .await + .expect_err("primary must not commit while its replica is unavailable"); + assert_eq!(error.code(), tonic::Code::Aborted); + + let blocked = primary_client + .ingest_records(IngestRecordsRequest { + records: vec![record(0)], + internal_protocol_version: DISTRIBUTED_PROTOCOL_VERSION, + }) + .await + .expect_err("an ambiguous replication failure must latch the primary closed"); + assert_eq!(blocked.code(), tonic::Code::FailedPrecondition); + let health = primary_client + .health_check(proto::HealthCheckRequest {}) + .await + .expect_err("latched primary must fail readiness"); + assert_eq!(health.code(), tonic::Code::FailedPrecondition); + + drop(primary_client); + stop_node(primary_node, primary_handle).await?; + stop_node(replica_node, replica_handle).await?; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn replica_pairing_rejects_wrong_credentials_and_mismatched_state() -> anyhow::Result<()> { + use proto::shard_service_server::ShardService; + + let _test_guard = PERSISTENT_TEST_LOCK.lock().await; + let volumes = TempDir::new()?; + let replica_node = node( + volumes.path().join("replica-data"), + volumes.path().join("replica-backup"), + )? + .into_replica(REPLICATION_TOKEN.into())?; + let (replica_addr, replica_handle) = spawn_node(replica_node.clone()).await?; + let replica_client = ShardServiceClient::connect(format!("http://{replica_addr}")).await?; + + let wrong_token_error = node( + volumes.path().join("wrong-token-primary-data"), + volumes.path().join("wrong-token-primary-backup"), + )? + .with_replica(replica_client.clone(), WRONG_REPLICATION_TOKEN.into()) + .await + .err() + .expect("replica pairing must authenticate the primary"); + assert_eq!(wrong_token_error.code(), tonic::Code::PermissionDenied); + + let mismatched_primary = node( + volumes.path().join("mismatched-primary-data"), + volumes.path().join("mismatched-primary-backup"), + )?; + ShardService::ingest_records( + &mismatched_primary, + Request::new(IngestRecordsRequest { + records: vec![record(0)], + internal_protocol_version: DISTRIBUTED_PROTOCOL_VERSION, + }), + ) + .await?; + let mismatch_error = mismatched_primary + .with_replica(replica_client, REPLICATION_TOKEN.into()) + .await + .err() + .expect("replica pairing must prove exact durable state equality"); + assert_eq!(mismatch_error.code(), tonic::Code::FailedPrecondition); + assert!( + mismatch_error.message().contains("durable state"), + "{}", + mismatch_error.message() + ); + + stop_node(replica_node, replica_handle).await?; + Ok(()) +} diff --git a/tests/temporal_evolution.rs b/tests/temporal_evolution.rs index cc89bfb..b1399c3 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: 3, + internal_protocol_version: 5, records: vec![record_c1.clone()], }) .await?; client_inc .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![record_c2.clone()], }) .await?; client_inc .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, 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: 3, + internal_protocol_version: 5, records: vec![record_c1, record_c2, record_c3], }) .await?; @@ -370,7 +370,7 @@ async fn ingestion_order_independence() -> anyhow::Result<()> { for &idx in order { client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![records[idx].clone()], }) .await?; @@ -447,19 +447,19 @@ async fn temporal_descriptor_evolution() -> anyhow::Result<()> { // Ingest over time client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![jan], }) .await?; client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![feb], }) .await?; client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![mar], }) .await?; @@ -572,7 +572,7 @@ async fn attribute_value_changes_over_time() -> anyhow::Result<()> { client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![initial, with_secondary], }) .await?; @@ -652,7 +652,7 @@ async fn multi_perspective_incremental_merge() -> anyhow::Result<()> { ); client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![crm_record], }) .await?; @@ -683,7 +683,7 @@ async fn multi_perspective_incremental_merge() -> anyhow::Result<()> { ); client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![erp_record], }) .await?; @@ -720,7 +720,7 @@ async fn multi_perspective_incremental_merge() -> anyhow::Result<()> { ); client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![hr_record], }) .await?; @@ -768,7 +768,7 @@ async fn idempotent_re_ingestion() -> anyhow::Result<()> { // First ingestion let resp1 = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![record.clone()], }) .await? @@ -778,7 +778,7 @@ async fn idempotent_re_ingestion() -> anyhow::Result<()> { // Second ingestion (identical) let resp2 = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![record.clone()], }) .await? @@ -794,7 +794,7 @@ async fn idempotent_re_ingestion() -> anyhow::Result<()> { // Third ingestion let resp3 = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![record], }) .await? @@ -847,13 +847,13 @@ async fn overlapping_temporal_ranges_extend_validity() -> anyhow::Result<()> { client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![jan_mar], }) .await?; client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![feb_apr], }) .await?; @@ -932,7 +932,7 @@ async fn temporal_gap_between_records() -> anyhow::Result<()> { client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![q1, q3], }) .await?; @@ -1084,7 +1084,7 @@ async fn same_cluster_regardless_of_insertion_order() -> anyhow::Result<()> { for &idx in order { let resp = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![records[idx].clone()], }) .await? @@ -1159,7 +1159,7 @@ async fn batch_ingestion_consistent_cluster() -> anyhow::Result<()> { let batch_resp = batch_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: records.clone(), }) .await? @@ -1186,7 +1186,7 @@ async fn batch_ingestion_consistent_cluster() -> anyhow::Result<()> { for record in &records { let resp = seq_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![record.clone()], }) .await? @@ -1271,7 +1271,7 @@ async fn independent_entities_stay_separate_any_order() -> anyhow::Result<()> { for &idx in order { client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![records[idx].clone()], }) .await?; @@ -1358,7 +1358,7 @@ async fn late_arrival_merges_into_existing_cluster() -> anyhow::Result<()> { let resp1 = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![rec1, rec2], }) .await? @@ -1385,7 +1385,7 @@ async fn late_arrival_merges_into_existing_cluster() -> anyhow::Result<()> { let resp2 = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 5, records: vec![late_rec], }) .await?