From 66357842c4a65224036f05d035b77b490a738818 Mon Sep 17 00:00:00 2001 From: Isaac Elbaz Date: Sat, 25 Jul 2026 12:31:54 -0400 Subject: [PATCH 1/7] Add verified off-host cluster backups --- Containerfile | 12 + README.md | 86 +++++- compose.yaml | 2 +- examples/unirust.toml | 12 + src/backup.rs | 512 +++++++++++++++++++++++++++++++ src/bin/unirust_backup.rs | 143 +++++++++ src/bin/unirust_rebalance.rs | 49 ++- src/bin/unirust_router.rs | 18 +- src/bin/unirust_shard.rs | 7 + src/lib.rs | 8 +- src/persistence.rs | 12 +- tests/backup_cli_safety.rs | 40 +++ tests/external_backup_restore.rs | 165 +++++++++- 13 files changed, 1036 insertions(+), 30 deletions(-) create mode 100644 src/backup.rs create mode 100644 src/bin/unirust_backup.rs create mode 100644 tests/backup_cli_safety.rs diff --git a/Containerfile b/Containerfile index 1ddd933..2901a1d 100644 --- a/Containerfile +++ b/Containerfile @@ -31,6 +31,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 +51,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 +86,14 @@ case "$1" in shift exec unirust_loadtest "$@" ;; + backup) + shift + exec unirust_backup "$@" + ;; + rebalance) + shift + exec unirust_rebalance "$@" + ;; *) exec "$@" ;; diff --git a/README.md b/README.md index 948f9ea..b6f7760 100644 --- a/README.md +++ b/README.md @@ -120,10 +120,19 @@ profile = "high-throughput" listen = "0.0.0.0:50061" id = 0 data_dir = "/var/lib/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 +149,19 @@ 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_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 @@ -381,23 +399,61 @@ 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: + +```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. + +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. +Unirust does not synchronously replicate live shard writes, so 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 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. +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. 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 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/unirust.toml b/examples/unirust.toml index 1a0be98..d956eb5 100644 --- a/examples/unirust.toml +++ b/examples/unirust.toml @@ -21,6 +21,10 @@ id = 0 # 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" # Run repair on startup (recovers from unclean shutdown) repair = false @@ -44,6 +48,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/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..4b5471f 100644 --- a/src/bin/unirust_router.rs +++ b/src/bin/unirust_router.rs @@ -55,11 +55,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" "# ); } diff --git a/src/bin/unirust_shard.rs b/src/bin/unirust_shard.rs index 4d4f11e..8bcf336 100644 --- a/src/bin/unirust_shard.rs +++ b/src/bin/unirust_shard.rs @@ -58,6 +58,10 @@ 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 CONFIG FILE (unirust.toml): profile = "billion-scale-high-performance" @@ -67,6 +71,9 @@ 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" "# ); } diff --git a/src/lib.rs b/src/lib.rs index 61f4036..c71802a 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}; diff --git a/src/persistence.rs b/src/persistence.rs index 375d068..b962293 100644 --- a/src/persistence.rs +++ b/src/persistence.rs @@ -296,6 +296,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 +344,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 +356,6 @@ pub fn restore_checkpoint_for_shard( } let source = source.canonicalize()?; - validate_rocksdb_checkpoint(&source)?; if destination.exists() { let metadata = fs::symlink_metadata(destination)?; 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/external_backup_restore.rs b/tests/external_backup_restore.rs index 4fd8fc5..58abeb6 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: 3, + }) + .await?; + client1 + .ingest_records(IngestRecordsRequest { + records: vec![shard1_record], + internal_protocol_version: 3, + }) + .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; From f98370d9e1d0d003f5e5c7ea3c8d5a2da2c6532a Mon Sep 17 00:00:00 2001 From: Isaac Elbaz Date: Sat, 25 Jul 2026 13:40:19 -0400 Subject: [PATCH 2/7] Add fail-closed synchronous shard replication --- README.md | 84 ++- examples/unirust.toml | 12 + proto/unirust.proto | 14 +- src/bin/unirust_shard.rs | 130 +++- src/config/mod.rs | 116 ++++ src/distributed.rs | 592 ++++++++++++++++-- src/persistence.rs | 47 ++ src/transport_security.rs | 18 + tests/distributed_apply_ontology.rs | 2 +- tests/distributed_conflicts_presets.rs | 2 +- .../distributed_cross_shard_reconciliation.rs | 54 +- ...d_cross_shard_reconciliation_persistent.rs | 38 +- tests/distributed_e2e.rs | 2 +- tests/distributed_ingest_stream.rs | 4 +- tests/distributed_metrics.rs | 2 +- tests/distributed_rebalance.rs | 6 +- tests/distributed_rebalance_stream.rs | 2 +- tests/distributed_reliability.rs | 12 +- tests/distributed_router_admin.rs | 2 +- tests/external_backup_restore.rs | 10 +- tests/process_crash_recovery.rs | 6 +- tests/shard_cli_safety.rs | 74 +++ tests/synchronous_replication.rs | 416 ++++++++++++ tests/temporal_evolution.rs | 48 +- 24 files changed, 1548 insertions(+), 145 deletions(-) create mode 100644 tests/synchronous_replication.rs diff --git a/README.md b/README.md index b6f7760..b3fc11e 100644 --- a/README.md +++ b/README.md @@ -152,6 +152,16 @@ write_buffer_mb = 256 | `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 | @@ -331,6 +341,67 @@ 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. +### 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. + +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 Configure each shard with a unique checkpoint root on storage outside its data @@ -431,10 +502,10 @@ 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. -Unirust does not synchronously replicate live shard writes, so 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. +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 @@ -455,6 +526,11 @@ 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`. An in-memory shard can only be started with the explicit `--ephemeral` flag and loses all records when the process exits. diff --git a/examples/unirust.toml b/examples/unirust.toml index d956eb5..22d3698 100644 --- a/examples/unirust.toml +++ b/examples/unirust.toml @@ -25,6 +25,18 @@ id = 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" +# 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 diff --git a/proto/unirust.proto b/proto/unirust.proto index 1c94ad9..3fe1ce7 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 --- diff --git a/src/bin/unirust_shard.rs b/src/bin/unirust_shard.rs index 8bcf336..28cceb3 100644 --- a/src/bin/unirust_shard.rs +++ b/src/bin/unirust_shard.rs @@ -1,9 +1,11 @@ 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}; fn parse_arg(flag: &str) -> Option { @@ -47,6 +49,18 @@ OPTIONS: --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 @@ -62,6 +76,19 @@ ENVIRONMENT: 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" @@ -74,6 +101,11 @@ CONFIG FILE (unirust.toml): 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" "# ); } @@ -173,6 +205,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() @@ -183,6 +236,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); } @@ -197,6 +257,17 @@ 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 \ @@ -229,7 +300,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, @@ -237,8 +308,57 @@ 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 {}", 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..8666092 100644 --- a/src/distributed.rs +++ b/src/distributed.rs @@ -34,6 +34,7 @@ use tokio::sync::{mpsc, oneshot}; use tokio::sync::{Mutex, RwLock}; 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 = 4; +const REPLICATION_TOKEN_HEADER: &str = "x-unirust-replication-token"; #[allow(clippy::result_large_err)] fn validate_distributed_protocol(protocol_version: u32) -> Result<(), Status> { @@ -1050,6 +1052,11 @@ pub struct ShardNode { ingest_worker_handles: Arc>>>, config_version: String, metrics: Arc, + shard_role: proto::ShardRole, + replica_client: Option>, + replication_token: Option, + replication_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. @@ -1100,6 +1107,109 @@ struct IngestJob { respond_to: oneshot::Sender, Status>>, } +struct ReplicationAttempt { + consistency: Arc, + armed: bool, +} + +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 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 + }; + let replication_attempt = ReplicationAttempt::new( + $self.replication_consistent.clone(), + replica_response.is_some(), + ); + (replication_serial, replica_response, replication_attempt) + }}; +} + impl ShardNode { pub fn new( shard_id: u32, @@ -1262,6 +1372,11 @@ impl ShardNode { 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)), + replication_gate: Arc::new(Mutex::new(())), allow_destructive_admin: false, cross_shard_conflicts: Arc::new(parking_lot::RwLock::new( recovered_cross_shard_conflicts, @@ -1322,6 +1437,11 @@ impl ShardNode { 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)), + replication_gate: Arc::new(Mutex::new(())), allow_destructive_admin: false, cross_shard_conflicts: Arc::new(parking_lot::RwLock::new(Vec::new())), }) @@ -1336,6 +1456,198 @@ 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: 0, + })) + .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.replication_consistent.load(Ordering::Acquire) { + Ok(()) + } else { + Err(Status::failed_precondition( + "primary and replica may have diverged; traffic is blocked until their complete \ + durable states are reconciled", + )) + } + } + + 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 +1679,7 @@ impl ShardNode { } fn ensure_no_pending_ingest(&self) -> Result<(), Status> { + self.ensure_replication_consistent()?; if self .ingest_wal .as_ref() @@ -1389,6 +1702,7 @@ impl ShardNode { } fn ensure_recovery_healthy(&self) -> Result<(), Status> { + self.ensure_replication_consistent()?; if self .ingest_wal .as_ref() @@ -1408,6 +1722,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!( @@ -2037,9 +2352,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,6 +2389,8 @@ 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 targets = { let mut unirust = write_unirust!(self); unirust @@ -2102,15 +2421,21 @@ impl proto::shard_service_server::ShardService for ShardNode { }, ) .collect(); - Ok(Response::new(proto::ReserveSourceRecordsResponse { - reservations, - })) + let response = proto::ReserveSourceRecordsResponse { reservations }; + 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 +2451,33 @@ 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 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 {}; + 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 +2498,8 @@ impl proto::shard_service_server::ShardService for ShardNode { ))); } + let (_replication_serial, replica_response, replication_attempt) = + replicate_to_replica!(self, set_ontology, replication_request); 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,13 +2563,21 @@ impl proto::shard_service_server::ShardService for ShardNode { } *config_guard = config; - Ok(Response::new(proto::ApplyOntologyResponse {})) + 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> { + self.authorize_mutation(&request)?; let _mutation_guard = self.mutation_gate.read().await; let start = Instant::now(); let request = request.into_inner(); @@ -2248,6 +2592,12 @@ impl proto::shard_service_server::ShardService for ShardNode { 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); // 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) @@ -2268,13 +2618,21 @@ impl proto::shard_service_server::ShardService for ShardNode { self.metrics .record_ingest(record_count, start.elapsed().as_micros() as u64); - Ok(Response::new(proto::IngestRecordsResponse { assignments })) + let response = proto::IngestRecordsResponse { assignments }; + replication_attempt.finish( + replica_response + .as_ref() + .is_none_or(|replica| replica == &response), + "record ingest", + )?; + Ok(Response::new(response)) } async fn ingest_records_stream( &self, request: Request>, ) -> Result, Status> { + self.authorize_mutation(&request)?; let _mutation_guard = self.mutation_gate.read().await; let start = Instant::now(); let mut stream = request.into_inner(); @@ -2299,6 +2657,12 @@ impl proto::shard_service_server::ShardService for ShardNode { self.ensure_store_healthy()?; self.ensure_ingest_payloads_idempotent(&chunk.records)?; + let replication_request = proto::IngestRecordsRequest { + records: chunk.records.clone(), + internal_protocol_version: DISTRIBUTED_PROTOCOL_VERSION, + }; + let (_replication_serial, replica_response, replication_attempt) = + replicate_to_replica!(self, ingest_records, replication_request); // 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) @@ -2317,6 +2681,12 @@ impl proto::shard_service_server::ShardService for ShardNode { dispatch_ingest_records(&self.ingest_txs, self.ingest_wal.as_deref(), chunk.records) .await? }; + replication_attempt.finish( + replica_response + .as_ref() + .is_none_or(|replica| replica.assignments == batch_assignments), + "streamed record ingest", + )?; assignments.extend(batch_assignments); } @@ -2499,6 +2869,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 +2891,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 +2928,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 +2964,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 +3001,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 +3025,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 +3206,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 +3214,30 @@ 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)?; 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, - })) + }; + 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 +3251,26 @@ 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)?; unirust .ingest_with_explicit_ids(records) .map_err(|err| Status::internal(err.to_string()))?; + let response = proto::ImportRecordsResponse { + imported: chunk.records.len() as u64, + }; + 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 +3280,13 @@ 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 mut summaries = if let Some(cache) = { let guard = read_unirust!(self); guard.cached_conflict_summaries() @@ -2917,13 +3361,20 @@ impl proto::shard_service_server::ShardService for ShardNode { .map(to_proto_conflict_summary) .collect(), }; + 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 +3383,12 @@ 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 config = self.ontology_config.lock().await.clone(); { let mut guard = write_unirust!(self); @@ -3031,9 +3488,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 @@ -3062,26 +3521,31 @@ impl proto::shard_service_server::ShardService for ShardNode { )); } + let (_replication_serial, replica_response, replication_attempt) = + replicate_to_replica!(self, apply_merge, replication_request); // 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(), + }, }; - - Ok(Response::new(proto::ApplyMergeResponse { - success: true, - records_updated: records_updated as u32, - error: String::new(), - })) + replication_attempt.finish( + replica_response + .as_ref() + .is_none_or(|replica| replica == &response), + "cross-shard merge", + )?; + Ok(Response::new(response)) } async fn get_dirty_boundary_keys( @@ -3153,11 +3617,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 +3637,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 +3649,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 +3722,8 @@ 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); // 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 +3745,14 @@ impl proto::shard_service_server::ShardService for ShardNode { *conflict_storage = updated; } - Ok(Response::new(proto::StoreCrossShardConflictsResponse { - stored_count, - })) + let response = proto::StoreCrossShardConflictsResponse { stored_count }; + replication_attempt.finish( + replica_response + .as_ref() + .is_none_or(|replica| replica == &response), + "cross-shard conflict storage", + )?; + Ok(Response::new(response)) } } @@ -3568,12 +4050,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, @@ -4913,6 +5419,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(), })) } @@ -5683,7 +6191,7 @@ mod wal_tests { ShardService::ingest_records( &shard, Request::new(proto::IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![original.clone()], }), ) @@ -5695,7 +6203,7 @@ mod wal_tests { let error = ShardService::ingest_records( &shard, Request::new(proto::IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![conflicting], }), ) diff --git a/src/persistence.rs b/src/persistence.rs index b962293..495edc2 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; @@ -417,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, @@ -2303,6 +2327,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()); 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/distributed_apply_ontology.rs b/tests/distributed_apply_ontology.rs index ae9775e..f98e376 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: 4, records: vec![record], }) .await?; diff --git a/tests/distributed_conflicts_presets.rs b/tests/distributed_conflicts_presets.rs index 6aa89d9..fd7ac39 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: 4, records: inputs.clone(), }) .await diff --git a/tests/distributed_cross_shard_reconciliation.rs b/tests/distributed_cross_shard_reconciliation.rs index 2ac957f..20f2c7d 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: 4, 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: 4, 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: 4, 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: 4, 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: 4, 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: 4, 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: 4, 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: 4, 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: 4, 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: 4, records: vec![record_input( 2, "person", @@ -797,7 +797,7 @@ async fn boundary_metadata_includes_perspective_strong_ids() -> anyhow::Result<( router_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![record1, record2], }) .await?; @@ -882,7 +882,7 @@ async fn dirty_boundary_keys_include_perspective_strong_ids() -> anyhow::Result< router_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![record1, record2], }) .await?; @@ -1014,7 +1014,7 @@ async fn transitive_cross_shard_conflict_detected() -> anyhow::Result<()> { ); shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![a1, a2], }) .await?; @@ -1046,7 +1046,7 @@ async fn transitive_cross_shard_conflict_detected() -> anyhow::Result<()> { ); shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![b1, b2], }) .await?; @@ -1076,7 +1076,7 @@ async fn transitive_cross_shard_conflict_detected() -> anyhow::Result<()> { ); shard2_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![c1, c2], }) .await?; @@ -1201,7 +1201,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: 4, records: vec![a1, a2], }) .await?; @@ -1232,7 +1232,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: 4, records: vec![b1, b2], }) .await?; @@ -1329,7 +1329,7 @@ async fn temporal_overlap_conflict_across_shards() -> anyhow::Result<()> { ); shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![a1, a2], }) .await?; @@ -1359,7 +1359,7 @@ async fn temporal_overlap_conflict_across_shards() -> anyhow::Result<()> { ); shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![b1, b2], }) .await?; @@ -1453,7 +1453,7 @@ async fn late_arriving_data_triggers_conflict() -> anyhow::Result<()> { ); shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![initial1, initial2], }) .await?; @@ -1498,7 +1498,7 @@ async fn late_arriving_data_triggers_conflict() -> anyhow::Result<()> { ); shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![late1, late2], }) .await?; @@ -1622,7 +1622,7 @@ async fn multi_hop_chain_conflict_across_four_shards() -> anyhow::Result<()> { ); shard0 .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![a1, a2], }) .await?; @@ -1652,7 +1652,7 @@ async fn multi_hop_chain_conflict_across_four_shards() -> anyhow::Result<()> { ); shard1 .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![b1, b2], }) .await?; @@ -1682,7 +1682,7 @@ async fn multi_hop_chain_conflict_across_four_shards() -> anyhow::Result<()> { ); shard2 .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![c1, c2], }) .await?; @@ -1707,7 +1707,7 @@ async fn multi_hop_chain_conflict_across_four_shards() -> anyhow::Result<()> { ); shard3 .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![d1, d2], }) .await?; @@ -1810,7 +1810,7 @@ async fn different_perspectives_no_false_positive_conflict() -> anyhow::Result<( ); shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![msci1, msci2], }) .await?; @@ -1840,7 +1840,7 @@ async fn different_perspectives_no_false_positive_conflict() -> anyhow::Result<( ); shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, 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..88cefc9 100644 --- a/tests/distributed_cross_shard_reconciliation_persistent.rs +++ b/tests/distributed_cross_shard_reconciliation_persistent.rs @@ -347,7 +347,7 @@ async fn cross_shard_merge_persistent_store() -> anyhow::Result<()> { ); shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![shard0_rec1, shard0_rec2], }) .await?; @@ -374,7 +374,7 @@ async fn cross_shard_merge_persistent_store() -> anyhow::Result<()> { ); shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![shard1_rec1, shard1_rec2], }) .await?; @@ -424,7 +424,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: 4, records: vec![record_input( shard_id as u32, "person", @@ -534,7 +534,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: 4, records: vec![original.clone()], }) .await? @@ -544,7 +544,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: 4, records: vec![changed.clone()], }) .await @@ -581,7 +581,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: 4, records: vec![changed], }) .await @@ -591,7 +591,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: 4, records: vec![original], }) .await? @@ -658,13 +658,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: 4, records: vec![original], }) .await?; @@ -673,7 +675,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 +690,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: 4, records: vec![changed], }) .await @@ -752,7 +756,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: 4, records: vec![record.clone()], }) .await @@ -790,7 +794,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: 4, records: vec![record.clone()], }) .await? @@ -803,7 +807,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: 4, records: vec![changed], }) .await @@ -836,7 +840,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: 4, records: vec![record_input( 0, "person", @@ -851,7 +855,7 @@ async fn partial_reconciliation_blocks_traffic_and_recovers_after_full_restart( .await?; client1 .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![record_input( 1, "person", @@ -968,7 +972,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: 4, records: vec![record_input( 0, "person", diff --git a/tests/distributed_e2e.rs b/tests/distributed_e2e.rs index e6e89e4..f612c85 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: 4, records: vec![record_a.clone(), record_b.clone()], }) .await? diff --git a/tests/distributed_ingest_stream.rs b/tests/distributed_ingest_stream.rs index 2024906..8eae107 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: 4, 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: 4, records: vec![record_input( 1, "person", diff --git a/tests/distributed_metrics.rs b/tests/distributed_metrics.rs index dc8eca9..5e11416 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: 4, records: vec![record_input( 0, "person", diff --git a/tests/distributed_rebalance.rs b/tests/distributed_rebalance.rs index 22141ef..6b84f19 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: 4, 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: 4, }) .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: 4, }) .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..33741bb 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: 4, records, }) .await?; diff --git a/tests/distributed_reliability.rs b/tests/distributed_reliability.rs index e074005..6d7bac4 100644 --- a/tests/distributed_reliability.rs +++ b/tests/distributed_reliability.rs @@ -277,7 +277,7 @@ async fn shard_recovery_after_restart() -> anyhow::Result<()> { .collect(); client .ingest_records(proto::IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records, }) .await?; @@ -342,7 +342,7 @@ async fn router_handles_partial_shard_availability() -> anyhow::Result<()> { ); let response = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![record], }) .await? @@ -523,7 +523,7 @@ async fn destructive_reset_is_disabled_and_preserves_records_by_default() -> any client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![record_input( 0, "person", @@ -603,7 +603,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: 4, records: vec![record_input( 0, "person", @@ -657,7 +657,7 @@ async fn ingest_operations_complete_before_shutdown() -> anyhow::Result<()> { let response = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records, }) .await? @@ -708,7 +708,7 @@ async fn cluster_stats_reflect_all_shards() -> anyhow::Result<()> { client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records, }) .await?; diff --git a/tests/distributed_router_admin.rs b/tests/distributed_router_admin.rs index 15f498b..ba800b7 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: 4, records, }) .await?; diff --git a/tests/external_backup_restore.rs b/tests/external_backup_restore.rs index 58abeb6..e680cec 100644 --- a/tests/external_backup_restore.rs +++ b/tests/external_backup_restore.rs @@ -216,13 +216,13 @@ async fn verified_export_survives_loss_of_checkpoint_volumes() -> anyhow::Result client0 .ingest_records(IngestRecordsRequest { records: vec![shard0_record], - internal_protocol_version: 3, + internal_protocol_version: 4, }) .await?; client1 .ingest_records(IngestRecordsRequest { records: vec![shard1_record], - internal_protocol_version: 3, + internal_protocol_version: 4, }) .await?; @@ -349,7 +349,7 @@ async fn external_checkpoint_restores_a_lost_shard_volume() -> anyhow::Result<() client .ingest_records(IngestRecordsRequest { records: vec![record(0, "backup@example.com")], - internal_protocol_version: 3, + internal_protocol_version: 4, }) .await?; let before = client.get_stats(StatsRequest {}).await?.into_inner(); @@ -391,7 +391,7 @@ async fn external_checkpoint_restores_a_lost_shard_volume() -> anyhow::Result<() client .ingest_records(IngestRecordsRequest { records: vec![record(1, "backup@example.com")], - internal_protocol_version: 3, + internal_protocol_version: 4, }) .await?; let retry_stats = client.get_stats(StatsRequest {}).await?.into_inner(); @@ -400,7 +400,7 @@ async fn external_checkpoint_restores_a_lost_shard_volume() -> anyhow::Result<() let error = client .ingest_records(IngestRecordsRequest { records: vec![record(2, "changed@example.com")], - internal_protocol_version: 3, + internal_protocol_version: 4, }) .await .expect_err("restored immutable source identity must reject a changed payload"); diff --git a/tests/process_crash_recovery.rs b/tests/process_crash_recovery.rs index c02e9d7..3826c5e 100644 --- a/tests/process_crash_recovery.rs +++ b/tests/process_crash_recovery.rs @@ -143,7 +143,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: 4, records, }) .await? @@ -191,7 +191,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: 4, records: vec![record( 128, "process-kill-after-restart".to_string(), @@ -225,7 +225,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: 4, records: vec![record( 129, "graceful-shutdown-after-restart".to_string(), diff --git a/tests/shard_cli_safety.rs b/tests/shard_cli_safety.rs index ddbada3..447e9bc 100644 --- a/tests/shard_cli_safety.rs +++ b/tests/shard_cli_safety.rs @@ -74,3 +74,77 @@ 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"), + "--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"), + "--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"), + "--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/synchronous_replication.rs b/tests/synchronous_replication.rs new file mode 100644 index 0000000..f8513f3 --- /dev/null +++ b/tests/synchronous_replication.rs @@ -0,0 +1,416 @@ +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, ConfigVersionRequest, IngestRecordsRequest, + RecordDescriptor, RecordIdentity, RecordInput, StatsRequest, +}; +use unirust_rs::distributed::{ + DistributedOntologyConfig, IdentityKeyConfig, RouterNode, ShardNode, + DISTRIBUTED_PROTOCOL_VERSION, +}; +use unirust_rs::{StreamingTuning, TuningProfile}; + +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 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, + }], + } +} + +#[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: vec![record(0)], + internal_protocol_version: 0, + }), + ) + .await? + .into_inner(); + assert_eq!(response.assignments.len(), 1); + + 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, 1); + 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![record(1)], + 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, 1); + + 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 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..7929a3e 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: 4, records: vec![record_c1.clone()], }) .await?; client_inc .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![record_c2.clone()], }) .await?; client_inc .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, 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: 4, 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: 4, 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: 4, records: vec![jan], }) .await?; client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, records: vec![feb], }) .await?; client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, 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: 4, 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: 4, 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: 4, 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: 4, 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: 4, 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: 4, 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: 4, 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: 4, records: vec![jan_mar], }) .await?; client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 3, + internal_protocol_version: 4, 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: 4, 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: 4, 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: 4, 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: 4, 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: 4, 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: 4, 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: 4, records: vec![late_rec], }) .await? From cbee6d6b0c2c9074952c3f3ead2ad9480f33ff7d Mon Sep 17 00:00:00 2001 From: Isaac Elbaz Date: Sat, 25 Jul 2026 14:36:23 -0400 Subject: [PATCH 3/7] Harden restart recovery and production deployment --- .github/workflows/ci.yml | 14 ++ .github/workflows/release.yml | 2 + Cargo.toml | 4 +- Containerfile | 4 +- README.md | 33 ++- examples/cluster.rs | 12 +- examples/unirust.toml | 2 +- scripts/podman_cluster.sh | 9 +- src/bin/unirust_router.rs | 15 +- src/bin/unirust_shard.rs | 39 +++- src/distributed.rs | 215 ++++++++++++++---- tests/process_crash_recovery.rs | 18 +- tests/process_replication_recovery.rs | 301 ++++++++++++++++++++++++++ tests/shard_cli_safety.rs | 46 ++++ tests/synchronous_replication.rs | 29 ++- 15 files changed, 659 insertions(+), 84 deletions(-) create mode 100644 tests/process_replication_recovery.rs 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 2901a1d..0607074 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 @@ -103,4 +103,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 b3fc11e..cc9c523 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,6 +120,7 @@ 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" @@ -350,6 +351,8 @@ 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 @@ -531,9 +534,17 @@ 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`. An in-memory shard can only -be started with the explicit `--ephemeral` flag and loses all records when the -process exits. +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/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 22d3698..9450707 100644 --- a/examples/unirust.toml +++ b/examples/unirust.toml @@ -15,7 +15,7 @@ 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" 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/bin/unirust_router.rs b/src/bin/unirust_router.rs index 4b5471f..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() { @@ -255,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 28cceb3..1c2cb1a 100644 --- a/src/bin/unirust_shard.rs +++ b/src/bin/unirust_shard.rs @@ -8,6 +8,9 @@ use unirust_rs::distributed::{proto, DistributedOntologyConfig, ShardNode}; 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() { @@ -43,6 +46,8 @@ 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 @@ -274,6 +279,28 @@ async fn main() -> anyhow::Result<()> { 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 @@ -364,14 +391,18 @@ async fn main() -> anyhow::Result<()> { "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/distributed.rs b/src/distributed.rs index 8666092..cd3155e 100644 --- a/src/distributed.rs +++ b/src/distributed.rs @@ -1056,6 +1056,7 @@ pub struct ShardNode { 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, @@ -1112,6 +1113,29 @@ struct ReplicationAttempt { armed: bool, } +struct DurableMutationAttempt { + 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 ReplicationAttempt { fn new(consistency: Arc, armed: bool) -> Self { Self { consistency, armed } @@ -1376,6 +1400,7 @@ impl ShardNode { 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( @@ -1441,6 +1466,7 @@ impl ShardNode { 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())), @@ -1608,14 +1634,23 @@ impl ShardNode { } fn ensure_replication_consistent(&self) -> Result<(), Status> { - if self.replication_consistent.load(Ordering::Acquire) { - Ok(()) - } else { - Err(Status::failed_precondition( + 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()) } fn authorize_mutation(&self, request: &Request) -> Result<(), Status> { @@ -1885,19 +1920,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!( @@ -1943,6 +1990,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, @@ -1950,9 +2018,6 @@ impl ShardNode { &snapshot.descriptors, )?); } - unirust - .ensure_store_healthy() - .map_err(|err| Status::failed_precondition(err.to_string()))?; Ok(records) } @@ -2070,36 +2135,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()?; @@ -2500,6 +2547,7 @@ 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 @@ -2563,6 +2611,7 @@ impl proto::shard_service_server::ShardService for ShardNode { } *config_guard = config; + local_attempt.finish(); let response = proto::ApplyOntologyResponse {}; replication_attempt.finish( replica_response @@ -2598,6 +2647,7 @@ impl proto::shard_service_server::ShardService for ShardNode { }; let (_replication_serial, replica_response, replication_attempt) = replicate_to_replica!(self, ingest_records, replication_request); + let local_attempt = self.begin_durable_mutation(); // 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) @@ -2618,6 +2668,7 @@ impl proto::shard_service_server::ShardService for ShardNode { self.metrics .record_ingest(record_count, start.elapsed().as_micros() as u64); + local_attempt.finish(); let response = proto::IngestRecordsResponse { assignments }; replication_attempt.finish( replica_response @@ -2663,6 +2714,7 @@ impl proto::shard_service_server::ShardService for ShardNode { }; let (_replication_serial, replica_response, replication_attempt) = replicate_to_replica!(self, ingest_records, replication_request); + let local_attempt = self.begin_durable_mutation(); // 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) @@ -2687,6 +2739,7 @@ impl proto::shard_service_server::ShardService for ShardNode { .is_none_or(|replica| replica.assignments == batch_assignments), "streamed record ingest", )?; + local_attempt.finish(); assignments.extend(batch_assignments); } @@ -3217,13 +3270,17 @@ impl proto::shard_service_server::ShardService for ShardNode { 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()))?; let response = proto::ImportRecordsResponse { imported: request.records.len() as u64, }; + local_attempt.finish(); replication_attempt.finish( replica_response .as_ref() @@ -3258,13 +3315,17 @@ impl proto::shard_service_server::ShardService for ShardNode { 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() @@ -3287,6 +3348,7 @@ impl proto::shard_service_server::ShardService for ShardNode { 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() @@ -3361,6 +3423,7 @@ impl proto::shard_service_server::ShardService for ShardNode { .map(to_proto_conflict_summary) .collect(), }; + local_attempt.finish(); replication_attempt.finish( replica_response .as_ref() @@ -3389,6 +3452,7 @@ impl proto::shard_service_server::ShardService for ShardNode { 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); @@ -3412,6 +3476,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 {})) } @@ -3520,9 +3585,15 @@ impl proto::shard_service_server::ShardService for ShardNode { "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); @@ -3539,6 +3610,9 @@ impl proto::shard_service_server::ShardService for ShardNode { error: err.to_string(), }, }; + if response.success { + local_attempt.finish(); + } replication_attempt.finish( replica_response .as_ref() @@ -3724,6 +3798,7 @@ impl proto::shard_service_server::ShardService for ShardNode { 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(); @@ -3746,6 +3821,7 @@ impl proto::shard_service_server::ShardService for ShardNode { } let response = proto::StoreCrossShardConflictsResponse { stored_count }; + local_attempt.finish(); replication_attempt.finish( replica_response .as_ref() @@ -6223,6 +6299,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/tests/process_crash_recovery.rs b/tests/process_crash_recovery.rs index 3826c5e..adb155d 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) @@ -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(); @@ -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); 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 447e9bc..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")) @@ -83,6 +126,7 @@ fn shard_binary_rejects_replica_mode_without_a_token() { .args([ "--data-dir", data.path().to_str().expect("UTF-8 path"), + "--allow-colocated-checkpoints", "--replica-mode", ]) .output() @@ -107,6 +151,7 @@ fn shard_binary_rejects_short_replication_token() { .args([ "--data-dir", data.path().to_str().expect("UTF-8 path"), + "--allow-colocated-checkpoints", "--replica-mode", "--allow-insecure-replication", "--replication-token-file", @@ -134,6 +179,7 @@ fn shard_binary_requires_mtls_for_replication_by_default() { .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"), diff --git a/tests/synchronous_replication.rs b/tests/synchronous_replication.rs index f8513f3..3398540 100644 --- a/tests/synchronous_replication.rs +++ b/tests/synchronous_replication.rs @@ -165,6 +165,23 @@ fn record(index: u32) -> RecordInput { } } +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<()> { @@ -199,13 +216,15 @@ async fn acknowledged_ingest_survives_primary_volume_loss_and_manual_failover() let response = RouterService::ingest_records( router.as_ref(), Request::new(IngestRecordsRequest { - records: vec![record(0)], + records: (0..128) + .map(|source| distinct_record(source, source)) + .collect(), internal_protocol_version: 0, }), ) .await? .into_inner(); - assert_eq!(response.assignments.len(), 1); + assert_eq!(response.assignments.len(), 128); let mut primary_client = ShardServiceClient::connect(format!("http://{primary_addr}")).await?; let primary_stats = primary_client @@ -216,7 +235,7 @@ async fn acknowledged_ingest_survives_primary_volume_loss_and_manual_failover() .get_stats(StatsRequest {}) .await? .into_inner(); - assert_eq!(primary_stats.record_count, 1); + assert_eq!(primary_stats.record_count, 128); assert_eq!(primary_stats, replica_stats); let primary_config = primary_client @@ -294,7 +313,7 @@ async fn acknowledged_ingest_survives_primary_volume_loss_and_manual_failover() let retry = RouterService::ingest_records( promoted_router.as_ref(), Request::new(IngestRecordsRequest { - records: vec![record(1)], + records: vec![distinct_record(999, 0)], internal_protocol_version: 0, }), ) @@ -307,7 +326,7 @@ async fn acknowledged_ingest_survives_primary_volume_loss_and_manual_failover() .get_stats(StatsRequest {}) .await? .into_inner(); - assert_eq!(promoted_stats.record_count, 1); + assert_eq!(promoted_stats.record_count, 128); drop(promoted_router); drop(promoted_client); From 3b79944cedd29db6670cc83be3d60cfd6eb4a969 Mon Sep 17 00:00:00 2001 From: Isaac Elbaz Date: Sat, 25 Jul 2026 14:42:28 -0400 Subject: [PATCH 4/7] Install libclang in the container builder --- Containerfile | 1 + 1 file changed, 1 insertion(+) diff --git a/Containerfile b/Containerfile index 0607074..17944fc 100644 --- a/Containerfile +++ b/Containerfile @@ -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/* From cd3dda6a666eb51bb2f866e56d82cd9c81f234a3 Mon Sep 17 00:00:00 2001 From: Isaac Elbaz Date: Sat, 25 Jul 2026 16:19:24 -0400 Subject: [PATCH 5/7] Bound cross-shard reconciliation at scale --- proto/unirust.proto | 22 +- src/distributed.rs | 653 +++++++++++------- src/lib.rs | 85 ++- src/linker.rs | 154 +++-- src/model.rs | 2 +- src/partitioned.rs | 26 +- src/persistence.rs | 22 +- src/sharding.rs | 2 +- tests/distributed_apply_ontology.rs | 2 +- tests/distributed_conflicts_presets.rs | 2 +- .../distributed_cross_shard_reconciliation.rs | 194 +++++- ...d_cross_shard_reconciliation_persistent.rs | 34 +- tests/distributed_e2e.rs | 2 +- tests/distributed_ingest_stream.rs | 4 +- tests/distributed_metrics.rs | 2 +- tests/distributed_rebalance.rs | 6 +- tests/distributed_rebalance_stream.rs | 2 +- tests/distributed_reliability.rs | 12 +- tests/distributed_router_admin.rs | 2 +- tests/external_backup_restore.rs | 10 +- tests/process_crash_recovery.rs | 6 +- tests/temporal_evolution.rs | 48 +- 22 files changed, 864 insertions(+), 428 deletions(-) diff --git a/proto/unirust.proto b/proto/unirust.proto index 3fe1ce7..485a988 100644 --- a/proto/unirust.proto +++ b/proto/unirust.proto @@ -468,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; @@ -484,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 { @@ -523,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. @@ -592,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/src/distributed.rs b/src/distributed.rs index cd3155e..2294a40 100644 --- a/src/distributed.rs +++ b/src/distributed.rs @@ -68,7 +68,7 @@ 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 = 4; +pub const DISTRIBUTED_PROTOCOL_VERSION: u32 = 5; const REPLICATION_TOKEN_HEADER: &str = "x-unirust-replication-token"; #[allow(clippy::result_large_err)] @@ -928,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 { @@ -1069,6 +1083,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 @@ -1603,7 +1621,8 @@ impl ShardNode { } let metadata = replica .get_boundary_metadata(Request::new(proto::GetBoundaryMetadataRequest { - since_version: 0, + since_version: u64::MAX, + signatures: Vec::new(), })) .await .map_err(|error| { @@ -3486,17 +3505,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, @@ -3566,20 +3593,8 @@ 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", @@ -3622,68 +3637,160 @@ impl proto::shard_service_server::ShardService for ShardNode { Ok(Response::new(response)) } + 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, })) } @@ -3949,10 +4056,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(); @@ -4226,7 +4336,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()))? @@ -4385,60 +4496,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(()) } @@ -4742,17 +4810,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"))?; @@ -4769,24 +4856,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 ))); } @@ -4823,12 +5002,11 @@ impl RouterNode { &self, result: &crate::sharding::ReconciliationResult, ) -> Result<(), Status> { - for (primary, secondary) in &result.merged_clusters { + 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?; + self.apply_merges_to_shard(shard_id, merges).await?; } } @@ -4848,7 +5026,10 @@ 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(()) } @@ -4974,74 +5155,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(); @@ -5053,56 +5184,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); } } } @@ -5114,14 +5245,28 @@ impl RouterNode { ) -> Result { use crate::sharding::IncrementalReconciler; + let metadata_started = Instant::now(); let mut reconciler = IncrementalReconciler::new(); - for metadata in self.fetch_boundary_metadata().await? { - reconciler.add_shard_boundary(boundary_index_from_metadata(&metadata)?); + for chunk in keys.chunks(RECONCILIATION_KEY_CHUNK) { + for metadata in self.fetch_boundary_metadata(chunk).await? { + reconciler.add_shard_boundary(boundary_index_from_metadata(&metadata)?); + } } - - let key_set: std::collections::HashSet<_> = keys.iter().copied().collect(); + let metadata_elapsed = metadata_started.elapsed(); + let key_set = keys.iter().copied().collect(); + let reconcile_started = Instant::now(); let result = reconciler.reconcile_keys(&key_set); + let reconcile_elapsed = reconcile_started.elapsed(); + let apply_started = Instant::now(); self.apply_reconciliation_result(&result).await?; + tracing::info!( + dirty_keys = keys.len(), + merges = result.merges_performed, + 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) } @@ -5866,15 +6011,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( @@ -6267,7 +6406,7 @@ mod wal_tests { ShardService::ingest_records( &shard, Request::new(proto::IngestRecordsRequest { - internal_protocol_version: 4, + internal_protocol_version: 5, records: vec![original.clone()], }), ) @@ -6279,7 +6418,7 @@ mod wal_tests { let error = ShardService::ingest_records( &shard, Request::new(proto::IngestRecordsRequest { - internal_protocol_version: 4, + internal_protocol_version: 5, records: vec![conflicting], }), ) diff --git a/src/lib.rs b/src/lib.rs index c71802a..00adbdd 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1217,14 +1217,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 @@ -1266,41 +1290,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..f213438 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)] @@ -413,9 +430,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, @@ -538,8 +555,8 @@ 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, @@ -1327,10 +1344,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 +1388,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 +1428,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 +1458,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 +1473,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 +1498,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)]) + } + + /// 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)), + ); - // 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 + 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; } } @@ -2110,8 +2182,8 @@ fn reconcile_cluster_ids( fn reconcile_global_cluster_ids( global_cluster_ids: &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, @@ -2173,7 +2245,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 495edc2..ab831a3 100644 --- a/src/persistence.rs +++ b/src/persistence.rs @@ -3220,16 +3220,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..e1e9eb3 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 { diff --git a/tests/distributed_apply_ontology.rs b/tests/distributed_apply_ontology.rs index f98e376..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: 4, + internal_protocol_version: 5, records: vec![record], }) .await?; diff --git a/tests/distributed_conflicts_presets.rs b/tests/distributed_conflicts_presets.rs index fd7ac39..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: 4, + 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 20f2c7d..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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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 88cefc9..a0b5814 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)) @@ -347,7 +347,7 @@ async fn cross_shard_merge_persistent_store() -> anyhow::Result<()> { ); shard0_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 4, + internal_protocol_version: 5, records: vec![shard0_rec1, shard0_rec2], }) .await?; @@ -374,7 +374,7 @@ async fn cross_shard_merge_persistent_store() -> anyhow::Result<()> { ); shard1_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 4, + internal_protocol_version: 5, records: vec![shard1_rec1, shard1_rec2], }) .await?; @@ -424,7 +424,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: 4, + internal_protocol_version: 5, records: vec![record_input( shard_id as u32, "person", @@ -534,7 +534,7 @@ async fn source_identity_reservation_survives_routing_change_and_restart() -> an let response = router_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 4, + internal_protocol_version: 5, records: vec![original.clone()], }) .await? @@ -544,7 +544,7 @@ async fn source_identity_reservation_survives_routing_change_and_restart() -> an let error = router_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 4, + internal_protocol_version: 5, records: vec![changed.clone()], }) .await @@ -581,7 +581,7 @@ async fn source_identity_reservation_survives_routing_change_and_restart() -> an let error = router_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 4, + internal_protocol_version: 5, records: vec![changed], }) .await @@ -591,7 +591,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: 4, + internal_protocol_version: 5, records: vec![original], }) .await? @@ -666,7 +666,7 @@ async fn router_backfills_legacy_records_before_serving_ingest() -> anyhow::Resu assert_eq!(status_before.source_reservation_backfill_version, 0); original_target_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 4, + internal_protocol_version: 5, records: vec![original], }) .await?; @@ -690,7 +690,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: 4, + internal_protocol_version: 5, records: vec![changed], }) .await @@ -756,7 +756,7 @@ async fn reserved_ingest_retries_after_target_failure_and_full_restart() -> anyh let error = router_client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 4, + internal_protocol_version: 5, records: vec![record.clone()], }) .await @@ -794,7 +794,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: 4, + internal_protocol_version: 5, records: vec![record.clone()], }) .await? @@ -807,7 +807,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: 4, + internal_protocol_version: 5, records: vec![changed], }) .await @@ -840,7 +840,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: 4, + internal_protocol_version: 5, records: vec![record_input( 0, "person", @@ -855,7 +855,7 @@ async fn partial_reconciliation_blocks_traffic_and_recovers_after_full_restart( .await?; client1 .ingest_records(IngestRecordsRequest { - internal_protocol_version: 4, + internal_protocol_version: 5, records: vec![record_input( 1, "person", @@ -972,7 +972,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: 4, + internal_protocol_version: 5, records: vec![record_input( 0, "person", diff --git a/tests/distributed_e2e.rs b/tests/distributed_e2e.rs index f612c85..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: 4, + internal_protocol_version: 5, records: vec![record_a.clone(), record_b.clone()], }) .await? diff --git a/tests/distributed_ingest_stream.rs b/tests/distributed_ingest_stream.rs index 8eae107..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: 4, + 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: 4, + internal_protocol_version: 5, records: vec![record_input( 1, "person", diff --git a/tests/distributed_metrics.rs b/tests/distributed_metrics.rs index 5e11416..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: 4, + internal_protocol_version: 5, records: vec![record_input( 0, "person", diff --git a/tests/distributed_rebalance.rs b/tests/distributed_rebalance.rs index 6b84f19..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: 4, + 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: 4, + 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: 4, + 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 33741bb..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: 4, + internal_protocol_version: 5, records, }) .await?; diff --git a/tests/distributed_reliability.rs b/tests/distributed_reliability.rs index 6d7bac4..d38e264 100644 --- a/tests/distributed_reliability.rs +++ b/tests/distributed_reliability.rs @@ -277,7 +277,7 @@ async fn shard_recovery_after_restart() -> anyhow::Result<()> { .collect(); client .ingest_records(proto::IngestRecordsRequest { - internal_protocol_version: 4, + internal_protocol_version: 5, records, }) .await?; @@ -342,7 +342,7 @@ async fn router_handles_partial_shard_availability() -> anyhow::Result<()> { ); let response = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 4, + internal_protocol_version: 5, records: vec![record], }) .await? @@ -523,7 +523,7 @@ async fn destructive_reset_is_disabled_and_preserves_records_by_default() -> any client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 4, + internal_protocol_version: 5, records: vec![record_input( 0, "person", @@ -603,7 +603,7 @@ async fn incomplete_ontology_update_blocks_cluster_traffic() -> anyhow::Result<( let ingest_error = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 4, + internal_protocol_version: 5, records: vec![record_input( 0, "person", @@ -657,7 +657,7 @@ async fn ingest_operations_complete_before_shutdown() -> anyhow::Result<()> { let response = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 4, + internal_protocol_version: 5, records, }) .await? @@ -708,7 +708,7 @@ async fn cluster_stats_reflect_all_shards() -> anyhow::Result<()> { client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 4, + internal_protocol_version: 5, records, }) .await?; diff --git a/tests/distributed_router_admin.rs b/tests/distributed_router_admin.rs index ba800b7..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: 4, + internal_protocol_version: 5, records, }) .await?; diff --git a/tests/external_backup_restore.rs b/tests/external_backup_restore.rs index e680cec..ee89d92 100644 --- a/tests/external_backup_restore.rs +++ b/tests/external_backup_restore.rs @@ -216,13 +216,13 @@ async fn verified_export_survives_loss_of_checkpoint_volumes() -> anyhow::Result client0 .ingest_records(IngestRecordsRequest { records: vec![shard0_record], - internal_protocol_version: 4, + internal_protocol_version: 5, }) .await?; client1 .ingest_records(IngestRecordsRequest { records: vec![shard1_record], - internal_protocol_version: 4, + internal_protocol_version: 5, }) .await?; @@ -349,7 +349,7 @@ async fn external_checkpoint_restores_a_lost_shard_volume() -> anyhow::Result<() client .ingest_records(IngestRecordsRequest { records: vec![record(0, "backup@example.com")], - internal_protocol_version: 4, + internal_protocol_version: 5, }) .await?; let before = client.get_stats(StatsRequest {}).await?.into_inner(); @@ -391,7 +391,7 @@ async fn external_checkpoint_restores_a_lost_shard_volume() -> anyhow::Result<() client .ingest_records(IngestRecordsRequest { records: vec![record(1, "backup@example.com")], - internal_protocol_version: 4, + internal_protocol_version: 5, }) .await?; let retry_stats = client.get_stats(StatsRequest {}).await?.into_inner(); @@ -400,7 +400,7 @@ async fn external_checkpoint_restores_a_lost_shard_volume() -> anyhow::Result<() let error = client .ingest_records(IngestRecordsRequest { records: vec![record(2, "changed@example.com")], - internal_protocol_version: 4, + internal_protocol_version: 5, }) .await .expect_err("restored immutable source identity must reject a changed payload"); diff --git a/tests/process_crash_recovery.rs b/tests/process_crash_recovery.rs index adb155d..6273517 100644 --- a/tests/process_crash_recovery.rs +++ b/tests/process_crash_recovery.rs @@ -153,7 +153,7 @@ async fn acknowledged_ingest_survives_process_kill_and_restart() -> anyhow::Resu .collect(); let response = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 4, + internal_protocol_version: 5, records, }) .await? @@ -201,7 +201,7 @@ async fn acknowledged_ingest_survives_process_kill_and_restart() -> anyhow::Resu let linked_after_restart = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 4, + internal_protocol_version: 5, records: vec![record( 128, "process-kill-after-restart".to_string(), @@ -235,7 +235,7 @@ async fn acknowledged_ingest_survives_process_kill_and_restart() -> anyhow::Resu let linked_after_shutdown = client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 4, + internal_protocol_version: 5, records: vec![record( 129, "graceful-shutdown-after-restart".to_string(), diff --git a/tests/temporal_evolution.rs b/tests/temporal_evolution.rs index 7929a3e..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: 4, + internal_protocol_version: 5, records: vec![record_c1.clone()], }) .await?; client_inc .ingest_records(IngestRecordsRequest { - internal_protocol_version: 4, + internal_protocol_version: 5, records: vec![record_c2.clone()], }) .await?; client_inc .ingest_records(IngestRecordsRequest { - internal_protocol_version: 4, + 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: 4, + 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: 4, + 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: 4, + internal_protocol_version: 5, records: vec![jan], }) .await?; client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 4, + internal_protocol_version: 5, records: vec![feb], }) .await?; client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + internal_protocol_version: 5, records: vec![jan_mar], }) .await?; client .ingest_records(IngestRecordsRequest { - internal_protocol_version: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + 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: 4, + internal_protocol_version: 5, records: vec![late_rec], }) .await? From a28ce20fe8e7da254974b3f2a38e973cf78d1981 Mon Sep 17 00:00:00 2001 From: Isaac Elbaz Date: Sat, 25 Jul 2026 16:55:24 -0400 Subject: [PATCH 6/7] Complete ingests after client cancellation --- README.md | 9 +- src/distributed.rs | 206 ++++++++++++----------- src/sharding.rs | 100 +++++++++-- tests/distributed_ingest_cancellation.rs | 115 +++++++++++++ 4 files changed, 313 insertions(+), 117 deletions(-) create mode 100644 tests/distributed_ingest_cancellation.rs diff --git a/README.md b/README.md index cc9c523..1447c00 100644 --- a/README.md +++ b/README.md @@ -337,10 +337,11 @@ 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 diff --git a/src/distributed.rs b/src/distributed.rs index 2294a40..34cfabe 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,7 +31,7 @@ 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; @@ -1061,6 +1061,7 @@ 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>>>, @@ -1409,6 +1410,7 @@ 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)), @@ -1475,6 +1477,7 @@ 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)), @@ -1672,6 +1675,85 @@ impl ShardNode { 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() @@ -2646,56 +2728,12 @@ impl proto::shard_service_server::ShardService for ShardNode { request: Request, ) -> Result, Status> { self.authorize_mutation(&request)?; - let _mutation_guard = self.mutation_gate.read().await; - let start = Instant::now(); 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)?; - - 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(); - // 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); - local_attempt.finish(); - let response = proto::IngestRecordsResponse { assignments }; - replication_attempt.finish( - replica_response - .as_ref() - .is_none_or(|replica| replica == &response), - "record ingest", - )?; - Ok(Response::new(response)) + let assignments = self + .await_ingest_commit(request.records, "record ingest") + .await?; + Ok(Response::new(proto::IngestRecordsResponse { assignments })) } async fn ingest_records_stream( @@ -2703,11 +2741,8 @@ impl proto::shard_service_server::ShardService for ShardNode { request: Request>, ) -> Result, Status> { self.authorize_mutation(&request)?; - let _mutation_guard = self.mutation_gate.read().await; - let start = Instant::now(); let mut stream = request.into_inner(); let mut assignments = Vec::new(); - let mut record_count = 0usize; while let Some(chunk) = stream .message() @@ -2718,52 +2753,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)?; - - let replication_request = proto::IngestRecordsRequest { - records: chunk.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(); - // 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? - }; - replication_attempt.finish( - replica_response - .as_ref() - .is_none_or(|replica| replica.assignments == batch_assignments), - "streamed record ingest", - )?; - local_attempt.finish(); + 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 })) } @@ -5245,18 +5240,25 @@ impl RouterNode { ) -> Result { use crate::sharding::IncrementalReconciler; - let metadata_started = Instant::now(); - let mut reconciler = IncrementalReconciler::new(); + 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 metadata_elapsed = metadata_started.elapsed(); - let key_set = keys.iter().copied().collect(); let reconcile_started = Instant::now(); - let result = reconciler.reconcile_keys(&key_set); - let reconcile_elapsed = reconcile_started.elapsed(); + let result = candidates.finish(); + reconcile_elapsed += reconcile_started.elapsed(); let apply_started = Instant::now(); self.apply_reconciliation_result(&result).await?; tracing::info!( diff --git a/src/sharding.rs b/src/sharding.rs index e1e9eb3..e9cdaf7 100644 --- a/src/sharding.rs +++ b/src/sharding.rs @@ -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/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(()) +} From 2cd5f07a2eb9b98f13156ad720bd45f951175a62 Mon Sep 17 00:00:00 2001 From: Isaac Elbaz Date: Sat, 25 Jul 2026 18:12:36 -0400 Subject: [PATCH 7/7] Harden cancellation and restart recovery --- README.md | 35 +++-- src/distributed.rs | 97 ++++++++++---- src/lib.rs | 6 + src/linker.rs | 93 +++++++++++-- src/persistence.rs | 72 ++++++++++ src/store.rs | 13 ++ ...d_cross_shard_reconciliation_persistent.rs | 123 +++++++++++++++++ tests/distributed_reliability.rs | 125 ++++++++++++++++++ tests/linker_state_recovery.rs | 77 ++++++++++- tests/support/mod.rs | 105 +++++++++++++++ tests/synchronous_replication.rs | 110 ++++++++++++++- 11 files changed, 802 insertions(+), 54 deletions(-) diff --git a/README.md b/README.md index 1447c00..952d1f0 100644 --- a/README.md +++ b/README.md @@ -288,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 diff --git a/src/distributed.rs b/src/distributed.rs index 34cfabe..d4ee4b4 100644 --- a/src/distributed.rs +++ b/src/distributed.rs @@ -1137,6 +1137,11 @@ struct DurableMutationAttempt { armed: bool, } +struct ConsistencyAttempt { + consistency: Arc, + armed: bool, +} + impl DurableMutationAttempt { fn new(consistency: Arc, armed: bool) -> Self { Self { consistency, armed } @@ -1155,6 +1160,28 @@ impl Drop for DurableMutationAttempt { } } +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 } @@ -1215,6 +1242,10 @@ macro_rules! replicate_to_replica { 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") @@ -1245,10 +1276,6 @@ macro_rules! replicate_to_replica { } else { None }; - let replication_attempt = ReplicationAttempt::new( - $self.replication_consistent.clone(), - replica_response.is_some(), - ); (replication_serial, replica_response, replication_attempt) }}; } @@ -2539,25 +2566,28 @@ 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() @@ -2570,6 +2600,7 @@ impl proto::shard_service_server::ShardService for ShardNode { ) .collect(); let response = proto::ReserveSourceRecordsResponse { reservations }; + local_attempt.finish(); replication_attempt.finish( replica_response .as_ref() @@ -2601,12 +2632,14 @@ 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()))?; let response = proto::MarkSourceReservationsBackfilledResponse {}; + local_attempt.finish(); replication_attempt.finish( replica_response .as_ref() @@ -4986,22 +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> { + ) -> 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_merges_to_shard(shard_id, merges).await?; + records_updated = records_updated + .saturating_add(self.apply_merges_to_shard(shard_id, merges).await?); } } @@ -5026,7 +5063,7 @@ impl RouterNode { .await?; } } - Ok(()) + Ok(records_updated) } /// Get read access to the cluster locality index. @@ -5260,10 +5297,11 @@ impl RouterNode { let result = candidates.finish(); reconcile_elapsed += reconcile_started.elapsed(); let apply_started = Instant::now(); - self.apply_reconciliation_result(&result).await?; + 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(), @@ -5428,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(); @@ -5447,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 {})) } diff --git a/src/lib.rs b/src/lib.rs index 00adbdd..483eb14 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -316,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) diff --git a/src/linker.rs b/src/linker.rs index f213438..05afcc1 100644 --- a/src/linker.rs +++ b/src/linker.rs @@ -415,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) @@ -548,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(), @@ -563,14 +566,27 @@ impl StreamingLinker { }; 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)?; } } @@ -632,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. @@ -842,6 +859,7 @@ impl StreamingLinker { cluster_ids, next_cluster_id, global_cluster_ids, + global_cluster_anchors, strong_id_summaries, tainted_identity_keys, record_perspectives, @@ -934,6 +952,7 @@ impl StreamingLinker { ); reconcile_global_cluster_ids( global_cluster_ids, + global_cluster_anchors, cross_shard_merges, boundary_signatures, dirty_boundary_keys, @@ -1112,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 @@ -1158,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; @@ -1171,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, @@ -1204,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, @@ -1316,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 } @@ -1774,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) } @@ -2181,6 +2237,7 @@ 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 BTreeMap<(ShardingKeySignature, GlobalClusterId), BoundaryData>, dirty_boundary_keys: &mut BTreeSet, @@ -2204,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() { @@ -2221,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 diff --git a/src/persistence.rs b/src/persistence.rs index ab831a3..6e89b2f 100644 --- a/src/persistence.rs +++ b/src/persistence.rs @@ -1453,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, @@ -3131,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); @@ -3159,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 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/tests/distributed_cross_shard_reconciliation_persistent.rs b/tests/distributed_cross_shard_reconciliation_persistent.rs index a0b5814..703f25f 100644 --- a/tests/distributed_cross_shard_reconciliation_persistent.rs +++ b/tests/distributed_cross_shard_reconciliation_persistent.rs @@ -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, @@ -951,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; diff --git a/tests/distributed_reliability.rs b/tests/distributed_reliability.rs index d38e264..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, @@ -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. 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/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 index 3398540..9a7c814 100644 --- a/tests/synchronous_replication.rs +++ b/tests/synchronous_replication.rs @@ -14,8 +14,8 @@ use tonic::server::NamedService; use tonic::transport::Server; use tonic::Request; use unirust_rs::distributed::proto::{ - self, shard_service_client::ShardServiceClient, ConfigVersionRequest, IngestRecordsRequest, - RecordDescriptor, RecordIdentity, RecordInput, StatsRequest, + self, shard_service_client::ShardServiceClient, ApplyOntologyRequest, ConfigVersionRequest, + IngestRecordsRequest, RecordDescriptor, RecordIdentity, RecordInput, StatsRequest, }; use unirust_rs::distributed::{ DistributedOntologyConfig, IdentityKeyConfig, RouterNode, ShardNode, @@ -23,6 +23,8 @@ use unirust_rs::distributed::{ }; 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 = @@ -108,6 +110,27 @@ async fn spawn_failable_node( 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(); @@ -335,6 +358,89 @@ async fn acknowledged_ingest_survives_primary_volume_loss_and_manual_failover() 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;