diff --git a/.gitignore b/.gitignore index b91857fb..7c17a0b0 100644 --- a/.gitignore +++ b/.gitignore @@ -53,3 +53,7 @@ CLAUDE.md .claude/ .rules .github/copilot-instructions.md + +# MemPalace per-project files (issue #185) +mempalace.yaml +entities.json diff --git a/CHANGELOG.md b/CHANGELOG.md index 41c510ff..6015118f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,27 @@ All notable changes to this project will be documented in this file. --- +## [v0.2.5] - 2026-XX-XX + +### Added + +- **ReadActor fast path for Eventual/LeaseRead** (#392): Dedicated `ReadActor` task serves `EventualConsistency` and `LeaseRead` without entering the Raft loop, eliminating channel contention under high read concurrency. + - `Arc` is owned exclusively by ReadActor — guarantees RocksDB LOCK release before Raft shutdown on `stop()`, fixing `test_snapshot_recovery_embedded` + - `ReadLease.revoke()`: atomic lease invalidation on leader demotion (replaces `invalidate()`) + - `EmbeddedClient` routes Eventual/Lease → ReadActor; falls back to `cmd_tx` on `LeaseInvalid`/`SmStopped` + - New `[raft.read_actor]` config section: `channel_capacity` (default 512), `max_drain` (default 100) + - **Benchmark (local embedded, 100 concurrent clients)**: Lease Read +10.9%, Eventual Read +13.1% vs v0.2.4; write throughput unchanged (see `benches/reports/v0.2.5/`) + +### Fixed + +- **fix(ttl) #398**: Removed `lease.enabled` flag — TTL expiration is always active. Fixes fatal crash when calling `put_with_ttl` without setting the (now-removed) `lease.enabled = true`. + +### Changed + +- `[raft.read_actor]` replaces the previous flat `read_actor_channel_capacity` / `read_actor_max_drain` fields in `[raft]`. Update existing config files accordingly. + +--- + ## [v0.2.4] - 2026-05-23 ### Added diff --git a/MIGRATION_GUIDE.md b/MIGRATION_GUIDE.md index 7c136ac5..fb768c2c 100644 --- a/MIGRATION_GUIDE.md +++ b/MIGRATION_GUIDE.md @@ -134,6 +134,7 @@ grep "WAL" /var/log/d-engine.log | v0.2.0–v0.2.2 | Absolute expiration | Compatible | ✅ Yes (clear WAL from v0.1.x) | | v0.2.3 | Same as v0.2.0+ | **Incompatible** | ✅ Yes (protobuf enum changes + API changes) | | v0.2.4 | Same as v0.2.0+ | Compatible (additive)| ✅ Yes (delete `snapshot/` — format changed to CF export)| +| v0.2.5 | Same as v0.2.0+ | Compatible (additive)| ⚠️ Minor (remove `lease.enabled` from config if present) | --- @@ -493,4 +494,50 @@ See [Watch Feature Guide](https://docs.rs/d-engine/latest/d_engine/docs/server_g --- +--- + +## For v0.2.4 Users: TTL Config Change in v0.2.5 (#398) + +### What Changed + +The `enabled` flag under `[raft.state_machine.lease]` has been removed. TTL expiration is now **always active** — no opt-in required. + +In v0.2.4, omitting `enabled = true` caused a fatal crash on `put_with_ttl`. This bug is fixed in v0.2.5. + +### Migration + +If your config contains `enabled = true` or `enabled = false`, remove the line: + +```toml +# Old (v0.2.4) — remove this line +[raft.state_machine.lease] +enabled = true # ← delete + +# New (v0.2.5) — TTL always active, no flag needed +[raft.state_machine.lease] +cleanup_interval_ms = 1000 +``` + +**Impact**: None if you do not touch the config — unrecognised fields are ignored. The only behavioral change is that TTL expiration is now unconditionally enabled. + +--- + +## For v0.2.4 Users: New `[raft.read_actor]` Config Section in v0.2.5 (#392) + +This section is **optional** — both fields have defaults and existing configs work without changes. + +```toml +[raft.read_actor] +# mpsc channel buffer for Eventual/LeaseRead fast path. +# Rule of thumb: ≥ 2× peak concurrent readers. Default: 512. +channel_capacity = 512 + +# Max reads drained per wakeup. Default: 100. +max_drain = 100 +``` + +No migration action required unless you want to tune read concurrency. + +--- + **Last Updated:** May 2026 diff --git a/README.md b/README.md index 9573bfb8..b6a5f51a 100644 --- a/README.md +++ b/README.md @@ -54,7 +54,7 @@ use std::time::Duration; #[tokio::main] async fn main() { - let engine = EmbeddedEngine::start("./data").await.unwrap(); + let engine = DefaultEmbeddedEngine::start("./data").await.unwrap(); engine.wait_ready(Duration::from_secs(5)).await.unwrap(); let client = engine.client(); @@ -123,13 +123,13 @@ Implement the `StorageEngine` and `StateMachine` traits for custom backends: ## Performance -### d-engine v0.2.4 vs etcd +### d-engine v0.2.5 vs etcd -![d-engine vs etcd comparison](https://raw.githubusercontent.com/deventlab/d-engine/main/benches/reports/v0.2.4/d-engine_comparison_v0.2.4.png) +![d-engine vs etcd comparison](https://raw.githubusercontent.com/deventlab/d-engine/main/benches/reports/v0.2.5/d-engine_comparison_v0.2.5.png) -### d-engine v0.2.4 vs v0.2.3 +### d-engine v0.2.5 vs v0.2.4 -![d-engine v0.2.4 vs v0.2.3 comparison](https://raw.githubusercontent.com/deventlab/d-engine/main/benches/reports/v0.2.4/d-engine_v0.2.3_vs_v0.2.4_embedded_mode.png) +![d-engine v0.2.5 vs v0.2.4 comparison](https://raw.githubusercontent.com/deventlab/d-engine/main/benches/reports/v0.2.5/d-engine_v0.2.5_vs_v0.2.4_embedded_mode.png) ```bash open benches/reports/ diff --git a/benches/embedded-bench/config/n1.toml b/benches/embedded-bench/config/n1.toml index ab276307..f53c2249 100644 --- a/benches/embedded-bench/config/n1.toml +++ b/benches/embedded-bench/config/n1.toml @@ -9,6 +9,10 @@ initial_cluster = [ db_root_dir = "./data/n1" log_dir = "./logs" +[raft.read_actor] +channel_capacity = 2048 +max_drain = 1024 + [raft.batching] # Maximum number of commands to accumulate in a single batch during drain operations max_batch_size = 200 diff --git a/benches/embedded-bench/config/n2.toml b/benches/embedded-bench/config/n2.toml index ee6d95f0..0d26472c 100644 --- a/benches/embedded-bench/config/n2.toml +++ b/benches/embedded-bench/config/n2.toml @@ -9,6 +9,10 @@ initial_cluster = [ db_root_dir = "./data/n2" log_dir = "./logs" +[raft.read_actor] +channel_capacity = 2048 +max_drain = 1024 + [raft.batching] # Maximum number of commands to accumulate in a single batch during drain operations max_batch_size = 200 diff --git a/benches/embedded-bench/config/n3.toml b/benches/embedded-bench/config/n3.toml index 651d5ce9..98a48063 100644 --- a/benches/embedded-bench/config/n3.toml +++ b/benches/embedded-bench/config/n3.toml @@ -9,6 +9,10 @@ initial_cluster = [ db_root_dir = "./data/n3" log_dir = "./logs" +[raft.read_actor] +channel_capacity = 2048 +max_drain = 1024 + [raft.batching] # Maximum number of commands to accumulate in a single batch during drain operations max_batch_size = 200 diff --git a/benches/embedded-bench/src/main.rs b/benches/embedded-bench/src/main.rs index caa0d23a..fa123ff6 100644 --- a/benches/embedded-bench/src/main.rs +++ b/benches/embedded-bench/src/main.rs @@ -1,4 +1,4 @@ -use d_engine::EmbeddedEngine; +use d_engine::DefaultEmbeddedEngine; use d_engine::protocol::ReadConsistencyPolicy; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; @@ -202,7 +202,7 @@ async fn main() { /// Run a single benchmark test with specified parameters #[allow(clippy::too_many_arguments)] async fn run_benchmark_task( - engine: &Arc, + engine: &Arc, test_name: &str, command: Commands, total: u64, @@ -294,7 +294,7 @@ async fn run_benchmark_task( /// Run all benchmark tests in batch mode async fn run_batch_tests( - engine: &Arc, + engine: &Arc, cli: &Cli, ) { println!("\n╔════════════════════════════════════════╗"); @@ -420,7 +420,7 @@ async fn run_batch_tests( /// Helper function to run a single test in batch mode async fn run_single_batch_test( - engine: &Arc, + engine: &Arc, test_name: &str, command: Commands, total: u64, @@ -447,7 +447,7 @@ async fn run_local_benchmark(cli: Cli) { println!("Starting local benchmark mode..."); let engine = Arc::new( - EmbeddedEngine::start_with(&cli.config_path) + DefaultEmbeddedEngine::start_with(&cli.config_path) .await .expect("Failed to start engine"), ); @@ -673,7 +673,7 @@ async fn run_http_server(cli: Cli) { println!("Health check port: {}", cli.health_port); let engine = Arc::new( - EmbeddedEngine::start_with(&cli.config_path) + DefaultEmbeddedEngine::start_with(&cli.config_path) .await .expect("Failed to start engine"), ); @@ -701,7 +701,7 @@ async fn run_http_server(cli: Cli) { } async fn start_health_check_server( - engine: Arc, + engine: Arc, port: u16, ) { let app = Router::new() @@ -718,7 +718,7 @@ async fn start_health_check_server( axum::serve(listener, app).await.expect("Health check server failed"); } -async fn health_primary(State(engine): State>) -> StatusCode { +async fn health_primary(State(engine): State>) -> StatusCode { if engine.is_leader() { StatusCode::OK } else { @@ -726,7 +726,7 @@ async fn health_primary(State(engine): State>) -> StatusCode } } -async fn health_replica(State(engine): State>) -> StatusCode { +async fn health_replica(State(engine): State>) -> StatusCode { if !engine.is_leader() { StatusCode::OK } else { @@ -735,7 +735,7 @@ async fn health_replica(State(engine): State>) -> StatusCode } async fn start_business_server( - engine: Arc, + engine: Arc, port: u16, ) { let app = Router::new() @@ -753,7 +753,7 @@ async fn start_business_server( } async fn handle_put( - State(engine): State>, + State(engine): State>, Json(req): Json, ) -> StatusCode { match engine.client().put(req.key.into_bytes(), req.value.into_bytes()).await { @@ -763,7 +763,7 @@ async fn handle_put( } async fn handle_get( - State(engine): State>, + State(engine): State>, Path(key): Path, ) -> Result, StatusCode> { match engine.client().get_eventual(key.into_bytes()).await { diff --git a/benches/reports/v0.2.5/bench_report_v0.2.5.md b/benches/reports/v0.2.5/bench_report_v0.2.5.md new file mode 100644 index 00000000..d2898011 --- /dev/null +++ b/benches/reports/v0.2.5/bench_report_v0.2.5.md @@ -0,0 +1,218 @@ +# d-engine v0.2.5 Benchmark Report + +**Test Environments**: + +- **Local**: Apple M2 Mac mini (8-core, 16GB RAM, 3-node cluster on localhost) +- **AWS**: EC2 c5.2xlarge (8 vCPUs, 16GB RAM, 50GB SSD) × 3 nodes + +**Test Dates**: + +- **Local v0.2.5 vs v0.2.4**: June 2, 2026 (6-round average, embedded) +- **AWS v0.2.5**: June 2, 2026 (2 × 5-round average, 10 rounds total; embedded and standalone) + +**Key/Value**: 8 bytes / 256 bytes + +--- + +## Local 3-Node Cluster: d-engine v0.2.5 vs d-engine v0.2.4 + +### Embedded Mode: v0.2.5 vs v0.2.4 + +_(v0.2.5: 6-round average; v0.2.4: 4-round average; v0.2.3: 4-round average (Lease/Eventual re-measured on same machine). See Benchmark Configuration for settings.)_ + +| **Scenario** | **Metric** | **v0.2.3** | **v0.2.4** | **v0.2.5** | **Δ (v0.2.4→v0.2.5)** | +| ------------------- | ----------- | ------------- | ------------- | ------------- | --------------------- | +| Single Client Write | Throughput | 10,075 ops/s | 9,740 ops/s | 9,658 ops/s | -0.8% → | +| | Avg Latency | 0.099 ms | 0.102 ms | 0.103 ms | stable | +| | p99 Latency | 0.139 ms | 0.177 ms | 0.203 ms | +14.7% → | +| High Conc. Write | Throughput | 176,314 ops/s | 233,821 ops/s | 224,176 ops/s | -4.1% → | +| | Avg Latency | 0.566 ms | 0.426 ms | 0.447 ms | +4.9% → | +| | p99 Latency | 1.566 ms | 1.164 ms | 0.912 ms | **-21.6%** ✅ | +| Linearizable Read | Throughput | 508,264 ops/s | 630,789 ops/s | 586,767 ops/s | -7.0% → | +| | Avg Latency | 0.197 ms | 0.157 ms | 0.170 ms | +8.3% → | +| | p99 Latency | 0.710 ms | 0.367 ms | 0.483 ms | +31.6% ⚠️ | +| Lease Read | Throughput | 705,530 ops/s | 730,836 ops/s | 893,254 ops/s | **+22.2%** ✅ | +| | Avg Latency | 0.116 ms | 0.136 ms | 0.007 ms | **-94.9%** ✅ | +| | p99 Latency | 0.342 ms | 0.337 ms | 0.066 ms | **-80.4%** ✅ | +| Eventual Read | Throughput | 742,602 ops/s | 752,198 ops/s | 884,781 ops/s | **+17.6%** ✅ | +| | Avg Latency | 0.115 ms | 0.132 ms | 0.007 ms | **-94.7%** ✅ | +| | p99 Latency | 0.382 ms | 0.343 ms | 0.067 ms | **-80.5%** ✅ | +| Hot-Key (10 keys) | Throughput | 499,527 ops/s | 659,429 ops/s | 622,459 ops/s | -5.6% → | +| | Avg Latency | 0.205 ms | 0.153 ms | 0.160 ms | stable | +| | p99 Latency | 0.638 ms | 0.343 ms | 0.425 ms | +23.9% ⚠️ | + +**Notes**: + +- **Lease/Eventual Read: headline win vs v0.2.4** — ReadActor fast path (#392) reduces avg latency from ~0.130 ms to ~0.007 ms (**-94%+**), throughput +22.2%/+17.6%. Direct SM reads eliminate all Raft/channel hops for Eventual/LeaseRead under 100 concurrent clients. +- v0.2.3 Lease/Eventual throughput (705K/742K) is the re-measured baseline from the same machine; original v0.2.3 values (852K/859K) were not reproducible due to different system state at the time of that test. +- Lease Read run-to-run variance: 819K–1,001K across 6 rounds (6-round avg 893K). Round 5 reached 1,001K; round 4 dipped to 819K (OS scheduler noise on a loaded Mac mini). +- HC Write p99 **-21.6%** vs v0.2.4 — tail latency improvement despite slight avg regression; 6-round variance includes runs with max p99.9 up to 2,943 µs skewing avg higher. +- Linearizable Read p99 +31.6% and HC Write avg +4.9% are within local benchmark noise on a shared Mac mini; neither read-index nor write-path changed from v0.2.4. +- SC Write and Hot-Key throughput within ±6% of v0.2.4; consistent with normal run-to-run variance across 6 rounds. +- Default config (`read_actor_channel_capacity = 512`) yields Lease/Eventual ~790K/~806K (+8%/+7% vs v0.2.4). Results above use tuned settings (see configuration). + +--- + +### Standalone Mode: v0.2.5 vs v0.2.4 vs v0.2.3 + +_(v0.2.5: 5-round average; v0.2.4: 5-round average; v0.2.3: 5-round average. All manually collected.)_ + +| **Scenario** | **Metric** | **v0.2.3** | **v0.2.4** | **v0.2.5** | **Δ (v0.2.4→v0.2.5)** | +| ------------------- | ----------- | ------------ | ------------ | ------------ | --------------------- | +| Single Client Write | Throughput | 6,421 ops/s | 5,245 ops/s | 5,234 ops/s | stable | +| | Avg Latency | 0.155 ms | 0.190 ms | 0.190 ms | stable | +| | p99 Latency | 0.200 ms | 0.235 ms | 0.237 ms | stable | +| High Conc. Write | Throughput | 55,285 ops/s | 59,733 ops/s | 60,608 ops/s | +1.5% → | +| | Avg Latency | 3.610 ms | 3.346 ms | 3.297 ms | -1.5% → | +| | p99 Latency | 6.720 ms | 6.325 ms | 6.372 ms | stable | +| Linearizable Read | Throughput | 63,210 ops/s | 71,702 ops/s | 71,918 ops/s | stable | +| | Avg Latency | 3.160 ms | 2.791 ms | 2.781 ms | stable | +| | p99 Latency | 5.810 ms | 5.933 ms | 6.078 ms | stable | +| Lease Read | Throughput | 67,878 ops/s | 72,593 ops/s | 80,449 ops/s | **+10.8%** ✅ | +| | Avg Latency | 2.950 ms | 2.756 ms | 2.493 ms | **-9.5%** ✅ | +| | p99 Latency | 6.200 ms | 5.762 ms | 5.848 ms | stable | +| Eventual Read | Throughput | 91,174 ops/s | 94,956 ops/s | 97,188 ops/s | +2.4% → | +| | Avg Latency | 2.190 ms | 2.103 ms | 2.054 ms | -2.3% → | +| | p99 Latency | 13.970 ms | 9.762 ms | 10.084 ms | stable | +| Hot-Key (10 keys) | Throughput | 74,017 ops/s | 84,863 ops/s | 84,911 ops/s | stable | +| | Avg Latency | 2.700 ms | 2.360 ms | 2.356 ms | stable | +| | p99 Latency | 5.490 ms | 5.602 ms | 5.492 ms | -2.0% → | + +**Notes**: + +- **v0.2.5 Standalone is stable vs v0.2.4** — all write and non-Lease read scenarios within ±2%, confirming ReadActor fast path (#392) is additive and introduces no regressions in Standalone mode. +- **Lease Read: +10.8% throughput / -9.5% avg latency vs v0.2.4** — ReadActor shared infrastructure benefits Standalone mode; gRPC RTT masks the direct SM path, so improvement is smaller than Embedded. +- SC Write shows -18.3% vs v0.2.3 but is stable vs v0.2.4; the latency floor shift (155 µs → 190 µs) was introduced before v0.2.4, not a #392 regression. +- Eventual Read p99 has high run-to-run variance (~±1 ms) due to Mac mini OS scheduler noise under 1000 concurrent clients; the improvement trend from v0.2.3 (13.97 ms) to v0.2.4 (9.76 ms) is the meaningful signal. +- Lease Read round-to-run variance: ~75K–91K across 5 rounds (round 4 reached 91K; other rounds avg ~78K). 5-round avg 80K. + +--- + +## AWS 3-Node Cluster: d-engine v0.2.5 vs v0.2.4 / etcd 3.2.0 + +**Hardware**: AWS EC2 c5.2xlarge (8 vCPUs, 16GB RAM, 50GB SSD) × 3 nodes +**Date**: June 2, 2026 | 2 × 5-round average (10 rounds total) | Key/Value: 8 bytes / 256 bytes +**etcd reference**: Official etcd benchmark (GCE, 8 vCPUs + 16GB + SSD × 3 nodes, etcd 3.2.0)² + +![d-engine v0.2.5 vs v0.2.4 Embedded Mode](d-engine_v0.2.5_vs_v0.2.4_embedded_mode.png) + +![d-engine v0.2.5 vs v0.2.4 Standalone Mode](d-engine_v0.2.5_vs_v0.2.4_standalone_mode.png) + +![d-engine v0.2.5 vs etcd 3.2.0](d-engine_comparison_v0.2.5.png) + +### Embedded Mode: v0.2.5 vs v0.2.4 / etcd 3.2.0 + +| **Scenario** | **Metric** | **v0.2.4 (AWS)** | **v0.2.5 (AWS)** | **Δ vs v0.2.4** | **etcd 3.2.0²** | **Δ vs etcd** | +| ------------------- | ----------- | ---------------- | ---------------- | --------------- | --------------- | ------------- | +| Single Client Write | Throughput | 4,398 ops/s | 4,428 ops/s | +0.7% → | 583 ops/s | **+7.6x** ✅ | +| | Avg Latency | 0.227 ms | 0.225 ms | stable | 1.6 ms | **-86%** ✅ | +| | p99 Latency | 0.316 ms | 0.303 ms | -4.1% → | — | — | +| High Conc. Write | Throughput | 110,798 ops/s | 111,618 ops/s | +0.7% → | 44,341 ops/s | **+152%** ✅ | +| | Avg Latency | 0.902 ms | 0.897 ms | stable | 22.0 ms | **-95.9%** ✅ | +| | p99 Latency | 1.063 ms | 1.060 ms | stable | — | — | +| Linearizable Read | Throughput | 327,355 ops/s | 345,872 ops/s | +5.7% → | 141,578 ops/s | **+144%** ✅ | +| | Avg Latency | 0.305 ms | 0.289 ms | -5.2% → | 5.5 ms | **-94.7%** ✅ | +| | p99 Latency | 0.374 ms | 0.348 ms | -7.0% → | — | — | +| Lease Read | Throughput | 341,254 ops/s | 1,231,541 ops/s | **+260.9%** ✅ | —³ | — | +| | Avg Latency | 0.293 ms | 0.005 ms | **-98.3%** ✅ | — | — | +| | p99 Latency | 0.367 ms | 0.018 ms | **-95.1%** ✅ | — | — | +| Eventual Read | Throughput | 363,407 ops/s | 1,235,832 ops/s | **+240.1%** ✅ | 185,758 ops/s | **+6.7x** ✅ | +| | Avg Latency | 0.274 ms | 0.005 ms | **-98.2%** ✅ | 2.2 ms | **-99.8%** ✅ | +| | p99 Latency | 0.343 ms | 0.019 ms | **-94.5%** ✅ | — | — | +| Hot-Key (10 keys) | Throughput | 319,957 ops/s | 346,830 ops/s | **+8.4%** ✅ | —³ | — | +| | Avg Latency | 0.313 ms | 0.288 ms | -8.0% → | — | — | +| | p99 Latency | 0.360 ms | 0.331 ms | -8.1% → | — | — | + +**Key Findings**: + +- **Lease/Eventual Read: ~3.6x throughput improvement vs v0.2.4** — ReadActor fast path (#392) delivers 261%/240% throughput gain and 95%+ latency reduction. On AWS, this is the headline win: zero-channel direct SM reads eliminate all Raft/channel contention. +- **Write performance stable vs v0.2.4** — HC Write +0.7%, SC Write +0.7%, within AWS noise tolerance. +- **Linearizable Read +5.7% vs v0.2.4** — observed benchmark delta; Linearizable Read path remains on the Raft loop and is unchanged by #392. +- **vs etcd 3.2.0** — Embedded mode outperforms etcd across all comparable metrics: SC Write **+7.6x**, HC Write **+152%**, Lin Read **+144%**, Eventual Read **+6.7x** throughput; avg latency **-86~99%** lower. +- **v0.2.5 vs v0.2.3 baseline** — v0.2.4 had Lease/Eventual regressions vs v0.2.3 (378K/395K). v0.2.5 with #392 ReadActor achieves **+3.3x** vs v0.2.3. + +--- + +### Standalone Mode: v0.2.5 vs v0.2.4 / etcd 3.2.0 + +| **Scenario** | **Metric** | **v0.2.4 (AWS)** | **v0.2.5 (AWS)** | **Δ vs v0.2.4** | **etcd 3.2.0²** | **Δ vs etcd** | +| ------------------- | ----------- | ---------------- | ---------------- | --------------- | --------------- | ------------- | +| Single Client Write | Throughput | 2,305 ops/s | 2,095 ops/s | -9.1% ⚠️ | 583 ops/s | **+3.6x** ✅ | +| | Avg Latency | 0.433 ms | 0.480 ms | +10.9% ⚠️ | 1.6 ms | **-70%** ✅ | +| | p99 Latency | 0.316 ms | 0.600 ms | +89.9% ⚠️ | — | — | +| High Conc. Write | Throughput | 59,687 ops/s | 53,338 ops/s | -10.6% ⚠️ | 44,341 ops/s | +20% → | +| | Avg Latency | 0.40 ms | 4.353 ms | +988% ⚠️ | 22.0 ms | **-80.2%** ✅ | +| | p99 Latency | 1.20 ms | 12.923 ms | +977% ⚠️ | — | — | +| Linearizable Read | Throughput | 77,907 ops/s | 80,145 ops/s | +2.9% → | 141,578 ops/s | -43% ⚠️ | +| | Avg Latency | 2.79 ms | 2.495 ms | -10.6% → | 5.5 ms | **-54.6%** ✅ | +| | p99 Latency | 5.93 ms | 5.913 ms | stable | — | — | +| Lease Read | Throughput | 76,144 ops/s | 86,466 ops/s | **+13.6%** ✅ | —³ | — | +| | Avg Latency | 2.76 ms | 2.311 ms | -16.3% ✅ | — | — | +| | p99 Latency | 5.76 ms | 5.957 ms | stable | — | — | +| Eventual Read | Throughput | 170,000 ops/s | 157,358 ops/s | -7.4% ⚠️ | 185,758 ops/s | -15% ⚠️ | +| | Avg Latency | 2.10 ms | 1.539 ms | -26.7% ✅ | 2.2 ms | -30% → | +| | p99 Latency | 9.76 ms | 6.674 ms | **-31.6%** ✅ | — | — | +| Hot-Key (10 keys) | Throughput | 90,609 ops/s | 92,711 ops/s | +2.3% → | —³ | — | +| | Avg Latency | 2.36 ms | 2.154 ms | -8.7% → | — | — | +| | p99 Latency | 5.60 ms | 5.922 ms | +5.8% → | — | — | + +**Key Findings**: + +- **Lease Read +13.6% vs v0.2.4** — ReadActor fast path benefit carries to Standalone mode (ReadActor is shared infrastructure). +- **Eventual Read: Mixed results vs v0.2.4** — Throughput -7.4%, but latency **-26.7%** avg / **-31.6%** p99, indicating improved tail behavior. +- **HC Write regression severe** — Latency spike 0.4ms → 4.4ms; ConnectionTimeout errors observed in raw rounds, suggesting network saturation or TCP pressure on c5.2xlarge under high concurrency. +- **SC Write regression** — -9.1% throughput vs v0.2.4, correlated with HC Write contention overhead. +- **vs etcd 3.2.0** — SC Write **+3.6x** and all latency metrics significantly better; Lin Read throughput below etcd (-43%) but latency -54.6% lower; HC Write and Eventual Read throughput regressions (vs v0.2.4) pull both metrics near or below etcd level. + +² etcd data sourced from [etcd official benchmark documentation](https://etcd.io/docs/v3.6/op-guide/performance/), tested on GCE infrastructure. Different cloud platform; results are for reference only. +³ etcd does not have an equivalent mode. + +--- + +## Key Changes Driving Results + +| Change | Impact | +| ---------------------------------- | ------------------------------------------------------------------------------------------------------ | +| ReadActor fast path (#392) | Lease/Eventual Read avg latency -94%+ vs v0.2.4; throughput +22.2%/+17.6%; bypasses Raft loop entirely | +| ReadLease.revoke() (#392) | Atomic lease invalidation on leader demotion; replaces `invalidate()` | +| Configurable ReadActor (#392) | `read_actor_channel_capacity` and `read_actor_max_drain` in `[raft]` | +| Fix: RocksDB LOCK on stop() (#392) | ReadActor is sole `Arc` holder; LOCK released before Raft shutdown | + +--- + +## ReadActor Parameter Tuning (Local, Embedded Mode) + +Four configurations tested to characterize `read_actor_channel_capacity` and `read_actor_max_drain` sensitivity: + +| Config | cap / drain / batch | Lease Read | Eventual Read | HC Write | vs v0.2.3 (705K/742K) | +| ------ | ---------------------- | ---------- | ------------- | --------- | --------------------- | +| C1 | 1024 / 1000 / 200 | ~790K | ~806K | ~232K | +12% / +9% | +| C2 | 1024 / 1000 / 300 | ~785K | ~803K | ~231K | +11% / +8% | +| C3 | 1024 / 2000 / 200 | ~756K | ~789K | ~227K | +7% / +6% | +| **C4** | **10240 / 2000 / 200** | **~810K** | **~851K** | **~232K** | **+15% / +15%** | + +**Key findings**: + +- `channel_capacity` is the dominant knob: 512→1024 yields +30%, 1024→10240 yields +3–6%. +- `max_drain` > `channel_capacity` has no effect: drain loop exits early once channel is empty. +- `max_batch_size = 300` does not improve reads and introduces write p99 instability. +- Rule of thumb: `read_actor_channel_capacity = 2× peak_concurrent_readers`. + +--- + +## Benchmark Configuration + +All Local Embedded results above were collected with the following configuration: + +```toml +[raft] +read_actor_channel_capacity = 10240 +read_actor_max_drain = 2000 + +[raft.persistence] +strategy = "MemFirst" +flush_policy = { Batch = { idle_flush_interval_ms = 1000 } } + +[raft.batching] +max_batch_size = 200 +``` diff --git a/benches/reports/v0.2.5/d-engine_comparison_v0.2.5.png b/benches/reports/v0.2.5/d-engine_comparison_v0.2.5.png new file mode 100644 index 00000000..48d6aed8 Binary files /dev/null and b/benches/reports/v0.2.5/d-engine_comparison_v0.2.5.png differ diff --git a/benches/reports/v0.2.5/d-engine_v0.2.5_vs_v0.2.4_embedded_mode.png b/benches/reports/v0.2.5/d-engine_v0.2.5_vs_v0.2.4_embedded_mode.png new file mode 100644 index 00000000..a1851eda Binary files /dev/null and b/benches/reports/v0.2.5/d-engine_v0.2.5_vs_v0.2.4_embedded_mode.png differ diff --git a/benches/reports/v0.2.5/d-engine_v0.2.5_vs_v0.2.4_standalone_mode.png b/benches/reports/v0.2.5/d-engine_v0.2.5_vs_v0.2.4_standalone_mode.png new file mode 100644 index 00000000..bd60e8a4 Binary files /dev/null and b/benches/reports/v0.2.5/d-engine_v0.2.5_vs_v0.2.4_standalone_mode.png differ diff --git a/config/base/raft.toml b/config/base/raft.toml index e47abad7..4374172e 100644 --- a/config/base/raft.toml +++ b/config/base/raft.toml @@ -2,6 +2,19 @@ # DEVELOPER SETTINGS — safe to adjust for your application # ============================================================================= +[raft.read_actor] +# ReadActor fast path — serves Eventual and LeaseRead requests without entering the Raft loop. + +# mpsc channel buffer between the client layer and ReadActor. +# Rule of thumb: set to at least 2× your peak concurrent Eventual/LeaseRead clients. +# Default: 512. +channel_capacity = 512 + +# Max reads drained per ReadActor wakeup. +# Setting above channel_capacity has no effect (channel empties first). +# Default: 100. +max_drain = 100 + [raft.batching] # Maximum number of commands to accumulate in a single batch during drain operations. # A single value covers all drain loops (cmd_rx, role_rx, new_commit_rx) because diff --git a/d-engine-core/src/config/raft.rs b/d-engine-core/src/config/raft.rs index 4dce9b0b..03a826c5 100644 --- a/d-engine-core/src/config/raft.rs +++ b/d-engine-core/src/config/raft.rs @@ -87,6 +87,10 @@ pub struct RaftConfig { #[serde(default = "default_ordered_channel_capacity")] pub ordered_channel_capacity: usize, + /// ReadActor configuration — tuning for the dedicated Eventual/LeaseRead fast path. + #[serde(default)] + pub read_actor: ReadActorConfig, + /// Configuration settings for new node auto join feature #[serde(default)] pub auto_join: AutoJoinConfig, @@ -145,6 +149,7 @@ impl Default for RaftConfig { snapshot_rpc_timeout_ms: default_snapshot_rpc_timeout_ms(), cmd_channel_capacity: default_cmd_channel_capacity(), ordered_channel_capacity: default_ordered_channel_capacity(), + read_actor: ReadActorConfig::default(), read_consistency: ReadConsistencyConfig::default(), backpressure: BackpressureConfig::default(), rpc_compression: RpcCompressionConfig::default(), @@ -174,24 +179,11 @@ impl RaftConfig { self.membership.validate()?; self.state_machine.validate()?; self.snapshot.validate()?; - self.read_consistency.validate()?; + self.read_consistency.validate(self.election.election_timeout_min)?; + self.read_actor.validate()?; self.watch.validate()?; self.persistence.validate()?; - // Enforce lease_duration < election_timeout (load-bearing safety constraint). - // The lease fast path for linearizable reads relies on the invariant that - // "lease valid" and "new leader elected" are mutually exclusive on the timeline. - // This holds only when: lease_duration < election_timeout - max_clock_drift. - // At minimum, lease_duration must be strictly less than election_timeout_min. - if self.read_consistency.lease_duration_ms >= self.election.election_timeout_min { - return Err(Error::Config(ConfigError::Message(format!( - "read_consistency.lease_duration_ms ({}) must be strictly less than \ - election_timeout_min ({}ms) — required for lease-based linearizable reads \ - to be safe under partition", - self.read_consistency.lease_duration_ms, self.election.election_timeout_min, - )))); - } - Ok(()) } } @@ -221,6 +213,67 @@ fn default_ordered_channel_capacity() -> usize { 1024 } +/// Configuration for the ReadActor — the dedicated read task that serves +/// Eventual and LeaseRead requests without entering the Raft loop. +/// +/// Exposed as `[raft.read_actor]` in TOML configuration. +/// +/// # Tuning Guidelines +/// - `channel_capacity`: set to at least 2× peak concurrent Eventual/LeaseRead clients +/// - `max_drain`: rarely needs changing; must not exceed `channel_capacity` meaningfully +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct ReadActorConfig { + /// mpsc channel buffer between the client layer and ReadActor. + /// Larger values reduce backpressure under high Eventual/LeaseRead concurrency. + /// Default: 512. + #[serde(default = "default_read_actor_channel_capacity")] + pub channel_capacity: usize, + + /// Max reads drained per ReadActor wakeup (mirrors Raft::drain_client_cmds). + /// After the first recv().await fires, the actor drains up to this many + /// additional commands with try_recv() before yielding. + /// Default: 100. + #[serde(default = "default_read_actor_max_drain")] + pub max_drain: usize, +} + +impl Default for ReadActorConfig { + fn default() -> Self { + Self { + channel_capacity: default_read_actor_channel_capacity(), + max_drain: default_read_actor_max_drain(), + } + } +} + +impl ReadActorConfig { + pub(crate) fn validate(&self) -> Result<()> { + if self.channel_capacity == 0 { + return Err(Error::Config(ConfigError::Message( + "read_actor.channel_capacity must be at least 1 \ + (0 causes mpsc::channel to panic at startup)" + .into(), + ))); + } + if self.max_drain == 0 { + return Err(Error::Config(ConfigError::Message( + "read_actor.max_drain must be at least 1 \ + (0 disables post-wakeup batching — ReadActor drains no additional commands)" + .into(), + ))); + } + Ok(()) + } +} + +fn default_read_actor_channel_capacity() -> usize { + 512 +} + +fn default_read_actor_max_drain() -> usize { + 100 +} + #[derive(Debug, Serialize, Deserialize, Clone)] pub struct ReplicationConfig { /// Heartbeat interval (milliseconds): how often the leader sends AppendEntries RPCs. @@ -991,12 +1044,21 @@ pub struct ReadConsistencyConfig { #[serde(default)] pub default_policy: ReadConsistencyPolicy, - /// Lease duration in milliseconds for LeaseRead policy + /// Lease duration in milliseconds for LeaseRead policy. + /// + /// How long the leader treats its lease as valid after the last quorum ACK. /// - /// Only applicable when using the LeaseRead policy. The leader considers - /// itself valid for this duration after successfully heartbeating to a quorum. + /// **Safety constraint** (enforced by `RaftConfig::validate()`): + /// `lease_duration_ms + network_rtt_p99_ms / 2 < election_timeout_min` /// - /// **MUST be > 0**. Config validation will reject 0 or invalid values. + /// Full theoretical bound (Raft §6.4): + /// `lease_duration_ms < election_timeout_min - rtt_p99/2 - max_clock_drift` + /// + /// The lease deadline is anchored to heartbeat *send* time, but the follower's + /// election timer resets at *receive* time (~rtt/2 later). Without accounting for + /// rtt/2 the safety invariant is violated under network latency — confirmed by + /// Jepsen set workload (19 elements lost). etcd absorbs this via `electionTimeout × 2/3`; + /// d-engine makes it explicit via `network_rtt_p99_ms`. #[serde(default = "default_lease_duration_ms")] pub lease_duration_ms: u64, @@ -1014,6 +1076,26 @@ pub struct ReadConsistencyConfig { /// Default: 10ms (safe buffer for single-node local deployments) #[serde(default = "default_state_machine_sync_timeout_ms")] pub state_machine_sync_timeout_ms: u64, + + /// Estimated p99 round-trip network latency between leader and followers, in milliseconds. + /// + /// Used to tighten the lease safety constraint beyond the basic + /// `lease_duration_ms < election_timeout_min` check. + /// + /// **Why this matters**: the lease deadline is anchored to the heartbeat *send* time, + /// but a follower's election timer resets at heartbeat *receive* time (~RTT/2 later). + /// The precise safety condition is: + /// `lease_duration_ms + network_rtt_p99_ms / 2 < election_timeout_min` + /// + /// Typical values: + /// - Same host / loopback: 0–1 ms + /// - Same datacenter: 1–2 ms + /// - Cross-AZ (AWS): 1–3 ms + /// - Cross-region: 50+ ms (LeaseRead not recommended) + /// + /// Default: 2ms (safe for same-datacenter and typical cross-AZ deployments). + #[serde(default = "default_network_rtt_p99_ms")] + pub network_rtt_p99_ms: u64, } impl Default for ReadConsistencyConfig { @@ -1023,12 +1105,13 @@ impl Default for ReadConsistencyConfig { lease_duration_ms: default_lease_duration_ms(), allow_client_override: default_allow_client_override(), state_machine_sync_timeout_ms: default_state_machine_sync_timeout_ms(), + network_rtt_p99_ms: default_network_rtt_p99_ms(), } } } fn default_lease_duration_ms() -> u64 { - // Conservative default: half of a typical heartbeat interval (~300ms) + // 2.5× the default heartbeat interval (100ms); safely below election_timeout_min (500ms). 250 } @@ -1037,18 +1120,43 @@ fn default_allow_client_override() -> bool { true } +fn default_network_rtt_p99_ms() -> u64 { + // 2ms covers same-datacenter and typical AWS cross-AZ deployments. + // Cross-region deployments should increase this value and reconsider using LeaseRead. + 2 +} + fn default_state_machine_sync_timeout_ms() -> u64 { 10 // 10ms is safe for typical <1ms apply latency on local SSD } impl ReadConsistencyConfig { - fn validate(&self) -> Result<()> { - // Validate read consistency configuration + pub(super) fn validate( + &self, + election_timeout_min: u64, + ) -> Result<()> { if self.lease_duration_ms == 0 { return Err(Error::Config(ConfigError::Message( "read_consistency.lease_duration_ms must be greater than 0".into(), ))); } + // Safety constraint (Raft §6.4): + // lease_duration_ms + rtt_p99/2 < election_timeout_min + // + // The follower's election timer resets at heartbeat receive time (~rtt/2 after send). + // Without this margin the effective lease window can exceed election_timeout_min, + // allowing stale reads after a network partition that heals before lease expiry. + let rtt_half_ms = self.network_rtt_p99_ms / 2; + // Use saturating_add: if the sum overflows u64 it saturates to u64::MAX, + // which is guaranteed >= any election_timeout_min, so the config is correctly rejected. + if self.lease_duration_ms.saturating_add(rtt_half_ms) >= election_timeout_min { + return Err(Error::Config(ConfigError::Message(format!( + "read_consistency.lease_duration_ms ({}) + network_rtt_p99_ms/2 ({}) \ + must be strictly less than election_timeout_min ({}ms) — \ + required for lease safety under partition (see network_rtt_p99_ms config)", + self.lease_duration_ms, rtt_half_ms, election_timeout_min, + )))); + } Ok(()) } } diff --git a/d-engine-core/src/config/raft_test.rs b/d-engine-core/src/config/raft_test.rs index 9a9bb2c3..f89825db 100644 --- a/d-engine-core/src/config/raft_test.rs +++ b/d-engine-core/src/config/raft_test.rs @@ -3,6 +3,7 @@ use crate::ElectionConfig; use crate::RaftConfig; use crate::ReadConsistencyConfig; use crate::RpcCompressionConfig; +use crate::config::raft::ReadActorConfig; #[test] fn test_invalid_election_timeout() { let mut config = RaftConfig::default(); @@ -232,3 +233,147 @@ fn test_config_rejects_lease_duration_not_less_than_election_timeout() { "lease_duration_ms < election_timeout_min must be accepted" ); } + +/// Direct unit test for ReadConsistencyConfig::validate() — lease safety constraint. +/// +/// This pins the Raft §6.4 invariant at the config layer: +/// lease_duration_ms < election_timeout_min +/// +/// Tested directly on ReadConsistencyConfig (not via RaftConfig) so the constraint +/// is verified at the closest possible layer to the data. +#[test] +fn test_read_consistency_config_lease_duration_rejects_unsafe_values() { + let election_timeout_min = 500u64; + + // equal → unsafe, must reject + let cfg = ReadConsistencyConfig { + lease_duration_ms: 500, + ..Default::default() + }; + assert!( + cfg.validate(election_timeout_min).is_err(), + "lease_duration_ms == election_timeout_min violates safety invariant" + ); + + // greater → unsafe, must reject + let cfg = ReadConsistencyConfig { + lease_duration_ms: 600, + ..Default::default() + }; + assert!( + cfg.validate(election_timeout_min).is_err(), + "lease_duration_ms > election_timeout_min violates safety invariant" + ); + + // strictly less → safe, must accept + let cfg = ReadConsistencyConfig { + lease_duration_ms: 250, + ..Default::default() + }; + assert!( + cfg.validate(election_timeout_min).is_ok(), + "lease_duration_ms < election_timeout_min must be accepted" + ); + + // zero → must reject (separate guard) + let cfg = ReadConsistencyConfig { + lease_duration_ms: 0, + ..Default::default() + }; + assert!( + cfg.validate(election_timeout_min).is_err(), + "lease_duration_ms = 0 must be rejected" + ); +} + +#[test] +fn test_read_actor_channel_capacity_zero_is_invalid() { + let cfg = ReadActorConfig { + channel_capacity: 0, + max_drain: 100, + }; + assert!( + cfg.validate().is_err(), + "channel_capacity = 0 must be rejected (mpsc::channel(0) panics at startup)" + ); +} + +#[test] +fn test_read_actor_max_drain_zero_is_invalid() { + let cfg = ReadActorConfig { + channel_capacity: 512, + max_drain: 0, + }; + assert!( + cfg.validate().is_err(), + "max_drain = 0 must be rejected (drains 0 commands per wakeup — disables batching)" + ); +} + +#[test] +fn test_read_actor_default_config_is_valid() { + assert!( + ReadActorConfig::default().validate().is_ok(), + "default ReadActorConfig must be valid" + ); +} + +#[test] +fn test_raft_config_propagates_read_actor_validation() { + let mut config = RaftConfig::default(); + config.read_actor.channel_capacity = 0; + assert!( + config.validate().is_err(), + "RaftConfig::validate must propagate ReadActorConfig::validate failures" + ); +} + +/// lease_duration_ms + network_rtt_p99_ms/2 must account for RTT/2 in the safety bound. +/// +/// Invariant (Raft §6.4): lease_duration_ms + rtt_p99_ms/2 < election_timeout_min +/// Even if lease_duration_ms alone is safe, adding rtt/2 can push it over the bound. +#[test] +fn test_read_consistency_config_rtt_p99_tightens_lease_safety_bound() { + let election_timeout_min = 500u64; + + // lease=498, rtt=4 → rtt/2=2 → 498+2=500 >= 500 → must reject + let cfg = ReadConsistencyConfig { + lease_duration_ms: 498, + network_rtt_p99_ms: 4, + ..Default::default() + }; + assert!( + cfg.validate(election_timeout_min).is_err(), + "lease + rtt/2 == election_timeout_min must be rejected" + ); + + // lease=497, rtt=4 → rtt/2=2 → 497+2=499 < 500 → must accept + let cfg = ReadConsistencyConfig { + lease_duration_ms: 497, + network_rtt_p99_ms: 4, + ..Default::default() + }; + assert!( + cfg.validate(election_timeout_min).is_ok(), + "lease + rtt/2 < election_timeout_min must be accepted" + ); +} + +/// Overflow-safe: u64::MAX lease_duration_ms must be rejected, not wrap around. +/// +/// Without saturating_add, u64::MAX + rtt_half wraps to a small value in release +/// builds, allowing an invalid config to pass validate() and violate lease safety. +#[test] +fn test_read_consistency_config_lease_duration_overflow_is_rejected() { + let election_timeout_min = 500u64; + + let cfg = ReadConsistencyConfig { + lease_duration_ms: u64::MAX, + network_rtt_p99_ms: 4, + ..Default::default() + }; + assert!( + cfg.validate(election_timeout_min).is_err(), + "u64::MAX lease_duration_ms must be rejected (saturating_add prevents wrap-around)" + ); +} diff --git a/d-engine-core/src/raft.rs b/d-engine-core/src/raft.rs index 42297fe0..6fd1b1f1 100644 --- a/d-engine-core/src/raft.rs +++ b/d-engine-core/src/raft.rs @@ -714,6 +714,14 @@ where self.cmd_tx.clone() } + pub fn read_lease(&self) -> Arc { + Arc::clone(&self.role.state().shared_state().lease) + } + + pub fn current_term(&self) -> u64 { + self.role.state().current_term() + } + /// Returns a cloned role event sender for internal use. /// /// # Warning diff --git a/d-engine-core/src/raft_role/leader_state.rs b/d-engine-core/src/raft_role/leader_state.rs index 5b9d9267..10e7ec90 100644 --- a/d-engine-core/src/raft_role/leader_state.rs +++ b/d-engine-core/src/raft_role/leader_state.rs @@ -5,6 +5,7 @@ use super::StateSnapshot; use super::buffers::BatchBuffer; use super::buffers::ProposeBatchBuffer; use super::candidate_state::CandidateState; +use super::read_lease::now_ms; use super::role_state::RaftRoleState; use super::role_state::check_and_trigger_snapshot; use super::role_state::send_replay_raft_event; @@ -73,7 +74,6 @@ use std::collections::VecDeque; use std::fmt::Debug; use std::marker::PhantomData; use std::sync::Arc; -use std::sync::Mutex; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU32; use std::sync::atomic::Ordering; @@ -405,6 +405,17 @@ pub struct LeaderState { /// Last time we checked for learners pub last_learner_check: Instant, + /// Timestamp captured just before the first AppendEntries RPC is sent each round. + /// Used as the deadline base in update_lease_timestamp so the lease expires at + /// `send_ts + lease_duration_ms` rather than `ack_ts + lease_duration_ms`, + /// eliminating the RTT/2 window that allowed stale reads after partition. + /// + /// Known limitation: single shared value overwritten on every execute_and_process_raft_rpc + /// call. In rare cases where tick() and drain_client_cmds() both fire in the same loop + /// iteration, an old ACK may use a newer send_ts (microsecond-level error, negligible + /// relative to lease_duration_ms). Per-round correctness requires Method B (token-based). + pub(crate) last_heartbeat_send_ts: u64, + // -- Stale Learner Handling -- /// Earliest time the oldest pending learner becomes stale. /// @@ -420,11 +431,6 @@ pub struct LeaderState { /// Queue of learners that have caught up and are pending promotion to voter. pub pending_promotions: VecDeque, - /// Monotonic instant of the last quorum ACK, used to check lease validity. - /// `None` until the first ACK; replaced on every subsequent quorum confirmation. - /// Must be `Instant` (not `SystemTime`) — wall clock can move backward under NTP. - pub(super) lease_instant: Mutex>, - // -- Cluster Topology Cache -- /// Cached cluster metadata (updated on membership changes) /// Avoids repeated async calls in hot paths @@ -702,6 +708,8 @@ impl RaftRoleState for LeaderState { self.node_id(), self.current_term() ); + // Revoke lease so ReadActor immediately returns LeaseInvalid on the next read. + self.shared_state.lease.revoke(); Ok(RaftRole::Follower(Box::new(self.into()))) } @@ -1119,6 +1127,9 @@ impl RaftRoleState for LeaderState { let my_term = self.current_term(); if my_term < vote_request.term { self.update_current_term(vote_request.term); + // Revoke lease immediately — before the Raft loop processes BecomeFollower, + // a concurrent ReadActor could still see the old valid lease (window-period bug). + self.shared_state.lease.revoke(); // Step down as Follower self.send_become_follower_event(None, &role_tx)?; @@ -1222,6 +1233,8 @@ impl RaftRoleState for LeaderState { my_id ); //TODO: if there is a bug? self.update_current_term(vote_request.term); + // Revoke lease immediately — window-period fix (see VoteRequest branch). + self.shared_state.lease.revoke(); self.send_become_follower_event( Some(append_entries_request.leader_id), &role_tx, @@ -1705,7 +1718,11 @@ impl RaftRoleState for LeaderState { // Refresh lease timestamp after all post-commit work to eliminate the yield-point // race window introduced when drain_commit_actions became async. if self.cluster_metadata.single_voter { - self.update_lease_timestamp(); + // Single-voter: self is the entire quorum, no peer RTT to account for. + self.update_lease_timestamp( + now_ms(), + ctx.node_config().raft.read_consistency.lease_duration_ms, + ); self.drain_pending_lease_reads(ctx); } } @@ -1767,6 +1784,8 @@ impl RaftRoleState for LeaderState { ); self.update_current_term(response.term); self.drain_pending_writes_with_error(ErrorCode::TermOutdated); + // Revoke lease immediately — window-period fix (see VoteRequest branch). + self.shared_state.lease.revoke(); self.send_become_follower_event(None, role_tx)?; return Err(ReplicationError::HigherTerm(response.term).into()); } @@ -1794,6 +1813,8 @@ impl RaftRoleState for LeaderState { if term > leader_term { self.update_current_term(term); self.drain_pending_writes_with_error(ErrorCode::TermOutdated); + // Revoke lease immediately — window-period fix (see VoteRequest branch). + self.shared_state.lease.revoke(); self.send_become_follower_event(None, role_tx)?; return Err(ReplicationError::HigherTerm(term).into()); } @@ -1880,9 +1901,17 @@ impl RaftRoleState for LeaderState { ) .is_some(); if quorum_confirmed { - // Refresh after all post-commit work to eliminate the yield-point race window - // introduced when drain_commit_actions became async. - self.update_lease_timestamp(); + // Anchor deadline to send time (not ACK time) to eliminate the RTT/2 window. + // Falls back to now_ms() only in tests that bypass execute_and_process_raft_rpc. + let send_ts = if self.last_heartbeat_send_ts > 0 { + self.last_heartbeat_send_ts + } else { + now_ms() + }; + self.update_lease_timestamp( + send_ts, + ctx.node_config().raft.read_consistency.lease_duration_ms, + ); self.drain_pending_lease_reads(ctx); // Path A drain (Bug #381 fix): serve linearizable reads that have been // waiting for quorum confirmation. Pure-read batches never advance @@ -3115,10 +3144,10 @@ impl LeaderState { scheduled_purge_upto: None, last_purged_index: None, //TODO last_learner_check: Instant::now(), + last_heartbeat_send_ts: 0, snapshot_in_progress: AtomicBool::new(false), stale_check_deadline: None, pending_promotions: VecDeque::new(), - lease_instant: Mutex::new(None), linearizable_read_buffer: Box::new(BatchBuffer::new(batch_size).with_length_gauge( node_id, "linearizable", @@ -3289,7 +3318,7 @@ impl LeaderState { // Step 6: Reschedule if any pending promotions remain if !self.pending_promotions.is_empty() { debug!( - "[Leader {}] 🔁 Re-sending PromoteReadyLearners for remaining pending: {:?}", + "[Leader {}] Re-sending PromoteReadyLearners for remaining pending: {:?}", self.node_id(), self.pending_promotions.iter().map(|p| p.node_id).collect::>() ); @@ -3383,6 +3412,11 @@ impl LeaderState { payload_count = payloads.len(), ); + // Phase 0: Record send timestamp before any AppendEntries RPC goes out. + // handle_append_result uses this to anchor deadline to send time (not ACK time), + // ensuring lease expires at send_ts + lease_duration_ms rather than ack_ts + lease_duration_ms. + self.last_heartbeat_send_ts = now_ms(); + // Phase 1: write entries to local log + prepare per-peer requests (serial in Raft loop). let requests = match ctx .replication_handler() @@ -3467,7 +3501,7 @@ impl LeaderState { if let Some(read_batch) = read_batch { let read_index = self.calculate_read_index(); let last_applied = ctx.state_machine().last_applied().index; - if (self.cluster_metadata.single_voter || self.is_lease_valid(ctx)) + if (self.cluster_metadata.single_voter || self.is_lease_valid()) && last_applied >= read_index { self.execute_pending_reads(read_batch, ctx); @@ -3712,22 +3746,25 @@ impl LeaderState { /// Returns true if the leader's lease is still within its validity window. /// /// Uses a monotonic `Instant` so NTP clock adjustments cannot extend the lease. - /// Returns false if no quorum ACK has been received yet (lease_instant is None). - pub fn is_lease_valid( - &self, - ctx: &RaftContext, - ) -> bool { - let lease_duration_ms = ctx.node_config().raft.read_consistency.lease_duration_ms; - match *self.lease_instant.lock().expect("lease_instant poisoned") { - Some(t) => (t.elapsed().as_millis() as u64) < lease_duration_ms, - None => false, - } + /// Returns true iff this leader holds a valid quorum-backed lease. + /// Uses a single atomic load (~4 ns) — no mutex, no lock contention. + pub fn is_lease_valid(&self) -> bool { + self.shared_state.lease.is_valid_for_leader(self.current_term(), now_ms()) } - /// Records the current monotonic instant as the last quorum-confirmed timestamp. - /// Called after every quorum ACK (handle_append_result, handle_log_flushed). - fn update_lease_timestamp(&self) { - *self.lease_instant.lock().expect("lease_instant poisoned") = Some(Instant::now()); + /// Renews the lease after every quorum ACK (handle_append_result, handle_log_flushed). + /// + /// `send_ts` is the timestamp recorded just before the heartbeat was sent. + /// Using send_ts (not now_ms()) eliminates the RTT/2 window: + /// deadline = send_ts + lease_duration_ms ≤ follower's election timer start + election_timeout + /// Single-voter callers pass now_ms() directly (no peer RTT to account for). + fn update_lease_timestamp( + &self, + send_ts: u64, + lease_duration_ms: u64, + ) { + let deadline = send_ts.saturating_add(lease_duration_ms); + self.shared_state.lease.renew(self.current_term(), deadline); } /// Unified write + linearizable read: single RPC for both (2*RTT → 1*RTT). @@ -3805,7 +3842,18 @@ impl LeaderState { #[cfg(test)] pub(crate) fn test_update_lease_timestamp(&mut self) { - self.update_lease_timestamp(); + // Renew with a generous 60-second deadline so tests don't flake on slow CI. + let deadline = now_ms().saturating_add(60_000); + self.shared_state.lease.renew(self.current_term(), deadline); + } + + /// Test helper: renew lease from an explicit send_ts (RTT/2 fix path). + #[cfg(test)] + pub(crate) fn test_renew_lease_from_send_ts( + &mut self, + lease_duration_ms: u64, + ) { + self.update_lease_timestamp(self.last_heartbeat_send_ts, lease_duration_ms); } /// Returns the current match_index for `peer_id`, or 0 if not yet tracked. @@ -3886,7 +3934,7 @@ impl LeaderState { ctx: &RaftContext, role_tx: &mpsc::UnboundedSender, ) -> Result<()> { - if self.is_lease_valid(ctx) { + if self.is_lease_valid() { // Lease valid - serve immediately let results = ctx .handlers @@ -3898,8 +3946,11 @@ impl LeaderState { } else { // Lease expired - need to confirm leadership before serving. if self.cluster_metadata.single_voter { - // Single-voter: self is the entire quorum, refresh immediately and serve. - self.update_lease_timestamp(); + // Single-voter: self is the entire quorum, no peer RTT to account for. + self.update_lease_timestamp( + now_ms(), + ctx.node_config().raft.read_consistency.lease_duration_ms, + ); let results = ctx .handlers .state_machine_handler @@ -3987,6 +4038,7 @@ impl From<&CandidateState> for LeaderState { scheduled_purge_upto: None, last_purged_index: candidate.last_purged_index, last_learner_check: Instant::now(), + last_heartbeat_send_ts: 0, snapshot_in_progress: AtomicBool::new(false), stale_check_deadline: None, pending_promotions: VecDeque::new(), @@ -3995,7 +4047,6 @@ impl From<&CandidateState> for LeaderState { total_voters: 0, replication_targets: vec![], }, - lease_instant: Mutex::new(None), linearizable_read_buffer: Box::new( BatchBuffer::new(candidate.node_config.raft.batching.max_batch_size) .with_length_gauge( diff --git a/d-engine-core/src/raft_role/leader_state_test/become_follower_test.rs b/d-engine-core/src/raft_role/leader_state_test/become_follower_test.rs new file mode 100644 index 00000000..96c2b42e --- /dev/null +++ b/d-engine-core/src/raft_role/leader_state_test/become_follower_test.rs @@ -0,0 +1,354 @@ +//! Tests verifying that `become_follower()` revokes the read lease. +//! +//! Safety invariant: a stepped-down leader must not allow ReadActor to serve +//! stale lease reads. The shared `Arc` must be invalid immediately +//! after `become_follower()` returns so any concurrent ReadActor check fails. +//! +//! # Coverage +//! - Direct call: `become_follower()` revokes lease (unit pin) +//! - End-to-end: higher-term VoteRequest → BecomeFollower → become_follower() → lease invalid +//! - End-to-end: higher-term AppendEntries → BecomeFollower → become_follower() → lease invalid +//! - End-to-end: higher-term AppendResult → BecomeFollower → become_follower() → lease invalid + +use std::sync::Arc; + +use tokio::sync::{mpsc, watch}; + +use crate::RaftNodeConfig; +use crate::event::{RaftEvent, RoleEvent}; +use crate::maybe_clone_oneshot::{MaybeCloneOneshot, RaftOneshot}; +use crate::now_ms; +use crate::raft_role::leader_state::LeaderState; +use crate::raft_role::role_state::RaftRoleState; +use crate::test_utils::MockBuilder; +use crate::test_utils::mock::MockTypeConfig; +use d_engine_proto::server::election::VoteRequest; +use d_engine_proto::server::replication::{AppendEntriesRequest, AppendEntriesResponse}; + +/// become_follower() revokes the shared Arc so ReadActor immediately +/// returns LeaseInvalid on the very next is_valid() check. +/// +/// This pins the safety contract: after step-down the ReadActor fast path is +/// blocked regardless of when (before or after) its thread checks the lease. +#[tokio::test] +async fn test_become_follower_revokes_read_lease() { + let mut state = LeaderState::::new(1, RaftNodeConfig::default().into()); + + // Clone Arc before step-down to observe the shared lease from the outside + // (simulates what ReadActor holds). + let lease = Arc::clone(&state.shared_state.lease); + + // Renew with a generous 60-second deadline. + state.test_update_lease_timestamp(); + assert!( + state.is_lease_valid(), + "precondition: lease must be valid before become_follower()" + ); + assert!( + lease.is_valid(now_ms()), + "precondition: same Arc must also report valid" + ); + + // Transition to follower — must call lease.revoke() atomically. + let _ = state.become_follower().expect("become_follower must succeed"); + + // The shared Arc must now be invalid. + assert!( + !lease.is_valid(now_ms()), + "become_follower() must revoke the shared Arc" + ); +} + +/// Higher-term VoteRequest → BecomeFollower event → become_follower() → lease revoked. +/// +/// End-to-end pin for the path: +/// ReceiveVoteRequest(term > current) → send_become_follower_event() +/// → Raft loop calls become_follower() → lease.revoke() +/// +/// The ReadActor holds the same Arc. After this chain the fast path +/// must return LeaseInvalid on the very next is_valid() check. +#[tokio::test] +async fn test_receive_higher_term_vote_request_revokes_lease() { + let (_graceful_tx, graceful_rx) = watch::channel(()); + let context = MockBuilder::new(graceful_rx) + .with_db_path("/tmp/test_vote_request_revokes_lease") + .build_context(); + + let mut state = LeaderState::::new(1, context.node_config.clone()); + let lease = Arc::clone(&state.shared_state.lease); + + state.test_update_lease_timestamp(); + assert!( + lease.is_valid(now_ms()), + "precondition: lease must be valid" + ); + + let (resp_tx, _resp_rx) = >::new(); + let (role_tx, mut role_rx) = mpsc::unbounded_channel(); + let event = RaftEvent::ReceiveVoteRequest( + VoteRequest { + term: 999, + candidate_id: 2, + last_log_index: 0, + last_log_term: 0, + }, + resp_tx, + ); + state.handle_raft_event(event, &context, role_tx).await.ok(); + + assert!( + matches!(role_rx.try_recv(), Ok(RoleEvent::BecomeFollower(_))), + "higher-term VoteRequest must emit BecomeFollower" + ); + + // Simulate Raft loop processing BecomeFollower. + let _ = state.become_follower().expect("become_follower must succeed"); + + assert!( + !lease.is_valid(now_ms()), + "lease must be revoked after VoteRequest step-down — ReadActor must not serve stale reads" + ); +} + +/// Higher-term AppendEntries → BecomeFollower event → become_follower() → lease revoked. +/// +/// End-to-end pin for the path: +/// AppendEntries(term > current) → send_become_follower_event() +/// → Raft loop calls become_follower() → lease.revoke() +#[tokio::test] +async fn test_receive_higher_term_append_entries_revokes_lease() { + let (_graceful_tx, graceful_rx) = watch::channel(()); + let context = MockBuilder::new(graceful_rx) + .with_db_path("/tmp/test_append_entries_revokes_lease") + .build_context(); + + let mut state = LeaderState::::new(1, context.node_config.clone()); + state.update_current_term(10); + let lease = Arc::clone(&state.shared_state.lease); + + state.test_update_lease_timestamp(); + assert!( + lease.is_valid(now_ms()), + "precondition: lease must be valid" + ); + + let (resp_tx, _resp_rx) = >::new(); + let (role_tx, mut role_rx) = mpsc::unbounded_channel(); + let event = RaftEvent::AppendEntries( + AppendEntriesRequest { + term: 11, // higher than current term 10 + leader_id: 2, + prev_log_index: 0, + prev_log_term: 0, + entries: vec![], + leader_commit_index: 0, + }, + resp_tx, + ); + state.handle_raft_event(event, &context, role_tx).await.ok(); + + assert!( + matches!(role_rx.try_recv(), Ok(RoleEvent::BecomeFollower(_))), + "higher-term AppendEntries must emit BecomeFollower" + ); + + // Simulate Raft loop processing BecomeFollower. + let _ = state.become_follower().expect("become_follower must succeed"); + + assert!( + !lease.is_valid(now_ms()), + "lease must be revoked after AppendEntries step-down" + ); +} + +/// Higher-term VoteRequest → lease revoked BEFORE Raft loop calls become_follower(). +/// +/// Pins the window-period safety contract: the Raft loop sends BecomeFollower into +/// role_tx but may not process it immediately (async event-driven). During that gap a +/// concurrent ReadActor task can observe the old valid lease and serve a stale LeaseRead. +/// +/// This test verifies the fix: revoke() is called at the detection point itself, so +/// the lease is invalid as soon as handle_raft_event() returns — before become_follower() +/// is ever called by the Raft loop. +/// +/// Expected to FAIL before the fix (lease still valid in window), PASS after. +#[tokio::test] +async fn test_vote_request_higher_term_lease_revoked_before_become_follower() { + let (_graceful_tx, graceful_rx) = watch::channel(()); + let context = MockBuilder::new(graceful_rx) + .with_db_path("/tmp/test_vote_request_revokes_lease_early") + .build_context(); + + let mut state = LeaderState::::new(1, context.node_config.clone()); + let lease = Arc::clone(&state.shared_state.lease); + + state.test_update_lease_timestamp(); + assert!( + lease.is_valid(now_ms()), + "precondition: lease must be valid" + ); + + let (resp_tx, _resp_rx) = >::new(); + let (role_tx, mut role_rx) = mpsc::unbounded_channel(); + let event = RaftEvent::ReceiveVoteRequest( + VoteRequest { + term: 999, + candidate_id: 2, + last_log_index: 0, + last_log_term: 0, + }, + resp_tx, + ); + state.handle_raft_event(event, &context, role_tx).await.ok(); + + assert!( + matches!(role_rx.try_recv(), Ok(RoleEvent::BecomeFollower(_))), + "higher-term VoteRequest must emit BecomeFollower" + ); + + // become_follower() is intentionally NOT called here — the Raft loop has not yet + // processed the BecomeFollower event. The lease must already be invalid. + assert!( + !lease.is_valid(now_ms()), + "lease must be revoked at detection point, not waiting for become_follower() \ + — window-period fix required in handle_raft_event VoteRequest branch" + ); +} + +/// Higher-term AppendEntries → lease revoked BEFORE Raft loop calls become_follower(). +/// +/// Same window-period contract as test_vote_request_higher_term_lease_revoked_before_become_follower, +/// for the AppendEntries detection path. +#[tokio::test] +async fn test_append_entries_higher_term_lease_revoked_before_become_follower() { + let (_graceful_tx, graceful_rx) = watch::channel(()); + let context = MockBuilder::new(graceful_rx) + .with_db_path("/tmp/test_append_entries_revokes_lease_early") + .build_context(); + + let mut state = LeaderState::::new(1, context.node_config.clone()); + state.update_current_term(10); + let lease = Arc::clone(&state.shared_state.lease); + + state.test_update_lease_timestamp(); + assert!( + lease.is_valid(now_ms()), + "precondition: lease must be valid" + ); + + let (resp_tx, _resp_rx) = >::new(); + let (role_tx, mut role_rx) = mpsc::unbounded_channel(); + let event = RaftEvent::AppendEntries( + AppendEntriesRequest { + term: 11, + leader_id: 2, + prev_log_index: 0, + prev_log_term: 0, + entries: vec![], + leader_commit_index: 0, + }, + resp_tx, + ); + state.handle_raft_event(event, &context, role_tx).await.ok(); + + assert!( + matches!(role_rx.try_recv(), Ok(RoleEvent::BecomeFollower(_))), + "higher-term AppendEntries must emit BecomeFollower" + ); + + // become_follower() intentionally NOT called — pinning the window-period invariant. + assert!( + !lease.is_valid(now_ms()), + "lease must be revoked at detection point, not waiting for become_follower() \ + — window-period fix required in handle_raft_event AppendEntries branch" + ); +} + +/// Higher-term AppendResult → lease revoked BEFORE Raft loop calls become_follower(). +/// +/// Same window-period contract for the handle_append_result detection path. +#[tokio::test] +async fn test_append_result_higher_term_lease_revoked_before_become_follower() { + let (_graceful_tx, graceful_rx) = watch::channel(()); + let context = MockBuilder::new(graceful_rx) + .with_db_path("/tmp/test_append_result_revokes_lease_early") + .build_context(); + + let mut state = LeaderState::::new(1, context.node_config.clone()); + let lease = Arc::clone(&state.shared_state.lease); + + state.test_update_lease_timestamp(); + assert!( + lease.is_valid(now_ms()), + "precondition: lease must be valid" + ); + + let (role_tx, mut role_rx) = mpsc::unbounded_channel(); + let higher_term_response = AppendEntriesResponse { + node_id: 2, + term: 999, + result: None, + }; + let result = state + .handle_append_result(2, Ok(higher_term_response), &context, &role_tx) + .await; + + assert!(result.is_err(), "higher-term AppendResult must return Err"); + assert!( + matches!(role_rx.try_recv(), Ok(RoleEvent::BecomeFollower(_))), + "higher-term AppendResult must emit BecomeFollower" + ); + + // become_follower() intentionally NOT called — pinning the window-period invariant. + assert!( + !lease.is_valid(now_ms()), + "lease must be revoked at detection point, not waiting for become_follower() \ + — window-period fix required in handle_append_result HigherTerm branch" + ); +} + +/// Higher-term AppendResult → BecomeFollower event → become_follower() → lease revoked. +/// +/// End-to-end pin for the path: +/// handle_append_result(response.term > current) → send_become_follower_event() +/// → Raft loop calls become_follower() → lease.revoke() +#[tokio::test] +async fn test_append_result_higher_term_revokes_lease() { + let (_graceful_tx, graceful_rx) = watch::channel(()); + let context = MockBuilder::new(graceful_rx) + .with_db_path("/tmp/test_append_result_revokes_lease") + .build_context(); + + let mut state = LeaderState::::new(1, context.node_config.clone()); + let lease = Arc::clone(&state.shared_state.lease); + + state.test_update_lease_timestamp(); + assert!( + lease.is_valid(now_ms()), + "precondition: lease must be valid" + ); + + let (role_tx, mut role_rx) = mpsc::unbounded_channel(); + let higher_term_response = AppendEntriesResponse { + node_id: 2, + term: 999, // higher than leader term=1 + result: None, + }; + let result = state + .handle_append_result(2, Ok(higher_term_response), &context, &role_tx) + .await; + + assert!(result.is_err(), "higher-term AppendResult must return Err"); + assert!( + matches!(role_rx.try_recv(), Ok(RoleEvent::BecomeFollower(_))), + "higher-term AppendResult must emit BecomeFollower" + ); + + // Simulate Raft loop processing BecomeFollower. + let _ = state.become_follower().expect("become_follower must succeed"); + + assert!( + !lease.is_valid(now_ms()), + "lease must be revoked after AppendResult step-down" + ); +} diff --git a/d-engine-core/src/raft_role/leader_state_test/client_read_test.rs b/d-engine-core/src/raft_role/leader_state_test/client_read_test.rs index 67d972d1..f5b6621a 100644 --- a/d-engine-core/src/raft_role/leader_state_test/client_read_test.rs +++ b/d-engine-core/src/raft_role/leader_state_test/client_read_test.rs @@ -753,7 +753,7 @@ async fn test_expired_lease_detection() { // Don't update lease timestamp - lease should be expired by default assert!( - !state.is_lease_valid(&context), + !state.is_lease_valid(), "Lease should be expired when timestamp is never updated" ); } @@ -796,10 +796,7 @@ async fn test_expired_lease_single_voter_refreshed_immediately() { membership.expect_replication_peers().returning(Vec::new); state.init_cluster_metadata(&Arc::new(membership)).await.unwrap(); - assert!( - !state.is_lease_valid(&context), - "Precondition: lease expired" - ); + assert!(!state.is_lease_valid(), "Precondition: lease expired"); let client_read_request = ClientReadRequest { client_id: 1, @@ -815,7 +812,7 @@ async fn test_expired_lease_single_voter_refreshed_immediately() { // Lease must be refreshed after flush assert!( - state.is_lease_valid(&context), + state.is_lease_valid(), "Lease must be refreshed after single-voter flush" ); @@ -1187,10 +1184,7 @@ async fn test_lease_reuse_after_linearizable_read_refresh() { state.init_cluster_metadata(&Arc::new(membership)).await.unwrap(); // Verify: Initial lease is invalid - assert!( - !state.is_lease_valid(&ctx), - "Lease should be invalid initially" - ); + assert!(!state.is_lease_valid(), "Lease should be invalid initially"); let (role_tx, _role_rx) = mpsc::unbounded_channel(); @@ -1213,7 +1207,7 @@ async fn test_lease_reuse_after_linearizable_read_refresh() { state.handle_log_flushed(1, &ctx, &role_tx).await; assert!( - state.is_lease_valid(&ctx), + state.is_lease_valid(), "Lease should be valid after log flush (single-voter)" ); @@ -1282,7 +1276,7 @@ async fn test_eventual_consistency_ignores_stale_lease() { // Verify: Lease is invalid (simulates stale leader scenario) assert!( - !state.is_lease_valid(&ctx), + !state.is_lease_valid(), "Lease should be invalid (stale leader)" ); @@ -1306,7 +1300,7 @@ async fn test_eventual_consistency_ignores_stale_lease() { // Verify: Lease still invalid (not refreshed by EventualConsistency) assert!( - !state.is_lease_valid(&ctx), + !state.is_lease_valid(), "EventualConsistency should not refresh lease" ); @@ -2015,7 +2009,7 @@ async fn test_lease_read_empty_payload_verification_hangs_in_multi_node() { !state.cluster_metadata.single_voter, "precondition: multi-voter" ); - assert!(!state.is_lease_valid(&ctx), "precondition: lease expired"); + assert!(!state.is_lease_valid(), "precondition: lease expired"); let (role_tx, _role_rx) = mpsc::unbounded_channel(); let req = ClientReadRequest { diff --git a/d-engine-core/src/raft_role/leader_state_test/lease_refresh_on_log_flushed_test.rs b/d-engine-core/src/raft_role/leader_state_test/lease_refresh_on_log_flushed_test.rs index ac8faf6b..2adc6d43 100644 --- a/d-engine-core/src/raft_role/leader_state_test/lease_refresh_on_log_flushed_test.rs +++ b/d-engine-core/src/raft_role/leader_state_test/lease_refresh_on_log_flushed_test.rs @@ -155,7 +155,7 @@ async fn test_single_voter_lease_refreshed_on_log_flushed() { // Precondition: Lease is invalid initially assert!( - !state.is_lease_valid(&ctx), + !state.is_lease_valid(), "Lease should be invalid before any log flush" ); @@ -166,7 +166,7 @@ async fn test_single_voter_lease_refreshed_on_log_flushed() { // Verify: Lease is now valid assert!( - state.is_lease_valid(&ctx), + state.is_lease_valid(), "Lease should be valid after log flush in single_voter cluster" ); @@ -197,17 +197,17 @@ async fn test_single_voter_lease_expires_and_refreshes() { let (role_tx, _role_rx) = mpsc::unbounded_channel(); // T0: Lease invalid - assert!(!state.is_lease_valid(&ctx)); + assert!(!state.is_lease_valid()); // T1: Log flush → lease valid last_entry_id.store(1, Ordering::Relaxed); state.handle_log_flushed(1, &ctx, &role_tx).await; - assert!(state.is_lease_valid(&ctx)); + assert!(state.is_lease_valid()); // T2: Wait for lease to expire sleep(Duration::from_millis(lease_duration_ms + 50)).await; assert!( - !state.is_lease_valid(&ctx), + !state.is_lease_valid(), "Lease should expire after lease_duration_ms" ); @@ -215,7 +215,7 @@ async fn test_single_voter_lease_expires_and_refreshes() { last_entry_id.store(2, Ordering::Relaxed); state.handle_log_flushed(2, &ctx, &role_tx).await; assert!( - state.is_lease_valid(&ctx), + state.is_lease_valid(), "Lease should be refreshed by subsequent log flush" ); } @@ -244,13 +244,13 @@ async fn test_single_voter_lease_not_refreshed_when_commit_unchanged() { // Precondition: commit=0, lease invalid assert_eq!(state.commit_index(), 0); - assert!(!state.is_lease_valid(&ctx)); + assert!(!state.is_lease_valid()); // Action: Flush with durable=0, last_entry_id=0 (no new entries, commit unchanged) state.handle_log_flushed(0, &ctx, &role_tx).await; // Verify: Lease remains invalid (no leadership proof) - assert!(!state.is_lease_valid(&ctx)); + assert!(!state.is_lease_valid()); assert_eq!(state.commit_index(), 0); // Action: Flush with durable=1 (commit advances); set last_entry_id=1 to match @@ -258,7 +258,7 @@ async fn test_single_voter_lease_not_refreshed_when_commit_unchanged() { state.handle_log_flushed(1, &ctx, &role_tx).await; // Verify: Now lease is valid - assert!(state.is_lease_valid(&ctx)); + assert!(state.is_lease_valid()); assert_eq!(state.commit_index(), 1); } @@ -292,7 +292,7 @@ async fn test_single_voter_continuous_flushes_maintain_lease() { last_entry_id.store(i, Ordering::Relaxed); state.handle_log_flushed(i, &ctx, &role_tx).await; assert!( - state.is_lease_valid(&ctx), + state.is_lease_valid(), "Lease should remain valid after flush {i}" ); assert_eq!(state.commit_index(), i); @@ -302,7 +302,7 @@ async fn test_single_voter_continuous_flushes_maintain_lease() { } // Final verification: lease still valid - assert!(state.is_lease_valid(&ctx)); + assert!(state.is_lease_valid()); } // ============================================================================ @@ -331,7 +331,7 @@ async fn test_multi_voter_lease_not_refreshed_on_log_flushed() { // Precondition: Multi-voter cluster, lease invalid assert!(!state.cluster_metadata.single_voter); - assert!(!state.is_lease_valid(&ctx)); + assert!(!state.is_lease_valid()); // Action: Log flush with durable=1 // Note: In multi-voter, commit won't advance without follower ACKs, @@ -340,7 +340,7 @@ async fn test_multi_voter_lease_not_refreshed_on_log_flushed() { // Verify: Lease remains invalid (must wait for handle_append_result) assert!( - !state.is_lease_valid(&ctx), + !state.is_lease_valid(), "Lease should NOT be refreshed by log flush in multi-voter cluster" ); @@ -372,18 +372,18 @@ async fn test_multi_voter_lease_refresh_requires_append_result() { let (role_tx, _role_rx) = mpsc::unbounded_channel(); // Precondition: Lease invalid - assert!(!state.is_lease_valid(&ctx)); + assert!(!state.is_lease_valid()); // Action 1: Log flush → no lease refresh state.handle_log_flushed(1, &ctx, &role_tx).await; - assert!(!state.is_lease_valid(&ctx)); + assert!(!state.is_lease_valid()); // Action 2: Simulate append_result path (via test helper) state.test_update_lease_timestamp(); // Verify: Lease now valid (proves separation of paths) assert!( - state.is_lease_valid(&ctx), + state.is_lease_valid(), "Lease should be valid after explicit update (simulating append_result)" ); } @@ -413,7 +413,7 @@ async fn test_single_voter_durable_regression_does_not_refresh_lease() { last_entry_id.store(5, Ordering::Relaxed); state.handle_log_flushed(5, &ctx, &role_tx).await; assert_eq!(state.commit_index(), 5); - assert!(state.is_lease_valid(&ctx)); + assert!(state.is_lease_valid()); // Wait to distinguish timestamps sleep(Duration::from_millis(50)).await; @@ -424,10 +424,7 @@ async fn test_single_voter_durable_regression_does_not_refresh_lease() { // Verify: Commit unchanged, lease not updated (timestamp same as T1) assert_eq!(state.commit_index(), 5, "Commit should not regress"); - assert!( - state.is_lease_valid(&ctx), - "Lease should remain valid from T1" - ); + assert!(state.is_lease_valid(), "Lease should remain valid from T1"); } /// **Business Scenario**: Very short lease duration expires quickly @@ -452,13 +449,13 @@ async fn test_single_voter_very_short_lease_expires_quickly() { let (role_tx, _role_rx) = mpsc::unbounded_channel(); // T0: Lease invalid - assert!(!state.is_lease_valid(&ctx)); + assert!(!state.is_lease_valid()); // T1: Flush updates lease last_entry_id.store(1, Ordering::Relaxed); state.handle_log_flushed(1, &ctx, &role_tx).await; assert!( - state.is_lease_valid(&ctx), + state.is_lease_valid(), "Lease should be valid immediately after flush" ); @@ -467,7 +464,7 @@ async fn test_single_voter_very_short_lease_expires_quickly() { // T3: Lease expired assert!( - !state.is_lease_valid(&ctx), + !state.is_lease_valid(), "Lease should expire with lease_duration_ms = 50" ); } @@ -491,14 +488,14 @@ async fn test_single_voter_large_lease_duration() { // Action: Single flush last_entry_id.store(1, Ordering::Relaxed); state.handle_log_flushed(1, &ctx, &role_tx).await; - assert!(state.is_lease_valid(&ctx)); + assert!(state.is_lease_valid()); // Wait 1 second (much less than 10s lease) sleep(Duration::from_secs(1)).await; // Verify: Lease still valid assert!( - state.is_lease_valid(&ctx), + state.is_lease_valid(), "Lease should remain valid within duration" ); } diff --git a/d-engine-core/src/raft_role/leader_state_test/lease_send_ts_test.rs b/d-engine-core/src/raft_role/leader_state_test/lease_send_ts_test.rs new file mode 100644 index 00000000..7471118d --- /dev/null +++ b/d-engine-core/src/raft_role/leader_state_test/lease_send_ts_test.rs @@ -0,0 +1,140 @@ +//! Tests for lease deadline anchored to heartbeat send time (RTT/2 fix). +//! +//! ## Background +//! The safety invariant for lease-based reads requires: +//! `lease_duration_ms < election_timeout_min` +//! +//! However, the original implementation renewed the lease with `now_ms()` at the +//! time the quorum ACK was *received and processed*, not when the heartbeat was +//! *sent*. This extends the effective lease duration by ~RTT/2, violating the +//! invariant in practice. +//! +//! The fix: record `last_heartbeat_send_ts` before Phase 1 of +//! `execute_and_process_raft_rpc` and use it as the deadline base in +//! `update_lease_timestamp`, so the lease expires at `send_ts + lease_duration_ms` +//! rather than `ack_ts + lease_duration_ms`. + +use crate::RaftNodeConfig; +use crate::raft_role::leader_state::LeaderState; +use crate::raft_role::read_lease::now_ms; +use crate::raft_role::role_state::RaftRoleState; +use crate::test_utils::MockBuilder; +use crate::test_utils::mock::MockTypeConfig; +use std::time::Duration; +use tokio::sync::watch; + +// ── helpers ────────────────────────────────────────────────────────────────── + +async fn setup_leader( + lease_duration_ms: u64 +) -> ( + LeaderState, + crate::RaftContext, +) { + let (_shutdown_tx, shutdown_rx) = watch::channel(()); + let mut node_config = RaftNodeConfig::default(); + node_config.raft.read_consistency.lease_duration_ms = lease_duration_ms; + let ctx = MockBuilder::new(shutdown_rx).with_node_config(node_config).build_context(); + let state = LeaderState::::new(1, ctx.node_config.clone()); + (state, ctx) +} + +// ── tests ───────────────────────────────────────────────────────────────────── + +/// Deadline must be anchored to the heartbeat *send* time, not the ACK *receive* time. +/// +/// With a simulated RTT of 5 ms: +/// - new impl: deadline = send_ts + lease_duration_ms (correct) +/// - old impl: deadline = ack_ts + lease_duration_ms (too late by RTT) +/// +/// At `send_ts + lease_duration_ms` the new lease must already be expired, +/// while the old impl would still report valid — proving the old window was real. +#[tokio::test] +async fn test_lease_deadline_anchored_to_heartbeat_send_time_not_ack_time() { + let lease_duration_ms = 100u64; + let (mut state, ctx) = setup_leader(lease_duration_ms).await; + + let send_ts = now_ms(); + state.last_heartbeat_send_ts = send_ts; + + // Simulate RTT: 5 ms elapses before the ACK is processed. + std::thread::sleep(Duration::from_millis(5)); + + // Quorum ACK arrives — renew using send_ts (new impl). + state.test_renew_lease_from_send_ts(ctx.node_config().raft.read_consistency.lease_duration_ms); + + let term = state.current_term(); + let new_deadline = send_ts + lease_duration_ms; + + // Just before the send-anchored deadline: still valid. + assert!( + state.shared_state.lease.is_valid_for_leader(term, new_deadline - 1), + "lease must be valid just before send-anchored deadline" + ); + // At the send-anchored deadline: expired. + assert!( + !state.shared_state.lease.is_valid_for_leader(term, new_deadline), + "lease must be expired at send_ts + lease_duration_ms" + ); + + // Document that the old implementation would NOT have expired yet. + // old deadline ≈ send_ts + RTT + lease_duration_ms > new_deadline. + let old_deadline_approx = now_ms() + lease_duration_ms; // ≈ send_ts + 5 + 100 + assert!( + old_deadline_approx > new_deadline, + "old impl extends lease by RTT: old={old_deadline_approx} new={new_deadline}" + ); +} + +/// Without a quorum ACK, recording `last_heartbeat_send_ts` must NOT advance +/// the lease deadline. A stale leader that never gets ACKs must not serve reads. +#[tokio::test] +async fn test_no_quorum_ack_does_not_advance_lease_deadline() { + let (mut state, _ctx) = setup_leader(100).await; + + // Record send timestamp — but no ACK arrives, so update_lease_timestamp is never called. + state.last_heartbeat_send_ts = now_ms(); + + let term = state.current_term(); + assert!( + !state.shared_state.lease.is_valid_for_leader(term, now_ms()), + "lease must remain invalid when no quorum ACK has been received" + ); +} + +/// Demonstrates the invariant `lease_duration_ms < election_timeout_min` using +/// concrete numbers, showing the old implementation violated it and the new one preserves it. +/// +/// Config: election_timeout_min=150ms, lease_duration_ms=140ms, RTT=10ms. +/// +/// old impl: effective deadline = send_ts + RTT + 140 = send_ts + 150 +/// = earliest possible election time → invariant broken ❌ +/// +/// new impl: effective deadline = send_ts + 140 +/// < send_ts + 150 (earliest election) → invariant holds ✅ +#[test] +fn test_new_impl_preserves_election_timeout_invariant_old_impl_violates_it() { + let election_timeout_min = 150u64; + let lease_duration_ms = 140u64; + let simulated_rtt = 10u64; + + let send_ts = 1_000u64; // arbitrary base + let ack_ts = send_ts + simulated_rtt; + + let new_deadline = send_ts + lease_duration_ms; // 1_140 + let old_deadline = ack_ts + lease_duration_ms; // 1_150 + let earliest_election = send_ts + election_timeout_min; // 1_150 + + // New impl: lease expires strictly before any possible election. + assert!( + new_deadline < earliest_election, + "new impl: deadline {new_deadline} must be < earliest election {earliest_election}" + ); + + // Old impl: lease expires at the same time as the earliest election — invariant broken. + assert!( + old_deadline >= earliest_election, + "old impl: deadline {old_deadline} must be >= earliest election {earliest_election} \ + (documents the pre-fix violation)" + ); +} diff --git a/d-engine-core/src/raft_role/leader_state_test/mod.rs b/d-engine-core/src/raft_role/leader_state_test/mod.rs index 3b263fc8..344203ff 100644 --- a/d-engine-core/src/raft_role/leader_state_test/mod.rs +++ b/d-engine-core/src/raft_role/leader_state_test/mod.rs @@ -60,3 +60,9 @@ mod stale_learner_deadline_test; #[cfg(test)] mod pending_reads_test; + +#[cfg(test)] +mod become_follower_test; + +#[cfg(test)] +mod lease_send_ts_test; diff --git a/d-engine-core/src/raft_role/leader_state_test/pending_lease_reads_test.rs b/d-engine-core/src/raft_role/leader_state_test/pending_lease_reads_test.rs index 7a99420c..ba7e96dd 100644 --- a/d-engine-core/src/raft_role/leader_state_test/pending_lease_reads_test.rs +++ b/d-engine-core/src/raft_role/leader_state_test/pending_lease_reads_test.rs @@ -166,10 +166,7 @@ async fn test_lease_read_expired_pushes_to_pending_lease_reads() { ) .await; - assert!( - !state.is_lease_valid(&context), - "precondition: lease expired" - ); + assert!(!state.is_lease_valid(), "precondition: lease expired"); let (resp_tx, _resp_rx) = MaybeCloneOneshot::new(); state.push_client_cmd(make_lease_read_cmd(resp_tx), &context); @@ -292,10 +289,7 @@ async fn test_single_voter_lease_read_served_immediately_on_expired_lease() { state.cluster_metadata.single_voter, "precondition: single-voter" ); - assert!( - !state.is_lease_valid(&context), - "precondition: lease expired" - ); + assert!(!state.is_lease_valid(), "precondition: lease expired"); let (resp_tx, mut resp_rx) = MaybeCloneOneshot::new(); let (role_tx, _role_rx) = mpsc::unbounded_channel(); diff --git a/d-engine-core/src/raft_role/leader_state_test/replication_test.rs b/d-engine-core/src/raft_role/leader_state_test/replication_test.rs index aeb38c9c..38fd79f0 100644 --- a/d-engine-core/src/raft_role/leader_state_test/replication_test.rs +++ b/d-engine-core/src/raft_role/leader_state_test/replication_test.rs @@ -1398,6 +1398,23 @@ async fn test_execute_and_process_raft_rpc_multi_node_empty_peer_updates() { let (tx_write, rx_write) = >::new(); let batch = VecDeque::from(vec![mock_request(tx_write)]); + // Initialize the monotonic clock before process_batch so now_ms() returns > 0 when + // Phase 0 runs. Without this, the very first now_ms() call returns 0 (epoch just + // initialized, elapsed = 0ms), making the post-call assertion on last_heartbeat_send_ts + // indistinguishable from the default value of 0. + // Use a bounded spin instead of a fixed sleep — more deterministic on slow/coarse CI. + crate::raft_role::read_lease::init_clock(); + let clock_start = std::time::Instant::now(); + while crate::raft_role::read_lease::now_ms() == 0 + && clock_start.elapsed() < std::time::Duration::from_millis(50) + { + std::thread::yield_now(); + } + assert!( + crate::raft_role::read_lease::now_ms() > 0, + "clock failed to advance past 0 within 50ms" + ); + let (role_tx, mut role_rx) = mpsc::unbounded_channel(); let result = context.state.process_batch(batch, &role_tx, &context.raft_context).await; @@ -1418,6 +1435,13 @@ async fn test_execute_and_process_raft_rpc_multi_node_empty_peer_updates() { rx.try_recv().is_err(), "Write is pending — no peer responses" ); + + // Phase 0 of execute_and_process_raft_rpc must have recorded last_heartbeat_send_ts. + // If this is 0, the RTT/2 fix (lease deadline anchored to send time) is silently broken. + assert!( + context.state.last_heartbeat_send_ts > 0, + "last_heartbeat_send_ts must be set by Phase 0 before AppendEntries is sent" + ); } /// Test merge_batch_to_write_metadata with empty payload but non-empty senders diff --git a/d-engine-core/src/raft_role/mod.rs b/d-engine-core/src/raft_role/mod.rs index 6fea0855..5ed47b96 100644 --- a/d-engine-core/src/raft_role/mod.rs +++ b/d-engine-core/src/raft_role/mod.rs @@ -3,6 +3,7 @@ pub mod candidate_state; pub mod follower_state; pub mod leader_state; pub mod learner_state; +pub mod read_lease; pub mod role_state; #[cfg(test)] @@ -18,6 +19,7 @@ mod leader_state_test; mod learner_state_test; use std::collections::HashMap; +use std::sync::Arc; use std::sync::atomic::AtomicU32; use std::sync::atomic::Ordering; @@ -28,6 +30,7 @@ use follower_state::FollowerState; pub use leader_state::ClusterMetadata; use leader_state::LeaderState; use learner_state::LearnerState; +pub use read_lease::{ReadLease, init_clock, now_ms}; use role_state::RaftRoleState; use serde::Deserialize; use serde::Deserializer; @@ -84,6 +87,10 @@ pub struct SharedState { /// In-memory leader ID for hot-path reads (0 = no leader) /// Performance optimization: avoid RwLock on AppendEntries path current_leader_id: AtomicU32, + + /// Shared lease state between Raft loop (writer) and EmbeddedClient (reader). + /// Arc ensures the same allocation is shared across role transitions via clone(). + pub lease: Arc, } impl Clone for SharedState { @@ -93,6 +100,7 @@ impl Clone for SharedState { hard_state: self.hard_state, commit_index: self.commit_index, current_leader_id: AtomicU32::new(self.current_leader_id.load(Ordering::Acquire)), + lease: Arc::clone(&self.lease), } } } @@ -150,6 +158,7 @@ impl SharedState { hard_state, commit_index: last_applied_index_option.unwrap_or(0), current_leader_id: AtomicU32::new(0), + lease: Arc::new(ReadLease::new()), } } diff --git a/d-engine-core/src/raft_role/read_lease.rs b/d-engine-core/src/raft_role/read_lease.rs new file mode 100644 index 00000000..9e855549 --- /dev/null +++ b/d-engine-core/src/raft_role/read_lease.rs @@ -0,0 +1,145 @@ +use std::sync::OnceLock; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Instant; + +// Process-local monotonic epoch. Initialized on first call to now_ms(). +// Using Instant (not SystemTime) guarantees NTP-safe monotonic time. +static EPOCH: OnceLock = OnceLock::new(); + +/// Monotonic milliseconds since first call. Safe to use for deadline arithmetic. +#[must_use] +pub fn now_ms() -> u64 { + EPOCH.get_or_init(Instant::now).elapsed().as_millis() as u64 +} + +/// Initializes the monotonic clock epoch. Call once at engine startup to avoid +/// first-call jitter on the hot path. +pub fn init_clock() { + let _ = now_ms(); +} + +/// Lock-free Raft read lease shared between the Raft loop (writer) and EmbeddedClient (reader). +/// +/// Packs term (16 bits) and deadline_ms (48 bits) into a single AtomicU64 so +/// that readers always observe a consistent (term, deadline) pair from one atomic +/// load — eliminating the ABA race that two separate AtomicU64 fields would have. +/// +/// Layout: `[63..48] term | [47..0] deadline_ms` +/// +/// - `deadline_ms = 0` → no valid lease (sentinel) +/// - `deadline_ms > now_ms()` → lease is still live +/// - `term` wraps at 16 bits (max 65535); Raft clusters rarely exceed a few thousand terms. +/// +/// # Writers (Raft loop only) +/// - `renew(term, deadline_ms)` — called on every quorum ACK +/// - `invalidate(new_term)` — called when leader steps down +/// +/// # Readers (EmbeddedClient hot path) +/// - `is_valid_for_leader(term, now_ms)` — validates term + deadline (~4 ns on L1 hit) +#[derive(Debug)] +pub struct ReadLease { + packed: AtomicU64, +} + +impl ReadLease { + const DEADLINE_MASK: u64 = (1u64 << 48) - 1; + const TERM_SHIFT: u32 = 48; + + pub fn new() -> Self { + Self { + packed: AtomicU64::new(0), + } + } + + #[inline] + fn pack( + term: u64, + deadline_ms: u64, + ) -> u64 { + assert!( + deadline_ms <= Self::DEADLINE_MASK, + "deadline_ms overflows 48 bits: {deadline_ms}" + ); + ((term & 0xFFFF) << Self::TERM_SHIFT) | (deadline_ms & Self::DEADLINE_MASK) + } + + #[inline] + fn unpack(v: u64) -> (u64, u64) { + (v >> Self::TERM_SHIFT, v & Self::DEADLINE_MASK) + } + + /// Renew the lease. Called by the Raft loop after every quorum ACK. + /// `deadline_ms` should be `now_ms() + lease_duration_ms`. + #[inline] + pub fn renew( + &self, + term: u64, + deadline_ms: u64, + ) { + self.packed.store(Self::pack(term, deadline_ms), Ordering::Release); + } + + /// Invalidate the lease. Called when leader steps down. + /// Sets `deadline_ms = 0`; subsequent `is_valid_for_leader` calls return false. + #[inline] + pub fn invalidate( + &self, + new_term: u64, + ) { + self.packed.store(Self::pack(new_term, 0), Ordering::Release); + } + + /// Revoke the lease immediately. Single atomic store(0, Release). + /// + /// Called by the Raft loop on ANY term change (leader step-down, new election). + /// Sets packed = 0 → deadline = 0 → `is_valid()` returns false until next `renew()`. + /// + /// Safety invariant: every term-change path in the Raft loop MUST call this. + /// Missing a call causes silent stale reads (LeaseRead correctness violation). + #[inline] + pub fn revoke(&self) { + self.packed.store(0, Ordering::Release); + } + + /// ReadActor hot path: single atomic load, no external term required. + /// + /// Returns `true` iff `deadline_ms > now_ms`. Safety relies on `revoke()` being + /// called on every term change — see `revoke()` doc for the invariant. + #[must_use] + #[inline] + pub fn is_valid( + &self, + now_ms: u64, + ) -> bool { + let packed = self.packed.load(Ordering::Acquire); + (packed & Self::DEADLINE_MASK) > now_ms + } + + /// Check both term and deadline in one atomic load (~4 ns on L1 hit). + /// + /// Used by both `LeaderState` (internal) and `EmbeddedClient` (fast path). + /// `current_term` is masked to 16 bits before comparison, matching the storage layout. + #[must_use] + #[inline] + pub fn is_valid_for_leader( + &self, + current_term: u64, + now_ms: u64, + ) -> bool { + let (term, deadline) = Self::unpack(self.packed.load(Ordering::Acquire)); + term == (current_term & 0xFFFF) && deadline > now_ms + } +} + +impl Default for ReadLease { + fn default() -> Self { + Self::new() + } +} + +// ReadLease is Send + Sync because AtomicU64 is Send + Sync. +// The compiler derives this automatically. + +#[cfg(test)] +#[path = "read_lease_test.rs"] +mod tests; diff --git a/d-engine-core/src/raft_role/read_lease_test.rs b/d-engine-core/src/raft_role/read_lease_test.rs new file mode 100644 index 00000000..bb4e58d3 --- /dev/null +++ b/d-engine-core/src/raft_role/read_lease_test.rs @@ -0,0 +1,312 @@ +use std::sync::Arc; +use std::sync::Barrier; +use std::thread; + +use super::{ReadLease, now_ms}; + +// ── now_ms ──────────────────────────────────────────────────────────────────── + +#[test] +fn test_now_ms_is_monotonically_nondecreasing() { + let a = now_ms(); + let b = now_ms(); + assert!(b >= a, "now_ms must never go backwards: a={a}, b={b}"); +} + +#[test] +fn test_now_ms_advances_over_time() { + let before = now_ms(); + thread::sleep(std::time::Duration::from_millis(5)); + let after = now_ms(); + assert!( + after > before, + "now_ms must advance over 5ms: before={before}, after={after}" + ); +} + +// ── new / default ───────────────────────────────────────────────────────────── + +#[test] +fn test_new_starts_invalid() { + let lease = ReadLease::new(); + // deadline=0 → always false regardless of term + assert!( + !lease.is_valid_for_leader(1, now_ms()), + "fresh ReadLease must be invalid (deadline=0)" + ); +} + +#[test] +fn test_new_is_invalid_for_any_term() { + let lease = ReadLease::new(); + assert!(!lease.is_valid_for_leader(1, now_ms())); + assert!(!lease.is_valid_for_leader(u16::MAX as u64, now_ms())); +} + +// ── renew ───────────────────────────────────────────────────────────────────── + +#[test] +fn test_renew_makes_lease_valid_for_matching_term() { + let lease = ReadLease::new(); + let deadline = now_ms() + 5_000; + lease.renew(1, deadline); + assert!(lease.is_valid_for_leader(1, now_ms())); +} + +#[test] +fn test_renew_invalid_for_leader_with_wrong_term() { + let lease = ReadLease::new(); + let deadline = now_ms() + 5_000; + lease.renew(42, deadline); + assert!(!lease.is_valid_for_leader(41, now_ms())); + assert!(!lease.is_valid_for_leader(43, now_ms())); +} + +#[test] +fn test_renew_with_past_deadline_is_invalid() { + let lease = ReadLease::new(); + let past = now_ms().saturating_sub(1_000); + lease.renew(1, past); + assert!(!lease.is_valid_for_leader(1, now_ms())); +} + +#[test] +fn test_renew_zero_deadline_is_invalid() { + let lease = ReadLease::new(); + lease.renew(1, 0); + assert!(!lease.is_valid_for_leader(1, now_ms())); +} + +// ── invalidate ──────────────────────────────────────────────────────────────── + +#[test] +fn test_invalidate_clears_valid_lease() { + let lease = ReadLease::new(); + lease.renew(1, now_ms() + 5_000); + assert!( + lease.is_valid_for_leader(1, now_ms()), + "pre-condition: lease must be valid" + ); + lease.invalidate(2); + assert!( + !lease.is_valid_for_leader(1, now_ms()), + "lease must be invalid after invalidate" + ); +} + +#[test] +fn test_invalidate_updates_term() { + let lease = ReadLease::new(); + lease.renew(1, now_ms() + 5_000); + lease.invalidate(2); + // deadline=0 after invalidate → false for any term + assert!(!lease.is_valid_for_leader(1, now_ms())); + assert!(!lease.is_valid_for_leader(2, now_ms())); +} + +#[test] +fn test_invalidate_on_fresh_lease_stays_invalid() { + let lease = ReadLease::new(); + lease.invalidate(5); + assert!(!lease.is_valid_for_leader(5, now_ms())); +} + +// ── renew → invalidate → renew cycle ───────────────────────────────────────── + +#[test] +fn test_renew_after_invalidate_makes_lease_valid_again() { + let lease = ReadLease::new(); + lease.renew(1, now_ms() + 5_000); + lease.invalidate(2); + assert!( + !lease.is_valid_for_leader(1, now_ms()), + "must be invalid after invalidate" + ); + lease.renew(3, now_ms() + 5_000); + assert!( + lease.is_valid_for_leader(3, now_ms()), + "must be valid after second renew" + ); + assert!( + !lease.is_valid_for_leader(1, now_ms()), + "old term must not match" + ); +} + +// ── term packing (16-bit truncation) ───────────────────────────────────────── + +#[test] +fn test_term_truncation_wraps_at_16_bits() { + let lease = ReadLease::new(); + let term_17bit: u64 = 0x1_0001; // bit 16 set — truncated to 0x0001 in 16-bit slot + lease.renew(term_17bit, now_ms() + 5_000); + // stored term = term_17bit & 0xFFFF = 1 + // both argument and stored value are masked to 16 bits, so 0x10001 → 1 on both sides + assert!( + lease.is_valid_for_leader(1, now_ms()), + "low 16 bits of 0x10001 == 1" + ); + assert!( + lease.is_valid_for_leader(term_17bit, now_ms()), + "0x10001 & 0xFFFF == 1 — same as passing 1" + ); + assert!( + !lease.is_valid_for_leader(2, now_ms()), + "term 2 must not match stored term 1" + ); +} + +// ── revoke ──────────────────────────────────────────────────────────────────── + +#[test] +fn test_revoke_clears_valid_lease() { + let lease = ReadLease::new(); + lease.renew(1, now_ms() + 5_000); + assert!( + lease.is_valid_for_leader(1, now_ms()), + "pre-condition: lease must be valid before revoke" + ); + lease.revoke(); + assert!( + !lease.is_valid_for_leader(1, now_ms()), + "lease must be invalid after revoke" + ); +} + +#[test] +fn test_revoke_on_fresh_lease_stays_invalid() { + let lease = ReadLease::new(); + lease.revoke(); + assert!(!lease.is_valid_for_leader(0, now_ms())); + assert!(!lease.is_valid_for_leader(1, now_ms())); +} + +#[test] +fn test_revoke_sets_packed_to_zero() { + let lease = ReadLease::new(); + lease.renew(42, now_ms() + 10_000); + lease.revoke(); + // After revoke, neither term=42 nor any other term should be valid. + assert!(!lease.is_valid_for_leader(42, now_ms())); + assert!(!lease.is_valid_for_leader(0, now_ms())); +} + +// ── is_valid ────────────────────────────────────────────────────────────────── + +#[test] +fn test_is_valid_returns_true_when_deadline_in_future() { + let lease = ReadLease::new(); + lease.renew(1, now_ms() + 5_000); + assert!( + lease.is_valid(now_ms()), + "is_valid must return true when deadline is in the future" + ); +} + +#[test] +fn test_is_valid_returns_false_when_deadline_in_past() { + let lease = ReadLease::new(); + let past = now_ms().saturating_sub(1_000); + lease.renew(1, past); + assert!( + !lease.is_valid(now_ms()), + "is_valid must return false when deadline is in the past" + ); +} + +#[test] +fn test_is_valid_returns_false_on_fresh_lease() { + let lease = ReadLease::new(); + assert!( + !lease.is_valid(now_ms()), + "fresh lease (deadline=0) must be invalid" + ); +} + +#[test] +fn test_is_valid_returns_false_after_revoke() { + let lease = ReadLease::new(); + lease.renew(1, now_ms() + 5_000); + lease.revoke(); + assert!( + !lease.is_valid(now_ms()), + "is_valid must return false after revoke" + ); +} + +#[test] +fn test_is_valid_ignores_term() { + // is_valid() only checks deadline, not term — any term works if deadline is in future + let lease = ReadLease::new(); + lease.renew(99, now_ms() + 5_000); + // is_valid does NOT require knowing the term + assert!( + lease.is_valid(now_ms()), + "is_valid must not care about term" + ); +} + +#[test] +fn test_is_valid_consistent_with_is_valid_for_leader() { + // When both term and deadline match, is_valid and is_valid_for_leader agree. + let lease = ReadLease::new(); + let deadline = now_ms() + 5_000; + lease.renew(7, deadline); + let now = now_ms(); + assert_eq!( + lease.is_valid(now), + lease.is_valid_for_leader(7, now), + "is_valid and is_valid_for_leader must agree when term matches" + ); +} + +// ── Arc sharing semantics ───────────────────────────────────────────────────── + +#[test] +fn test_arc_read_lease_visible_across_clones() { + let lease = Arc::new(ReadLease::new()); + let reader: Arc = Arc::clone(&lease); + + lease.renew(1, now_ms() + 5_000); + assert!( + reader.is_valid_for_leader(1, now_ms()), + "renew on original must be visible via clone" + ); + + lease.invalidate(2); + assert!( + !reader.is_valid_for_leader(1, now_ms()), + "invalidate on original must be visible via clone" + ); +} + +// ── concurrent safety (basic) ───────────────────────────────────────────────── + +#[test] +fn test_concurrent_renew_does_not_panic() { + let lease = Arc::new(ReadLease::new()); + let barrier = Arc::new(Barrier::new(4)); + let mut handles = Vec::new(); + + for term in 1u64..=3 { + let l: Arc = Arc::clone(&lease); + let b = Arc::clone(&barrier); + handles.push(thread::spawn(move || { + b.wait(); + let deadline = now_ms() + 5_000; + l.renew(term, deadline); + })); + } + // reader thread + { + let l: Arc = Arc::clone(&lease); + let b = Arc::clone(&barrier); + handles.push(thread::spawn(move || { + b.wait(); + let _ = l.is_valid_for_leader(1, now_ms()); + })); + } + for h in handles { + h.join().expect("thread panicked"); + } +} diff --git a/d-engine-core/src/state_machine_handler/snapshot_assembler_test.rs b/d-engine-core/src/state_machine_handler/snapshot_assembler_test.rs index c620d5f3..939f1475 100644 --- a/d-engine-core/src/state_machine_handler/snapshot_assembler_test.rs +++ b/d-engine-core/src/state_machine_handler/snapshot_assembler_test.rs @@ -197,12 +197,17 @@ async fn test_assembler_on_stale_temp_file_does_not_append_old_data() { "snapshot-".to_string(), )); - // First transfer: write 3 chunks then drop (simulates timeout / leader change) + // First transfer: write 3 chunks, flush to OS buffer, then drop without finalize. + // Explicit flush is required: tokio::fs::File has an internal write buffer that is + // NOT guaranteed to reach the OS kernel on drop. Without flush, the stale file size + // is non-deterministic (CI flakiness). flush_to_disk() models the production scenario + // where some chunks were flushed to the OS before the transfer timed out. { let mut assembler = SnapshotAssembler::new(path_mgr.clone()).await.unwrap(); for i in 0..3u32 { assembler.write_chunk(i, Bytes::from(vec![0xAAu8; 1024])).await.unwrap(); } + assembler.flush_to_disk().await.unwrap(); // Drop without finalize — stale temp file stays on disk with 3 * 1024 bytes } diff --git a/d-engine-core/src/storage/state_machine.rs b/d-engine-core/src/storage/state_machine.rs index 42613c4b..95eeef46 100644 --- a/d-engine-core/src/storage/state_machine.rs +++ b/d-engine-core/src/storage/state_machine.rs @@ -85,6 +85,17 @@ pub trait StateMachine: Send + Sync + 'static { /// This is typically a sync operation for state management. fn stop(&self) -> Result<(), Error>; + /// Permanently close underlying storage resources (e.g. DB handle, file descriptors). + /// + /// Called by `EmbeddedEngine::stop()` before the Raft loop exits to release + /// OS-level resources (e.g. RocksDB LOCK file) without waiting for all + /// `Arc` clones to drop. + /// + /// Default is a no-op — custom state machine implementations do not need to + /// override this unless they hold exclusive OS resources that must be released + /// before the process exits or a new engine instance is started. + fn close_storage(&self) {} + /// Checks if the state machine is currently running. /// Sync operation as it just checks an atomic boolean. fn is_running(&self) -> bool; @@ -96,6 +107,46 @@ pub trait StateMachine: Send + Sync + 'static { key_buffer: &[u8], ) -> Result, Error>; + /// Retrieves multiple values by key from the state machine in a single call. + /// + /// # Why this belongs in the protocol layer + /// + /// `get()` is already part of the protocol interface because the ReadActor hot path + /// requires direct SM access. `get_multi` is its natural batch extension — the same + /// reasoning applies. An Iterator's `fold` has a default impl; `Vec` overrides it for + /// performance. Users don't need to know which path executes. + /// + /// # Read coherence requirement + /// + /// The consistency unit is the **request**, not the key. All returned values must + /// come from the same applied state. A torn read — key A from apply index 100, + /// key B from apply index 105 — is a correctness violation in coordinator workloads: + /// + /// - Service registry: addr=v2 + version=v2.2 routes traffic to the wrong instance + /// - Leader election: leader=node-3 + term=43 is a phantom state that never existed + /// - Quota management: used=850 with a reset window_start blocks valid requests + /// + /// etcd (Range + MVCC), TiKV RawKV (BatchGet + RocksDB snapshot), and Consul stale + /// reads all guarantee snapshot consistency even in eventual/stale modes. + /// + /// # Default implementation + /// + /// Calls `get()` sequentially. Correct when the caller holds no write locks, but + /// does not guarantee snapshot isolation under concurrent writes. Override in + /// implementations that support it (RocksDB `db.snapshot()`, FileStateMachine + /// read-lock held for the full batch). + /// + /// # Position contract + /// + /// `result[i]` corresponds to `keys[i]`. Missing keys produce `None` at their + /// position; the result length always equals `keys.len()`. + fn get_multi( + &self, + keys: &[Bytes], + ) -> Result>, Error> { + keys.iter().map(|k| self.get(k)).collect() + } + /// Returns the term of a specific log entry by its ID. /// Sync operation as it queries in-memory data. fn entry_term( diff --git a/d-engine-core/src/storage/state_machine_test.rs b/d-engine-core/src/storage/state_machine_test.rs index e51e8ee3..c265b1bc 100644 --- a/d-engine-core/src/storage/state_machine_test.rs +++ b/d-engine-core/src/storage/state_machine_test.rs @@ -48,6 +48,14 @@ impl StateMachineTestSuite { Self::test_ungraceful_shutdown_recovery(&builder).await?; Self::test_reset_operation(builder.build().await?).await?; + Self::test_get_multi_returns_values_in_key_order(builder.build().await?).await?; + Self::test_get_multi_absent_keys_return_none(builder.build().await?).await?; + Self::test_get_multi_empty_keys_returns_empty(builder.build().await?).await?; + Self::test_get_multi_partial_hit_some_keys_missing(builder.build().await?).await?; + Self::test_get_multi_service_registry_coherence(builder.build().await?).await?; + Self::test_get_multi_election_state_coherence(builder.build().await?).await?; + Self::test_get_multi_quota_management_coherence(builder.build().await?).await?; + builder.cleanup().await?; Ok(()) } @@ -834,6 +842,326 @@ impl StateMachineTestSuite { Ok(()) } + + // ── get_multi tests ─────────────────────────────────────────────────────── + + /// Verify that get_multi preserves input key order in its output. + /// + /// The caller maps result[i] back to keys[i] without a key lookup. + /// Any reordering silently breaks that mapping. + pub async fn test_get_multi_returns_values_in_key_order( + sm: Arc + ) -> Result<(), Error> { + sm.apply_chunk(&[ + create_insert_entry(1, Bytes::from("key_a"), Bytes::from("value_a")), + create_insert_entry(2, Bytes::from("key_b"), Bytes::from("value_b")), + create_insert_entry(3, Bytes::from("key_c"), Bytes::from("value_c")), + ]) + .await?; + + // Request in reverse alphabetical order to catch any accidental sorting. + let keys = vec![ + Bytes::from("key_c"), + Bytes::from("key_a"), + Bytes::from("key_b"), + ]; + let values = sm.get_multi(&keys)?; + + assert_eq!(values.len(), 3, "result length must match key count"); + assert_eq!( + values[0], + Some(Bytes::from("value_c")), + "position 0 = key_c" + ); + assert_eq!( + values[1], + Some(Bytes::from("value_a")), + "position 1 = key_a" + ); + assert_eq!( + values[2], + Some(Bytes::from("value_b")), + "position 2 = key_b" + ); + Ok(()) + } + + /// Verify that absent keys produce None entries, not errors or panics. + /// + /// Callers rely on None to distinguish "key not found" from a storage failure. + /// Returning Err for a missing key would force callers to treat a normal business + /// state (no registration yet) as an error. + pub async fn test_get_multi_absent_keys_return_none( + sm: Arc + ) -> Result<(), Error> { + let keys = vec![Bytes::from("ghost_1"), Bytes::from("ghost_2")]; + let values = sm.get_multi(&keys)?; + + assert_eq!(values.len(), 2); + assert!(values[0].is_none(), "absent key must return None, not Err"); + assert!(values[1].is_none(), "absent key must return None, not Err"); + Ok(()) + } + + /// Verify that an empty key slice returns an empty result without error. + /// + /// Callers may pass an empty slice when the coordinator has no keys to fetch + /// (e.g., an empty service namespace). The implementation must not panic or + /// return an error in this case. + pub async fn test_get_multi_empty_keys_returns_empty( + sm: Arc + ) -> Result<(), Error> { + let values = sm.get_multi(&[])?; + assert!( + values.is_empty(), + "empty key list must produce empty result" + ); + Ok(()) + } + + /// Verify that a mix of present and absent keys returns Some/None correctly. + /// + /// The position contract must hold even when some slots are None. + /// A naive implementation that skips missing keys and compacts the result + /// would shift positions, silently misrouting values to wrong fields. + pub async fn test_get_multi_partial_hit_some_keys_missing( + sm: Arc + ) -> Result<(), Error> { + sm.apply_chunk(&[create_insert_entry( + 1, + Bytes::from("present"), + Bytes::from("v"), + )]) + .await?; + + let keys = vec![ + Bytes::from("missing_before"), + Bytes::from("present"), + Bytes::from("missing_after"), + ]; + let values = sm.get_multi(&keys)?; + + assert_eq!( + values.len(), + 3, + "result length must equal key count even with gaps" + ); + assert!( + values[0].is_none(), + "missing_before must be None at position 0" + ); + assert_eq!( + values[1], + Some(Bytes::from("v")), + "present must have value at position 1" + ); + assert!( + values[2].is_none(), + "missing_after must be None at position 2" + ); + Ok(()) + } + + /// Service registry coordinator scenario: batch read must not produce a torn state. + /// + /// In distributed service routing, (addr, version, health) must all come from + /// the same applied state. A torn read — addr from apply index 100, version from + /// index 105 — would cause traffic to route to the old instance while the client + /// believes it runs the new version, breaking gray-release correctness. + /// + /// This test verifies that after a full v2 deployment, get_multi returns all three + /// fields from v2 with no v1 values mixed in. + pub async fn test_get_multi_service_registry_coherence( + sm: Arc + ) -> Result<(), Error> { + // v1: initial deployment + sm.apply_chunk(&[ + create_insert_entry( + 1, + Bytes::from("/service/payment/addr"), + Bytes::from("10.0.0.1:8080"), + ), + create_insert_entry( + 2, + Bytes::from("/service/payment/version"), + Bytes::from("v2.2"), + ), + create_insert_entry( + 3, + Bytes::from("/service/payment/health"), + Bytes::from("healthy"), + ), + ]) + .await?; + + // v2: new deployment — all three fields transition together + sm.apply_chunk(&[ + create_insert_entry( + 4, + Bytes::from("/service/payment/addr"), + Bytes::from("10.0.0.2:8080"), + ), + create_insert_entry( + 5, + Bytes::from("/service/payment/version"), + Bytes::from("v2.3"), + ), + create_insert_entry( + 6, + Bytes::from("/service/payment/health"), + Bytes::from("starting"), + ), + ]) + .await?; + + let keys = vec![ + Bytes::from("/service/payment/addr"), + Bytes::from("/service/payment/version"), + Bytes::from("/service/payment/health"), + ]; + let values = sm.get_multi(&keys)?; + + assert_eq!(values.len(), 3); + // All three must reflect the v2 state consistently. + // Any mix (e.g., addr=v2 + version=v2.2) is a torn read. + assert_eq!( + values[0], + Some(Bytes::from("10.0.0.2:8080")), + "addr must be v2.3 instance" + ); + assert_eq!(values[1], Some(Bytes::from("v2.3")), "version must be v2.3"); + assert_eq!( + values[2], + Some(Bytes::from("starting")), + "health must be v2.3 startup value" + ); + Ok(()) + } + + /// Leader election coordinator scenario: batch read must not produce a phantom state. + /// + /// (leader, term, lease_expire) must be internally consistent. A torn read can + /// produce a combination that never existed — e.g., leader=node-3 (old term) with + /// term=43 (new term). A coordinator observing this phantom state may incorrectly + /// conclude that node-3 holds a lease in term 43 and take split-brain actions. + /// + /// This test verifies that after a term transition, get_multi returns all three + /// fields from the new term with no old-term values mixed in. + pub async fn test_get_multi_election_state_coherence( + sm: Arc + ) -> Result<(), Error> { + // Term 42: node-3 is leader + sm.apply_chunk(&[ + create_insert_entry(1, Bytes::from("/election/leader"), Bytes::from("node-3")), + create_insert_entry(2, Bytes::from("/election/term"), Bytes::from("42")), + create_insert_entry( + 3, + Bytes::from("/election/lease_expire"), + Bytes::from("1748600000"), + ), + ]) + .await?; + + // Term 43: node-5 wins new election + sm.apply_chunk(&[ + create_insert_entry(4, Bytes::from("/election/leader"), Bytes::from("node-5")), + create_insert_entry(5, Bytes::from("/election/term"), Bytes::from("43")), + create_insert_entry( + 6, + Bytes::from("/election/lease_expire"), + Bytes::from("1748603600"), + ), + ]) + .await?; + + let keys = vec![ + Bytes::from("/election/leader"), + Bytes::from("/election/term"), + Bytes::from("/election/lease_expire"), + ]; + let values = sm.get_multi(&keys)?; + + assert_eq!(values.len(), 3); + // Must reflect term-43 state consistently. + // leader=node-3 + term=43 is a phantom state that never existed. + assert_eq!( + values[0], + Some(Bytes::from("node-5")), + "leader must be node-5 (term 43)" + ); + assert_eq!(values[1], Some(Bytes::from("43")), "term must be 43"); + assert_eq!( + values[2], + Some(Bytes::from("1748603600")), + "lease_expire must be term-43 value" + ); + Ok(()) + } + + /// Quota management coordinator scenario: batch read must return coherent counters. + /// + /// (limit, used, window_start) from different apply indexes produce incorrect + /// rate-limit calculations — a pure business correctness failure that is independent + /// of Linearizable semantics. Reading used=850 with a new window_start makes the + /// rate limiter reject requests that should be allowed in the new window. + /// + /// This test verifies that after a window reset, get_multi returns all three + /// counters from the new window with no old-window values mixed in. + pub async fn test_get_multi_quota_management_coherence( + sm: Arc + ) -> Result<(), Error> { + // Window 1: tenant-A has consumed 850 of their 1000 quota + sm.apply_chunk(&[ + create_insert_entry(1, Bytes::from("/quota/tenant-A/limit"), Bytes::from("1000")), + create_insert_entry(2, Bytes::from("/quota/tenant-A/used"), Bytes::from("850")), + create_insert_entry( + 3, + Bytes::from("/quota/tenant-A/window_start"), + Bytes::from("1748599900"), + ), + ]) + .await?; + + // Window 2: quota window resets — used counter goes back to 0 + sm.apply_chunk(&[ + create_insert_entry(4, Bytes::from("/quota/tenant-A/limit"), Bytes::from("1000")), + create_insert_entry(5, Bytes::from("/quota/tenant-A/used"), Bytes::from("0")), + create_insert_entry( + 6, + Bytes::from("/quota/tenant-A/window_start"), + Bytes::from("1748603500"), + ), + ]) + .await?; + + let keys = vec![ + Bytes::from("/quota/tenant-A/limit"), + Bytes::from("/quota/tenant-A/used"), + Bytes::from("/quota/tenant-A/window_start"), + ]; + let values = sm.get_multi(&keys)?; + + assert_eq!(values.len(), 3); + // Must read window-2 consistently. + // used=850 + window_start from window-2 is a torn read that wrongly + // blocks requests that should be free in the new window. + assert_eq!( + values[0], + Some(Bytes::from("1000")), + "limit is unchanged across windows" + ); + assert_eq!( + values[1], + Some(Bytes::from("0")), + "used must be 0 after window reset" + ); + assert_eq!( + values[2], + Some(Bytes::from("1748603500")), + "window_start must be from window-2" + ); + Ok(()) + } } /// Helper function to create an Insert ApplyEntry diff --git a/d-engine-server/README.md b/d-engine-server/README.md index 0ba57f76..608a01db 100644 --- a/d-engine-server/README.md +++ b/d-engine-server/README.md @@ -44,12 +44,12 @@ d-engine-server = "0.2" ### Embedded Mode (zero-overhead local client) ```rust -use d_engine_server::EmbeddedEngine; +use d_engine_server::DefaultEmbeddedEngine; use std::time::Duration; #[tokio::main] async fn main() -> Result<(), Box> { - let engine = EmbeddedEngine::start_with("config.toml").await?; + let engine = DefaultEmbeddedEngine::start_with("config.toml").await?; engine.wait_ready(Duration::from_secs(5)).await?; let client = engine.client(); @@ -184,13 +184,13 @@ See the [Server Guide](https://docs.rs/d-engine/latest/d_engine/docs/server_guid High-level API for embedding d-engine in Rust applications: ```rust -use d_engine_server::EmbeddedEngine; +use d_engine_server::DefaultEmbeddedEngine; use std::time::Duration; #[tokio::main] async fn main() -> Result<(), Box> { // Start embedded engine with config file - let engine = EmbeddedEngine::start_with("d-engine.toml").await?; + let engine = DefaultEmbeddedEngine::start_with("d-engine.toml").await?; // Wait for leader election let leader = engine.wait_ready(Duration::from_secs(5)).await?; diff --git a/d-engine-server/benches/lease_performance.rs b/d-engine-server/benches/lease_performance.rs index fac06eac..bc07b70a 100644 --- a/d-engine-server/benches/lease_performance.rs +++ b/d-engine-server/benches/lease_performance.rs @@ -19,8 +19,8 @@ use criterion::black_box; use criterion::criterion_group; use criterion::criterion_main; use d_engine_core::{ApplyEntry, Command, StateMachine}; -use d_engine_server::storage::DefaultLease; use d_engine_server::storage::FileStateMachine; +use d_engine_server::storage::TtlLease; use tempfile::TempDir; /// Create FileStateMachine with Background strategy (lease enabled) @@ -35,7 +35,7 @@ async fn create_sm_background() -> (FileStateMachine, TempDir) { .await .expect("Failed to create state machine"); - let lease = Arc::new(DefaultLease::new(lease_config)); + let lease = Arc::new(TtlLease::new(lease_config)); sm.set_lease(lease); (sm, temp_dir) diff --git a/d-engine-server/benches/state_machine.rs b/d-engine-server/benches/state_machine.rs index f5836212..87d8a80c 100644 --- a/d-engine-server/benches/state_machine.rs +++ b/d-engine-server/benches/state_machine.rs @@ -40,7 +40,7 @@ async fn create_test_state_machine() -> (FileStateMachine, TempDir) { cleanup_interval_ms: 1000, max_cleanup_duration_ms: 1, }; - let lease = Arc::new(d_engine_server::storage::DefaultLease::new(lease_config)); + let lease = Arc::new(d_engine_server::storage::TtlLease::new(lease_config)); sm.set_lease(lease); (sm, temp_dir) diff --git a/d-engine-server/benches/ttl.rs b/d-engine-server/benches/ttl.rs index c01f9130..87a7ccdf 100644 --- a/d-engine-server/benches/ttl.rs +++ b/d-engine-server/benches/ttl.rs @@ -23,7 +23,7 @@ use tempfile::TempDir; /// Helper to create a temporary state machine for benchmarking async fn create_test_state_machine() -> (FileStateMachine, TempDir) { - use d_engine_server::storage::DefaultLease; + use d_engine_server::storage::TtlLease; let temp_dir = TempDir::new().expect("Failed to create temp dir"); // For TTL benchmarks, we need lease enabled @@ -37,7 +37,7 @@ async fn create_test_state_machine() -> (FileStateMachine, TempDir) { .expect("Failed to create state machine"); // Manually inject lease for benchmarking - let lease = std::sync::Arc::new(DefaultLease::new(lease_config)); + let lease = std::sync::Arc::new(TtlLease::new(lease_config)); sm.set_lease(lease); sm.load_lease_data().await.expect("Failed to load lease data"); @@ -80,7 +80,7 @@ fn create_noop_entry(index: u64) -> ApplyEntry { /// Target: < 1ms for typical workload (100 expired keys) /// This directly benchmarks lease.on_apply() to isolate cleanup overhead fn bench_piggyback_cleanup(c: &mut Criterion) { - use d_engine_server::storage::DefaultLease; + use d_engine_server::storage::TtlLease; let mut group = c.benchmark_group("piggyback_cleanup"); @@ -90,7 +90,7 @@ fn bench_piggyback_cleanup(c: &mut Criterion) { cleanup_interval_ms: 1000, max_cleanup_duration_ms: 1, }; - let lease = DefaultLease::new(lease_config); + let lease = TtlLease::new(lease_config); // Register keys with very short TTL for i in 0..*expired_count { @@ -121,13 +121,13 @@ fn bench_piggyback_cleanup(c: &mut Criterion) { /// Measures the cost of registering a TTL entry in the TTL manager /// Target: < 100ns per registration fn bench_ttl_registration(c: &mut Criterion) { - use d_engine_server::storage::DefaultLease; + use d_engine_server::storage::TtlLease; let lease_config = d_engine_core::config::LeaseConfig { cleanup_interval_ms: 1000, max_cleanup_duration_ms: 1, }; - let lease = DefaultLease::new(lease_config); + let lease = TtlLease::new(lease_config); c.bench_function("ttl_registration", |b| { let mut counter = 0u64; diff --git a/d-engine-server/benches/watch_overhead.rs b/d-engine-server/benches/watch_overhead.rs index b9fa7a59..ca69b4cc 100644 --- a/d-engine-server/benches/watch_overhead.rs +++ b/d-engine-server/benches/watch_overhead.rs @@ -22,7 +22,7 @@ use criterion::criterion_group; use criterion::criterion_main; use d_engine_core::{ApplyEntry, Command, StateMachine}; use d_engine_server::RocksDBUnifiedEngine; -use d_engine_server::api::EmbeddedEngine; +use d_engine_server::api::DefaultEmbeddedEngine; use tempfile::TempDir; use tokio::time::sleep; @@ -71,7 +71,8 @@ fn create_test_entries( } /// Create a temporary EmbeddedEngine for benchmarking -async fn create_embedded_engine() -> Result<(EmbeddedEngine, TempDir), Box> { +async fn create_embedded_engine() +-> Result<(DefaultEmbeddedEngine, TempDir), Box> { let temp_dir = TempDir::new()?; let db_path = temp_dir.path().join("db"); @@ -101,9 +102,12 @@ watcher_buffer_size = 100 let storage = Arc::new(storage); let state_machine = Arc::new(state_machine); - let engine = - EmbeddedEngine::start_custom(storage, state_machine, Some(config_path.to_str().unwrap())) - .await?; + let engine = DefaultEmbeddedEngine::start_custom( + storage, + state_machine, + Some(config_path.to_str().unwrap()), + ) + .await?; // Wait for ready engine.wait_ready(Duration::from_secs(5)).await?; diff --git a/d-engine-server/src/api/embedded.rs b/d-engine-server/src/api/embedded.rs index 2599d439..f04c4775 100644 --- a/d-engine-server/src/api/embedded.rs +++ b/d-engine-server/src/api/embedded.rs @@ -122,6 +122,8 @@ //! //! For auto-forwarding with gRPC overhead, use standalone mode with `GrpcClient`. +use super::embedded_client::EmbeddedClient; +use super::embedded_read_handle::EmbeddedReadHandle; use crate::Result; #[cfg(feature = "rocksdb")] use crate::RocksDBStateMachine; @@ -131,8 +133,10 @@ use crate::RocksDBStorageEngine; use crate::RocksDBUnifiedEngine; use crate::StateMachine; use crate::StorageEngine; -use crate::api::EmbeddedClient; use crate::node::NodeBuilder; +use crate::node::RaftTypeConfig; +use d_engine_core::TypeConfig; +use std::fmt::Debug; use std::sync::Arc; use std::time::Duration; use tokio::sync::Mutex; @@ -141,10 +145,13 @@ use tokio::task::JoinHandle; use tracing::error; use tracing::info; -struct Inner { +struct Inner { + /// SM reference for shutdown: close_storage() releases OS resources (e.g. RocksDB LOCK) + /// before the Raft loop exits, without waiting for all Arc clones to drop. + sm: Arc, node_handle: Mutex>>>, shutdown_tx: watch::Sender<()>, - client: Arc, + client: Arc>, leader_elected_rx: watch::Receiver>, membership_rx: watch::Receiver, is_stopped: Mutex, @@ -175,12 +182,22 @@ struct Inner { /// /// engine.stop().await?; /// ``` -#[derive(Clone)] -pub struct EmbeddedEngine { - inner: Arc, +pub struct EmbeddedEngine { + inner: Arc>, } -impl EmbeddedEngine { +impl Clone for EmbeddedEngine { + fn clone(&self) -> Self { + Self { + inner: Arc::clone(&self.inner), + } + } +} + +// ─── RocksDB default start ─────────────────────────────────────────────────── + +#[cfg(feature = "rocksdb")] +impl EmbeddedEngine> { /// Start engine with an explicit data directory. /// /// `data_dir` has highest priority and always overrides `cluster.db_root_dir` from @@ -199,7 +216,6 @@ impl EmbeddedEngine { /// // Works with any AsRef /// let engine = EmbeddedEngine::start(std::path::Path::new("/var/lib/my-app")).await?; /// ``` - #[cfg(feature = "rocksdb")] pub async fn start(data_dir: impl AsRef) -> Result { let mut config = d_engine_core::RaftNodeConfig::new()?; config.cluster.db_root_dir = data_dir.as_ref().to_path_buf(); @@ -227,7 +243,7 @@ impl EmbeddedEngine { (storage, sm) }; - let lease = Arc::new(crate::storage::DefaultLease::new( + let lease = Arc::new(crate::storage::TtlLease::new( config.raft.state_machine.lease.clone(), )); sm.set_lease(lease); @@ -248,7 +264,6 @@ impl EmbeddedEngine { /// let engine = EmbeddedEngine::start_with("config/node1.toml").await?; /// engine.wait_ready(Duration::from_secs(5)).await?; /// ``` - #[cfg(feature = "rocksdb")] pub async fn start_with(config_path: &str) -> Result { let config = d_engine_core::RaftNodeConfig::new()? .with_override_config(config_path)? @@ -276,14 +291,22 @@ impl EmbeddedEngine { (storage, sm) }; - let lease = Arc::new(crate::storage::DefaultLease::new( + let lease = Arc::new(crate::storage::TtlLease::new( config.raft.state_machine.lease.clone(), )); sm.set_lease(lease); Self::start_custom(Arc::new(storage), Arc::new(sm), Some(config_path)).await } +} + +// ─── Custom SE/SM: start_custom + start_node ───────────────────────────────── +impl EmbeddedEngine> +where + SE: StorageEngine + Debug + 'static, + SM: StateMachine + Debug + 'static, +{ /// Start engine with custom storage and state machine. /// /// Advanced API for users providing custom storage implementations. @@ -299,15 +322,11 @@ impl EmbeddedEngine { /// let sm = Arc::new(MyCustomStateMachine::new()?); /// let engine = EmbeddedEngine::start_custom(storage, sm, None).await?; /// ``` - pub async fn start_custom( + pub async fn start_custom( storage_engine: Arc, state_machine: Arc, config_path: Option<&str>, - ) -> Result - where - SE: StorageEngine + std::fmt::Debug + 'static, - SM: StateMachine + std::fmt::Debug + 'static, - { + ) -> Result { let node_config = if let Some(path) = config_path { d_engine_core::RaftNodeConfig::default() .with_override_config(path)? @@ -320,19 +339,23 @@ impl EmbeddedEngine { } /// Build and launch the node from a validated config and pre-built storage. - async fn start_node( + async fn start_node( node_config: d_engine_core::RaftNodeConfig, storage_engine: Arc, state_machine: Arc, - ) -> Result - where - SE: StorageEngine + std::fmt::Debug + 'static, - SM: StateMachine + std::fmt::Debug + 'static, - { + ) -> Result { info!("Starting embedded d-engine"); + d_engine_core::init_clock(); let (shutdown_tx, shutdown_rx) = watch::channel(()); + // Clone SM before NodeBuilder consumes it — EmbeddedReadHandle holds this + // for the direct-SM fast path. LOCK release is decoupled from Arc counting + // (via close_db()), so holding multiple Arc clones is safe. + let sm_for_client = Arc::clone(&state_machine); + + let sm_for_engine = Arc::clone(&state_machine); + let node = NodeBuilder::init(node_config, shutdown_rx) .storage_engine(storage_engine) .state_machine(state_machine) @@ -342,27 +365,27 @@ impl EmbeddedEngine { let leader_elected_rx = node.leader_change_notifier(); let membership_rx = node.membership_change_notifier(); - #[cfg(not(feature = "watch"))] - let client = Arc::new(EmbeddedClient::new_internal( - node.event_tx.clone(), - node.cmd_tx.clone(), - node.node_id, - Duration::from_millis(node.node_config.raft.general_raft_timeout_duration_in_ms), - )); + let read_handle = + EmbeddedReadHandle::new(sm_for_client, node.read_lease(), node.cmd_tx.clone()); - #[cfg(feature = "watch")] let client = { - let watch_registry = node.watch_registry.clone(); - let mut client = EmbeddedClient::new_internal( + let base = EmbeddedClient::new_internal( node.event_tx.clone(), - node.cmd_tx.clone(), + read_handle, node.node_id, Duration::from_millis(node.node_config.raft.general_raft_timeout_duration_in_ms), ); - if let Some(registry) = &watch_registry { - client = client.with_watch_registry(registry.clone()); - } - Arc::new(client) + + #[cfg(feature = "watch")] + let base = { + let mut c = base; + if let Some(registry) = &node.watch_registry { + c = c.with_watch_registry(registry.clone()); + } + c + }; + + Arc::new(base) }; let node_id = node.node_id(); @@ -380,6 +403,7 @@ impl EmbeddedEngine { Ok(Self { inner: Arc::new(Inner { + sm: sm_for_engine, node_handle: Mutex::new(Some(node_handle)), shutdown_tx, client, @@ -390,7 +414,11 @@ impl EmbeddedEngine { }), }) } +} + +// ─── General API (all TypeConfig) ──────────────────────────────────────────── +impl EmbeddedEngine { /// Wait until the cluster is ready to serve requests. /// /// Blocks until a leader has been elected **and** its no-op entry is committed @@ -590,7 +618,7 @@ impl EmbeddedEngine { /// let client = engine.client(); /// client.put(b"key", b"value").await?; /// ``` - pub fn client(&self) -> Arc { + pub fn client(&self) -> Arc> { Arc::clone(&self.inner.client) } @@ -619,7 +647,12 @@ impl EmbeddedEngine { info!("Stopping embedded d-engine"); - // Send shutdown signal + // Close SM storage resources first (releases RocksDB LOCK immediately). + // Must happen before the Raft loop exits so the LOCK is freed before any + // subsequent engine restart on the same data directory. + self.inner.sm.close_storage(); + + // Send shutdown signal to Raft node let _ = self.inner.shutdown_tx.send(()); // Wait for node task to complete @@ -661,7 +694,7 @@ impl EmbeddedEngine { } } -impl Drop for EmbeddedEngine { +impl Drop for EmbeddedEngine { fn drop(&mut self) { // Warn if stop() was not called if let Ok(handle) = self.inner.node_handle.try_lock() @@ -672,3 +705,7 @@ impl Drop for EmbeddedEngine { } } } + +#[cfg(test)] +#[path = "embedded_test/mod.rs"] +mod tests; diff --git a/d-engine-server/src/api/embedded_client.rs b/d-engine-server/src/api/embedded_client.rs index c00ac7e9..c9bf3c56 100644 --- a/d-engine-server/src/api/embedded_client.rs +++ b/d-engine-server/src/api/embedded_client.rs @@ -15,129 +15,72 @@ //! client.put(b"key", b"value").await?; //! ``` -use std::time::Duration; - #[cfg(feature = "watch")] use std::sync::Arc; +use std::time::Duration; use bytes::Bytes; use d_engine_core::MaybeCloneOneshot; use d_engine_core::RaftEvent; use d_engine_core::RaftOneshot; use d_engine_core::ScanResult; +use d_engine_core::TypeConfig; use d_engine_core::client::{ - ClientApi, ClientApiError, ClientApiResult, ClientReadRequest, ClientResponsePayload, - ClientWriteRequest, ErrorCode, LeaderHint, ReadResults, WriteOperation, + ClientApi, ClientApiResult, ClientResponsePayload, ClientWriteRequest, ErrorCode, + WriteOperation, }; use d_engine_core::config::ReadConsistencyPolicy; use tokio::sync::mpsc; +use super::embedded_read_handle::EmbeddedReadHandle; +pub(crate) use super::standalone_read_handle::{ + channel_closed_error, map_error_response, server_error, timeout_error, +}; + +#[cfg(feature = "watch")] +use d_engine_core::client::ClientApiError; #[cfg(feature = "watch")] use d_engine_core::watch::WatchRegistry; -// ============================================================================ -// Error helpers - simplify ClientApiError construction for embedded client -// ============================================================================ - -fn channel_closed_error() -> ClientApiError { - ClientApiError::Network { - code: ErrorCode::ConnectionTimeout, - message: "Channel closed, node may be shutting down".to_string(), - retry_after_ms: None, - leader_hint: None, - } -} - -fn timeout_error(duration: Duration) -> ClientApiError { - ClientApiError::Network { - code: ErrorCode::ConnectionTimeout, - message: format!("Operation timed out after {duration:?}"), - retry_after_ms: Some(1000), - leader_hint: None, - } -} - -fn not_leader_error( - leader_id: Option, - leader_address: Option, - retry_after_ms: Option, -) -> ClientApiError { - let message = match (&leader_address, &leader_id) { - (Some(addr), _) => format!("Not leader, try leader at: {addr}"), - (None, Some(id)) => format!("Not leader, leader_id: {id}"), - (None, None) => "Not leader".to_string(), - }; - - let leader_hint = match (&leader_id, &leader_address) { - (Some(id_str), Some(addr)) => id_str.parse::().ok().map(|id| LeaderHint { - leader_id: id, - address: addr.clone(), - }), - _ => None, - }; - - ClientApiError::Network { - code: ErrorCode::NotLeader, - message, - retry_after_ms: retry_after_ms.or(Some(100)), - leader_hint, - } -} - -fn server_error(msg: String) -> ClientApiError { - ClientApiError::Business { - code: ErrorCode::Uncategorized, - message: msg, - required_action: None, - } -} - -/// Unwrap a `ClientResponsePayload` as a `ReadResults`, returning a -/// `Protocol { InvalidResponse }` error for any other variant. -/// -/// Centralises the match so both `get_with_consistency` and -/// `get_multi_with_consistency` share identical error semantics, and so -/// the logic can be unit-tested without standing up a full Raft channel. -fn extract_read_payload(result: Option) -> ClientApiResult { - match result { - Some(ClientResponsePayload::Read(r)) => Ok(r), - Some(ClientResponsePayload::Write(_)) => Err(ClientApiError::Protocol { - code: ErrorCode::InvalidResponse, - message: "expected ReadData payload, got WriteResult".to_string(), - supported_versions: None, - }), - None => Err(ClientApiError::Protocol { - code: ErrorCode::InvalidResponse, - message: "expected ReadData payload, got None".to_string(), - supported_versions: None, - }), - } -} - /// Zero-overhead KV client for embedded mode. Obtained via `EmbeddedEngine::client()`. /// +/// `T` is the [`TypeConfig`] that carries the concrete `StateMachine` type, enabling +/// direct monomorphized calls to `T::SM::get()` on the fast path — no vtable dispatch. +/// /// For standalone/gRPC mode use `GrpcClient` instead. Both implement `ClientApi`. -#[derive(Clone)] -pub struct EmbeddedClient { +pub struct EmbeddedClient { event_tx: mpsc::Sender, - cmd_tx: mpsc::Sender, + read_handle: EmbeddedReadHandle, client_id: u32, timeout: Duration, #[cfg(feature = "watch")] watch_registry: Option>, } -impl EmbeddedClient { - /// Internal constructor (used by EmbeddedEngine) +impl Clone for EmbeddedClient { + fn clone(&self) -> Self { + Self { + event_tx: self.event_tx.clone(), + read_handle: self.read_handle.clone(), + client_id: self.client_id, + timeout: self.timeout, + #[cfg(feature = "watch")] + watch_registry: self.watch_registry.clone(), + } + } +} + +impl EmbeddedClient { + /// Internal constructor (used by EmbeddedEngine). pub(crate) fn new_internal( event_tx: mpsc::Sender, - cmd_tx: mpsc::Sender, + read_handle: EmbeddedReadHandle, client_id: u32, timeout: Duration, ) -> Self { Self { event_tx, - cmd_tx, + read_handle, client_id, timeout, #[cfg(feature = "watch")] @@ -155,24 +98,6 @@ impl EmbeddedClient { self } - fn map_error_response( - error: ErrorCode, - leader_hint: Option, - retry_after_ms: Option, - ) -> ClientApiError { - match error { - ErrorCode::NotLeader => { - let (leader_id, leader_address) = if let Some(hint) = leader_hint { - (Some(hint.leader_id.to_string()), Some(hint.address)) - } else { - (None, None) - }; - not_leader_error(leader_id, leader_address, retry_after_ms) - } - _ => server_error(format!("Error code: {error:?}")), - } - } - /// Store a key-value pair with strong consistency. /// /// # Errors @@ -194,7 +119,8 @@ impl EmbeddedClient { let (resp_tx, resp_rx) = MaybeCloneOneshot::new(); - self.cmd_tx + self.read_handle + .cmd_tx .send(d_engine_core::ClientCmd::Propose(request, resp_tx)) .await .map_err(|_| channel_closed_error())?; @@ -208,7 +134,7 @@ impl EmbeddedClient { result.map_err(|status| server_error(format!("RPC error: {}", status.message())))?; if response.error != ErrorCode::Success { - return Err(Self::map_error_response( + return Err(map_error_response( response.error, response.leader_hint, response.retry_after_ms, @@ -291,37 +217,9 @@ impl EmbeddedClient { key: impl AsRef<[u8]>, consistency: ReadConsistencyPolicy, ) -> ClientApiResult> { - let request = ClientReadRequest { - client_id: self.client_id, - keys: vec![Bytes::copy_from_slice(key.as_ref())], - consistency_policy: Some(consistency), - }; - - let (resp_tx, resp_rx) = MaybeCloneOneshot::new(); - - self.cmd_tx - .send(d_engine_core::ClientCmd::Read(request, resp_tx)) + self.read_handle + .get(key.as_ref(), consistency, self.client_id, self.timeout) .await - .map_err(|_| channel_closed_error())?; - - let result = tokio::time::timeout(self.timeout, resp_rx) - .await - .map_err(|_| timeout_error(self.timeout))? - .map_err(|_| channel_closed_error())?; - - let response = - result.map_err(|status| server_error(format!("RPC error: {}", status.message())))?; - - if response.error != ErrorCode::Success { - return Err(Self::map_error_response( - response.error, - response.leader_hint, - response.retry_after_ms, - )); - } - - let read_results = extract_read_payload(response.result)?; - Ok(read_results.entries.first().map(|e| e.value.clone())) } /// Get multiple keys with linearizable consistency. @@ -359,47 +257,17 @@ impl EmbeddedClient { } /// Advanced: Get multiple keys with explicit consistency policy. + /// + /// Eventual and LeaseRead policies use the ReadActor fast path (no Raft round-trip). + /// LinearizableRead always goes through the Raft readIndex protocol via cmd_tx. pub async fn get_multi_with_consistency( &self, keys: &[Bytes], consistency: ReadConsistencyPolicy, ) -> ClientApiResult>> { - let request = ClientReadRequest { - client_id: self.client_id, - keys: keys.to_vec(), - consistency_policy: Some(consistency), - }; - - let (resp_tx, resp_rx) = MaybeCloneOneshot::new(); - - self.cmd_tx - .send(d_engine_core::ClientCmd::Read(request, resp_tx)) + self.read_handle + .get_batch(keys, consistency, self.client_id, self.timeout) .await - .map_err(|_| channel_closed_error())?; - - let result = tokio::time::timeout(self.timeout, resp_rx) - .await - .map_err(|_| timeout_error(self.timeout))? - .map_err(|_| channel_closed_error())?; - - let response = - result.map_err(|status| server_error(format!("RPC error: {}", status.message())))?; - - if response.error != ErrorCode::Success { - return Err(Self::map_error_response( - response.error, - response.leader_hint, - response.retry_after_ms, - )); - } - - let read_results = extract_read_payload(response.result)?; - // Reconstruct result vector in requested key order. - // Server only returns results for keys that exist, so we must - // map by key to preserve positional correspondence with input. - let results_by_key: std::collections::HashMap<_, _> = - read_results.entries.into_iter().map(|e| (e.key, e.value)).collect(); - Ok(keys.iter().map(|k| results_by_key.get(k).cloned()).collect()) } /// Delete a key-value pair with strong consistency. @@ -420,7 +288,8 @@ impl EmbeddedClient { let (resp_tx, resp_rx) = MaybeCloneOneshot::new(); - self.cmd_tx + self.read_handle + .cmd_tx .send(d_engine_core::ClientCmd::Propose(request, resp_tx)) .await .map_err(|_| channel_closed_error())?; @@ -434,7 +303,7 @@ impl EmbeddedClient { result.map_err(|status| server_error(format!("RPC error: {}", status.message())))?; if response.error != ErrorCode::Success { - return Err(Self::map_error_response( + return Err(map_error_response( response.error, response.leader_hint, response.retry_after_ms, @@ -636,7 +505,8 @@ impl EmbeddedClient { ) -> ClientApiResult { let (resp_tx, resp_rx) = MaybeCloneOneshot::new(); - self.cmd_tx + self.read_handle + .cmd_tx .send(d_engine_core::ClientCmd::Scan( Bytes::copy_from_slice(prefix.as_ref()), resp_tx, @@ -674,7 +544,7 @@ impl EmbeddedClient { } } -impl std::fmt::Debug for EmbeddedClient { +impl std::fmt::Debug for EmbeddedClient { fn fmt( &self, f: &mut std::fmt::Formatter<'_>, @@ -688,7 +558,7 @@ impl std::fmt::Debug for EmbeddedClient { // Implement ClientApi trait #[async_trait::async_trait] -impl ClientApi for EmbeddedClient { +impl ClientApi for EmbeddedClient { async fn put( &self, key: impl AsRef<[u8]> + Send, @@ -714,7 +584,8 @@ impl ClientApi for EmbeddedClient { let (resp_tx, resp_rx) = MaybeCloneOneshot::new(); - self.cmd_tx + self.read_handle + .cmd_tx .send(d_engine_core::ClientCmd::Propose(request, resp_tx)) .await .map_err(|_| channel_closed_error())?; @@ -728,7 +599,7 @@ impl ClientApi for EmbeddedClient { result.map_err(|status| server_error(format!("RPC error: {}", status.message())))?; if response.error != ErrorCode::Success { - return Err(Self::map_error_response( + return Err(map_error_response( response.error, response.leader_hint, response.retry_after_ms, @@ -776,7 +647,8 @@ impl ClientApi for EmbeddedClient { let (resp_tx, resp_rx) = MaybeCloneOneshot::new(); - self.cmd_tx + self.read_handle + .cmd_tx .send(d_engine_core::ClientCmd::Propose(request, resp_tx)) .await .map_err(|_| channel_closed_error())?; @@ -790,7 +662,7 @@ impl ClientApi for EmbeddedClient { result.map_err(|status| server_error(format!("RPC error: {}", status.message())))?; if response.error != ErrorCode::Success { - return Err(Self::map_error_response( + return Err(map_error_response( response.error, response.leader_hint, response.retry_after_ms, @@ -858,172 +730,5 @@ impl ClientApi for EmbeddedClient { } #[cfg(test)] -mod error_helper_tests { - use d_engine_core::client::KvEntry; - - use super::*; - - // ─── not_leader_error ──────────────────────────────────────────────────── - - #[test] - fn test_not_leader_uses_server_retry_after_ms_when_provided() { - let err = not_leader_error( - Some("1".to_string()), - Some("127.0.0.1:5001".to_string()), - Some(500), - ); - match err { - ClientApiError::Network { retry_after_ms, .. } => { - assert_eq!(retry_after_ms, Some(500)); - } - _ => panic!("expected Network error"), - } - } - - #[test] - fn test_not_leader_falls_back_to_100ms_when_server_provides_none() { - let err = not_leader_error(None, None, None); - match err { - ClientApiError::Network { retry_after_ms, .. } => { - assert_eq!(retry_after_ms, Some(100)); - } - _ => panic!("expected Network error"), - } - } - - #[test] - fn test_not_leader_zero_is_not_treated_as_none() { - // Some(0) is an explicit server instruction, not absent — must be preserved - let err = not_leader_error(None, None, Some(0)); - match err { - ClientApiError::Network { retry_after_ms, .. } => { - assert_eq!(retry_after_ms, Some(0)); - } - _ => panic!("expected Network error"), - } - } - - // ─── map_error_response ────────────────────────────────────────────────── - - #[test] - fn test_map_error_response_not_leader_forwards_retry_after_ms() { - let err = EmbeddedClient::map_error_response( - ErrorCode::NotLeader, - Some(LeaderHint { - leader_id: 2, - address: "10.0.0.2:5002".into(), - }), - Some(250), - ); - match err { - ClientApiError::Network { - code, - retry_after_ms, - leader_hint, - .. - } => { - assert_eq!(code, ErrorCode::NotLeader); - assert_eq!(retry_after_ms, Some(250)); - let h = leader_hint.unwrap(); - assert_eq!(h.leader_id, 2); - } - _ => panic!("expected Network error"), - } - } - - #[test] - fn test_map_error_response_not_leader_falls_back_to_100ms_when_none() { - let err = EmbeddedClient::map_error_response(ErrorCode::NotLeader, None, None); - match err { - ClientApiError::Network { retry_after_ms, .. } => { - assert_eq!(retry_after_ms, Some(100)); - } - _ => panic!("expected Network error"), - } - } - - // ─── extract_read_payload ──────────────────────────────────────────────── - // - // These tests call `extract_read_payload` directly — the actual function - // used by both `get_with_consistency` and `get_multi_with_consistency`. - // This ensures the error behaviour is tested at the implementation site, - // not via a copy of the match logic that could silently drift. - - #[test] - fn test_extract_read_payload_returns_read_results_on_success() { - // Happy path: a well-formed Read payload must be unwrapped without error. - let entries = vec![KvEntry { - key: Bytes::from("k"), - value: Bytes::from("v"), - }]; - let payload = Some(ClientResponsePayload::Read(ReadResults { - entries: entries.clone(), - })); - - let result = extract_read_payload(payload).unwrap(); - assert_eq!(result.entries.len(), 1); - assert_eq!(result.entries[0].key, Bytes::from("k")); - assert_eq!(result.entries[0].value, Bytes::from("v")); - } - - #[test] - fn test_extract_read_payload_rejects_write_result_payload() { - // A WriteResult inside a read response is a server protocol violation. - // It must surface as InvalidResponse, not silently become "key not found". - let payload = Some(ClientResponsePayload::Write( - d_engine_core::client::WriteResult { succeeded: true }, - )); - - let err = extract_read_payload(payload).unwrap_err(); - assert_eq!(err.code(), ErrorCode::InvalidResponse); - assert!( - err.message().contains("WriteResult"), - "error message should identify the unexpected variant; got: {}", - err.message() - ); - } - - #[test] - fn test_extract_read_payload_rejects_none_payload() { - // A Success response with no payload is a protocol violation. - // It must surface as InvalidResponse, not silently become "key not found". - let err = extract_read_payload(None).unwrap_err(); - assert_eq!(err.code(), ErrorCode::InvalidResponse); - assert!( - err.message().contains("None"), - "error message should identify missing payload; got: {}", - err.message() - ); - } - - #[test] - fn test_extract_read_payload_empty_entries_is_valid() { - // An empty ReadResults is a legitimate response (no keys matched). - // It must not be treated as an error. - let payload = Some(ClientResponsePayload::Read(ReadResults { entries: vec![] })); - - let result = extract_read_payload(payload).unwrap(); - assert!(result.entries.is_empty()); - } - - #[test] - fn test_extract_read_payload_multiple_entries_are_preserved() { - // All entries in the ReadResults must be passed through unchanged. - let payload = Some(ClientResponsePayload::Read(ReadResults { - entries: vec![ - KvEntry { - key: Bytes::from("k1"), - value: Bytes::from("v1"), - }, - KvEntry { - key: Bytes::from("k2"), - value: Bytes::from("v2"), - }, - ], - })); - - let result = extract_read_payload(payload).unwrap(); - assert_eq!(result.entries.len(), 2); - assert_eq!(result.entries[1].key, Bytes::from("k2")); - } -} +#[path = "embedded_client_test/embedded_client_test.rs"] +mod tests; diff --git a/d-engine-server/src/api/embedded_client_test.rs b/d-engine-server/src/api/embedded_client_test/embedded_client_test.rs similarity index 72% rename from d-engine-server/src/api/embedded_client_test.rs rename to d-engine-server/src/api/embedded_client_test/embedded_client_test.rs index f82d9e6d..6cf68037 100644 --- a/d-engine-server/src/api/embedded_client_test.rs +++ b/d-engine-server/src/api/embedded_client_test/embedded_client_test.rs @@ -1,105 +1,101 @@ -#[cfg(all(test, feature = "rocksdb"))] -mod tests { - use std::collections::HashMap; - - use bytes::Bytes; - - /// Test that get_multi reconstructs results in correct key order. - /// - /// Verifies the critical fix: when server returns only results for - /// existing keys (sparse), we must map by key to preserve positional - /// correspondence with the input key vector. - #[test] - fn test_get_multi_result_reconstruction() { - // Simulate server response scenario: - // Request: [key1, key2, key3] - // Exists: [key1, key3] (key2 is missing) - // Server returns: [key1→value1, key3→value3] - - let requested_keys: Vec<_> = [ - Bytes::from("key1"), - Bytes::from("key2"), - Bytes::from("key3"), - ] - .to_vec(); - let server_results: Vec<_> = vec![ - (Bytes::from("key1"), Bytes::from("value1")), - (Bytes::from("key3"), Bytes::from("value3")), - ]; - - // Simulate the fixed reconstruction logic - let results_by_key: HashMap<_, _> = server_results.into_iter().collect(); - let reconstructed: Vec> = - requested_keys.iter().map(|k| results_by_key.get(k).cloned()).collect(); - - // Expected: [Some(value1), None, Some(value3)] - assert_eq!( - reconstructed.len(), - 3, - "Result count must match request count" - ); - assert_eq!( - reconstructed[0], - Some(Bytes::from("value1")), - "Position 0 is key1" - ); - assert_eq!(reconstructed[1], None, "Position 1 is key2 (missing)"); - assert_eq!( - reconstructed[2], - Some(Bytes::from("value3")), - "Position 2 is key3" - ); - } +use std::collections::HashMap; + +use bytes::Bytes; + +/// Test that get_multi reconstructs results in correct key order. +/// +/// Verifies the critical fix: when server returns only results for +/// existing keys (sparse), we must map by key to preserve positional +/// correspondence with the input key vector. +#[test] +fn test_get_multi_result_reconstruction() { + // Simulate server response scenario: + // Request: [key1, key2, key3] + // Exists: [key1, key3] (key2 is missing) + // Server returns: [key1→value1, key3→value3] + + let requested_keys: Vec<_> = [ + Bytes::from("key1"), + Bytes::from("key2"), + Bytes::from("key3"), + ] + .to_vec(); + let server_results: Vec<_> = vec![ + (Bytes::from("key1"), Bytes::from("value1")), + (Bytes::from("key3"), Bytes::from("value3")), + ]; + + // Simulate the fixed reconstruction logic + let results_by_key: HashMap<_, _> = server_results.into_iter().collect(); + let reconstructed: Vec> = + requested_keys.iter().map(|k| results_by_key.get(k).cloned()).collect(); + + // Expected: [Some(value1), None, Some(value3)] + assert_eq!( + reconstructed.len(), + 3, + "Result count must match request count" + ); + assert_eq!( + reconstructed[0], + Some(Bytes::from("value1")), + "Position 0 is key1" + ); + assert_eq!(reconstructed[1], None, "Position 1 is key2 (missing)"); + assert_eq!( + reconstructed[2], + Some(Bytes::from("value3")), + "Position 2 is key3" + ); +} - /// Test edge case: all requested keys exist. - #[test] - fn test_get_multi_all_keys_exist() { - let requested_keys: Vec<_> = [Bytes::from("a"), Bytes::from("b")].to_vec(); - let server_results: Vec<_> = vec![ - (Bytes::from("a"), Bytes::from("1")), - (Bytes::from("b"), Bytes::from("2")), - ]; - - let results_by_key: HashMap<_, _> = server_results.into_iter().collect(); - let reconstructed: Vec> = - requested_keys.iter().map(|k| results_by_key.get(k).cloned()).collect(); - - assert_eq!(reconstructed.len(), 2); - assert_eq!(reconstructed[0], Some(Bytes::from("1"))); - assert_eq!(reconstructed[1], Some(Bytes::from("2"))); - } +/// Test edge case: all requested keys exist. +#[test] +fn test_get_multi_all_keys_exist() { + let requested_keys: Vec<_> = [Bytes::from("a"), Bytes::from("b")].to_vec(); + let server_results: Vec<_> = vec![ + (Bytes::from("a"), Bytes::from("1")), + (Bytes::from("b"), Bytes::from("2")), + ]; + + let results_by_key: HashMap<_, _> = server_results.into_iter().collect(); + let reconstructed: Vec> = + requested_keys.iter().map(|k| results_by_key.get(k).cloned()).collect(); + + assert_eq!(reconstructed.len(), 2); + assert_eq!(reconstructed[0], Some(Bytes::from("1"))); + assert_eq!(reconstructed[1], Some(Bytes::from("2"))); +} - /// Test edge case: no requested keys exist. - #[test] - fn test_get_multi_no_keys_exist() { - let requested_keys: Vec<_> = - [Bytes::from("x"), Bytes::from("y"), Bytes::from("z")].to_vec(); - let server_results: Vec<(Bytes, Bytes)> = Vec::new(); +/// Test edge case: no requested keys exist. +#[test] +fn test_get_multi_no_keys_exist() { + let requested_keys: Vec<_> = [Bytes::from("x"), Bytes::from("y"), Bytes::from("z")].to_vec(); + let server_results: Vec<(Bytes, Bytes)> = Vec::new(); - let results_by_key: HashMap<_, _> = server_results.into_iter().collect(); - let reconstructed: Vec> = - requested_keys.iter().map(|k| results_by_key.get(k).cloned()).collect(); + let results_by_key: HashMap<_, _> = server_results.into_iter().collect(); + let reconstructed: Vec> = + requested_keys.iter().map(|k| results_by_key.get(k).cloned()).collect(); - assert_eq!(reconstructed.len(), 3); - assert!(reconstructed.iter().all(|r| r.is_none())); - } + assert_eq!(reconstructed.len(), 3); + assert!(reconstructed.iter().all(|r| r.is_none())); +} - /// Test that empty byte values are preserved correctly. - #[test] - fn test_get_multi_preserves_empty_values() { - let requested_keys: Vec<_> = [Bytes::from("empty"), Bytes::from("nonempty")].to_vec(); - let server_results: Vec<_> = vec![ - (Bytes::from("empty"), Bytes::new()), - (Bytes::from("nonempty"), Bytes::from("v")), - ]; - - let results_by_key: HashMap<_, _> = server_results.into_iter().collect(); - let reconstructed: Vec> = - requested_keys.iter().map(|k| results_by_key.get(k).cloned()).collect(); - - assert_eq!(reconstructed[0], Some(Bytes::new())); - assert_eq!(reconstructed[1], Some(Bytes::from("v"))); - } +/// Test that empty byte values are preserved correctly. +#[test] +fn test_get_multi_preserves_empty_values() { + let requested_keys: Vec<_> = [Bytes::from("empty"), Bytes::from("nonempty")].to_vec(); + let server_results: Vec<_> = vec![ + (Bytes::from("empty"), Bytes::new()), + (Bytes::from("nonempty"), Bytes::from("v")), + ]; + + let results_by_key: HashMap<_, _> = server_results.into_iter().collect(); + let reconstructed: Vec> = + requested_keys.iter().map(|k| results_by_key.get(k).cloned()).collect(); + + assert_eq!(reconstructed[0], Some(Bytes::new())); + assert_eq!(reconstructed[1], Some(Bytes::from("v"))); } // ============================================================================= @@ -113,10 +109,10 @@ mod integration_tests { use d_engine_core::ClientApi; use tempfile::TempDir; - use crate::api::EmbeddedEngine; + use crate::api::DefaultEmbeddedEngine; /// Helper to create a test EmbeddedEngine - async fn create_test_engine() -> (EmbeddedEngine, TempDir) { + async fn create_test_engine() -> (DefaultEmbeddedEngine, TempDir) { let temp_dir = TempDir::new().expect("Failed to create temp dir"); let db_path = temp_dir.path().join("db"); @@ -145,7 +141,7 @@ max_batch_size = 1 ); std::fs::write(&config_path, config_content).expect("Failed to write config"); - let engine = EmbeddedEngine::start_with(config_path.to_str().unwrap()) + let engine = DefaultEmbeddedEngine::start_with(config_path.to_str().unwrap()) .await .expect("Failed to start engine"); @@ -319,7 +315,7 @@ max_batch_size = 1 engine.stop().await.expect("Failed to stop engine"); } - } + } // cas_operations mod scan_prefix_operations { use super::*; @@ -382,5 +378,5 @@ max_batch_size = 1 engine.stop().await.expect("Failed to stop engine"); } - } -} + } // scan_prefix_operations +} // integration_tests diff --git a/d-engine-server/src/api/embedded_client_test/mod.rs b/d-engine-server/src/api/embedded_client_test/mod.rs new file mode 100644 index 00000000..6548344e --- /dev/null +++ b/d-engine-server/src/api/embedded_client_test/mod.rs @@ -0,0 +1,2 @@ +#[cfg(test)] +mod embedded_client_test; diff --git a/d-engine-server/src/api/embedded_read_handle.rs b/d-engine-server/src/api/embedded_read_handle.rs new file mode 100644 index 00000000..b7c9cfe3 --- /dev/null +++ b/d-engine-server/src/api/embedded_read_handle.rs @@ -0,0 +1,176 @@ +//! EmbeddedReadHandle — zero-channel read path for embedded mode. +//! +//! Replaces `ReadHandle` (which routes through the ReadActor channel) for +//! `EmbeddedClient`. Eventual and LeaseRead call `SM::get_multi()` directly — +//! no oneshot allocation, no task context switch, no channel round-trip. +//! +//! # Routing +//! +//! ```text +//! Eventual / LeaseRead (valid) → sm.get_multi() (direct, zero channel) +//! Eventual / LeaseRead (error) → cmd_tx fallback (SM stopped or error) +//! LeaseRead (invalid lease) → cmd_tx fallback +//! LinearizableRead → cmd_tx always (Raft readIndex) +//! ``` +//! +//! # Shutdown safety +//! +//! After `EmbeddedEngine::stop()` calls `sm.close_db()`, `sm.get_multi()` returns +//! `NotServing`. The handle falls through to `cmd_tx` which is also closed at +//! that point, so callers receive a `channel_closed_error` — the same behaviour +//! as any other shutdown scenario. + +use std::marker::PhantomData; +use std::time::Duration; + +use bytes::Bytes; +use d_engine_core::MaybeCloneOneshot; +use d_engine_core::RaftOneshot; +use d_engine_core::ReadLease; +use d_engine_core::StateMachine; +use d_engine_core::TypeConfig; +use d_engine_core::client::{ClientApiResult, ClientReadRequest, ErrorCode}; +use d_engine_core::config::ReadConsistencyPolicy; +use d_engine_core::now_ms; +use std::sync::Arc; +use tokio::sync::mpsc; + +use super::standalone_read_handle::{ + channel_closed_error, extract_read_payload, map_error_response, server_error, timeout_error, +}; + +// ── EmbeddedReadHandle ──────────────────────────────────────────────────────── + +/// Direct-SM read handle for embedded mode. +/// +/// Cloneable — `EmbeddedClient` and its clones each hold their own instance. +/// Cloning is cheap: all fields are `Arc` or `mpsc::Sender` (reference-counted). +pub(crate) struct EmbeddedReadHandle { + sm: Arc, + lease: Arc, + pub(crate) cmd_tx: mpsc::Sender, + _phantom: PhantomData T>, +} + +impl Clone for EmbeddedReadHandle { + fn clone(&self) -> Self { + Self { + sm: Arc::clone(&self.sm), + lease: Arc::clone(&self.lease), + cmd_tx: self.cmd_tx.clone(), + _phantom: PhantomData, + } + } +} + +impl EmbeddedReadHandle { + pub(crate) fn new( + sm: Arc, + lease: Arc, + cmd_tx: mpsc::Sender, + ) -> Self { + Self { + sm, + lease, + cmd_tx, + _phantom: PhantomData, + } + } + + /// Single-key read. Convenience wrapper around [`Self::get_batch`]. + pub async fn get( + &self, + key: &[u8], + consistency: ReadConsistencyPolicy, + client_id: u32, + timeout: Duration, + ) -> ClientApiResult> { + let key_bytes = Bytes::copy_from_slice(key); + let mut results = self.get_batch(&[key_bytes], consistency, client_id, timeout).await?; + Ok(results.pop().flatten()) + } + + /// Multi-key read with consistency routing. + /// + /// For `Eventual` and `LeaseRead` (when the lease is valid), all keys are + /// served from a single `SM::get_multi()` snapshot — no Raft round-trip. + /// On any SM error, or when the lease is invalid, the request falls through + /// to `cmd_tx` (Raft readIndex protocol). + pub async fn get_batch( + &self, + keys: &[Bytes], + consistency: ReadConsistencyPolicy, + client_id: u32, + timeout: Duration, + ) -> ClientApiResult>> { + match consistency { + ReadConsistencyPolicy::EventualConsistency => { + if let Ok(values) = self.sm.get_multi(keys) { + return Ok(values); + } + // SM stopped or hard error → fall through to cmd_tx + } + ReadConsistencyPolicy::LeaseRead => { + if self.lease.is_valid(now_ms()) + && let Ok(values) = self.sm.get_multi(keys) + { + return Ok(values); + + // SM stopped or hard error → fall through + } + // invalid lease → fall through to cmd_tx + } + ReadConsistencyPolicy::LinearizableRead => { + // always cmd_tx — Raft readIndex protocol + } + } + + self.cmd_tx_path(keys, consistency, client_id, timeout).await + } + + // ── cmd_tx fallback ─────────────────────────────────────────────────────── + + async fn cmd_tx_path( + &self, + keys: &[Bytes], + consistency: ReadConsistencyPolicy, + client_id: u32, + timeout: Duration, + ) -> ClientApiResult>> { + let request = ClientReadRequest { + client_id, + keys: keys.to_vec(), + consistency_policy: Some(consistency), + }; + let (resp_tx, resp_rx) = MaybeCloneOneshot::new(); + self.cmd_tx + .send(d_engine_core::ClientCmd::Read(request, resp_tx)) + .await + .map_err(|_| channel_closed_error())?; + + let result = tokio::time::timeout(timeout, resp_rx) + .await + .map_err(|_| timeout_error(timeout))? + .map_err(|_| channel_closed_error())?; + + let response = + result.map_err(|status| server_error(format!("RPC error: {}", status.message())))?; + + if response.error != ErrorCode::Success { + return Err(map_error_response( + response.error, + response.leader_hint, + response.retry_after_ms, + )); + } + + let read_results = extract_read_payload(response.result)?; + let entry_map: std::collections::HashMap = + read_results.entries.into_iter().map(|e| (e.key, e.value)).collect(); + Ok(keys.iter().map(|k| entry_map.get(k).cloned()).collect()) + } +} + +#[cfg(test)] +#[path = "embedded_read_handle_test.rs"] +mod tests; diff --git a/d-engine-server/src/api/embedded_read_handle_test.rs b/d-engine-server/src/api/embedded_read_handle_test.rs new file mode 100644 index 00000000..c8ec8461 --- /dev/null +++ b/d-engine-server/src/api/embedded_read_handle_test.rs @@ -0,0 +1,199 @@ +/// Tests for EmbeddedReadHandle — the zero-channel direct-SM read path. +/// +/// Proof technique (same as embedded_client_fast_path_test): +/// cmd_rx is dropped so any fallback to cmd_tx returns a channel-closed error. +/// Success proves the SM was called directly; error proves fallback was triggered. +#[cfg(test)] +mod embedded_read_handle_tests { + use std::sync::Arc; + use std::time::Duration; + + use bytes::Bytes; + use d_engine_core::config::ReadConsistencyPolicy; + use d_engine_core::{MockStateMachine, MockTypeConfig, ReadLease, now_ms}; + use tokio::sync::mpsc; + + use super::super::EmbeddedReadHandle; + + // ── Helpers ─────────────────────────────────────────────────────────────── + + fn valid_lease() -> Arc { + let lease = Arc::new(ReadLease::new()); + lease.renew(1, now_ms() + 60_000); + lease + } + + fn invalid_lease() -> Arc { + let lease = Arc::new(ReadLease::new()); + lease.revoke(); + lease + } + + fn sm_with_value(value: Bytes) -> MockStateMachine { + let mut sm = MockStateMachine::new(); + sm.expect_get_multi() + .returning(move |keys| Ok(keys.iter().map(|_| Some(value.clone())).collect())); + sm + } + + fn sm_stopped() -> MockStateMachine { + // SM stopped: get_multi returns NotServing error + let mut sm = MockStateMachine::new(); + sm.expect_get_multi().returning(|_| { + Err(d_engine_core::StorageError::NotServing("state machine stopped".into()).into()) + }); + sm + } + + /// Build a handle. `drop_cmd_rx=true` closes the cmd channel so any fallback errors. + fn make_handle( + sm: MockStateMachine, + lease: Arc, + drop_cmd_rx: bool, + ) -> EmbeddedReadHandle { + let (cmd_tx, cmd_rx) = mpsc::channel(1); + if drop_cmd_rx { + drop(cmd_rx); + } + EmbeddedReadHandle::new(Arc::new(sm), lease, cmd_tx) + } + + // ── Eventual fast path ──────────────────────────────────────────────────── + + /// Eventual read goes directly to SM; cmd_rx dropped proves cmd_tx was not used. + #[tokio::test] + async fn test_eventual_read_bypasses_cmd_tx() { + let handle = make_handle(sm_with_value(Bytes::from("v")), valid_lease(), true); + let result = handle + .get( + b"k", + ReadConsistencyPolicy::EventualConsistency, + 1, + Duration::from_millis(100), + ) + .await; + assert_eq!(result.unwrap(), Some(Bytes::from("v"))); + } + + /// Eventual read when SM is stopped falls back to cmd_tx. + /// cmd_rx dropped → fallback returns error (proves fallback was taken, not panic). + #[tokio::test] + async fn test_eventual_sm_stopped_falls_back_to_cmd_tx() { + let handle = make_handle(sm_stopped(), valid_lease(), true); + let result = handle + .get( + b"k", + ReadConsistencyPolicy::EventualConsistency, + 1, + Duration::from_millis(100), + ) + .await; + assert!(result.is_err(), "stopped SM must fall back to cmd_tx"); + } + + // ── LeaseRead fast path ─────────────────────────────────────────────────── + + /// LeaseRead with valid lease goes directly to SM. + #[tokio::test] + async fn test_lease_read_valid_lease_bypasses_cmd_tx() { + let handle = make_handle(sm_with_value(Bytes::from("lease_val")), valid_lease(), true); + let result = handle + .get( + b"k", + ReadConsistencyPolicy::LeaseRead, + 1, + Duration::from_millis(100), + ) + .await; + assert_eq!(result.unwrap(), Some(Bytes::from("lease_val"))); + } + + /// LeaseRead with invalid (revoked) lease falls back to cmd_tx. + #[tokio::test] + async fn test_lease_read_invalid_lease_falls_back_to_cmd_tx() { + let mut sm = MockStateMachine::new(); + // get_multi must NOT be called — lease is invalid + sm.expect_get_multi().never(); + + let handle = make_handle(sm, invalid_lease(), true); + let result = handle + .get( + b"k", + ReadConsistencyPolicy::LeaseRead, + 1, + Duration::from_millis(100), + ) + .await; + assert!(result.is_err(), "invalid lease must fall back to cmd_tx"); + } + + /// LeaseRead with valid lease but stopped SM falls back to cmd_tx. + #[tokio::test] + async fn test_lease_read_sm_stopped_falls_back_to_cmd_tx() { + let handle = make_handle(sm_stopped(), valid_lease(), true); + let result = handle + .get( + b"k", + ReadConsistencyPolicy::LeaseRead, + 1, + Duration::from_millis(100), + ) + .await; + assert!(result.is_err(), "stopped SM must fall back to cmd_tx"); + } + + // ── LinearizableRead ───────────────────────────────────────────────────── + + /// Linearizable always uses cmd_tx. cmd_rx dropped → error proves cmd_tx was used. + #[tokio::test] + async fn test_linearizable_always_uses_cmd_tx() { + let mut sm = MockStateMachine::new(); + sm.expect_get_multi().never(); // must not be called + + let handle = make_handle(sm, valid_lease(), true); + let result = handle + .get( + b"k", + ReadConsistencyPolicy::LinearizableRead, + 1, + Duration::from_millis(100), + ) + .await; + assert!( + result.is_err(), + "linearizable must always go through cmd_tx" + ); + } + + // ── get_batch multi-key ─────────────────────────────────────────────────── + + /// get_batch returns values for all keys in order. + #[tokio::test] + async fn test_get_batch_returns_all_keys_in_order() { + let mut sm = MockStateMachine::new(); + sm.expect_get_multi().returning(|keys| { + Ok(keys + .iter() + .enumerate() + .map(|(i, _)| Some(Bytes::from(format!("v{i}")))) + .collect()) + }); + + let handle = make_handle(sm, valid_lease(), true); + let keys = vec![Bytes::from("k0"), Bytes::from("k1"), Bytes::from("k2")]; + let result = handle + .get_batch( + &keys, + ReadConsistencyPolicy::EventualConsistency, + 1, + Duration::from_millis(100), + ) + .await + .unwrap(); + + assert_eq!(result.len(), 3); + assert_eq!(result[0], Some(Bytes::from("v0"))); + assert_eq!(result[1], Some(Bytes::from("v1"))); + assert_eq!(result[2], Some(Bytes::from("v2"))); + } +} diff --git a/d-engine-server/src/api/embedded_env_test.rs b/d-engine-server/src/api/embedded_test/embedded_env_test.rs similarity index 100% rename from d-engine-server/src/api/embedded_env_test.rs rename to d-engine-server/src/api/embedded_test/embedded_env_test.rs diff --git a/d-engine-server/src/api/embedded_test.rs b/d-engine-server/src/api/embedded_test/embedded_test.rs similarity index 94% rename from d-engine-server/src/api/embedded_test.rs rename to d-engine-server/src/api/embedded_test/embedded_test.rs index 8324be3d..60e498e5 100644 --- a/d-engine-server/src/api/embedded_test.rs +++ b/d-engine-server/src/api/embedded_test/embedded_test.rs @@ -9,6 +9,8 @@ mod embedded_engine_tests { use crate::storage::FileStateMachine; use crate::storage::FileStorageEngine; + type TestEngine = EmbeddedEngine; + async fn create_test_storage_and_sm() -> ( Arc, Arc, @@ -39,7 +41,7 @@ mod embedded_engine_tests { let (storage, sm, _temp_dir) = create_test_storage_and_sm().await; // Start embedded engine (single node) - let engine = EmbeddedEngine::start_custom(storage, sm, None) + let engine = TestEngine::start_custom(storage, sm, None) .await .expect("Failed to start engine"); @@ -65,7 +67,7 @@ mod embedded_engine_tests { async fn test_wait_ready_timeout() { let (storage, sm, _temp_dir) = create_test_storage_and_sm().await; - let engine = EmbeddedEngine::start_custom(storage, sm, None) + let engine = TestEngine::start_custom(storage, sm, None) .await .expect("Failed to start engine"); @@ -85,7 +87,7 @@ mod embedded_engine_tests { async fn test_leader_change_notifier_basic() { let (storage, sm, _temp_dir) = create_test_storage_and_sm().await; - let engine = EmbeddedEngine::start_custom(storage, sm, None) + let engine = TestEngine::start_custom(storage, sm, None) .await .expect("Failed to start engine"); @@ -119,7 +121,7 @@ mod embedded_engine_tests { async fn test_ready_and_wait_ready_sequence() { let (storage, sm, _temp_dir) = create_test_storage_and_sm().await; - let engine = EmbeddedEngine::start_custom(storage, sm, None) + let engine = TestEngine::start_custom(storage, sm, None) .await .expect("Failed to start engine"); @@ -148,7 +150,7 @@ mod embedded_engine_tests { async fn test_client_available_after_wait_ready() { let (storage, sm, _temp_dir) = create_test_storage_and_sm().await; - let engine = EmbeddedEngine::start_custom(storage, sm, None) + let engine = TestEngine::start_custom(storage, sm, None) .await .expect("Failed to start engine"); @@ -174,7 +176,7 @@ mod embedded_engine_tests { async fn test_multiple_leader_change_notifier_subscribers() { let (storage, sm, _temp_dir) = create_test_storage_and_sm().await; - let engine = EmbeddedEngine::start_custom(storage, sm, None) + let engine = TestEngine::start_custom(storage, sm, None) .await .expect("Failed to start engine"); @@ -211,7 +213,7 @@ mod embedded_engine_tests { async fn test_engine_stop_cleans_up() { let (storage, sm, _temp_dir) = create_test_storage_and_sm().await; - let engine = EmbeddedEngine::start_custom(storage, sm, None) + let engine = TestEngine::start_custom(storage, sm, None) .await .expect("Failed to start engine"); @@ -224,7 +226,7 @@ mod embedded_engine_tests { async fn test_wait_ready_race_condition_already_elected() { let (storage, sm, _temp_dir) = create_test_storage_and_sm().await; - let engine = EmbeddedEngine::start_custom(storage, sm, None) + let engine = TestEngine::start_custom(storage, sm, None) .await .expect("Failed to start engine"); @@ -259,7 +261,7 @@ mod embedded_engine_tests { let (storage, sm, _temp_dir) = create_test_storage_and_sm().await; let engine = Arc::new( - EmbeddedEngine::start_custom(storage, sm, None) + TestEngine::start_custom(storage, sm, None) .await .expect("Failed to start engine"), ); @@ -305,7 +307,7 @@ mod embedded_engine_tests { async fn test_wait_ready_check_current_value_first() { let (storage, sm, _temp_dir) = create_test_storage_and_sm().await; - let engine = EmbeddedEngine::start_custom(storage, sm, None) + let engine = TestEngine::start_custom(storage, sm, None) .await .expect("Failed to start engine"); @@ -513,7 +515,7 @@ listen_addr = "127.0.0.1:0" async fn test_drop_without_stop_warning() { let (storage, sm, _temp_dir) = create_test_storage_and_sm().await; - let engine = EmbeddedEngine::start_custom(storage, sm, None) + let engine = TestEngine::start_custom(storage, sm, None) .await .expect("Failed to start engine"); @@ -522,6 +524,29 @@ listen_addr = "127.0.0.1:0" // verifies the code path doesn't panic drop(engine); } + + /// stop() must release the RocksDB LOCK before returning. + /// + /// If ReadActor still held Arc after stop(), RocksDB would refuse the + /// second open() with "lock file already held by process" — no sleep needed. + #[tokio::test] + #[serial(tmp_db)] + async fn test_stop_releases_sm_lock_immediately() { + let data_dir = std::path::PathBuf::from("/tmp/d-engine-lock-release-test"); + let _ = std::fs::remove_dir_all(&data_dir); + + let engine = + EmbeddedEngine::start(&data_dir).await.expect("First start should succeed"); + engine.stop().await.expect("stop() should succeed"); + + // Immediately reopen — no sleep. If the LOCK is still held this fails. + let engine2 = EmbeddedEngine::start(&data_dir).await.expect( + "Second start must succeed immediately after stop() — LOCK was not released", + ); + engine2.stop().await.ok(); + + let _ = std::fs::remove_dir_all(&data_dir); + } } #[cfg(feature = "watch")] @@ -535,7 +560,7 @@ listen_addr = "127.0.0.1:0" async fn test_watch_registers_successfully() { let (storage, sm, _temp_dir) = create_test_storage_and_sm().await; - let engine = EmbeddedEngine::start_custom(storage, sm, None) + let engine = TestEngine::start_custom(storage, sm, None) .await .expect("Failed to start engine"); @@ -612,7 +637,7 @@ listen_addr = "127.0.0.1:0" // temp_dir dropped here — underlying paths are now invalid }; - let engine = EmbeddedEngine::start_custom(storage, sm, None) + let engine = TestEngine::start_custom(storage, sm, None) .await .expect("Failed to start engine"); @@ -677,10 +702,9 @@ listen_addr = "127.0.0.1:0" ) .unwrap(); - let engine = - EmbeddedEngine::start_custom(storage, sm, Some(config_path.to_str().unwrap())) - .await - .expect("Failed to start engine"); + let engine = TestEngine::start_custom(storage, sm, Some(config_path.to_str().unwrap())) + .await + .expect("Failed to start engine"); engine .wait_ready(Duration::from_secs(5)) @@ -751,7 +775,7 @@ listen_addr = "127.0.0.1:0" async fn test_is_leader_single_node() { let (storage, sm, _temp_dir) = create_test_storage_and_sm().await; - let engine = EmbeddedEngine::start_custom(storage, sm, None) + let engine = TestEngine::start_custom(storage, sm, None) .await .expect("Failed to start engine"); @@ -791,7 +815,7 @@ listen_addr = "127.0.0.1:0" async fn test_leader_info_before_election() { let (storage, sm, _temp_dir) = create_test_storage_and_sm().await; - let engine = EmbeddedEngine::start_custom(storage, sm, None) + let engine = TestEngine::start_custom(storage, sm, None) .await .expect("Failed to start engine"); @@ -842,7 +866,7 @@ listen_addr = "127.0.0.1:0" let (storage, sm, _temp_dir) = create_test_storage_and_sm().await; let engine = Arc::new( - EmbeddedEngine::start_custom(storage, sm, None) + TestEngine::start_custom(storage, sm, None) .await .expect("Failed to start engine"), ); @@ -925,7 +949,7 @@ listen_addr = "127.0.0.1:0" async fn test_leader_info_consistency() { let (storage, sm, _temp_dir) = create_test_storage_and_sm().await; - let engine = EmbeddedEngine::start_custom(storage, sm, None) + let engine = TestEngine::start_custom(storage, sm, None) .await .expect("Failed to start engine"); diff --git a/d-engine-server/src/api/embedded_test/mod.rs b/d-engine-server/src/api/embedded_test/mod.rs new file mode 100644 index 00000000..c9c4cfc0 --- /dev/null +++ b/d-engine-server/src/api/embedded_test/mod.rs @@ -0,0 +1,5 @@ +#[cfg(test)] +mod embedded_test; + +#[cfg(test)] +mod embedded_env_test; diff --git a/d-engine-server/src/api/mod.rs b/d-engine-server/src/api/mod.rs index dc9166ec..e78e38ab 100644 --- a/d-engine-server/src/api/mod.rs +++ b/d-engine-server/src/api/mod.rs @@ -2,20 +2,26 @@ mod embedded; mod embedded_client; +mod embedded_read_handle; mod standalone; +mod standalone_read_handle; -pub use embedded::EmbeddedEngine; -pub use embedded_client::EmbeddedClient; pub use standalone::StandaloneEngine; +pub(crate) use standalone_read_handle::StandaloneReadHandle; -#[cfg(test)] -mod embedded_client_test; +/// Embedded engine generic over any `(SE, SM)` pair. +pub type EmbeddedEngine = embedded::EmbeddedEngine>; -#[cfg(test)] -mod embedded_test; +/// Embedded client generic over any `(SE, SM)` pair. +pub type EmbeddedClient = + embedded_client::EmbeddedClient>; -#[cfg(test)] -mod embedded_env_test; +/// Zero-param embedded engine alias using the default RocksDB backend. +#[cfg(feature = "rocksdb")] +pub type DefaultEmbeddedEngine = + EmbeddedEngine; -#[cfg(test)] -mod standalone_test; +/// Zero-param embedded client alias using the default RocksDB backend. +#[cfg(feature = "rocksdb")] +pub type DefaultEmbeddedClient = + EmbeddedClient; diff --git a/d-engine-server/src/api/standalone.rs b/d-engine-server/src/api/standalone.rs index f01958f2..17f64d40 100644 --- a/d-engine-server/src/api/standalone.rs +++ b/d-engine-server/src/api/standalone.rs @@ -68,7 +68,7 @@ impl StandaloneEngine { (storage, sm) }; - let lease = Arc::new(crate::storage::DefaultLease::new( + let lease = Arc::new(crate::storage::TtlLease::new( config.raft.state_machine.lease.clone(), )); sm.set_lease(lease); @@ -122,7 +122,7 @@ impl StandaloneEngine { (storage, sm) }; - let lease = Arc::new(crate::storage::DefaultLease::new( + let lease = Arc::new(crate::storage::TtlLease::new( config.raft.state_machine.lease.clone(), )); sm.set_lease(lease); @@ -187,3 +187,7 @@ impl StandaloneEngine { node.run().await } } + +#[cfg(test)] +#[path = "standalone_test.rs"] +mod tests; diff --git a/d-engine-server/src/api/standalone_read_handle.rs b/d-engine-server/src/api/standalone_read_handle.rs new file mode 100644 index 00000000..80ea65e2 --- /dev/null +++ b/d-engine-server/src/api/standalone_read_handle.rs @@ -0,0 +1,242 @@ +//! StandaloneReadHandle — read routing for standalone (gRPC) mode. +//! +//! Routes reads through the ReadActor channel for Eventual/LeaseRead fast path, +//! falling back to `cmd_tx` (Raft readIndex) for LinearizableRead or on error. +//! +//! For embedded mode, use `EmbeddedReadHandle` which calls the SM directly +//! without any channel overhead. + +use std::time::Duration; + +use bytes::Bytes; +use d_engine_core::MaybeCloneOneshot; +use d_engine_core::RaftOneshot; +use d_engine_core::client::{ + ClientApiError, ClientApiResult, ClientReadRequest, ClientResponsePayload, ErrorCode, + LeaderHint, ReadResults, +}; +use d_engine_core::config::ReadConsistencyPolicy; +use tokio::sync::{mpsc, oneshot}; + +use crate::read_actor::{ReadActorError, ReadCmd}; + +// ── Error helpers ───────────────────────────────────────────────────────────── + +pub(crate) fn channel_closed_error() -> ClientApiError { + ClientApiError::Network { + code: ErrorCode::ConnectionTimeout, + message: "Channel closed, node may be shutting down".to_string(), + retry_after_ms: None, + leader_hint: None, + } +} + +pub(crate) fn timeout_error(duration: Duration) -> ClientApiError { + ClientApiError::Network { + code: ErrorCode::ConnectionTimeout, + message: format!("Operation timed out after {duration:?}"), + retry_after_ms: Some(1000), + leader_hint: None, + } +} + +pub(crate) fn server_error(msg: String) -> ClientApiError { + ClientApiError::Business { + code: ErrorCode::Uncategorized, + message: msg, + required_action: None, + } +} + +pub(crate) fn not_leader_error( + leader_id: Option, + leader_address: Option, + retry_after_ms: Option, +) -> ClientApiError { + let message = match (&leader_address, &leader_id) { + (Some(addr), _) => format!("Not leader, try leader at: {addr}"), + (None, Some(id)) => format!("Not leader, leader_id: {id}"), + (None, None) => "Not leader".to_string(), + }; + let leader_hint = match (&leader_id, &leader_address) { + (Some(id_str), Some(addr)) => id_str.parse::().ok().map(|id| LeaderHint { + leader_id: id, + address: addr.clone(), + }), + _ => None, + }; + ClientApiError::Network { + code: ErrorCode::NotLeader, + message, + retry_after_ms: retry_after_ms.or(Some(100)), + leader_hint, + } +} + +pub(crate) fn map_error_response( + error: ErrorCode, + leader_hint: Option, + retry_after_ms: Option, +) -> ClientApiError { + match error { + ErrorCode::NotLeader => { + let (leader_id, leader_address) = if let Some(hint) = leader_hint { + (Some(hint.leader_id.to_string()), Some(hint.address)) + } else { + (None, None) + }; + not_leader_error(leader_id, leader_address, retry_after_ms) + } + _ => server_error(format!("Error code: {error:?}")), + } +} + +/// Unwrap a `ClientResponsePayload` as `ReadResults`. +pub(crate) fn extract_read_payload( + result: Option +) -> ClientApiResult { + match result { + Some(ClientResponsePayload::Read(r)) => Ok(r), + Some(ClientResponsePayload::Write(_)) => Err(ClientApiError::Protocol { + code: ErrorCode::InvalidResponse, + message: "expected ReadData payload, got WriteResult".to_string(), + supported_versions: None, + }), + None => Err(ClientApiError::Protocol { + code: ErrorCode::InvalidResponse, + message: "expected ReadData payload, got None".to_string(), + supported_versions: None, + }), + } +} + +// ── StandaloneReadHandle ────────────────────────────────────────────────────── + +/// Routes read requests to the ReadActor fast path or `cmd_tx` for standalone gRPC mode. +/// +/// Cloneable — the gRPC `Node` handler holds a clone per request. +/// The underlying channels are shared (`mpsc::Sender` is cheaply cloneable). +#[derive(Clone)] +pub(crate) struct StandaloneReadHandle { + /// ReadActor channel. `None` if node was started without a ReadActor. + pub(crate) read_tx: Option>, + /// Raft command channel for the cmd_tx fallback path. + pub(crate) cmd_tx: mpsc::Sender, +} + +impl StandaloneReadHandle { + /// Create a `StandaloneReadHandle` wiring the fast-path sender and the Raft cmd channel. + pub(crate) fn new( + read_tx: Option>, + cmd_tx: mpsc::Sender, + ) -> Self { + Self { read_tx, cmd_tx } + } + + /// Route a single-key read. Convenience wrapper around [`Self::get_batch`]. + /// + /// # Routing rules + /// | Policy | Path | On failure | + /// |----------------|----------------------------|-----------------------| + /// | Eventual/Lease | ReadActor (fast path) | fall through to cmd_tx| + /// | Linearizable | cmd_tx always | — | + /// | SmError | direct return | no fallback | + /// + /// Note: Not called internally; exposed as public API for callers to use. + #[allow(dead_code)] + pub async fn get( + &self, + key: &[u8], + consistency: ReadConsistencyPolicy, + client_id: u32, + timeout: Duration, + ) -> ClientApiResult> { + let key_bytes = Bytes::copy_from_slice(key); + let mut results = self.get_batch(&[key_bytes], consistency, client_id, timeout).await?; + Ok(results.pop().flatten()) + } + + /// Route a multi-key read atomically. + /// + /// For `Eventual` / `LeaseRead`, all keys are sent to ReadActor in a single + /// `ReadCmd` and read from one consistent snapshot. On fallback or + /// `Linearizable`, the full key list is forwarded to `cmd_tx` as a single + /// `ClientCmd::Read`. + /// + /// The returned `Vec` is the same length as `keys` and positionally ordered. + pub async fn get_batch( + &self, + keys: &[Bytes], + consistency: ReadConsistencyPolicy, + client_id: u32, + timeout: Duration, + ) -> ClientApiResult>> { + // Fast path: Eventual and LeaseRead bypass cmd_tx via ReadActor. + if let Some(read_tx) = &self.read_tx + && matches!( + consistency, + ReadConsistencyPolicy::EventualConsistency | ReadConsistencyPolicy::LeaseRead + ) + { + let (reply_tx, reply_rx) = oneshot::channel(); + if read_tx + .send(ReadCmd { + keys: keys.to_vec(), + consistency: consistency.clone(), + reply: reply_tx, + }) + .await + .is_ok() + { + match tokio::time::timeout(timeout, reply_rx).await { + Ok(Ok(Ok(values))) => return Ok(values), + Ok(Ok(Err(ReadActorError::SmError(e)))) => return Err(server_error(e)), + // LeaseInvalid / SmStopped → fall through to cmd_tx + Ok(Ok(Err(ReadActorError::LeaseInvalid | ReadActorError::SmStopped))) => {} + // ReadActor exited without replying → fall through + Ok(Err(_)) => {} + // ReadActor stalled beyond timeout → fall through to cmd_tx + Err(_timeout) => {} + } + } + // send() failed (channel closed) → fall through to cmd_tx + } + + // cmd_tx path (Raft readIndex — always correct, Linearizable always arrives here) + let request = ClientReadRequest { + client_id, + keys: keys.to_vec(), + consistency_policy: Some(consistency), + }; + let (resp_tx, resp_rx) = MaybeCloneOneshot::new(); + self.cmd_tx + .send(d_engine_core::ClientCmd::Read(request, resp_tx)) + .await + .map_err(|_| channel_closed_error())?; + + let result = tokio::time::timeout(timeout, resp_rx) + .await + .map_err(|_| timeout_error(timeout))? + .map_err(|_| channel_closed_error())?; + + let response = + result.map_err(|status| server_error(format!("RPC error: {}", status.message())))?; + + if response.error != ErrorCode::Success { + return Err(map_error_response( + response.error, + response.leader_hint, + response.retry_after_ms, + )); + } + + let read_results = extract_read_payload(response.result)?; + let entry_map: std::collections::HashMap = + read_results.entries.into_iter().map(|e| (e.key, e.value)).collect(); + Ok(keys.iter().map(|k| entry_map.get(k).cloned()).collect()) + } +} + +#[cfg(test)] +#[path = "standalone_read_handle_test.rs"] +mod tests; diff --git a/d-engine-server/src/api/standalone_read_handle_test.rs b/d-engine-server/src/api/standalone_read_handle_test.rs new file mode 100644 index 00000000..c57a3713 --- /dev/null +++ b/d-engine-server/src/api/standalone_read_handle_test.rs @@ -0,0 +1,752 @@ +/// Unit tests for StandaloneReadHandle routing logic. +/// +/// StandaloneReadHandle is the single source of truth for read routing: +/// Eventual/LeaseRead → ReadActor fast path; fallback to cmd_tx on LeaseInvalid/SmStopped +/// LinearizableRead → cmd_tx always +/// SmError → returned directly, NO cmd_tx fallback +/// +/// Proof technique (same as embedded_client_fast_path_test.rs): +/// cmd_rx is dropped so any fallback to cmd_tx returns a Network/channel-closed error. +/// A Business error proves SmError was returned directly. +/// An Ok(value) with cmd_rx dropped proves the fast path was taken. +#[cfg(test)] +mod read_handle_tests { + use std::sync::Arc; + use std::time::Duration; + + use crate::read_actor::run_read_actor; + use bytes::Bytes; + use d_engine_core::client::{ClientResponse, KvEntry}; + use d_engine_core::config::ReadConsistencyPolicy; + use d_engine_core::{ClientApiError, ClientCmd, Error, MockStateMachine, ReadLease, now_ms}; + use tokio::sync::mpsc; + use tokio::task::JoinHandle; + + use super::super::StandaloneReadHandle; + + // ── Helpers ─────────────────────────────────────────────────────────────── + + fn valid_lease() -> Arc { + let l = Arc::new(ReadLease::new()); + l.renew(1, now_ms() + 60_000); + l + } + + fn revoked_lease() -> Arc { + let l = Arc::new(ReadLease::new()); + l.revoke(); + l + } + + fn expired_lease() -> Arc { + let l = Arc::new(ReadLease::new()); + l.renew(1, now_ms().saturating_sub(1)); + l + } + + fn sm_running_with(value: Bytes) -> MockStateMachine { + let mut sm = MockStateMachine::new(); + sm.expect_is_running().returning(|| true); + sm.expect_get_multi() + .returning(move |keys| Ok(keys.iter().map(|_| Some(value.clone())).collect())); + sm + } + + fn sm_running_empty() -> MockStateMachine { + let mut sm = MockStateMachine::new(); + sm.expect_is_running().returning(|| true); + sm.expect_get_multi().returning(|keys| Ok(keys.iter().map(|_| None).collect())); + sm + } + + fn sm_stopped() -> MockStateMachine { + let mut sm = MockStateMachine::new(); + sm.expect_is_running().returning(|| false); + sm + } + + fn sm_with_error() -> MockStateMachine { + let mut sm = MockStateMachine::new(); + sm.expect_is_running().returning(|| true); + sm.expect_get_multi().returning(|_| Err(Error::Fatal("disk failure".into()))); + sm + } + + /// Build a StandaloneReadHandle with a live ReadActor. + /// + /// `drop_cmd_rx=true` closes the cmd_tx receiver — any cmd_tx fallback + /// returns a channel-closed Network error (proof the fast path was NOT taken). + async fn make_handle( + sm: MockStateMachine, + lease: Arc, + drop_cmd_rx: bool, + ) -> (StandaloneReadHandle, JoinHandle<()>) { + let (read_tx, read_rx) = mpsc::channel(8); + let (cmd_tx, cmd_rx) = mpsc::channel(1); + if drop_cmd_rx { + drop(cmd_rx); + } + let handle = tokio::spawn(run_read_actor(read_rx, lease, Arc::new(sm), 64)); + let rh = StandaloneReadHandle::new(Some(read_tx), cmd_tx); + (rh, handle) + } + + const TIMEOUT: Duration = Duration::from_millis(100); + + // ── Eventual fast path ──────────────────────────────────────────────────── + + /// Eventual, SM running, key exists → fast path returns value. + /// cmd_rx dropped: Ok(value) proves fast path was taken. + #[tokio::test] + async fn test_eventual_fast_path_returns_value() { + let (rh, handle) = + make_handle(sm_running_with(Bytes::from("v1")), valid_lease(), true).await; + let result = rh.get(b"k", ReadConsistencyPolicy::EventualConsistency, 1, TIMEOUT).await; + assert_eq!(result.unwrap(), Some(Bytes::from("v1"))); + drop(rh); + handle.await.unwrap(); + } + + /// Eventual, SM running, key missing → fast path returns None. + #[tokio::test] + async fn test_eventual_fast_path_returns_none_for_missing_key() { + let (rh, handle) = make_handle(sm_running_empty(), valid_lease(), true).await; + let result = rh + .get( + b"missing", + ReadConsistencyPolicy::EventualConsistency, + 1, + TIMEOUT, + ) + .await; + assert_eq!(result.unwrap(), None); + drop(rh); + handle.await.unwrap(); + } + + /// Eventual, SM stopped → SmStopped → fallback to cmd_tx. + /// cmd_rx dropped: error is Network (channel-closed), proving cmd_tx fallback triggered. + #[tokio::test] + async fn test_eventual_sm_stopped_falls_back_to_cmd_tx() { + let (rh, handle) = make_handle(sm_stopped(), valid_lease(), true).await; + let result = rh.get(b"k", ReadConsistencyPolicy::EventualConsistency, 1, TIMEOUT).await; + assert!( + matches!(result, Err(ClientApiError::Network { .. })), + "SmStopped must fall back to cmd_tx (Network error), got: {:?}", + result + ); + drop(rh); + handle.await.unwrap(); + } + + // ── LeaseRead fast path ─────────────────────────────────────────────────── + + /// LeaseRead, valid lease, SM running → fast path returns value. + #[tokio::test] + async fn test_lease_read_fast_path_returns_value() { + let (rh, handle) = make_handle( + sm_running_with(Bytes::from("lease_val")), + valid_lease(), + true, + ) + .await; + let result = rh.get(b"k", ReadConsistencyPolicy::LeaseRead, 1, TIMEOUT).await; + assert_eq!(result.unwrap(), Some(Bytes::from("lease_val"))); + drop(rh); + handle.await.unwrap(); + } + + /// LeaseRead, revoked lease → LeaseInvalid → fallback to cmd_tx. + #[tokio::test] + async fn test_lease_read_revoked_falls_back_to_cmd_tx() { + let (rh, handle) = + make_handle(sm_running_with(Bytes::from("x")), revoked_lease(), true).await; + let result = rh.get(b"k", ReadConsistencyPolicy::LeaseRead, 1, TIMEOUT).await; + assert!( + matches!(result, Err(ClientApiError::Network { .. })), + "revoked lease must fall back to cmd_tx, got: {:?}", + result + ); + drop(rh); + handle.await.unwrap(); + } + + /// LeaseRead, expired lease → LeaseInvalid → fallback to cmd_tx. + #[tokio::test] + async fn test_lease_read_expired_falls_back_to_cmd_tx() { + let (rh, handle) = + make_handle(sm_running_with(Bytes::from("x")), expired_lease(), true).await; + let result = rh.get(b"k", ReadConsistencyPolicy::LeaseRead, 1, TIMEOUT).await; + assert!( + matches!(result, Err(ClientApiError::Network { .. })), + "expired lease must fall back to cmd_tx, got: {:?}", + result + ); + drop(rh); + handle.await.unwrap(); + } + + /// LeaseRead, SM stopped (even with valid lease) → LeaseInvalid → fallback. + #[tokio::test] + async fn test_lease_read_sm_stopped_falls_back_to_cmd_tx() { + let (rh, handle) = make_handle(sm_stopped(), valid_lease(), true).await; + let result = rh.get(b"k", ReadConsistencyPolicy::LeaseRead, 1, TIMEOUT).await; + assert!( + matches!(result, Err(ClientApiError::Network { .. })), + "SM stopped must fall back to cmd_tx, got: {:?}", + result + ); + drop(rh); + handle.await.unwrap(); + } + + // ── SmError — direct return, no cmd_tx fallback ─────────────────────────── + + /// sm.get() returns Err → SmError → Business error returned DIRECTLY. + /// cmd_rx dropped: if fallback were taken, error would be Network. + /// Business error proves SmError was returned without touching cmd_tx. + #[tokio::test] + async fn test_sm_error_returned_directly_no_cmd_tx_fallback() { + let (rh, handle) = make_handle(sm_with_error(), valid_lease(), true).await; + let result = rh.get(b"k", ReadConsistencyPolicy::EventualConsistency, 1, TIMEOUT).await; + assert!( + matches!(result, Err(ClientApiError::Business { .. })), + "SmError must produce Business error (direct), not Network (fallback), got: {:?}", + result + ); + drop(rh); + handle.await.unwrap(); + } + + // ── Linearizable — always cmd_tx ───────────────────────────────────────── + + /// LinearizableRead must never touch ReadActor — always routed to cmd_tx. + /// cmd_rx dropped: Network error proves cmd_tx was used (not fast path). + #[tokio::test] + async fn test_linearizable_always_uses_cmd_tx() { + // SM has no get() expectation — mockall panics if get() is called (implicit safety net). + let mut sm = MockStateMachine::new(); + sm.expect_is_running().returning(|| true); + let (rh, handle) = make_handle(sm, valid_lease(), true).await; + + let result = rh.get(b"k", ReadConsistencyPolicy::LinearizableRead, 1, TIMEOUT).await; + assert!( + matches!(result, Err(ClientApiError::Network { .. })), + "LinearizableRead must always use cmd_tx, got: {:?}", + result + ); + drop(rh); + handle.await.unwrap(); + } + + // ── No read_tx — always cmd_tx ──────────────────────────────────────────── + + /// StandaloneReadHandle without read_tx (None) always falls through to cmd_tx for any policy. + #[tokio::test] + async fn test_no_read_tx_all_policies_use_cmd_tx() { + let (cmd_tx, cmd_rx) = mpsc::channel::(1); + drop(cmd_rx); + let rh = StandaloneReadHandle::new(None, cmd_tx); + + for policy in [ + ReadConsistencyPolicy::EventualConsistency, + ReadConsistencyPolicy::LeaseRead, + ReadConsistencyPolicy::LinearizableRead, + ] { + let result = rh.get(b"k", policy.clone(), 1, TIMEOUT).await; + assert!( + matches!(result, Err(ClientApiError::Network { .. })), + "no read_tx must always use cmd_tx, policy={policy:?}, got: {:?}", + result + ); + } + } + + // ── ReadActor channel closed ────────────────────────────────────────────── + + /// ReadActor already stopped (read_rx dropped) → send fails → fallback to cmd_tx. + /// cmd_rx also dropped: Network error proves fallback was taken. + #[tokio::test] + async fn test_read_actor_gone_falls_back_to_cmd_tx() { + let (read_tx, read_rx) = mpsc::channel::(1); + let (cmd_tx, cmd_rx) = mpsc::channel::(1); + drop(read_rx); // ReadActor is gone + drop(cmd_rx); // cmd_tx also closed — proves fallback was attempted + + let rh = StandaloneReadHandle::new(Some(read_tx), cmd_tx); + let result = rh.get(b"k", ReadConsistencyPolicy::EventualConsistency, 1, TIMEOUT).await; + assert!( + matches!(result, Err(ClientApiError::Network { .. })), + "closed ReadActor must fall back to cmd_tx, got: {:?}", + result + ); + } + + // ── cmd_tx path: alignment helper ──────────────────────────────────────── + + /// Builds a StandaloneReadHandle wired to a one-shot mock Raft responder. + /// + /// The responder returns the given sparse `entries` exactly as + /// `read_from_state_machine` would — caller is responsible for verifying + /// that `get_batch` re-aligns them to the input key positions. + async fn make_cmd_tx_handle(entries: Vec) -> StandaloneReadHandle { + let (cmd_tx, mut cmd_rx) = mpsc::channel(1); + tokio::spawn(async move { + if let Some(ClientCmd::Read(_, resp_tx)) = cmd_rx.recv().await { + let _ = resp_tx.send(Ok(ClientResponse::read_results(entries))); + } + }); + StandaloneReadHandle::new(None, cmd_tx) + } + + // ── cmd_tx path: HashMap re-alignment correctness ───────────────────────── + + /// cmd_tx returns sparse KvEntry with middle key missing. + /// get_batch must produce [Some, None, Some] — length 3, not 2. + /// This is the core regression for the HashMap re-alignment fix. + #[tokio::test] + async fn test_get_batch_cmd_tx_middle_key_missing_aligned() { + let entries = vec![ + KvEntry { + key: Bytes::from_static(b"k1"), + value: Bytes::from_static(b"v1"), + }, + KvEntry { + key: Bytes::from_static(b"k3"), + value: Bytes::from_static(b"v3"), + }, + ]; + let rh = make_cmd_tx_handle(entries).await; + let keys = vec![ + Bytes::from_static(b"k1"), + Bytes::from_static(b"k2"), + Bytes::from_static(b"k3"), + ]; + let result = rh + .get_batch(&keys, ReadConsistencyPolicy::LinearizableRead, 1, TIMEOUT) + .await + .unwrap(); + assert_eq!(result.len(), 3, "length must equal keys.len()"); + assert_eq!(result[0], Some(Bytes::from_static(b"v1"))); + assert_eq!( + result[1], None, + "missing key must be None at correct position" + ); + assert_eq!(result[2], Some(Bytes::from_static(b"v3"))); + } + + /// cmd_tx returns empty entries (all keys missing). + /// get_batch must produce [None, None, None] — length 3, not 0. + #[tokio::test] + async fn test_get_batch_cmd_tx_all_keys_missing_returns_all_none() { + let rh = make_cmd_tx_handle(vec![]).await; + let keys = vec![ + Bytes::from_static(b"x"), + Bytes::from_static(b"y"), + Bytes::from_static(b"z"), + ]; + let result = rh + .get_batch(&keys, ReadConsistencyPolicy::LinearizableRead, 1, TIMEOUT) + .await + .unwrap(); + assert_eq!(result.len(), 3); + assert!( + result.iter().all(|v| v.is_none()), + "all missing keys must be None: {:?}", + result + ); + } + + /// cmd_tx returns sparse entries with first key missing. + /// get_batch must produce [None, Some, Some]. + #[tokio::test] + async fn test_get_batch_cmd_tx_first_key_missing_aligned() { + let entries = vec![ + KvEntry { + key: Bytes::from_static(b"k2"), + value: Bytes::from_static(b"v2"), + }, + KvEntry { + key: Bytes::from_static(b"k3"), + value: Bytes::from_static(b"v3"), + }, + ]; + let rh = make_cmd_tx_handle(entries).await; + let keys = vec![ + Bytes::from_static(b"missing"), + Bytes::from_static(b"k2"), + Bytes::from_static(b"k3"), + ]; + let result = rh + .get_batch(&keys, ReadConsistencyPolicy::LinearizableRead, 1, TIMEOUT) + .await + .unwrap(); + assert_eq!(result.len(), 3); + assert_eq!(result[0], None); + assert_eq!(result[1], Some(Bytes::from_static(b"v2"))); + assert_eq!(result[2], Some(Bytes::from_static(b"v3"))); + } + + /// cmd_tx returns sparse entries with last key missing. + /// get_batch must produce [Some, Some, None]. + #[tokio::test] + async fn test_get_batch_cmd_tx_last_key_missing_aligned() { + let entries = vec![ + KvEntry { + key: Bytes::from_static(b"k1"), + value: Bytes::from_static(b"v1"), + }, + KvEntry { + key: Bytes::from_static(b"k2"), + value: Bytes::from_static(b"v2"), + }, + ]; + let rh = make_cmd_tx_handle(entries).await; + let keys = vec![ + Bytes::from_static(b"k1"), + Bytes::from_static(b"k2"), + Bytes::from_static(b"missing"), + ]; + let result = rh + .get_batch(&keys, ReadConsistencyPolicy::LinearizableRead, 1, TIMEOUT) + .await + .unwrap(); + assert_eq!(result.len(), 3); + assert_eq!(result[0], Some(Bytes::from_static(b"v1"))); + assert_eq!(result[1], Some(Bytes::from_static(b"v2"))); + assert_eq!(result[2], None); + } + + /// cmd_tx returns all keys — positive case, no alignment work needed. + #[tokio::test] + async fn test_get_batch_cmd_tx_all_keys_exist_aligned() { + let entries = vec![ + KvEntry { + key: Bytes::from_static(b"k1"), + value: Bytes::from_static(b"v1"), + }, + KvEntry { + key: Bytes::from_static(b"k2"), + value: Bytes::from_static(b"v2"), + }, + ]; + let rh = make_cmd_tx_handle(entries).await; + let keys = vec![Bytes::from_static(b"k1"), Bytes::from_static(b"k2")]; + let result = rh + .get_batch(&keys, ReadConsistencyPolicy::LinearizableRead, 1, TIMEOUT) + .await + .unwrap(); + assert_eq!(result.len(), 2); + assert_eq!(result[0], Some(Bytes::from_static(b"v1"))); + assert_eq!(result[1], Some(Bytes::from_static(b"v2"))); + } + + // ── get_batch() multi-key fast path ────────────────────────────────────── + + fn sm_running_with_multi(vals: Vec<(&'static [u8], &'static [u8])>) -> MockStateMachine { + let mut sm = MockStateMachine::new(); + sm.expect_is_running().returning(|| true); + sm.expect_get_multi().returning(move |keys| { + Ok(keys + .iter() + .map(|k| { + for (key, val) in &vals { + if &k[..] == *key { + return Some(Bytes::from_static(val)); + } + } + None + }) + .collect()) + }); + sm + } + + /// Multi-key Eventual batch: all keys found → Vec positionally ordered. + /// cmd_rx dropped: Ok proves fast path was taken. + #[tokio::test] + async fn test_get_batch_eventual_fast_path_returns_ordered_values() { + let sm = sm_running_with_multi(vec![(b"k1", b"v1"), (b"k2", b"v2")]); + let (rh, handle) = make_handle(sm, valid_lease(), true).await; + + let keys = vec![Bytes::from_static(b"k1"), Bytes::from_static(b"k2")]; + let result = rh + .get_batch( + &keys, + ReadConsistencyPolicy::EventualConsistency, + 1, + TIMEOUT, + ) + .await; + + let values = result.unwrap(); + assert_eq!(values.len(), 2); + assert_eq!(values[0], Some(Bytes::from_static(b"v1"))); + assert_eq!(values[1], Some(Bytes::from_static(b"v2"))); + drop(rh); + handle.await.unwrap(); + } + + /// Multi-key Eventual batch: some keys missing → None at correct positions. + #[tokio::test] + async fn test_get_batch_eventual_partial_missing_keys() { + let sm = sm_running_with_multi(vec![(b"k1", b"v1")]); + let (rh, handle) = make_handle(sm, valid_lease(), true).await; + + let keys = vec![Bytes::from_static(b"k1"), Bytes::from_static(b"missing")]; + let result = rh + .get_batch( + &keys, + ReadConsistencyPolicy::EventualConsistency, + 1, + TIMEOUT, + ) + .await; + + let values = result.unwrap(); + assert_eq!(values.len(), 2); + assert_eq!(values[0], Some(Bytes::from_static(b"v1"))); + assert_eq!(values[1], None); + drop(rh); + handle.await.unwrap(); + } + + /// Multi-key LeaseRead: valid lease → fast path returns all values. + #[tokio::test] + async fn test_get_batch_lease_read_valid_lease_fast_path() { + let sm = sm_running_with_multi(vec![(b"k1", b"v1"), (b"k2", b"v2")]); + let (rh, handle) = make_handle(sm, valid_lease(), true).await; + + let keys = vec![Bytes::from_static(b"k1"), Bytes::from_static(b"k2")]; + let result = rh.get_batch(&keys, ReadConsistencyPolicy::LeaseRead, 1, TIMEOUT).await; + + assert!( + result.is_ok(), + "multi-key LeaseRead fast path must succeed: {:?}", + result.err() + ); + drop(rh); + handle.await.unwrap(); + } + + /// Multi-key LeaseRead: revoked lease → LeaseInvalid → fallback to cmd_tx. + #[tokio::test] + async fn test_get_batch_lease_revoked_falls_back_to_cmd_tx() { + let sm = sm_running_with_multi(vec![(b"k1", b"v1")]); + let (rh, handle) = make_handle(sm, revoked_lease(), true).await; + + let keys = vec![Bytes::from_static(b"k1"), Bytes::from_static(b"k2")]; + let result = rh.get_batch(&keys, ReadConsistencyPolicy::LeaseRead, 1, TIMEOUT).await; + + assert!( + matches!(result, Err(ClientApiError::Network { .. })), + "revoked lease must fall back to cmd_tx, got: {:?}", + result + ); + drop(rh); + handle.await.unwrap(); + } + + /// Multi-key Linearizable: always goes through cmd_tx (no ReadActor). + /// cmd_rx dropped: Network error proves cmd_tx was used. + #[tokio::test] + async fn test_get_batch_linearizable_always_uses_cmd_tx() { + let mut sm = MockStateMachine::new(); + sm.expect_is_running().returning(|| true); + // no expect_get() — mockall panics if get() is called + let (rh, handle) = make_handle(sm, valid_lease(), true).await; + + let keys = vec![Bytes::from_static(b"k1"), Bytes::from_static(b"k2")]; + let result = rh.get_batch(&keys, ReadConsistencyPolicy::LinearizableRead, 1, TIMEOUT).await; + + assert!( + matches!(result, Err(ClientApiError::Network { .. })), + "Linearizable get_batch must always use cmd_tx, got: {:?}", + result + ); + drop(rh); + handle.await.unwrap(); + } + + /// ReadActor stalls (never replies) → timeout must fire and fall through to cmd_tx. + /// Without timeout on reply_rx.await, this test would hang indefinitely. + #[tokio::test] + async fn test_read_actor_stall_respects_timeout() { + let (read_tx, _read_rx) = mpsc::channel::(1); + // _read_rx is held alive (ReadActor never processes requests) but cmd_rx dropped + let (cmd_tx, cmd_rx) = mpsc::channel::(1); + drop(cmd_rx); + + let rh = StandaloneReadHandle::new(Some(read_tx), cmd_tx); + + let short_timeout = Duration::from_millis(50); + let start = tokio::time::Instant::now(); + let result = rh + .get( + b"k", + ReadConsistencyPolicy::EventualConsistency, + 1, + short_timeout, + ) + .await; + + assert!( + start.elapsed() < Duration::from_millis(500), + "stalled ReadActor must not block beyond timeout, took {:?}", + start.elapsed() + ); + assert!( + result.is_err(), + "stalled ReadActor + closed cmd_tx must return error, got: {:?}", + result + ); + } +} + +// ── Error helper tests ──────────────────────────────────────────────────────── +// +// These functions live in standalone_read_handle.rs and are shared by both +// StandaloneReadHandle and EmbeddedReadHandle (via re-export). Testing them +// here keeps the coverage co-located with the implementation. + +#[cfg(test)] +mod error_helper_tests { + use bytes::Bytes; + use d_engine_core::client::{ClientApiError, ClientResponsePayload, ErrorCode}; + use d_engine_core::client::{KvEntry, LeaderHint, ReadResults}; + + use super::super::{extract_read_payload, map_error_response, not_leader_error}; + + // ─── not_leader_error ──────────────────────────────────────────────────── + + #[test] + fn test_not_leader_uses_server_retry_after_ms_when_provided() { + let err = not_leader_error( + Some("1".to_string()), + Some("127.0.0.1:5001".to_string()), + Some(500), + ); + match err { + ClientApiError::Network { retry_after_ms, .. } => { + assert_eq!(retry_after_ms, Some(500)); + } + _ => panic!("expected Network error"), + } + } + + #[test] + fn test_not_leader_falls_back_to_100ms_when_server_provides_none() { + let err = not_leader_error(None, None, None); + match err { + ClientApiError::Network { retry_after_ms, .. } => { + assert_eq!(retry_after_ms, Some(100)); + } + _ => panic!("expected Network error"), + } + } + + #[test] + fn test_not_leader_zero_is_not_treated_as_none() { + let err = not_leader_error(None, None, Some(0)); + match err { + ClientApiError::Network { retry_after_ms, .. } => { + assert_eq!(retry_after_ms, Some(0)); + } + _ => panic!("expected Network error"), + } + } + + // ─── map_error_response ────────────────────────────────────────────────── + + #[test] + fn test_map_error_response_not_leader_forwards_retry_after_ms() { + let err = map_error_response( + ErrorCode::NotLeader, + Some(LeaderHint { + leader_id: 2, + address: "10.0.0.2:5002".into(), + }), + Some(250), + ); + match err { + ClientApiError::Network { + code, + retry_after_ms, + leader_hint, + .. + } => { + assert_eq!(code, ErrorCode::NotLeader); + assert_eq!(retry_after_ms, Some(250)); + assert_eq!(leader_hint.unwrap().leader_id, 2); + } + _ => panic!("expected Network error"), + } + } + + #[test] + fn test_map_error_response_not_leader_falls_back_to_100ms_when_none() { + let err = map_error_response(ErrorCode::NotLeader, None, None); + match err { + ClientApiError::Network { retry_after_ms, .. } => { + assert_eq!(retry_after_ms, Some(100)); + } + _ => panic!("expected Network error"), + } + } + + // ─── extract_read_payload ──────────────────────────────────────────────── + + #[test] + fn test_extract_read_payload_returns_read_results_on_success() { + let payload = Some(ClientResponsePayload::Read(ReadResults { + entries: vec![KvEntry { + key: Bytes::from("k"), + value: Bytes::from("v"), + }], + })); + let result = extract_read_payload(payload).unwrap(); + assert_eq!(result.entries[0].key, Bytes::from("k")); + } + + #[test] + fn test_extract_read_payload_rejects_write_result_payload() { + let payload = Some(ClientResponsePayload::Write( + d_engine_core::client::WriteResult { succeeded: true }, + )); + let err = extract_read_payload(payload).unwrap_err(); + assert_eq!(err.code(), ErrorCode::InvalidResponse); + assert!(err.message().contains("WriteResult")); + } + + #[test] + fn test_extract_read_payload_rejects_none_payload() { + let err = extract_read_payload(None).unwrap_err(); + assert_eq!(err.code(), ErrorCode::InvalidResponse); + assert!(err.message().contains("None")); + } + + #[test] + fn test_extract_read_payload_empty_entries_is_valid() { + let payload = Some(ClientResponsePayload::Read(ReadResults { entries: vec![] })); + assert!(extract_read_payload(payload).unwrap().entries.is_empty()); + } + + #[test] + fn test_extract_read_payload_multiple_entries_are_preserved() { + let payload = Some(ClientResponsePayload::Read(ReadResults { + entries: vec![ + KvEntry { + key: Bytes::from("k1"), + value: Bytes::from("v1"), + }, + KvEntry { + key: Bytes::from("k2"), + value: Bytes::from("v2"), + }, + ], + })); + let result = extract_read_payload(payload).unwrap(); + assert_eq!(result.entries.len(), 2); + assert_eq!(result.entries[1].key, Bytes::from("k2")); + } +} diff --git a/d-engine-server/src/lib.rs b/d-engine-server/src/lib.rs index 2e14bf30..7967bd48 100644 --- a/d-engine-server/src/lib.rs +++ b/d-engine-server/src/lib.rs @@ -98,6 +98,12 @@ mod proto_convert; #[cfg(test)] mod proto_convert_test; +/// ReadActor — fast-path read task for Eventual and LeaseRead. +/// Kept in server (not core) because it is a performance optimization, not a Raft protocol component. +pub(crate) mod read_actor; +#[cfg(test)] +mod read_actor_test; + /// Node lifecycle management /// /// Contains [`Node`] and [`NodeBuilder`] for server setup. @@ -117,6 +123,10 @@ pub use d_engine_core::LeaderInfo; pub mod storage; // -------------------- Primary Entry Points -------------------- +#[cfg(feature = "rocksdb")] +pub use api::DefaultEmbeddedClient; +#[cfg(feature = "rocksdb")] +pub use api::DefaultEmbeddedEngine; pub use api::EmbeddedEngine; pub use api::StandaloneEngine; pub use membership::MembershipSnapshot; diff --git a/d-engine-server/src/membership/membership_snapshot.rs b/d-engine-server/src/membership/membership_snapshot.rs index 315f696f..e0c60799 100644 --- a/d-engine-server/src/membership/membership_snapshot.rs +++ b/d-engine-server/src/membership/membership_snapshot.rs @@ -2,7 +2,7 @@ use std::collections::BTreeSet; /// A point-in-time snapshot of committed cluster membership. /// -/// Delivered via [`crate::EmbeddedEngine::watch_membership`] whenever a `ConfChange` +/// Delivered via `EmbeddedEngine::watch_membership` whenever a `ConfChange` /// entry commits. The snapshot reflects the membership state **after** the /// change has been applied, so `borrow()` always returns a consistent view. /// diff --git a/d-engine-server/src/network/grpc/grpc_raft_service.rs b/d-engine-server/src/network/grpc/grpc_raft_service.rs index 467d11c3..3e673064 100644 --- a/d-engine-server/src/network/grpc/grpc_raft_service.rs +++ b/d-engine-server/src/network/grpc/grpc_raft_service.rs @@ -467,15 +467,51 @@ where "Unknown consistency_policy value received, degrading to cluster default" ); } + + let timeout_duration = + Duration::from_millis(self.node_config.raft.general_raft_timeout_duration_in_ms); let core_req = proto_convert::to_core_read_req(proto_req); + + // Fast path: Eventual/LeaseRead → ReadHandle (ReadActor + cmd_tx fallback). + { + use d_engine_core::client::ClientApiError; + use d_engine_core::config::ReadConsistencyPolicy; + if let Some(ref policy) = core_req.consistency_policy + && matches!( + policy, + ReadConsistencyPolicy::EventualConsistency | ReadConsistencyPolicy::LeaseRead + ) + { + return match self + .read_handle + .get_batch( + &core_req.keys, + policy.clone(), + core_req.client_id, + timeout_duration, + ) + .await + { + Ok(values) => Ok(tonic::Response::new( + proto_convert::fast_path_batch_read_response(&core_req.keys, values), + )), + Err(ClientApiError::Business { message, .. }) => Err(Status::internal(message)), + Err(ClientApiError::Network { message, .. }) => { + Err(Status::unavailable(message)) + } + Err(other) => Err(Status::internal(format!("{other:?}"))), + }; + } + } + + // cmd_tx path: Linearizable or unrecognized policy. let (resp_tx, resp_rx) = MaybeCloneOneshot::new(); - self.cmd_tx + self.read_handle + .cmd_tx .send(d_engine_core::ClientCmd::Read(core_req, resp_tx)) .await .map_err(|_| Status::internal("Command channel closed"))?; - let timeout_duration = - Duration::from_millis(self.node_config.raft.general_raft_timeout_duration_in_ms); handle_rpc_timeout(resp_rx, timeout_duration, "handle_client_read") .await .map(|resp| resp.map(proto_convert::to_proto_response)) diff --git a/d-engine-server/src/network/grpc/grpc_raft_service_fast_path_test.rs b/d-engine-server/src/network/grpc/grpc_raft_service_fast_path_test.rs new file mode 100644 index 00000000..ed7d50c0 --- /dev/null +++ b/d-engine-server/src/network/grpc/grpc_raft_service_fast_path_test.rs @@ -0,0 +1,357 @@ +/// Fast-path tests for `handle_client_read` via `ReadHandle`. +/// +/// Routing rules under test: +/// Eventual/LeaseRead (any key count) → ReadActor fast path; fallback on LeaseInvalid/SmStopped +/// LinearizableRead → cmd_tx always +/// SmError → error returned directly (no cmd_tx fallback) +/// +/// Proof technique (identical to read_handle_test.rs): +/// The ReadHandle's cmd_rx is dropped. Any fallback to cmd_tx produces a tonic +/// Status error (channel closed). An Ok response with cmd_rx dropped proves the +/// fast path was taken. +#[cfg(test)] +mod grpc_fast_path_tests { + use std::sync::Arc; + use std::sync::atomic::Ordering; + + use crate::read_actor::run_read_actor; + use bytes::Bytes; + use d_engine_core::{MockStateMachine, ReadLease, now_ms}; + use d_engine_proto::client::ClientReadRequest; + use d_engine_proto::client::ReadConsistencyPolicy as ProtoPolicy; + use d_engine_proto::client::client_response::SuccessResult; + use d_engine_proto::client::raft_client_service_server::RaftClientService; + use tokio::sync::mpsc; + use tokio::task::JoinHandle; + use tonic::Request; + + use crate::Node; + use crate::api::StandaloneReadHandle; + use crate::test_utils::mock_node; + use d_engine_core::MockTypeConfig; + + // ── Helpers ─────────────────────────────────────────────────────────────── + + fn valid_lease() -> Arc { + let l = Arc::new(ReadLease::new()); + l.renew(1, now_ms() + 60_000); + l + } + + fn revoked_lease() -> Arc { + let l = Arc::new(ReadLease::new()); + l.revoke(); + l + } + + fn sm_running_with(value: Bytes) -> MockStateMachine { + let mut sm = MockStateMachine::new(); + sm.expect_is_running().returning(|| true); + sm.expect_get_multi() + .returning(move |keys| Ok(keys.iter().map(|_| Some(value.clone())).collect())); + sm + } + + fn sm_stopped() -> MockStateMachine { + let mut sm = MockStateMachine::new(); + sm.expect_is_running().returning(|| false); + sm + } + + fn sm_with_error() -> MockStateMachine { + let mut sm = MockStateMachine::new(); + sm.expect_is_running().returning(|| true); + sm.expect_get_multi() + .returning(|_| Err(d_engine_core::Error::Fatal("disk failure".into()))); + sm + } + + /// Build a Node wired to a live ReadActor. + /// + /// `drop_cmd_rx=true` closes the ReadHandle's cmd_rx — any fallback to cmd_tx + /// returns a tonic Status error (channel closed), proving the fast path was NOT taken. + async fn make_node( + sm: MockStateMachine, + lease: Arc, + drop_cmd_rx: bool, + ) -> (Node, JoinHandle<()>) { + let (read_tx, read_rx) = mpsc::channel(8); + let (cmd_tx, cmd_rx) = mpsc::channel(1); + if drop_cmd_rx { + drop(cmd_rx); + } + let handle = tokio::spawn(run_read_actor(read_rx, lease, Arc::new(sm), 64)); + let rh = StandaloneReadHandle::new(Some(read_tx), cmd_tx); + + let (_, graceful_rx) = tokio::sync::watch::channel(()); + let mut node = mock_node("/tmp/grpc_fast_path_test", graceful_rx, None); + node.read_handle = rh; + node.ready.store(true, Ordering::SeqCst); + (node, handle) + } + + fn req( + key: &[u8], + policy: ProtoPolicy, + ) -> ClientReadRequest { + ClientReadRequest { + client_id: 1, + keys: vec![Bytes::copy_from_slice(key)], + consistency_policy: Some(policy as i32), + } + } + + fn multi_key_req( + keys: &[&[u8]], + policy: ProtoPolicy, + ) -> ClientReadRequest { + ClientReadRequest { + client_id: 1, + keys: keys.iter().map(|k| Bytes::copy_from_slice(k)).collect(), + consistency_policy: Some(policy as i32), + } + } + + // ── Eventual fast path ──────────────────────────────────────────────────── + + /// Eventual, SM running, key exists → fast path returns Ok with the value. + /// cmd_rx dropped: Ok proves fast path was taken, not cmd_tx. + #[tokio::test] + async fn test_eventual_single_key_fast_path_returns_value() { + let (node, handle) = + make_node(sm_running_with(Bytes::from("fast_v")), valid_lease(), true).await; + + let result = node + .handle_client_read(Request::new(req(b"k", ProtoPolicy::EventualConsistency))) + .await; + + let resp = result.expect("fast path must succeed with cmd_rx dropped"); + let cr = resp.into_inner(); + assert_eq!(cr.error, 0, "error must be Success"); + if let Some(SuccessResult::ReadData(read_data)) = cr.success_result { + assert_eq!(read_data.results.len(), 1); + assert_eq!(read_data.results[0].value, Bytes::from("fast_v")); + } else { + panic!("expected ReadData payload"); + } + drop(node); + handle.await.unwrap(); + } + + /// Eventual, SM running, key missing → fast path returns Ok with empty entries. + #[tokio::test] + async fn test_eventual_single_key_fast_path_returns_empty_for_missing_key() { + let mut sm = MockStateMachine::new(); + sm.expect_is_running().returning(|| true); + sm.expect_get_multi().returning(|keys| Ok(keys.iter().map(|_| None).collect())); + + let (node, handle) = make_node(sm, valid_lease(), true).await; + let result = node + .handle_client_read(Request::new(req( + b"missing", + ProtoPolicy::EventualConsistency, + ))) + .await; + + let resp = result.expect("fast path must succeed"); + let cr = resp.into_inner(); + assert_eq!(cr.error, 0); + if let Some(SuccessResult::ReadData(read_data)) = cr.success_result { + assert!(read_data.results.is_empty(), "missing key → empty results"); + } else { + panic!("expected ReadData payload"); + } + drop(node); + handle.await.unwrap(); + } + + /// Eventual, SM stopped → SmStopped → fallback to cmd_tx. + /// cmd_rx dropped: error proves fallback was triggered. + #[tokio::test] + async fn test_eventual_sm_stopped_falls_back_to_cmd_tx() { + let (node, handle) = make_node(sm_stopped(), valid_lease(), true).await; + + let result = node + .handle_client_read(Request::new(req(b"k", ProtoPolicy::EventualConsistency))) + .await; + + assert!( + result.is_err(), + "SM stopped must fall back to closed cmd_tx" + ); + drop(node); + handle.await.unwrap(); + } + + // ── LeaseRead fast path ─────────────────────────────────────────────────── + + /// LeaseRead, valid lease, SM running → fast path returns Ok. + #[tokio::test] + async fn test_lease_read_valid_lease_fast_path_returns_ok() { + let (node, handle) = + make_node(sm_running_with(Bytes::from("lease_v")), valid_lease(), true).await; + + let result = node.handle_client_read(Request::new(req(b"k", ProtoPolicy::LeaseRead))).await; + + assert!(result.is_ok(), "fast path must succeed: {:?}", result.err()); + drop(node); + handle.await.unwrap(); + } + + /// LeaseRead, revoked lease → LeaseInvalid → fallback to cmd_tx. + #[tokio::test] + async fn test_lease_read_revoked_falls_back_to_cmd_tx() { + let (node, handle) = + make_node(sm_running_with(Bytes::from("x")), revoked_lease(), true).await; + + let result = node.handle_client_read(Request::new(req(b"k", ProtoPolicy::LeaseRead))).await; + + assert!( + result.is_err(), + "revoked lease must fall back to closed cmd_tx" + ); + drop(node); + handle.await.unwrap(); + } + + /// LeaseRead, SM stopped (even valid lease) → LeaseInvalid → fallback. + #[tokio::test] + async fn test_lease_read_sm_stopped_falls_back_to_cmd_tx() { + let (node, handle) = make_node(sm_stopped(), valid_lease(), true).await; + + let result = node.handle_client_read(Request::new(req(b"k", ProtoPolicy::LeaseRead))).await; + + assert!( + result.is_err(), + "SM stopped must fall back to closed cmd_tx" + ); + drop(node); + handle.await.unwrap(); + } + + // ── SmError — direct return, no cmd_tx fallback ─────────────────────────── + + /// sm.get() returns Err → SmError → error returned directly without cmd_tx fallback. + /// cmd_rx dropped: if fallback were taken, the error would come from channel-closed. + /// Either way the result is Err, but the important invariant (no fallback) is + /// proven at the ReadHandle layer; here we verify the error propagates to gRPC. + #[tokio::test] + async fn test_sm_error_propagates_as_grpc_error() { + let (node, handle) = make_node(sm_with_error(), valid_lease(), true).await; + + let result = node + .handle_client_read(Request::new(req(b"k", ProtoPolicy::EventualConsistency))) + .await; + + assert!( + result.is_err(), + "SmError must surface as tonic Status error" + ); + drop(node); + handle.await.unwrap(); + } + + // ── Linearizable — always cmd_tx ───────────────────────────────────────── + + /// LinearizableRead must never touch ReadActor — always routed to cmd_tx. + /// No expect_get_multi() on SM: mockall panics if get_multi() is called (implicit safety net). + /// cmd_rx dropped: error proves cmd_tx was used (not fast path). + #[tokio::test] + async fn test_linearizable_always_uses_cmd_tx() { + let mut sm = MockStateMachine::new(); + sm.expect_is_running().returning(|| true); + // no expect_get_multi() — any get_multi() call panics + + let (node, handle) = make_node(sm, valid_lease(), true).await; + + let result = node + .handle_client_read(Request::new(req(b"k", ProtoPolicy::LinearizableRead))) + .await; + + assert!(result.is_err(), "LinearizableRead must always use cmd_tx"); + drop(node); + handle.await.unwrap(); + } + + // ── Multi-key — ReadActor fast path ────────────────────────────────────── + + /// Multi-key Eventual read routes through ReadActor fast path (same as single-key). + /// cmd_rx dropped: Ok proves ReadActor was used, not cmd_tx. + #[tokio::test] + async fn test_multi_key_eventual_fast_path_returns_values() { + let mut sm = MockStateMachine::new(); + sm.expect_is_running().returning(|| true); + sm.expect_get_multi() + .returning(|keys| Ok(keys.iter().map(|_| Some(Bytes::from_static(b"val"))).collect())); + + let (node, handle) = make_node(sm, valid_lease(), true).await; + + let result = node + .handle_client_read(Request::new(multi_key_req( + &[b"k1", b"k2"], + ProtoPolicy::EventualConsistency, + ))) + .await; + + assert!( + result.is_ok(), + "multi-key Eventual must use ReadActor fast path: {:?}", + result.err() + ); + drop(node); + handle.await.unwrap(); + } + + /// Multi-key LeaseRead with valid lease routes through ReadActor fast path. + /// cmd_rx dropped: Ok proves ReadActor was used, not cmd_tx. + #[tokio::test] + async fn test_multi_key_lease_read_fast_path_returns_values() { + let mut sm = MockStateMachine::new(); + sm.expect_is_running().returning(|| true); + sm.expect_get_multi() + .returning(|keys| Ok(keys.iter().map(|_| Some(Bytes::from_static(b"val"))).collect())); + + let (node, handle) = make_node(sm, valid_lease(), true).await; + + let result = node + .handle_client_read(Request::new(multi_key_req( + &[b"k1", b"k2"], + ProtoPolicy::LeaseRead, + ))) + .await; + + assert!( + result.is_ok(), + "multi-key LeaseRead must use ReadActor fast path: {:?}", + result.err() + ); + drop(node); + handle.await.unwrap(); + } + + // ── No read_tx — always cmd_tx ──────────────────────────────────────────── + + /// ReadHandle without read_tx (no ReadActor wired) falls through to cmd_tx for all policies. + #[tokio::test] + async fn test_no_read_tx_all_policies_use_cmd_tx() { + let (cmd_tx, cmd_rx) = mpsc::channel::(1); + drop(cmd_rx); + + let (_, graceful_rx) = tokio::sync::watch::channel(()); + let mut node = mock_node("/tmp/grpc_fast_path_test_no_rt", graceful_rx, None); + node.read_handle = StandaloneReadHandle::new(None, cmd_tx); + node.ready.store(true, Ordering::SeqCst); + + for policy in [ + ProtoPolicy::EventualConsistency, + ProtoPolicy::LeaseRead, + ProtoPolicy::LinearizableRead, + ] { + let result = node.handle_client_read(Request::new(req(b"k", policy))).await; + assert!( + result.is_err(), + "no read_tx must always use cmd_tx (closed → error), policy={policy:?}" + ); + } + } +} diff --git a/d-engine-server/src/network/grpc/mod.rs b/d-engine-server/src/network/grpc/mod.rs index bbec6533..0a7a22fb 100644 --- a/d-engine-server/src/network/grpc/mod.rs +++ b/d-engine-server/src/network/grpc/mod.rs @@ -13,6 +13,9 @@ pub(crate) mod grpc_transport; #[cfg(test)] mod grpc_raft_service_test; +#[cfg(test)] +mod grpc_raft_service_fast_path_test; + #[cfg(test)] mod grpc_transport_test; diff --git a/d-engine-server/src/node/builder.rs b/d-engine-server/src/node/builder.rs index 2a35ce8e..46d89b65 100644 --- a/d-engine-server/src/node/builder.rs +++ b/d-engine-server/src/node/builder.rs @@ -32,6 +32,7 @@ use std::fmt::Debug; use std::sync::Arc; use std::sync::atomic::AtomicBool; +use crate::read_actor::run_read_actor; use d_engine_core::ClusterConfig; use d_engine_core::CommitHandler; use d_engine_core::CommitHandlerDependencies; @@ -566,13 +567,26 @@ where let commit_handler_handle = Self::spawn_state_machine_commit_listener(commit_handler); let event_tx = raft_core.event_sender(); + let read_lease = raft_core.read_lease(); let (rpc_ready_tx, _rpc_ready_rx) = watch::channel(false); + // Spawn ReadActor — sole owner of Arc on the read fast path. + let (read_tx, read_rx) = mpsc::channel(node_config_arc.raft.read_actor.channel_capacity); + let max_drain = node_config_arc.raft.read_actor.max_drain; + let read_actor_handle = tokio::spawn(run_read_actor( + read_rx, + Arc::clone(&read_lease), + state_machine, + max_drain, + )); + let read_handle = crate::api::StandaloneReadHandle::new(Some(read_tx), cmd_tx.clone()); + let node = Node::> { node_id, raft_core: Arc::new(Mutex::new(raft_core)), membership, event_tx: event_tx.clone(), + read_handle, cmd_tx, ready: AtomicBool::new(false), rpc_ready_tx, @@ -584,9 +598,11 @@ where #[cfg(feature = "watch")] _watch_dispatcher_handle: watch_system.map(|(_, _, handle, _)| handle), sm_worker_handle: std::sync::Mutex::new(Some(sm_worker_handle)), + read_actor_handle: std::sync::Mutex::new(Some(read_actor_handle)), _commit_handler_handle: Some(commit_handler_handle), _lease_cleanup_handle: lease_cleanup_handle, shutdown_signal: self.shutdown_signal.clone(), + read_lease, }; self.node = Some(Arc::new(node)); diff --git a/d-engine-server/src/node/mod.rs b/d-engine-server/src/node/mod.rs index d09cbe26..4bd53766 100644 --- a/d-engine-server/src/node/mod.rs +++ b/d-engine-server/src/node/mod.rs @@ -46,6 +46,7 @@ use d_engine_core::Membership; use d_engine_core::Raft; use d_engine_core::RaftEvent; use d_engine_core::RaftNodeConfig; +use d_engine_core::ReadLease; use d_engine_core::Result; use d_engine_core::TypeConfig; use d_engine_core::alias::MOF; @@ -113,6 +114,10 @@ where /// ensuring `Arc` is released before run() returns. pub(crate) sm_worker_handle: std::sync::Mutex>>, + /// ReadActor task handle. Aborted in run() after sm_worker exits so that the + /// `Arc` (and the RocksDB LOCK) is released before run() returns. + pub(crate) read_actor_handle: std::sync::Mutex>>, + /// Commit handler task handle (background log application) pub(crate) _commit_handler_handle: Option>, @@ -121,6 +126,12 @@ where /// Shutdown signal for graceful termination pub(crate) shutdown_signal: watch::Receiver<()>, + + /// Read lease — same Arc as the one inside SharedState, exposed for EmbeddedClient. + pub(crate) read_lease: Arc, + + /// Fast-path read routing. Always set; `read_tx = None` means no ReadActor. + pub(crate) read_handle: crate::api::StandaloneReadHandle, } impl Debug for Node @@ -183,6 +194,14 @@ where .ok(); } + // Abort ReadActor after sm_worker exits so Arc ref-count goes to zero + // and RocksDB LOCK is released before run() returns. + let ra_handle = self.read_actor_handle.lock().unwrap().take(); + if let Some(ra_handle) = ra_handle { + ra_handle.abort(); + let _ = ra_handle.await; + } + Ok(()) } @@ -327,4 +346,19 @@ where pub fn node_id(&self) -> u32 { self.node_id } + + /// Returns a clone of the read lease handle. + /// + /// The returned `Arc` is the same object held inside the Raft loop's + /// `SharedState`. EmbeddedClient uses it to check lease validity without going + /// through `cmd_tx`. + pub fn read_lease(&self) -> Arc { + Arc::clone(&self.read_lease) + } + + /// Returns the current Raft term from the hard state. + /// Used to seed `EmbeddedClient::known_term` at startup. + pub async fn current_term(&self) -> u64 { + self.raft_core.lock().await.current_term() + } } diff --git a/d-engine-server/src/proto_convert.rs b/d-engine-server/src/proto_convert.rs index a2d675f6..dbb00f05 100644 --- a/d-engine-server/src/proto_convert.rs +++ b/d-engine-server/src/proto_convert.rs @@ -211,6 +211,33 @@ pub(crate) fn to_proto_response(r: ClientResponse) -> proto_client::ClientRespon } } +/// Build a fast-path read response from a batch of keys and their optional values. +/// +/// Keys with `None` values are omitted from results (key not found). +/// The result slice is positionally ordered and corresponds 1:1 to `keys`. +pub(crate) fn fast_path_batch_read_response( + keys: &[bytes::Bytes], + values: Vec>, +) -> proto_client::ClientResponse { + let results = keys + .iter() + .zip(values) + .filter_map(|(key, value)| { + value.map(|v| ClientResult { + key: key.clone(), + value: v, + }) + }) + .collect(); + proto_client::ClientResponse { + error: 0, // ErrorCode::Success + metadata: None, + success_result: Some(proto_client::client_response::SuccessResult::ReadData( + ProtoReadResults { results }, + )), + } +} + // ─── Helpers ────────────────────────────────────────────────────────────────── #[inline] diff --git a/d-engine-server/src/read_actor.rs b/d-engine-server/src/read_actor.rs new file mode 100644 index 00000000..246f6668 --- /dev/null +++ b/d-engine-server/src/read_actor.rs @@ -0,0 +1,131 @@ +//! ReadActor — dedicated read task for Eventual and LeaseRead fast path. +//! +//! Lives in `d-engine-server`, not `d-engine-core`, because it is a +//! performance optimisation component, not a Raft protocol component. +//! Removing it does not affect correctness of consensus, election, or +//! log replication. (Contrast: TiKV's LocalReader lives in tikv/src/server/, +//! not in raft-rs.) +//! +//! # Routing contract +//! +//! ```text +//! Eventual / LeaseRead ──► ReadActor ──► SM.get_multi() (no Raft loop) +//! Linearizable ──► cmd_tx ──► Raft loop (readIndex) +//! ``` +//! +//! # Lifecycle guarantee +//! +//! The caller drops `read_tx` to signal shutdown. `run_read_actor` exits when +//! `read_rx` is drained and closed, at which point `Arc` is dropped and the +//! RocksDB LOCK file is released — before `stop()` signals the Raft loop. + +use std::sync::Arc; + +use bytes::Bytes; +use d_engine_core::ReadLease; +use d_engine_core::StateMachine; +use d_engine_core::config::ReadConsistencyPolicy; +use d_engine_core::now_ms; +use tokio::sync::{mpsc, oneshot}; + +// ── ReadCmd ──────────────────────────────────────────────────────────────────── + +/// A read request dispatched from `ReadHandle` to `ReadActor`. +/// +/// Only `EventualConsistency` and `LeaseRead` are valid; `LinearizableRead` +/// must go through `cmd_tx` (Raft readIndex protocol). +/// +/// # Ordering guarantee +/// +/// The reply `Vec` is guaranteed to be the same length as `keys` and +/// positionally corresponds to the input key order. +pub(crate) struct ReadCmd { + pub keys: Vec, + pub consistency: ReadConsistencyPolicy, + pub reply: oneshot::Sender>, ReadActorError>>, +} + +// ── ReadActorError ───────────────────────────────────────────────────────────── + +/// Error returned by the ReadActor to the caller. +/// +/// `LeaseInvalid` and `SmStopped` are recoverable: fall back to `cmd_tx`. +/// `SmError` is a hard failure; do NOT retry via cmd_tx. +#[derive(Debug)] +pub(crate) enum ReadActorError { + /// Lease has been revoked or expired. Retry via cmd_tx. + LeaseInvalid, + /// State machine is not running (node is stopping). Retry via cmd_tx. + SmStopped, + /// SM returned a hard error. + SmError(String), +} + +// ── run_read_actor ───────────────────────────────────────────────────────────── + +/// Run the ReadActor loop until `read_rx` is closed. +/// +/// Designed to be spawned as an independent task: +/// ```ignore +/// tokio::spawn(run_read_actor(read_rx, lease, sm, max_drain)); +/// ``` +/// +/// `max_drain` caps how many commands are drained per wake-up (mirrors +/// `batching.max_batch_size` in the Raft loop). After the first `recv().await` +/// unblocks, the actor drains up to `max_drain` pending commands with +/// `try_recv()` before yielding — amortising the task-switch cost across +/// request batches. +pub(crate) async fn run_read_actor( + mut read_rx: mpsc::Receiver, + lease: Arc, + sm: Arc, + max_drain: usize, +) where + SM: StateMachine, +{ + while let Some(cmd) = read_rx.recv().await { + let reply = serve_read(&cmd, &lease, &sm); + let _ = cmd.reply.send(reply); + + // Drain commands that arrived while processing, without going back to + // sleep — same pattern as Raft::drain_client_cmds(). + let mut count = 1; + while count < max_drain { + match read_rx.try_recv() { + Ok(cmd) => { + let reply = serve_read(&cmd, &lease, &sm); + let _ = cmd.reply.send(reply); + count += 1; + } + Err(_) => break, + } + } + } + // read_rx closed → Arc drops here, releasing RocksDB LOCK +} + +fn serve_read( + cmd: &ReadCmd, + lease: &Arc, + sm: &Arc, +) -> Result>, ReadActorError> +where + SM: StateMachine, +{ + match cmd.consistency { + ReadConsistencyPolicy::EventualConsistency => { + if !sm.is_running() { + return Err(ReadActorError::SmStopped); + } + } + ReadConsistencyPolicy::LeaseRead => { + if !lease.is_valid(now_ms()) || !sm.is_running() { + return Err(ReadActorError::LeaseInvalid); + } + } + // LinearizableRead must not reach ReadActor — guard defensively + _ => return Err(ReadActorError::LeaseInvalid), + } + + sm.get_multi(&cmd.keys).map_err(|e| ReadActorError::SmError(e.to_string())) +} diff --git a/d-engine-server/src/read_actor_test.rs b/d-engine-server/src/read_actor_test.rs new file mode 100644 index 00000000..f3704589 --- /dev/null +++ b/d-engine-server/src/read_actor_test.rs @@ -0,0 +1,330 @@ +use std::sync::Arc; + +use bytes::Bytes; +use d_engine_core::config::ReadConsistencyPolicy; +use d_engine_core::{Error, MockStateMachine, ReadLease, now_ms}; +use tokio::sync::{mpsc, oneshot}; + +use crate::read_actor::{ReadActorError, ReadCmd, run_read_actor}; + +fn make_sm_running() -> MockStateMachine { + let mut sm = MockStateMachine::new(); + sm.expect_is_running().returning(|| true); + sm +} + +fn make_sm_stopped() -> MockStateMachine { + let mut sm = MockStateMachine::new(); + sm.expect_is_running().returning(|| false); + sm +} + +fn valid_lease() -> Arc { + let lease = Arc::new(ReadLease::new()); + lease.renew(1, now_ms() + 60_000); + lease +} + +fn revoked_lease() -> Arc { + let lease = Arc::new(ReadLease::new()); + lease.revoke(); + lease +} + +fn send_read( + tx: &mpsc::Sender, + key: &[u8], + consistency: ReadConsistencyPolicy, +) -> oneshot::Receiver>, ReadActorError>> { + let (reply_tx, reply_rx) = oneshot::channel(); + tx.try_send(ReadCmd { + keys: vec![Bytes::copy_from_slice(key)], + consistency, + reply: reply_tx, + }) + .expect("channel should have capacity"); + reply_rx +} + +fn send_read_multi( + tx: &mpsc::Sender, + keys: Vec, + consistency: ReadConsistencyPolicy, +) -> oneshot::Receiver>, ReadActorError>> { + let (reply_tx, reply_rx) = oneshot::channel(); + tx.try_send(ReadCmd { + keys, + consistency, + reply: reply_tx, + }) + .expect("channel should have capacity"); + reply_rx +} + +/// Eventual read succeeds when SM is running and key exists +#[tokio::test] +async fn test_read_actor_eventual_read_returns_value() { + let mut sm = make_sm_running(); + sm.expect_get_multi().returning(|keys| { + Ok(keys + .iter() + .map(|k| { + if &k[..] == b"k1" { + Some(Bytes::from_static(b"v1")) + } else { + None + } + }) + .collect()) + }); + + let lease = valid_lease(); + let (tx, rx) = mpsc::channel(8); + let handle = tokio::spawn(run_read_actor(rx, lease, Arc::new(sm), 64)); + + let reply_rx = send_read(&tx, b"k1", ReadConsistencyPolicy::EventualConsistency); + drop(tx); + + let values = reply_rx.await.expect("reply sent").unwrap(); + assert_eq!(values.len(), 1); + assert_eq!(values[0], Some(Bytes::from_static(b"v1"))); + handle.await.unwrap(); +} + +/// Eventual read returns None for missing key +#[tokio::test] +async fn test_read_actor_eventual_read_missing_key_returns_none() { + let mut sm = make_sm_running(); + sm.expect_get_multi().returning(|keys| Ok(keys.iter().map(|_| None).collect())); + + let (tx, rx) = mpsc::channel(8); + let handle = tokio::spawn(run_read_actor(rx, valid_lease(), Arc::new(sm), 64)); + + let reply_rx = send_read(&tx, b"missing", ReadConsistencyPolicy::EventualConsistency); + drop(tx); + + let values = reply_rx.await.unwrap().unwrap(); + assert_eq!(values.len(), 1); + assert_eq!(values[0], None); + handle.await.unwrap(); +} + +/// Multi-key Eventual read returns all values in positional order +#[tokio::test] +async fn test_read_actor_eventual_multi_key_returns_ordered_values() { + let mut sm = make_sm_running(); + sm.expect_get_multi().returning(|keys| { + Ok(keys + .iter() + .map(|k| { + if &k[..] == b"k1" { + Some(Bytes::from_static(b"v1")) + } else { + None + } + }) + .collect()) + }); + + let (tx, rx) = mpsc::channel(8); + let handle = tokio::spawn(run_read_actor(rx, valid_lease(), Arc::new(sm), 64)); + + let keys = vec![Bytes::from_static(b"k1"), Bytes::from_static(b"k2")]; + let reply_rx = send_read_multi(&tx, keys, ReadConsistencyPolicy::EventualConsistency); + drop(tx); + + let values = reply_rx.await.unwrap().unwrap(); + assert_eq!(values.len(), 2); + assert_eq!(values[0], Some(Bytes::from_static(b"v1"))); + assert_eq!(values[1], None); + handle.await.unwrap(); +} + +/// Eventual read fails when SM is stopped +#[tokio::test] +async fn test_read_actor_eventual_sm_stopped_returns_error() { + let sm = make_sm_stopped(); + let (tx, rx) = mpsc::channel(8); + let handle = tokio::spawn(run_read_actor(rx, valid_lease(), Arc::new(sm), 64)); + + let reply_rx = send_read(&tx, b"k", ReadConsistencyPolicy::EventualConsistency); + drop(tx); + + let err = reply_rx.await.unwrap().unwrap_err(); + assert!(matches!(err, ReadActorError::SmStopped)); + handle.await.unwrap(); +} + +/// LeaseRead succeeds when lease is valid and SM is running +#[tokio::test] +async fn test_read_actor_lease_read_valid_lease_returns_value() { + let mut sm = make_sm_running(); + sm.expect_get_multi() + .returning(|keys| Ok(keys.iter().map(|_| Some(Bytes::from_static(b"val"))).collect())); + + let (tx, rx) = mpsc::channel(8); + let handle = tokio::spawn(run_read_actor(rx, valid_lease(), Arc::new(sm), 64)); + + let reply_rx = send_read(&tx, b"k", ReadConsistencyPolicy::LeaseRead); + drop(tx); + + assert!(reply_rx.await.unwrap().is_ok()); + handle.await.unwrap(); +} + +/// LeaseRead returns LeaseInvalid when lease is revoked +#[tokio::test] +async fn test_read_actor_lease_revoked_returns_lease_invalid() { + let sm = make_sm_running(); + let (tx, rx) = mpsc::channel(8); + let handle = tokio::spawn(run_read_actor(rx, revoked_lease(), Arc::new(sm), 64)); + + let reply_rx = send_read(&tx, b"k", ReadConsistencyPolicy::LeaseRead); + drop(tx); + + let err = reply_rx.await.unwrap().unwrap_err(); + assert!(matches!(err, ReadActorError::LeaseInvalid)); + handle.await.unwrap(); +} + +/// LeaseRead returns LeaseInvalid when lease has expired (deadline in the past) +#[tokio::test] +async fn test_read_actor_lease_expired_returns_lease_invalid() { + let sm = make_sm_running(); + let lease = Arc::new(ReadLease::new()); + lease.renew(1, now_ms().saturating_sub(1)); + + let (tx, rx) = mpsc::channel(8); + let handle = tokio::spawn(run_read_actor(rx, lease, Arc::new(sm), 64)); + + let reply_rx = send_read(&tx, b"k", ReadConsistencyPolicy::LeaseRead); + drop(tx); + + let err = reply_rx.await.unwrap().unwrap_err(); + assert!(matches!(err, ReadActorError::LeaseInvalid)); + handle.await.unwrap(); +} + +/// LeaseRead returns LeaseInvalid when SM is stopped (even if lease is valid) +#[tokio::test] +async fn test_read_actor_lease_valid_but_sm_stopped_returns_error() { + let sm = make_sm_stopped(); + let (tx, rx) = mpsc::channel(8); + let handle = tokio::spawn(run_read_actor(rx, valid_lease(), Arc::new(sm), 64)); + + let reply_rx = send_read(&tx, b"k", ReadConsistencyPolicy::LeaseRead); + drop(tx); + + let err = reply_rx.await.unwrap().unwrap_err(); + assert!(matches!(err, ReadActorError::LeaseInvalid)); + handle.await.unwrap(); +} + +/// ReadActor exits cleanly when all senders are dropped (lifecycle guarantee) +#[tokio::test] +async fn test_read_actor_exits_when_channel_closed() { + let sm = make_sm_running(); + let sm_arc = Arc::new(sm); + let sm_weak = Arc::downgrade(&sm_arc); + + let (tx, rx) = mpsc::channel(8); + let handle = tokio::spawn(run_read_actor(rx, valid_lease(), sm_arc, 64)); + + drop(tx); + + handle.await.unwrap(); + + assert!( + sm_weak.upgrade().is_none(), + "Arc should be dropped when ReadActor exits" + ); +} + +/// serve_read calls sm.get_multi() — verified by explicit expect_get_multi() expectation. +/// +/// Tests that only set up expect_get() rely on MockStateMachine's default-impl delegation +/// (get_multi → get). This test sets up expect_get_multi() directly to verify that +/// serve_read → sm.get_multi() is the actual call path, not just an indirect get() chain. +#[tokio::test] +async fn test_read_actor_eventual_multi_key_uses_get_multi() { + let mut sm = make_sm_running(); + sm.expect_get_multi().returning(|keys| { + Ok(keys + .iter() + .map(|k| { + if &k[..] == b"k1" { + Some(Bytes::from_static(b"v1")) + } else { + None + } + }) + .collect()) + }); + + let (tx, rx) = mpsc::channel(8); + let handle = tokio::spawn(run_read_actor(rx, valid_lease(), Arc::new(sm), 64)); + + let keys = vec![Bytes::from_static(b"k1"), Bytes::from_static(b"k2")]; + let reply_rx = send_read_multi(&tx, keys, ReadConsistencyPolicy::EventualConsistency); + drop(tx); + + let values = reply_rx.await.unwrap().unwrap(); + assert_eq!(values.len(), 2); + assert_eq!(values[0], Some(Bytes::from_static(b"v1"))); + assert_eq!(values[1], None); + handle.await.unwrap(); +} + +/// sm.get_multi() returning Err propagates as SmError through serve_read. +#[tokio::test] +async fn test_read_actor_sm_get_multi_error_returns_sm_error() { + let mut sm = make_sm_running(); + sm.expect_get_multi().returning(|_| Err(Error::Fatal("disk failure".into()))); + + let (tx, rx) = mpsc::channel(8); + let handle = tokio::spawn(run_read_actor(rx, valid_lease(), Arc::new(sm), 64)); + + let reply_rx = send_read(&tx, b"k", ReadConsistencyPolicy::EventualConsistency); + drop(tx); + + let err = reply_rx.await.unwrap().unwrap_err(); + assert!(matches!(err, ReadActorError::SmError(_))); + handle.await.unwrap(); +} + +/// LinearizableRead must never reach ReadActor — defensive guard returns LeaseInvalid. +/// If get_multi() were called, mockall would panic (no expectation set) — implicit safety net. +#[tokio::test] +async fn test_read_actor_linearizable_read_returns_lease_invalid() { + let sm = make_sm_running(); // no expect_get_multi() — any get_multi() call panics + + let (tx, rx) = mpsc::channel(8); + let handle = tokio::spawn(run_read_actor(rx, valid_lease(), Arc::new(sm), 64)); + + let reply_rx = send_read(&tx, b"k", ReadConsistencyPolicy::LinearizableRead); + drop(tx); + + let err = reply_rx.await.unwrap().unwrap_err(); + assert!(matches!(err, ReadActorError::LeaseInvalid)); + handle.await.unwrap(); +} + +/// revoke() immediately invalidates an in-flight lease +#[tokio::test] +async fn test_read_actor_revoke_invalidates_in_flight_reads() { + let sm = make_sm_running(); + let lease = valid_lease(); + let lease_clone = Arc::clone(&lease); + + let (tx, rx) = mpsc::channel(8); + let handle = tokio::spawn(run_read_actor(rx, lease, Arc::new(sm), 64)); + + lease_clone.revoke(); + + let reply_rx = send_read(&tx, b"k", ReadConsistencyPolicy::LeaseRead); + drop(tx); + + let err = reply_rx.await.unwrap().unwrap_err(); + assert!(matches!(err, ReadActorError::LeaseInvalid)); + handle.await.unwrap(); +} diff --git a/d-engine-server/src/storage/adaptors/file/file_state_machine.rs b/d-engine-server/src/storage/adaptors/file/file_state_machine.rs index 9b4952c3..c5f15d11 100644 --- a/d-engine-server/src/storage/adaptors/file/file_state_machine.rs +++ b/d-engine-server/src/storage/adaptors/file/file_state_machine.rs @@ -115,7 +115,7 @@ use tracing::error; use tracing::info; use tracing::warn; -use crate::storage::DefaultLease; +use crate::storage::TtlLease; type FileStateMachineDataType = RwLock>; @@ -242,9 +242,9 @@ pub struct FileStateMachine { data: FileStateMachineDataType, // (value, term) // Lease management for automatic key expiration - // DefaultLease is thread-safe internally (uses DashMap + Mutex) + // TtlLease is thread-safe internally (uses DashMap + Mutex) // Injected by NodeBuilder after construction - lease: Option>, + lease: Option>, // Raft state with disk persistence last_applied_index: AtomicU64, @@ -305,7 +305,7 @@ impl FileStateMachine { /// Also available for testing and benchmarks. pub fn set_lease( &mut self, - lease: Arc, + lease: Arc, ) { self.lease = Some(lease); } @@ -1017,6 +1017,12 @@ impl StateMachine for FileStateMachine { Ok(()) } + fn close_storage(&self) { + self.running.store(false, Ordering::SeqCst); + // FileStateMachine has no exclusive OS lock file. + // Marking not-running is sufficient; Drop handles final flush. + } + fn stop(&self) -> Result<(), Error> { // Ensure all data is flushed to disk before stopping self.running.store(false, Ordering::SeqCst); @@ -1581,4 +1587,14 @@ impl StateMachine for FileStateMachine { let revision = self.last_applied_index.load(Ordering::SeqCst); Ok(ScanResult { entries, revision }) } + + fn get_multi( + &self, + keys: &[Bytes], + ) -> d_engine_core::Result>> { + // Hold read lock for the full batch so all keys come from the same consistent view. + // apply_chunk holds the write lock; parking_lot RwLock ensures they don't interleave. + let data = self.data.read(); + Ok(keys.iter().map(|k| data.get(k).map(|(v, _)| v.clone())).collect()) + } } diff --git a/d-engine-server/src/storage/adaptors/file/file_state_machine_test.rs b/d-engine-server/src/storage/adaptors/file/file_state_machine_test.rs index fb40cd35..dcf49be1 100644 --- a/d-engine-server/src/storage/adaptors/file/file_state_machine_test.rs +++ b/d-engine-server/src/storage/adaptors/file/file_state_machine_test.rs @@ -9,6 +9,32 @@ use d_engine_core::Error; use d_engine_core::storage::state_machine_test::{StateMachineBuilder, StateMachineTestSuite}; use tempfile::TempDir; +// ── close_storage() tests ───────────────────────────────────────────────────── + +/// close_storage() must mark the SM as not running. +/// FileStateMachine has no exclusive OS lock, so close_storage() is a lifecycle +/// signal — subsequent reads must fail with NotServing. +#[tokio::test] +async fn test_file_sm_close_storage_marks_not_running() { + let dir = TempDir::new().unwrap(); + let sm = FileStateMachine::new(dir.path().to_path_buf()).await.unwrap(); + sm.start().await.unwrap(); + assert!(sm.is_running()); + + sm.close_storage(); + + assert!(!sm.is_running(), "close_storage must mark SM not running"); +} + +/// stop() after close_storage() must not panic (idempotent shutdown). +#[tokio::test] +async fn test_file_sm_stop_after_close_storage_is_safe() { + let dir = TempDir::new().unwrap(); + let sm = FileStateMachine::new(dir.path().to_path_buf()).await.unwrap(); + sm.close_storage(); + assert!(sm.stop().is_ok(), "stop after close_storage must not error"); +} + /// Builder for FileStateMachine test instances struct FileStateMachineBuilder { temp_dir: TempDir, @@ -26,7 +52,7 @@ impl FileStateMachineBuilder { impl StateMachineBuilder for FileStateMachineBuilder { async fn build(&self) -> Result, Error> { // Use fixed path to support restart recovery testing - let path = self.temp_dir.path().join("file_sm"); + let path = self.temp_dir.path().to_path_buf().join("file_sm"); let sm = FileStateMachine::new(path).await?; Ok(Arc::new(sm)) } @@ -60,7 +86,7 @@ async fn test_file_state_machine_performance() { #[tokio::test] async fn test_wal_replay_after_crash() { let temp_dir = tempfile::tempdir().unwrap(); - let data_dir = temp_dir.path().to_path_buf(); + let data_dir = temp_dir.path().to_path_buf().to_path_buf(); let sm = FileStateMachine::new(data_dir.clone()).await.unwrap(); @@ -136,7 +162,7 @@ async fn test_wal_replay_after_crash() { #[tokio::test] async fn test_cas_failure_wal_replay_does_not_corrupt_data() { let temp_dir = tempfile::tempdir().unwrap(); - let data_dir = temp_dir.path().to_path_buf(); + let data_dir = temp_dir.path().to_path_buf().to_path_buf(); let sm = FileStateMachine::new(data_dir.clone()).await.unwrap(); // Establish initial value @@ -190,7 +216,9 @@ async fn test_cas_failure_wal_replay_does_not_corrupt_data() { #[tokio::test] async fn test_cas_in_same_chunk_as_preceding_insert() { let temp_dir = tempfile::tempdir().unwrap(); - let sm = FileStateMachine::new(temp_dir.path().to_path_buf()).await.unwrap(); + let sm = FileStateMachine::new(temp_dir.path().to_path_buf().to_path_buf()) + .await + .unwrap(); let results = sm .apply_chunk(&[ @@ -224,7 +252,7 @@ async fn test_cas_in_same_chunk_as_preceding_insert() { #[tokio::test] async fn test_cas_success_wal_replay_applies_new_value() { let temp_dir = tempfile::tempdir().unwrap(); - let data_dir = temp_dir.path().to_path_buf(); + let data_dir = temp_dir.path().to_path_buf().to_path_buf(); let sm = FileStateMachine::new(data_dir.clone()).await.unwrap(); sm.apply_chunk(&[ApplyEntry { @@ -278,7 +306,7 @@ async fn test_cas_success_wal_replay_applies_new_value() { #[tokio::test] async fn test_replay_wal_unknown_opcode_returns_error() { let temp_dir = tempfile::tempdir().unwrap(); - let data_dir = temp_dir.path().to_path_buf(); + let data_dir = temp_dir.path().to_path_buf().to_path_buf(); // Write one valid entry so the WAL file exists and has a known-good prefix. let sm = FileStateMachine::new(data_dir.clone()).await.unwrap(); @@ -325,7 +353,9 @@ async fn test_replay_wal_unknown_opcode_returns_error() { #[tokio::test] async fn test_file_sm_scan_prefix_returns_matching_keys() { let temp_dir = tempfile::tempdir().unwrap(); - let sm = FileStateMachine::new(temp_dir.path().to_path_buf()).await.unwrap(); + let sm = FileStateMachine::new(temp_dir.path().to_path_buf().to_path_buf()) + .await + .unwrap(); let entries = vec![ ApplyEntry { @@ -375,7 +405,9 @@ async fn test_file_sm_scan_prefix_returns_matching_keys() { #[tokio::test] async fn test_file_sm_scan_prefix_empty_namespace() { let temp_dir = tempfile::tempdir().unwrap(); - let sm = FileStateMachine::new(temp_dir.path().to_path_buf()).await.unwrap(); + let sm = FileStateMachine::new(temp_dir.path().to_path_buf().to_path_buf()) + .await + .unwrap(); sm.apply_chunk(&[ApplyEntry { index: 1, @@ -397,11 +429,115 @@ async fn test_file_sm_scan_prefix_empty_namespace() { ); } +// ── get_multi tests ─────────────────────────────────────────────────────────── + +/// FileStateMachine get_multi holds a read lock for the entire batch, ensuring +/// that state transitions are atomic from the reader's perspective. +/// +/// Because apply_chunk acquires a write lock when updating self.data, a concurrent +/// write cannot interleave between the individual key reads inside get_multi. +/// The result is always either the full pre-write state or the full post-write state — +/// never a mix of values from different apply indexes. +/// +/// This test verifies sequential coherence: after each state transition, all keys +/// in the batch reflect the same version with no cross-version contamination. +#[tokio::test] +async fn test_get_multi_file_sm_atomic_state_transitions() { + let temp_dir = tempfile::tempdir().unwrap(); + let sm = FileStateMachine::new(temp_dir.path().to_path_buf().to_path_buf()) + .await + .unwrap(); + + // v1: initial state + sm.apply_chunk(&[ + ApplyEntry { + index: 1, + term: 1, + command: Command::Insert { + key: Bytes::from("/svc/addr"), + value: Bytes::from("10.0.0.1"), + ttl_secs: None, + }, + }, + ApplyEntry { + index: 2, + term: 1, + command: Command::Insert { + key: Bytes::from("/svc/version"), + value: Bytes::from("v1"), + ttl_secs: None, + }, + }, + ]) + .await + .unwrap(); + + let keys = vec![Bytes::from("/svc/addr"), Bytes::from("/svc/version")]; + + // Read after v1 — both fields must be from v1 + let v1_result = sm.get_multi(&keys).unwrap(); + assert_eq!( + v1_result[0], + Some(Bytes::from("10.0.0.1")), + "addr must be v1" + ); + assert_eq!(v1_result[1], Some(Bytes::from("v1")), "version must be v1"); + + // Apply v2: both fields transition together + sm.apply_chunk(&[ + ApplyEntry { + index: 3, + term: 1, + command: Command::Insert { + key: Bytes::from("/svc/addr"), + value: Bytes::from("10.0.0.2"), + ttl_secs: None, + }, + }, + ApplyEntry { + index: 4, + term: 1, + command: Command::Insert { + key: Bytes::from("/svc/version"), + value: Bytes::from("v2"), + ttl_secs: None, + }, + }, + ]) + .await + .unwrap(); + + // Read after v2 — both fields must be from v2 + let v2_result = sm.get_multi(&keys).unwrap(); + assert_eq!( + v2_result[0], + Some(Bytes::from("10.0.0.2")), + "addr must be v2" + ); + assert_eq!(v2_result[1], Some(Bytes::from("v2")), "version must be v2"); + + // Verify internal consistency of each result: addr and version must come from + // the same version. A torn read (addr=v2 + version=v1) would indicate that + // the read lock is not held across the full batch, which is a correctness bug. + for (label, result) in [("v1_result", &v1_result), ("v2_result", &v2_result)] { + let is_v1 = + result[0] == Some(Bytes::from("10.0.0.1")) && result[1] == Some(Bytes::from("v1")); + let is_v2 = + result[0] == Some(Bytes::from("10.0.0.2")) && result[1] == Some(Bytes::from("v2")); + assert!( + is_v1 || is_v2, + "{label}: torn read — result is neither pure v1 nor pure v2: {result:?}" + ); + } +} + /// scan_prefix revision equals last_applied_index at scan time. #[tokio::test] async fn test_file_sm_scan_prefix_revision_reflects_applied_index() { let temp_dir = tempfile::tempdir().unwrap(); - let sm = FileStateMachine::new(temp_dir.path().to_path_buf()).await.unwrap(); + let sm = FileStateMachine::new(temp_dir.path().to_path_buf().to_path_buf()) + .await + .unwrap(); let entries: Vec = (1u64..=3) .map(|i| ApplyEntry { diff --git a/d-engine-server/src/storage/adaptors/rocksdb/rocksdb_state_machine.rs b/d-engine-server/src/storage/adaptors/rocksdb/rocksdb_state_machine.rs index 72975850..39f9158b 100644 --- a/d-engine-server/src/storage/adaptors/rocksdb/rocksdb_state_machine.rs +++ b/d-engine-server/src/storage/adaptors/rocksdb/rocksdb_state_machine.rs @@ -4,7 +4,7 @@ use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; use std::time::SystemTime; -use arc_swap::ArcSwap; +use arc_swap::ArcSwapOption; use bytes::Bytes; use d_engine_core::ApplyEntry; use d_engine_core::ApplyResult; @@ -40,7 +40,7 @@ use tracing::info; use tracing::instrument; use tracing::warn; -use crate::storage::DefaultLease; +use crate::storage::TtlLease; use super::STATE_MACHINE_CF; use super::STATE_MACHINE_META_CF; @@ -76,16 +76,18 @@ struct CfExportFile { /// RocksDB-based state machine implementation with lease support #[derive(Debug)] pub struct RocksDBStateMachine { - db: Arc>, + // ArcSwapOption: close_db() stores None to release the RocksDB LOCK file immediately, + // decoupled from Arc reference count reaching zero. + db: ArcSwapOption, is_serving: AtomicBool, last_applied_index: AtomicU64, last_applied_term: AtomicU64, last_snapshot_metadata: RwLock>, // Lease management for automatic key expiration - // DefaultLease is thread-safe internally (uses DashMap + Mutex) + // TtlLease is thread-safe internally (uses DashMap + Mutex) // Injected by NodeBuilder after construction - lease: Option>, + lease: Option>, } /// Returns the lexicographic successor of `prefix` for use as an iterator upper bound. @@ -130,7 +132,7 @@ impl RocksDBStateMachine { let last_snapshot_metadata = Self::load_snapshot_metadata(&db_arc)?; Ok(Self { - db: Arc::new(ArcSwap::new(db_arc)), + db: ArcSwapOption::new(Some(db_arc)), is_serving: AtomicBool::new(true), last_applied_index: AtomicU64::new(last_applied_index), last_applied_term: AtomicU64::new(last_applied_term), @@ -148,7 +150,7 @@ impl RocksDBStateMachine { let last_snapshot_metadata = Self::load_snapshot_metadata(&db)?; Ok(Self { - db: Arc::new(ArcSwap::new(db)), + db: ArcSwapOption::new(Some(db)), is_serving: AtomicBool::new(true), last_applied_index: AtomicU64::new(last_applied_index), last_applied_term: AtomicU64::new(last_applied_term), @@ -164,7 +166,7 @@ impl RocksDBStateMachine { /// Also available for testing and benchmarks. pub fn set_lease( &mut self, - lease: Arc, + lease: Arc, ) { self.lease = Some(lease); } @@ -176,15 +178,49 @@ impl RocksDBStateMachine { &self, new_db: DB, ) { - self.db.store(Arc::new(new_db)); + self.db.store(Some(Arc::new(new_db))); } - // Injects lease configuration into this state machine. - // - // Framework-internal method: called by NodeBuilder::build() during initialization. - // Opens RocksDB with the standard configuration // ========== Private helper methods ========== + /// Run `f` with a borrowed `&DB`, returning `NotServing` if `close_db()` was called. + /// + /// Uses `ArcSwapOption::load()` — a seqlock read with no Arc clone. The Guard never + /// escapes this call, so the hot-path cost is identical to the pre-refactor code. + /// + /// **When to use**: any operation whose entire DB access fits inside a single closure + /// body: reads, metadata writes, batch commits, etc. + /// + /// **When NOT to use**: if the DB handle must cross an `await` point or if complex + /// lifetime dependencies prevent the whole operation from fitting in one closure + /// (e.g. `apply_chunk`, where `batch` borrows `cf` which borrows `db` across the + /// full function body). In those cases use the inline Guard pattern directly, or + /// `load_db()` when the `Arc` must be moved (e.g. `spawn_blocking`). + fn with_db( + &self, + f: F, + ) -> Result + where + F: FnOnce(&DB) -> Result, + { + let guard = self.db.load(); // Guard — no Arc clone + match guard.as_deref() { + Some(db) => f(db), + None => Err(StorageError::NotServing("state machine stopped".into()).into()), + } + } + + /// Clone the live `Arc` for paths that must *move* the handle (e.g. `spawn_blocking`). + /// + /// Costs 2 atomic ops (Arc increment + future decrement). Only use when the handle + /// must be moved into a closure or across an await point. For all other cases prefer + /// `with_db()` (zero Arc clone) or the inline Guard pattern. + fn load_db(&self) -> Result, Error> { + self.db + .load_full() + .ok_or_else(|| StorageError::NotServing("state machine stopped".into()).into()) + } + fn load_state_machine_metadata(db: &Arc) -> Result<(u64, u64), Error> { let cf = db .cf_handle(STATE_MACHINE_META_CF) @@ -231,51 +267,55 @@ impl RocksDBStateMachine { } fn persist_state_machine_metadata(&self) -> Result<(), Error> { - let db = self.db.load(); - let cf = db - .cf_handle(STATE_MACHINE_META_CF) - .ok_or_else(|| StorageError::DbError("State machine meta CF not found".to_string()))?; - let index = self.last_applied_index.load(Ordering::SeqCst); let term = self.last_applied_term.load(Ordering::SeqCst); - - db.put_cf(&cf, LAST_APPLIED_INDEX_KEY, index.to_be_bytes()) - .map_err(|e| StorageError::DbError(e.to_string()))?; - db.put_cf(&cf, LAST_APPLIED_TERM_KEY, term.to_be_bytes()) - .map_err(|e| StorageError::DbError(e.to_string()))?; - - Ok(()) + self.with_db(|db| { + let cf = db.cf_handle(STATE_MACHINE_META_CF).ok_or_else(|| { + Error::System(d_engine_core::SystemError::Storage(StorageError::DbError( + "State machine meta CF not found".to_string(), + ))) + })?; + db.put_cf(&cf, LAST_APPLIED_INDEX_KEY, index.to_be_bytes()) + .map_err(|e| StorageError::DbError(e.to_string()))?; + db.put_cf(&cf, LAST_APPLIED_TERM_KEY, term.to_be_bytes()) + .map_err(|e| StorageError::DbError(e.to_string()))?; + Ok(()) + }) } fn persist_snapshot_metadata(&self) -> Result<(), Error> { - let db = self.db.load(); - let cf = db - .cf_handle(STATE_MACHINE_META_CF) - .ok_or_else(|| StorageError::DbError("State machine meta CF not found".to_string()))?; - - if let Some(metadata) = self.last_snapshot_metadata.read().clone() { - let bytes = bincode::serialize(&metadata).map_err(StorageError::BincodeError)?; - db.put_cf(&cf, SNAPSHOT_METADATA_KEY, bytes) - .map_err(|e| StorageError::DbError(e.to_string()))?; - } - Ok(()) + let snapshot = self.last_snapshot_metadata.read().clone(); + self.with_db(|db| { + let cf = db.cf_handle(STATE_MACHINE_META_CF).ok_or_else(|| { + Error::System(d_engine_core::SystemError::Storage(StorageError::DbError( + "State machine meta CF not found".to_string(), + ))) + })?; + if let Some(metadata) = &snapshot { + let bytes = bincode::serialize(metadata).map_err(StorageError::BincodeError)?; + db.put_cf(&cf, SNAPSHOT_METADATA_KEY, bytes) + .map_err(|e| StorageError::DbError(e.to_string()))?; + } + Ok(()) + }) } fn persist_ttl_metadata(&self) -> Result<(), Error> { - if let Some(ref lease) = self.lease { - let db = self.db.load(); + let Some(ref lease) = self.lease else { + return Ok(()); + }; + let ttl_snapshot = lease.to_snapshot(); + self.with_db(|db| { let cf = db.cf_handle(STATE_MACHINE_META_CF).ok_or_else(|| { - StorageError::DbError("State machine meta CF not found".to_string()) + Error::System(d_engine_core::SystemError::Storage(StorageError::DbError( + "State machine meta CF not found".to_string(), + ))) })?; - - let ttl_snapshot = lease.to_snapshot(); - - db.put_cf(&cf, TTL_STATE_KEY, ttl_snapshot) + db.put_cf(&cf, TTL_STATE_KEY, &ttl_snapshot) .map_err(|e| StorageError::DbError(e.to_string()))?; - debug!("Persisted TTL state to RocksDB"); - } - Ok(()) + Ok(()) + }) } /// Loads TTL state from RocksDB metadata after lease injection. @@ -286,26 +326,26 @@ impl RocksDBStateMachine { let Some(ref lease) = self.lease else { return Ok(()); // No lease configured }; - - let db = self.db.load(); - let cf = db - .cf_handle(STATE_MACHINE_META_CF) - .ok_or_else(|| StorageError::DbError("State machine meta CF not found".to_string()))?; - - match db - .get_cf(&cf, TTL_STATE_KEY) - .map_err(|e| StorageError::DbError(e.to_string()))? - { - Some(ttl_data) => { - lease.reload(&ttl_data)?; - debug!("Loaded TTL state from RocksDB: {} active TTLs", lease.len()); - } - None => { - debug!("No TTL state found in RocksDB"); + self.with_db(|db| { + let cf = db.cf_handle(STATE_MACHINE_META_CF).ok_or_else(|| { + Error::System(d_engine_core::SystemError::Storage(StorageError::DbError( + "State machine meta CF not found".to_string(), + ))) + })?; + match db + .get_cf(&cf, TTL_STATE_KEY) + .map_err(|e| StorageError::DbError(e.to_string()))? + { + Some(ttl_data) => { + lease.reload(&ttl_data)?; + debug!("Loaded TTL state from RocksDB: {} active TTLs", lease.len()); + } + None => { + debug!("No TTL state found in RocksDB"); + } } - } - - Ok(()) + Ok(()) + }) } /// Piggyback cleanup: Remove expired keys with time budget @@ -347,13 +387,13 @@ impl RocksDBStateMachine { } // Get database handle - let db = self.db.load(); - let cf = match db.cf_handle(STATE_MACHINE_CF) { - Some(cf) => cf, - None => { - error!("State machine CF not found during TTL cleanup"); - return 0; - } + let guard = self.db.load(); + let Some(db) = guard.as_deref() else { + return 0; + }; + let Some(cf) = db.cf_handle(STATE_MACHINE_CF) else { + error!("State machine CF not found during TTL cleanup"); + return 0; }; // Cleanup expired keys with time budget @@ -410,8 +450,7 @@ impl RocksDBStateMachine { &self, batch: WriteBatch, ) -> Result<(), Error> { - self.db.load().write(&batch).map_err(|e| StorageError::DbError(e.to_string()))?; - Ok(()) + self.with_db(|db| db.write(&batch).map_err(|e| StorageError::DbError(e.to_string()).into())) } // ===== Snapshot restore helpers ===== @@ -424,8 +463,7 @@ impl RocksDBStateMachine { metadata: &SnapshotMetadata, snapshot_dir: &std::path::Path, ) -> Result<(), Error> { - let db = self.db.load(); - Self::restore_from_cf_export(&db, snapshot_dir)?; + self.with_db(|db| Self::restore_from_cf_export(db, snapshot_dir))?; info!("Snapshot restore complete"); @@ -615,7 +653,10 @@ impl RocksDBStateMachine { opts.set_iterate_upper_bound(upper); } - let db = self.db.load(); + let guard = self.db.load(); + let Some(db) = guard.as_deref() else { + return Err(StorageError::NotServing("state machine stopped".into()).into()); + }; let cf = db .cf_handle(STATE_MACHINE_CF) .ok_or_else(|| StorageError::DbError("STATE_MACHINE_CF not found".into()))?; @@ -633,6 +674,38 @@ impl RocksDBStateMachine { let revision = self.last_applied_index.load(Ordering::SeqCst); Ok(ScanResult { entries, revision }) } + + /// Permanently close the RocksDB handle and release the LOCK file. + /// + /// Unlike `stop()`, this is not reversible with `start()`. + /// Called by `EmbeddedEngine::stop()` before the Raft loop exits so the LOCK + /// is released without waiting for `Arc` reference counts to reach zero. + /// + /// Sequence: flush → save hard state → cancel background work → store None. + /// After this, `load_db()` returns `NotServing`; `stop()` afterwards is safe. + pub(crate) fn close_db(&self) { + self.is_serving.store(false, Ordering::SeqCst); + + if let Err(e) = self.persist_ttl_metadata() { + error!("close_db: failed to persist TTL metadata: {e:?}"); + } + if let Err(e) = self.save_hard_state() { + error!("close_db: failed to save hard state: {e:?}"); + } + if let Err(e) = self.flush() { + error!("close_db: failed to flush: {e:?}"); + } + { + let guard = self.db.load(); + if let Some(db) = guard.as_deref() { + db.cancel_all_background_work(true); + } + } + + // Release Arc: refcount decrements here, RocksDB LOCK file released. + self.db.store(None); + info!("RocksDB state machine: DB closed, LOCK released"); + } } #[async_trait] @@ -650,12 +723,20 @@ impl StateMachine for RocksDBStateMachine { Ok(()) } + fn close_storage(&self) { + self.close_db(); + } + fn stop(&self) -> Result<(), Error> { self.is_serving.store(false, Ordering::SeqCst); - // Graceful shutdown: persist TTL state to disk - // This ensures lease data survives across restarts - if let Err(e) = self.persist_ttl_metadata() { + // Skip TTL persist if close_db() was already called (DB is None). + // Node::stop() calls sm.stop() after EmbeddedEngine has already called close_db(), + // so this path must be safe to call on a closed DB. + // Skip TTL persist if close_db() has already set DB to None. + if self.db.load().is_some() + && let Err(e) = self.persist_ttl_metadata() + { error!("Failed to persist TTL metadata on shutdown: {:?}", e); return Err(e); } @@ -682,15 +763,17 @@ impl StateMachine for RocksDBStateMachine { .into()); } - let db = self.db.load(); - let cf = db - .cf_handle(STATE_MACHINE_CF) - .ok_or_else(|| StorageError::DbError("State machine CF not found".to_string()))?; - - match db.get_cf(&cf, key_buffer).map_err(|e| StorageError::DbError(e.to_string()))? { - Some(value) => Ok(Some(Bytes::copy_from_slice(&value))), - None => Ok(None), - } + self.with_db(|db| { + let cf = db.cf_handle(STATE_MACHINE_CF).ok_or_else(|| { + Error::System(d_engine_core::SystemError::Storage(StorageError::DbError( + "State machine CF not found".to_string(), + ))) + })?; + match db.get_cf(&cf, key_buffer).map_err(|e| StorageError::DbError(e.to_string()))? { + Some(value) => Ok(Some(Bytes::copy_from_slice(&value))), + None => Ok(None), + } + }) } fn entry_term( @@ -708,7 +791,13 @@ impl StateMachine for RocksDBStateMachine { &self, chunk: &[ApplyEntry], ) -> Result, Error> { - let db = self.db.load(); + // Inline guard: db and cf must share a lifetime across the entire function body + // (batch borrows cf, CAS reads borrow db). Using with_db() closure would require + // moving the whole async function body into a sync closure — use Guard directly. + let guard = self.db.load(); + let Some(db) = guard.as_deref() else { + return Err(StorageError::NotServing("state machine stopped".into()).into()); + }; let cf = db .cf_handle(STATE_MACHINE_CF) .ok_or_else(|| StorageError::DbError("State machine CF not found".to_string()))?; @@ -767,7 +856,7 @@ impl StateMachine for RocksDBStateMachine { // Read through WriteBatchWithIndex so earlier writes in this batch // are visible, preventing stale-read linearizability violations. let current_value = batch - .get_from_batch_and_db_cf(&*db, &cf, key, &ReadOptions::default()) + .get_from_batch_and_db_cf(db, &cf, key, &ReadOptions::default()) .map_err(|e| StorageError::DbError(format!("CAS read failed: {e}")))?; let cas_success = match (current_value, expected) { @@ -796,10 +885,7 @@ impl StateMachine for RocksDBStateMachine { } } - self.db - .load() - .write_wbwi(&batch) - .map_err(|e| StorageError::DbError(e.to_string()))?; + db.write_wbwi(&batch).map_err(|e| StorageError::DbError(e.to_string()))?; if let Some(highest) = highest_index_entry { self.update_last_applied(highest); @@ -809,15 +895,15 @@ impl StateMachine for RocksDBStateMachine { } fn len(&self) -> usize { - let db = self.db.load(); - let cf = match db.cf_handle(STATE_MACHINE_CF) { - Some(cf) => cf, - None => return 0, + let guard = self.db.load(); + let Some(db) = guard.as_deref() else { + return 0; + }; + let Some(cf) = db.cf_handle(STATE_MACHINE_CF) else { + return 0; }; - // Note: This is an expensive operation because it iterates over all keys. - let iter = db.iterator_cf(&cf, IteratorMode::Start); - iter.count() + db.iterator_cf(&cf, IteratorMode::Start).count() } fn update_last_applied( @@ -904,7 +990,7 @@ impl StateMachine for RocksDBStateMachine { // run on tokio's blocking thread pool via spawn_blocking, not on an async worker thread. // This prevents snapshot disk I/O from starving the Raft event loop under concurrent // writes (#315). - let db = self.db.load_full(); // Arc: Send, safe to move into spawn_blocking + let db = self.load_db()?; // Arc: Send, safe to move into spawn_blocking let dir = new_snapshot_dir.clone(); tokio::task::spawn_blocking(move || -> Result<(), Error> { std::fs::create_dir_all(&dir)?; @@ -970,18 +1056,12 @@ impl StateMachine for RocksDBStateMachine { } fn flush(&self) -> Result<(), Error> { - let db = self.db.load(); - - // Step 1: Sync WAL to disk (critical!) - // true = sync to disk - db.flush_wal(true).map_err(|e| StorageError::DbError(e.to_string()))?; - // Step 2: Flush memtables to SST files - db.flush().map_err(|e| StorageError::DbError(e.to_string()))?; - - // Persist state machine metadata (last_applied_index, last_applied_term, snapshot_metadata) - self.persist_state_machine_metadata()?; - - Ok(()) + self.with_db(|db| { + db.flush_wal(true).map_err(|e| StorageError::DbError(e.to_string()))?; + db.flush().map_err(|e| StorageError::DbError(e.to_string()))?; + Ok(()) + })?; + self.persist_state_machine_metadata() } async fn flush_async(&self) -> Result<(), Error> { @@ -990,29 +1070,26 @@ impl StateMachine for RocksDBStateMachine { #[instrument(skip(self))] async fn reset(&self) -> Result<(), Error> { - let db = self.db.load(); - let cf = db - .cf_handle(STATE_MACHINE_CF) - .ok_or_else(|| StorageError::DbError("State machine CF not found".to_string()))?; - - // Delete all keys in the state machine - let mut batch = WriteBatch::default(); - let iter = db.iterator_cf(&cf, IteratorMode::Start); - - for item in iter { - let (key, _) = item.map_err(|e| StorageError::DbError(e.to_string()))?; - batch.delete_cf(&cf, &key); - } - - db.write(&batch).map_err(|e| StorageError::DbError(e.to_string()))?; + self.with_db(|db| { + let cf = db.cf_handle(STATE_MACHINE_CF).ok_or_else(|| { + Error::System(d_engine_core::SystemError::Storage(StorageError::DbError( + "State machine CF not found".to_string(), + ))) + })?; + let mut batch = WriteBatch::default(); + let iter = db.iterator_cf(&cf, IteratorMode::Start); + for item in iter { + let (key, _) = item.map_err(|e| StorageError::DbError(e.to_string()))?; + batch.delete_cf(&cf, &key); + } + db.write(&batch).map_err(|e| StorageError::DbError(e.to_string()))?; + Ok(()) + })?; - // Reset metadata self.last_applied_index.store(0, Ordering::SeqCst); self.last_applied_term.store(0, Ordering::SeqCst); *self.last_snapshot_metadata.write() = None; - // Note: Lease is managed by NodeBuilder and doesn't need reset - self.persist_state_machine_metadata()?; self.persist_snapshot_metadata()?; @@ -1049,16 +1126,18 @@ impl StateMachine for RocksDBStateMachine { ); // Delete expired keys from RocksDB - let db = self.db.load(); - let cf = db - .cf_handle(STATE_MACHINE_CF) - .ok_or_else(|| StorageError::DbError("State machine CF not found".to_string()))?; - - let mut batch = WriteBatch::default(); - for key in &expired_keys { - batch.delete_cf(&cf, key); - } - + let batch = self.with_db(|db| { + let cf = db.cf_handle(STATE_MACHINE_CF).ok_or_else(|| { + Error::System(d_engine_core::SystemError::Storage(StorageError::DbError( + "State machine CF not found".to_string(), + ))) + })?; + let mut batch = WriteBatch::default(); + for key in &expired_keys { + batch.delete_cf(&cf, key); + } + Ok(batch) + })?; self.apply_batch(batch)?; info!( @@ -1075,9 +1154,43 @@ impl StateMachine for RocksDBStateMachine { ) -> Result { self.scan_prefix(prefix) } + + fn get_multi( + &self, + keys: &[Bytes], + ) -> Result>, Error> { + if !self.is_serving.load(Ordering::SeqCst) { + return Err(StorageError::NotServing( + "State machine is restoring from snapshot".to_string(), + ) + .into()); + } + self.with_db(|db| { + let cf = db.cf_handle(STATE_MACHINE_CF).ok_or_else(|| { + Error::System(d_engine_core::SystemError::Storage(StorageError::DbError( + "State machine CF not found".to_string(), + ))) + })?; + // One snapshot covers the full batch: all keys read from the same point-in-time. + // Snapshot lives only within this call — no lifetime escaping, no unsafe. + let snap = db.snapshot(); + keys.iter() + .map(|k| { + snap.get_cf(&cf, k) + .map(|v| v.map(|b| Bytes::copy_from_slice(&b))) + .map_err(|e| StorageError::DbError(e.to_string()).into()) + }) + .collect() + }) + } } impl Drop for RocksDBStateMachine { fn drop(&mut self) { + // If close_db() was already called, DB is None — nothing to flush or cancel. + if self.db.load().is_none() { + return; + } + // save_hard_state() persists last_applied metadata before flush // This is critical to prevent replay of already-applied entries on restart if let Err(e) = self.save_hard_state() { @@ -1091,8 +1204,11 @@ impl Drop for RocksDBStateMachine { debug!("RocksDBStateMachine flushed successfully on drop"); } - // This ensures flush operations are truly finished - self.db.load().cancel_all_background_work(true); // true = wait for completion - debug!("RocksDB background work cancelled on drop"); + // Ensure flush operations are truly finished (no Arc clone needed — just &DB) + let guard = self.db.load(); + if let Some(db) = guard.as_deref() { + db.cancel_all_background_work(true); + debug!("RocksDB background work cancelled on drop"); + } } } diff --git a/d-engine-server/src/storage/adaptors/rocksdb/rocksdb_state_machine_test.rs b/d-engine-server/src/storage/adaptors/rocksdb/rocksdb_state_machine_test.rs index a4fd99c9..43b69e54 100644 --- a/d-engine-server/src/storage/adaptors/rocksdb/rocksdb_state_machine_test.rs +++ b/d-engine-server/src/storage/adaptors/rocksdb/rocksdb_state_machine_test.rs @@ -343,6 +343,153 @@ async fn test_scan_prefix_all_0xff_no_upper_bound() { assert_eq!(result.entries[0].0, Bytes::copy_from_slice(key)); } +// ── get_multi tests ─────────────────────────────────────────────────────────── + +/// RocksDB get_multi uses db.snapshot() to guarantee snapshot isolation. +/// +/// A snapshot taken at the start of the batch read freezes the visible state: +/// writes committed after the snapshot is taken are invisible to that batch. +/// This means every (key, value) pair in the result comes from the same +/// point-in-time state, preventing torn reads across apply indexes. +/// +/// This test verifies two things: +/// 1. get_multi returns the current state accurately after each apply. +/// 2. Each result set is internally consistent — no v1/v2 mix is possible. +#[tokio::test] +async fn test_get_multi_rocksdb_snapshot_reads_consistent_state() { + let dir = tempfile::TempDir::new().unwrap(); + let (_storage, sm) = RocksDBUnifiedEngine::open(dir.path()).unwrap(); + + // Apply v1: initial deployment state + sm.apply_chunk(&[ + insert_at(b"/svc/addr", b"10.0.0.1:8080", 1), + insert_at(b"/svc/version", b"v1", 2), + insert_at(b"/svc/health", b"healthy", 3), + ]) + .await + .unwrap(); + + let keys = vec![ + Bytes::from_static(b"/svc/addr"), + Bytes::from_static(b"/svc/version"), + Bytes::from_static(b"/svc/health"), + ]; + + // Read after v1 — must see full v1 state + let v1_result = sm.get_multi(&keys).unwrap(); + assert_eq!( + v1_result[0], + Some(Bytes::from_static(b"10.0.0.1:8080")), + "addr must be v1" + ); + assert_eq!( + v1_result[1], + Some(Bytes::from_static(b"v1")), + "version must be v1" + ); + assert_eq!( + v1_result[2], + Some(Bytes::from_static(b"healthy")), + "health must be v1" + ); + + // Apply v2: new deployment — all three fields transition + sm.apply_chunk(&[ + insert_at(b"/svc/addr", b"10.0.0.2:8080", 4), + insert_at(b"/svc/version", b"v2", 5), + insert_at(b"/svc/health", b"starting", 6), + ]) + .await + .unwrap(); + + // Read after v2 — must see full v2 state, no v1 values mixed in + let v2_result = sm.get_multi(&keys).unwrap(); + assert_eq!( + v2_result[0], + Some(Bytes::from_static(b"10.0.0.2:8080")), + "addr must be v2" + ); + assert_eq!( + v2_result[1], + Some(Bytes::from_static(b"v2")), + "version must be v2" + ); + assert_eq!( + v2_result[2], + Some(Bytes::from_static(b"starting")), + "health must be v2" + ); + + // Verify each result set is internally consistent (snapshot isolation guarantee). + // A torn read — e.g., addr=v2 + version=v1 — would indicate the snapshot is not + // held for the full batch, which would be a correctness violation. + for (label, result) in [("v1_result", &v1_result), ("v2_result", &v2_result)] { + let is_v1 = result[0] == Some(Bytes::from_static(b"10.0.0.1:8080")) + && result[1] == Some(Bytes::from_static(b"v1")) + && result[2] == Some(Bytes::from_static(b"healthy")); + let is_v2 = result[0] == Some(Bytes::from_static(b"10.0.0.2:8080")) + && result[1] == Some(Bytes::from_static(b"v2")) + && result[2] == Some(Bytes::from_static(b"starting")); + assert!( + is_v1 || is_v2, + "{label}: torn read detected — result is neither pure v1 nor pure v2: {result:?}" + ); + } +} + +// ── close_db() tests ────────────────────────────────────────────────────────── +// +// close_db() is a permanent, one-way shutdown that releases the RocksDB LOCK file +// immediately, without waiting for Arc to be the sole owner. It is separate +// from stop() because stop()/start() must remain a reversible cycle (used during +// snapshot restoration). + +/// close_db() marks SM not-running and makes reads return an error. +#[tokio::test] +async fn test_close_db_prevents_reads() { + let dir = TempDir::new().unwrap(); + let (_storage, sm) = RocksDBUnifiedEngine::open(dir.path()).unwrap(); + sm.start().await.unwrap(); + sm.apply_chunk(&[insert_at(b"k", b"v", 1)]).await.unwrap(); + assert_eq!(sm.get(b"k").unwrap(), Some(Bytes::from_static(b"v"))); + + sm.close_db(); + + assert!(!sm.is_running(), "is_running must be false after close_db"); + assert!(sm.get(b"k").is_err(), "get must fail after close_db"); + assert!( + sm.get_multi(&[Bytes::from_static(b"k")]).is_err(), + "get_multi must fail after close_db" + ); +} + +/// stop() after close_db() must succeed without panicking. +/// Node::stop() calls sm.stop() after EmbeddedEngine has already called close_db(); +/// that sequence must be safe. +#[test] +fn test_stop_after_close_db_is_safe() { + let dir = TempDir::new().unwrap(); + let (_storage, sm) = RocksDBUnifiedEngine::open(dir.path()).unwrap(); + sm.close_db(); + assert!(sm.stop().is_ok(), "stop after close_db must not error"); +} + +/// Drop after close_db() must not panic or double-flush. +/// EmbeddedEngine::stop() calls close_db() before the Raft loop exits; the SM +/// is then dropped — Drop must detect the closed state and skip the flush. +#[tokio::test] +async fn test_drop_after_close_db_does_not_panic() { + let dir = TempDir::new().unwrap(); + { + let (_storage, sm) = RocksDBUnifiedEngine::open(dir.path()).unwrap(); + sm.start().await.unwrap(); + sm.apply_chunk(&[insert_at(b"k", b"v", 1)]).await.unwrap(); + sm.close_db(); + // sm drops here — Drop must be a no-op + } + // reaching here without panic means the test passes +} + // ── Helpers ─────────────────────────────────────────────────────────────────── fn insert_at( diff --git a/d-engine-server/src/storage/adaptors/rocksdb/rocksdb_unified_engine_test.rs b/d-engine-server/src/storage/adaptors/rocksdb/rocksdb_unified_engine_test.rs index d38d60d6..2f7332e0 100644 --- a/d-engine-server/src/storage/adaptors/rocksdb/rocksdb_unified_engine_test.rs +++ b/d-engine-server/src/storage/adaptors/rocksdb/rocksdb_unified_engine_test.rs @@ -15,7 +15,7 @@ use super::RocksDBStateMachine; use super::RocksDBStorageEngine; use super::RocksDBUnifiedEngine; use crate::StateMachine; -use crate::storage::DefaultLease; +use crate::storage::TtlLease; // Both test suites call build() twice within a single persistence test (write data → drop → // build() again → verify data survived). Both builders therefore: @@ -169,8 +169,8 @@ fn test_concurrent_open_same_path_fails() { // ── Lease round-trip tests ──────────────────────────────────────────────────── -fn make_lease() -> Arc { - Arc::new(DefaultLease::new(LeaseConfig { +fn make_lease() -> Arc { + Arc::new(TtlLease::new(LeaseConfig { cleanup_interval_ms: 1000, max_cleanup_duration_ms: 10, })) diff --git a/d-engine-server/src/storage/lease.rs b/d-engine-server/src/storage/lease.rs index 459afb6b..0fbd5f4e 100644 --- a/d-engine-server/src/storage/lease.rs +++ b/d-engine-server/src/storage/lease.rs @@ -55,7 +55,7 @@ use crate::Result; /// - Per-key overhead: ~50 bytes (single DashMap entry) /// - Expired keys are removed automatically during cleanup #[derive(Debug)] -pub struct DefaultLease { +pub struct TtlLease { /// Lease cleanup configuration (immutable after creation) config: d_engine_core::config::LeaseConfig, @@ -72,7 +72,7 @@ pub struct DefaultLease { has_keys: AtomicBool, } -impl DefaultLease { +impl TtlLease { /// Creates a new default lease manager with the given configuration. /// /// # Arguments @@ -179,7 +179,7 @@ impl DefaultLease { } } -impl Lease for DefaultLease { +impl Lease for TtlLease { /// Register a key with TTL (Time-To-Live). /// /// # TTL Semantics diff --git a/d-engine-server/src/storage/lease_integration_test.rs b/d-engine-server/src/storage/lease_integration_test.rs index bb5e1cd8..3730276c 100644 --- a/d-engine-server/src/storage/lease_integration_test.rs +++ b/d-engine-server/src/storage/lease_integration_test.rs @@ -4,7 +4,7 @@ //! Integration tests for Lease functionality across the full stack //! //! This module contains comprehensive tests for: -//! - DefaultLease unit tests (moved from ttl_manager.rs) +//! - TtlLease unit tests (moved from ttl_manager.rs) //! - FileStateMachine Lease integration tests //! - RocksDBStateMachine Lease integration tests @@ -15,12 +15,12 @@ mod lease_tests { use bytes::Bytes; use d_engine_core::Lease; - use crate::storage::DefaultLease; + use crate::storage::TtlLease; #[test] fn test_register_and_get_expired() { let config = d_engine_core::config::LeaseConfig::default(); - let manager = DefaultLease::new(config); + let manager = TtlLease::new(config); // Register keys with 1 second TTL manager.register(Bytes::from("key1"), 1); @@ -43,7 +43,7 @@ mod lease_tests { #[test] fn test_unregister() { let config = d_engine_core::config::LeaseConfig::default(); - let manager = DefaultLease::new(config); + let manager = TtlLease::new(config); manager.register(Bytes::from("key1"), 10); assert_eq!(manager.len(), 1); @@ -55,7 +55,7 @@ mod lease_tests { #[test] fn test_update_ttl() { let config = d_engine_core::config::LeaseConfig::default(); - let manager = DefaultLease::new(config); + let manager = TtlLease::new(config); // Register with 10 seconds manager.register(Bytes::from("key1"), 10); @@ -70,13 +70,13 @@ mod lease_tests { #[test] fn test_snapshot_roundtrip() { let config = d_engine_core::config::LeaseConfig::default(); - let manager = DefaultLease::new(config.clone()); + let manager = TtlLease::new(config.clone()); manager.register(Bytes::from("key1"), 3600); manager.register(Bytes::from("key2"), 7200); let snapshot = manager.to_snapshot(); - let restored = DefaultLease::from_snapshot(&snapshot, config); + let restored = TtlLease::from_snapshot(&snapshot, config); assert_eq!(restored.len(), 2); } @@ -84,7 +84,7 @@ mod lease_tests { #[test] fn test_snapshot_filters_expired() { let config = d_engine_core::config::LeaseConfig::default(); - let manager = DefaultLease::new(config.clone()); + let manager = TtlLease::new(config.clone()); // Register key with 1 second TTL manager.register(Bytes::from("key1"), 1); @@ -93,7 +93,7 @@ mod lease_tests { sleep(Duration::from_secs(2)); let snapshot = manager.to_snapshot(); - let restored = DefaultLease::from_snapshot(&snapshot, config); + let restored = TtlLease::from_snapshot(&snapshot, config); // Only key2 should be restored assert_eq!(restored.len(), 1); @@ -108,8 +108,8 @@ mod file_state_machine_tests { use tempfile::TempDir; use tokio::time::sleep; - use crate::storage::DefaultLease; use crate::storage::FileStateMachine; + use crate::storage::TtlLease; /// Helper to create a FileStateMachine with lease injected for testing async fn create_file_state_machine_with_lease( @@ -117,7 +117,7 @@ mod file_state_machine_tests { lease_config: d_engine_core::config::LeaseConfig, ) -> FileStateMachine { let mut sm = FileStateMachine::new(path).await.unwrap(); - let lease = std::sync::Arc::new(DefaultLease::new(lease_config)); + let lease = std::sync::Arc::new(TtlLease::new(lease_config)); sm.set_lease(lease); sm.load_lease_data().await.unwrap(); sm @@ -779,7 +779,7 @@ mod rocksdb_state_machine_tests { lease_config: d_engine_core::config::LeaseConfig, ) -> (RocksDBStorageEngine, RocksDBStateMachine) { let (storage, mut sm) = RocksDBUnifiedEngine::open(&path).unwrap(); - let lease = std::sync::Arc::new(crate::storage::DefaultLease::new(lease_config)); + let lease = std::sync::Arc::new(crate::storage::TtlLease::new(lease_config)); sm.set_lease(lease); sm.load_lease_data().await.unwrap(); (storage, sm) diff --git a/d-engine-server/src/storage/lease_unit_test.rs b/d-engine-server/src/storage/lease_unit_test.rs index e83469b2..550ea720 100644 --- a/d-engine-server/src/storage/lease_unit_test.rs +++ b/d-engine-server/src/storage/lease_unit_test.rs @@ -1,4 +1,4 @@ -//! Comprehensive unit tests for DefaultLease implementation +//! Comprehensive unit tests for TtlLease implementation //! //! Tests cover: //! - Basic registration, expiration, and cleanup @@ -14,7 +14,7 @@ use std::time::SystemTime; use bytes::Bytes; use d_engine_core::Lease; -use crate::storage::lease::DefaultLease; +use crate::storage::lease::TtlLease; fn default_config() -> d_engine_core::config::LeaseConfig { d_engine_core::config::LeaseConfig { @@ -29,7 +29,7 @@ fn default_config() -> d_engine_core::config::LeaseConfig { #[test] fn test_register_and_is_expired() { - let lease = DefaultLease::new(default_config()); + let lease = TtlLease::new(default_config()); lease.register(Bytes::from("key1"), 1); assert!(!lease.is_expired(b"key1")); @@ -40,7 +40,7 @@ fn test_register_and_is_expired() { #[test] fn test_unregister() { - let lease = DefaultLease::new(default_config()); + let lease = TtlLease::new(default_config()); lease.register(Bytes::from("key1"), 10); assert_eq!(lease.len(), 1); @@ -52,7 +52,7 @@ fn test_unregister() { #[test] fn test_unregister_nonexistent_key() { - let lease = DefaultLease::new(default_config()); + let lease = TtlLease::new(default_config()); // Should not panic on unregistering non-existent key lease.unregister(b"nonexistent"); @@ -61,7 +61,7 @@ fn test_unregister_nonexistent_key() { #[test] fn test_get_expired_keys() { - let lease = DefaultLease::new(default_config()); + let lease = TtlLease::new(default_config()); lease.register(Bytes::from("key1"), 1); lease.register(Bytes::from("key2"), 1); @@ -75,7 +75,7 @@ fn test_get_expired_keys() { #[test] fn test_is_empty() { - let lease = DefaultLease::new(default_config()); + let lease = TtlLease::new(default_config()); assert!(lease.is_empty()); lease.register(Bytes::from("key1"), 10); @@ -91,7 +91,7 @@ fn test_is_empty() { #[test] fn test_has_lease_keys() { - let lease = DefaultLease::new(default_config()); + let lease = TtlLease::new(default_config()); assert!(!lease.has_lease_keys()); lease.register(Bytes::from("key1"), 10); @@ -103,7 +103,7 @@ fn test_has_lease_keys() { #[test] fn test_may_have_expired_keys_false_when_no_keys() { - let lease = DefaultLease::new(default_config()); + let lease = TtlLease::new(default_config()); // No keys registered, should return false immediately assert!(!lease.may_have_expired_keys(SystemTime::now())); @@ -111,7 +111,7 @@ fn test_may_have_expired_keys_false_when_no_keys() { #[test] fn test_may_have_expired_keys_false_when_all_valid() { - let lease = DefaultLease::new(default_config()); + let lease = TtlLease::new(default_config()); lease.register(Bytes::from("key1"), 3600); lease.register(Bytes::from("key2"), 7200); @@ -122,7 +122,7 @@ fn test_may_have_expired_keys_false_when_all_valid() { #[test] fn test_may_have_expired_keys_true_when_has_expired() { - let lease = DefaultLease::new(default_config()); + let lease = TtlLease::new(default_config()); lease.register(Bytes::from("key1"), 1); sleep(Duration::from_secs(2)); @@ -137,7 +137,7 @@ fn test_may_have_expired_keys_true_when_has_expired() { #[test] fn test_get_expiration() { - let lease = DefaultLease::new(default_config()); + let lease = TtlLease::new(default_config()); lease.register(Bytes::from("key1"), 3600); @@ -148,7 +148,7 @@ fn test_get_expiration() { #[test] fn test_get_expiration_not_found() { - let lease = DefaultLease::new(default_config()); + let lease = TtlLease::new(default_config()); let expiration = lease.get_expiration(b"nonexistent"); assert!(expiration.is_none()); @@ -160,7 +160,7 @@ fn test_get_expiration_not_found() { #[test] fn test_register_updates_existing_key() { - let lease = DefaultLease::new(default_config()); + let lease = TtlLease::new(default_config()); // Register with 10s TTL lease.register(Bytes::from("key1"), 10); @@ -181,7 +181,7 @@ fn test_register_updates_existing_key() { #[test] fn test_multiple_keys_same_expiration() { - let lease = DefaultLease::new(default_config()); + let lease = TtlLease::new(default_config()); let ttl = 5u64; lease.register(Bytes::from("key1"), ttl); @@ -199,7 +199,7 @@ fn test_multiple_keys_same_expiration() { #[test] fn test_mixed_expired_and_valid_keys() { - let lease = DefaultLease::new(default_config()); + let lease = TtlLease::new(default_config()); lease.register(Bytes::from("expire_soon"), 1); lease.register(Bytes::from("expire_later"), 3600); @@ -222,13 +222,13 @@ fn test_mixed_expired_and_valid_keys() { #[test] fn test_snapshot_roundtrip() { let config = default_config(); - let lease1 = DefaultLease::new(config.clone()); + let lease1 = TtlLease::new(config.clone()); lease1.register(Bytes::from("key1"), 3600); lease1.register(Bytes::from("key2"), 7200); let snapshot = lease1.to_snapshot(); - let lease2 = DefaultLease::from_snapshot(&snapshot, config); + let lease2 = TtlLease::from_snapshot(&snapshot, config); assert_eq!(lease2.len(), 2); assert!(!lease2.is_expired(b"key1")); @@ -238,7 +238,7 @@ fn test_snapshot_roundtrip() { #[test] fn test_snapshot_roundtrip_filters_expired_keys() { let config = default_config(); - let lease1 = DefaultLease::new(config.clone()); + let lease1 = TtlLease::new(config.clone()); lease1.register(Bytes::from("valid_key"), 3600); lease1.register(Bytes::from("expired_key"), 1); @@ -248,7 +248,7 @@ fn test_snapshot_roundtrip_filters_expired_keys() { let snapshot = lease1.to_snapshot(); // from_snapshot should filter out expired keys - let lease2 = DefaultLease::from_snapshot(&snapshot, config); + let lease2 = TtlLease::from_snapshot(&snapshot, config); // Only valid_key should remain assert_eq!(lease2.len(), 1); @@ -260,14 +260,14 @@ fn test_from_snapshot_with_invalid_data() { let config = default_config(); // from_snapshot with invalid data should return empty lease - let lease = DefaultLease::from_snapshot(&[0xFF, 0xFE], config); + let lease = TtlLease::from_snapshot(&[0xFF, 0xFE], config); assert_eq!(lease.len(), 0); assert!(!lease.has_lease_keys()); } #[test] fn test_reload() { - let lease = DefaultLease::new(default_config()); + let lease = TtlLease::new(default_config()); lease.register(Bytes::from("key1"), 3600); let snapshot = lease.to_snapshot(); @@ -283,7 +283,7 @@ fn test_reload() { #[test] fn test_reload_invalid_data() { - let lease = DefaultLease::new(default_config()); + let lease = TtlLease::new(default_config()); lease.register(Bytes::from("key1"), 3600); assert_eq!(lease.len(), 1); @@ -306,7 +306,7 @@ fn test_reload_clears_apply_counter() { max_cleanup_duration_ms: 1, }; - let lease = DefaultLease::new(config.clone()); + let lease = TtlLease::new(config.clone()); lease.register(Bytes::from("key1"), 3600); @@ -333,7 +333,7 @@ fn test_reload_clears_apply_counter() { fn test_on_apply_piggyback_removed() { let config = default_config(); - let lease = DefaultLease::new(config); + let lease = TtlLease::new(config); lease.register(Bytes::from("key1"), 1); sleep(Duration::from_secs(2)); @@ -354,6 +354,6 @@ fn test_custom_cleanup_interval() { max_cleanup_duration_ms: 1, }; - let lease = DefaultLease::new(config); + let lease = TtlLease::new(config); assert_eq!(lease.config().cleanup_interval_ms, 5000); } diff --git a/d-engine-server/src/storage/mod.rs b/d-engine-server/src/storage/mod.rs index ba98c842..432eb9e1 100644 --- a/d-engine-server/src/storage/mod.rs +++ b/d-engine-server/src/storage/mod.rs @@ -26,7 +26,7 @@ pub use adaptors::*; pub use buffered::*; // Re-export Lease trait from core for convenience pub use d_engine_core::Lease; -pub use lease::DefaultLease; +pub use lease::TtlLease; #[cfg(test)] mod lease_integration_test; diff --git a/d-engine-server/src/test_utils/mock/mock_node_builder.rs b/d-engine-server/src/test_utils/mock/mock_node_builder.rs index aeeb550d..8cbeb150 100644 --- a/d-engine-server/src/test_utils/mock/mock_node_builder.rs +++ b/d-engine-server/src/test_utils/mock/mock_node_builder.rs @@ -316,6 +316,7 @@ impl MockBuilder { let cmd_tx = raft.cmd_sender(); let node_config = raft.ctx.node_config.clone(); let membership = raft.ctx.membership.clone(); + let read_lease = raft.read_lease(); let (rpc_ready_tx, _rpc_ready_rx) = watch::channel(false); let leader_notifier = LeaderNotifier::new(); let (_membership_tx, membership_rx) = watch::channel(MembershipSnapshot::default()); @@ -325,6 +326,7 @@ impl MockBuilder { raft_core: Arc::new(Mutex::new(raft)), membership, event_tx, + read_handle: crate::api::StandaloneReadHandle::new(None, cmd_tx.clone()), cmd_tx, ready: AtomicBool::new(false), rpc_ready_tx, @@ -336,9 +338,11 @@ impl MockBuilder { #[cfg(feature = "watch")] _watch_dispatcher_handle: None, sm_worker_handle: std::sync::Mutex::new(None), + read_actor_handle: std::sync::Mutex::new(None), _commit_handler_handle: None, _lease_cleanup_handle: None, shutdown_signal, + read_lease, } } @@ -360,6 +364,7 @@ impl MockBuilder { .expect("Should succeed to validate RaftNodeConfig") }); let membership = raft.ctx.membership.clone(); + let read_lease = raft.read_lease(); trace!( node_config.raft.election.election_timeout_min, "build_node_with_rpc_server" @@ -374,6 +379,7 @@ impl MockBuilder { raft_core: Arc::new(Mutex::new(raft)), membership, event_tx, + read_handle: crate::api::StandaloneReadHandle::new(None, cmd_tx.clone()), cmd_tx, ready: AtomicBool::new(false), rpc_ready_tx, @@ -385,9 +391,11 @@ impl MockBuilder { #[cfg(feature = "watch")] _watch_dispatcher_handle: None, sm_worker_handle: std::sync::Mutex::new(None), + read_actor_handle: std::sync::Mutex::new(None), _commit_handler_handle: None, _lease_cleanup_handle: None, shutdown_signal: shutdown.clone(), + read_lease, }); let node_clone = node.clone(); let listen_address = node_config_arc.cluster.listen_address; diff --git a/d-engine-server/tests/cas_operations/bank_transfer_invariant_embedded.rs b/d-engine-server/tests/cas_operations/bank_transfer_invariant_embedded.rs index d7e66711..838b3e77 100644 --- a/d-engine-server/tests/cas_operations/bank_transfer_invariant_embedded.rs +++ b/d-engine-server/tests/cas_operations/bank_transfer_invariant_embedded.rs @@ -3,7 +3,7 @@ use d_engine_core::ClientApi; use d_engine_core::ClientApiError; use d_engine_server::RocksDBUnifiedEngine; -use d_engine_server::api::EmbeddedEngine; +use d_engine_server::api::DefaultEmbeddedEngine; use serial_test::serial; use std::sync::Arc; use std::time::Duration; @@ -81,7 +81,7 @@ async fn test_concurrent_transfers_preserve_bank_invariant() let config_path = format!("/tmp/d-engine-bank-node{node_id}.toml"); tokio::fs::write(&config_path, &config_str).await?; - let engine = EmbeddedEngine::start_custom( + let engine = DefaultEmbeddedEngine::start_custom( Arc::new(storage), Arc::new(state_machine), Some(&config_path), diff --git a/d-engine-server/tests/cas_operations/distributed_lock_embedded.rs b/d-engine-server/tests/cas_operations/distributed_lock_embedded.rs index 02fee19e..217729ca 100644 --- a/d-engine-server/tests/cas_operations/distributed_lock_embedded.rs +++ b/d-engine-server/tests/cas_operations/distributed_lock_embedded.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use std::time::Duration; use d_engine_core::ClientApi; -use d_engine_server::api::EmbeddedEngine; +use d_engine_server::api::DefaultEmbeddedEngine; use tracing::info; use tracing_test::traced_test; @@ -58,7 +58,7 @@ async fn test_distributed_lock_embedded() -> Result<(), Box Result<(), Box Result<(), Box Result<(), Box Result<(), Box> { let temp_dir = tempfile::tempdir()?; @@ -16,7 +16,7 @@ async fn test_single_node_lifecycle() -> Result<(), Box> } // Start embedded engine — data_dir takes highest priority - let engine = EmbeddedEngine::start(&data_dir).await?; + let engine = DefaultEmbeddedEngine::start(&data_dir).await?; // Clean up environment variables immediately unsafe { @@ -80,7 +80,7 @@ async fn test_leader_notification() -> Result<(), Box> { std::env::set_var("RAFT__CLUSTER__LISTEN_ADDRESS", "127.0.0.1:9002"); } - let engine = EmbeddedEngine::start(&data_dir).await?; + let engine = DefaultEmbeddedEngine::start(&data_dir).await?; unsafe { std::env::remove_var("RAFT__CLUSTER__NODE_ID"); @@ -112,7 +112,7 @@ async fn test_leader_notification() -> Result<(), Box> { async fn test_data_persistence() -> Result<(), Box> { use std::time::Duration; - use d_engine_server::EmbeddedEngine; + use d_engine_server::DefaultEmbeddedEngine; let temp_dir = tempfile::tempdir()?; let data_dir = temp_dir.path().join("db"); @@ -125,7 +125,7 @@ async fn test_data_persistence() -> Result<(), Box> { // First session: write data { - let engine = EmbeddedEngine::start(&data_dir).await?; + let engine = DefaultEmbeddedEngine::start(&data_dir).await?; engine.wait_ready(Duration::from_secs(5)).await?; engine.client().put(b"persist-key".to_vec(), b"persist-value".to_vec()).await?; @@ -138,7 +138,7 @@ async fn test_data_persistence() -> Result<(), Box> { // Second session: verify data still exists { - let engine = EmbeddedEngine::start(&data_dir).await?; + let engine = DefaultEmbeddedEngine::start(&data_dir).await?; engine.wait_ready(Duration::from_secs(5)).await?; let value = engine.client().get_linearizable(b"persist-key".to_vec()).await?; diff --git a/d-engine-server/tests/cluster_state_and_metadata/cluster_state_consensus_embedded.rs b/d-engine-server/tests/cluster_state_and_metadata/cluster_state_consensus_embedded.rs index 3999243e..40d9c416 100644 --- a/d-engine-server/tests/cluster_state_and_metadata/cluster_state_consensus_embedded.rs +++ b/d-engine-server/tests/cluster_state_and_metadata/cluster_state_consensus_embedded.rs @@ -1,4 +1,4 @@ -//! Integration tests for EmbeddedEngine cluster state APIs (Ticket #234) +//! Integration tests for DefaultEmbeddedEngine cluster state APIs (Ticket #234) //! //! These tests verify is_leader() and leader_info() behavior in multi-node scenarios: //! - Multi-node leader election @@ -7,7 +7,7 @@ use crate::common::get_available_ports; use crate::common::node_config; -use d_engine_server::EmbeddedEngine; +use d_engine_server::DefaultEmbeddedEngine; use std::sync::Arc; use std::time::Duration; use tracing::info; @@ -18,7 +18,7 @@ use d_engine_server::RocksDBUnifiedEngine; /// Test: 3-node cluster should elect exactly one leader /// /// Setup: -/// - Start 3 EmbeddedEngine instances as a cluster +/// - Start 3 DefaultEmbeddedEngine instances as a cluster /// - Wait for leader election /// /// Verification: @@ -105,25 +105,34 @@ listen_address = '127.0.0.1:{}' let db_path1 = config1.cluster.db_root_dir.join("node1/db"); tokio::fs::create_dir_all(&db_path1).await?; let (storage1, sm1) = RocksDBUnifiedEngine::open(&db_path1)?; - let engine1 = - EmbeddedEngine::start_custom(Arc::new(storage1), Arc::new(sm1), Some(node1_config_path)) - .await?; + let engine1 = DefaultEmbeddedEngine::start_custom( + Arc::new(storage1), + Arc::new(sm1), + Some(node1_config_path), + ) + .await?; let config2 = node_config(&node2_config); let db_path2 = config2.cluster.db_root_dir.join("node2/db"); tokio::fs::create_dir_all(&db_path2).await?; let (storage2, sm2) = RocksDBUnifiedEngine::open(&db_path2)?; - let engine2 = - EmbeddedEngine::start_custom(Arc::new(storage2), Arc::new(sm2), Some(node2_config_path)) - .await?; + let engine2 = DefaultEmbeddedEngine::start_custom( + Arc::new(storage2), + Arc::new(sm2), + Some(node2_config_path), + ) + .await?; let config3 = node_config(&node3_config); let db_path3 = config3.cluster.db_root_dir.join("node3/db"); tokio::fs::create_dir_all(&db_path3).await?; let (storage3, sm3) = RocksDBUnifiedEngine::open(&db_path3)?; - let engine3 = - EmbeddedEngine::start_custom(Arc::new(storage3), Arc::new(sm3), Some(node3_config_path)) - .await?; + let engine3 = DefaultEmbeddedEngine::start_custom( + Arc::new(storage3), + Arc::new(sm3), + Some(node3_config_path), + ) + .await?; // Wait for leader election info!("Waiting for leader election..."); @@ -243,25 +252,34 @@ general_raft_timeout_duration_in_ms = 3000 let db_path1 = config1.cluster.db_root_dir.join("node1/db"); tokio::fs::create_dir_all(&db_path1).await?; let (storage1, sm1) = RocksDBUnifiedEngine::open(&db_path1)?; - let engine1 = - EmbeddedEngine::start_custom(Arc::new(storage1), Arc::new(sm1), Some(node1_config_path)) - .await?; + let engine1 = DefaultEmbeddedEngine::start_custom( + Arc::new(storage1), + Arc::new(sm1), + Some(node1_config_path), + ) + .await?; let config2 = node_config(&node2_config); let db_path2 = config2.cluster.db_root_dir.join("node2/db"); tokio::fs::create_dir_all(&db_path2).await?; let (storage2, sm2) = RocksDBUnifiedEngine::open(&db_path2)?; - let engine2 = - EmbeddedEngine::start_custom(Arc::new(storage2), Arc::new(sm2), Some(node2_config_path)) - .await?; + let engine2 = DefaultEmbeddedEngine::start_custom( + Arc::new(storage2), + Arc::new(sm2), + Some(node2_config_path), + ) + .await?; let config3 = node_config(&node3_config); let db_path3 = config3.cluster.db_root_dir.join("node3/db"); tokio::fs::create_dir_all(&db_path3).await?; let (storage3, sm3) = RocksDBUnifiedEngine::open(&db_path3)?; - let engine3 = - EmbeddedEngine::start_custom(Arc::new(storage3), Arc::new(sm3), Some(node3_config_path)) - .await?; + let engine3 = DefaultEmbeddedEngine::start_custom( + Arc::new(storage3), + Arc::new(sm3), + Some(node3_config_path), + ) + .await?; // Wait for leader election info!("Waiting for leader election..."); @@ -413,25 +431,34 @@ general_raft_timeout_duration_in_ms = 3000 let db_path1 = config1.cluster.db_root_dir.join("node1/db"); tokio::fs::create_dir_all(&db_path1).await?; let (storage1, sm1) = RocksDBUnifiedEngine::open(&db_path1)?; - let engine1 = - EmbeddedEngine::start_custom(Arc::new(storage1), Arc::new(sm1), Some(node1_config_path)) - .await?; + let engine1 = DefaultEmbeddedEngine::start_custom( + Arc::new(storage1), + Arc::new(sm1), + Some(node1_config_path), + ) + .await?; let config2 = node_config(&node2_config); let db_path2 = config2.cluster.db_root_dir.join("node2/db"); tokio::fs::create_dir_all(&db_path2).await?; let (storage2, sm2) = RocksDBUnifiedEngine::open(&db_path2)?; - let engine2 = - EmbeddedEngine::start_custom(Arc::new(storage2), Arc::new(sm2), Some(node2_config_path)) - .await?; + let engine2 = DefaultEmbeddedEngine::start_custom( + Arc::new(storage2), + Arc::new(sm2), + Some(node2_config_path), + ) + .await?; let config3 = node_config(&node3_config); let db_path3 = config3.cluster.db_root_dir.join("node3/db"); tokio::fs::create_dir_all(&db_path3).await?; let (storage3, sm3) = RocksDBUnifiedEngine::open(&db_path3)?; - let engine3 = - EmbeddedEngine::start_custom(Arc::new(storage3), Arc::new(sm3), Some(node3_config_path)) - .await?; + let engine3 = DefaultEmbeddedEngine::start_custom( + Arc::new(storage3), + Arc::new(sm3), + Some(node3_config_path), + ) + .await?; // Wait for initial leader election tokio::time::sleep(Duration::from_secs(5)).await; diff --git a/d-engine-server/tests/consistent_reads/lease_read_embedded.rs b/d-engine-server/tests/consistent_reads/lease_read_embedded.rs index 1c70afa0..4ddb6cb2 100644 --- a/d-engine-server/tests/consistent_reads/lease_read_embedded.rs +++ b/d-engine-server/tests/consistent_reads/lease_read_embedded.rs @@ -9,14 +9,14 @@ use std::time::Duration; use d_engine_core::ClientApi; -use d_engine_server::EmbeddedEngine; +use d_engine_server::DefaultEmbeddedEngine; use tempfile::TempDir; use tracing_test::traced_test; use crate::common::get_available_ports; -/// Helper to create a test EmbeddedEngine with lease configuration -async fn create_test_engine_with_lease(test_name: &str) -> (EmbeddedEngine, TempDir) { +/// Helper to create a test DefaultEmbeddedEngine with lease configuration +async fn create_test_engine_with_lease(test_name: &str) -> (DefaultEmbeddedEngine, TempDir) { let temp_dir = tempfile::tempdir().expect("Failed to create temp dir"); let db_path = temp_dir.path().join(test_name); @@ -39,7 +39,7 @@ state_machine_sync_timeout_ms = 2000 ); std::fs::write(&config_path, config_content).expect("Failed to write config"); - let engine = EmbeddedEngine::start_with(config_path.to_str().unwrap()) + let engine = DefaultEmbeddedEngine::start_with(config_path.to_str().unwrap()) .await .expect("Failed to start engine"); diff --git a/d-engine-server/tests/consistent_reads/linearizable_read_batching_embedded.rs b/d-engine-server/tests/consistent_reads/linearizable_read_batching_embedded.rs index 6dde2d23..411f338b 100644 --- a/d-engine-server/tests/consistent_reads/linearizable_read_batching_embedded.rs +++ b/d-engine-server/tests/consistent_reads/linearizable_read_batching_embedded.rs @@ -3,12 +3,12 @@ //! These tests verify that batching multiple linearizable read requests //! improves throughput by sharing a single verify_leadership() call. //! -//! Uses embedded mode with EmbeddedEngine for production-ready testing. +//! Uses embedded mode with DefaultEmbeddedEngine for production-ready testing. use std::time::Duration; use bytes::Bytes; -use d_engine_server::EmbeddedEngine; +use d_engine_server::DefaultEmbeddedEngine; use tempfile::TempDir; use tokio::time::Instant; @@ -16,11 +16,11 @@ use tracing_test::traced_test; use crate::common::get_available_ports; -/// Helper to create a test EmbeddedEngine with batching config +/// Helper to create a test DefaultEmbeddedEngine with batching config async fn create_engine_with_batching( test_name: &str, size_threshold: usize, -) -> (EmbeddedEngine, TempDir) { +) -> (DefaultEmbeddedEngine, TempDir) { let temp_dir = tempfile::tempdir().expect("Failed to create temp dir"); let db_path = temp_dir.path().join(test_name); @@ -47,7 +47,7 @@ size_threshold = {} ); std::fs::write(&config_path, config_content).expect("Failed to write config"); - let engine = EmbeddedEngine::start_with(config_path.to_str().unwrap()) + let engine = DefaultEmbeddedEngine::start_with(config_path.to_str().unwrap()) .await .expect("Failed to start engine"); diff --git a/d-engine-server/tests/consistent_reads/linearizable_read_consistency_embedded.rs b/d-engine-server/tests/consistent_reads/linearizable_read_consistency_embedded.rs index 336a874b..6617fef3 100644 --- a/d-engine-server/tests/consistent_reads/linearizable_read_consistency_embedded.rs +++ b/d-engine-server/tests/consistent_reads/linearizable_read_consistency_embedded.rs @@ -3,10 +3,10 @@ //! These tests verify that linearizable reads use a fixed read_index calculated //! at request arrival time, preventing unnecessary waiting for concurrent writes. //! -//! Uses embedded mode with EmbeddedEngine for production-ready testing. +//! Uses embedded mode with DefaultEmbeddedEngine for production-ready testing. use crate::common::{create_node_config, get_available_ports, node_config}; -use d_engine_server::EmbeddedEngine; +use d_engine_server::DefaultEmbeddedEngine; use d_engine_server::RocksDBUnifiedEngine; use std::sync::Arc; use std::time::Duration; @@ -15,7 +15,7 @@ use tokio::time::Instant; use tracing_test::traced_test; /// Helper to create a test EmbeddedEngine -async fn create_test_engine(test_name: &str) -> (EmbeddedEngine, TempDir) { +async fn create_test_engine(test_name: &str) -> (DefaultEmbeddedEngine, TempDir) { let temp_dir = tempfile::tempdir().expect("Failed to create temp dir"); let db_path = temp_dir.path().join(test_name); @@ -38,7 +38,7 @@ state_machine_sync_timeout_ms = 2000 ); std::fs::write(&config_path, config_content).expect("Failed to write config"); - let engine = EmbeddedEngine::start_with(config_path.to_str().unwrap()) + let engine = DefaultEmbeddedEngine::start_with(config_path.to_str().unwrap()) .await .expect("Failed to start engine"); @@ -330,7 +330,7 @@ async fn test_read_index_fixed_with_concurrent_writes_multi_node() let config_path = format!("/tmp/d-engine-test-linear-read-node{node_id}.toml"); tokio::fs::write(&config_path, &config_str).await?; - let engine = EmbeddedEngine::start_custom( + let engine = DefaultEmbeddedEngine::start_custom( Arc::new(storage), Arc::new(state_machine), Some(&config_path), diff --git a/d-engine-server/tests/drain_batching/select_fairness_embedded.rs b/d-engine-server/tests/drain_batching/select_fairness_embedded.rs index e2b9e295..8b928d83 100644 --- a/d-engine-server/tests/drain_batching/select_fairness_embedded.rs +++ b/d-engine-server/tests/drain_batching/select_fairness_embedded.rs @@ -11,7 +11,7 @@ use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Duration; -use d_engine_server::EmbeddedEngine; +use d_engine_server::DefaultEmbeddedEngine; use tempfile::TempDir; use tokio::sync::Semaphore; use tokio::time::Instant; @@ -19,7 +19,7 @@ use tokio::time::Instant; use crate::common::get_available_ports; /// Helper to create a single-node test engine with embedded mode -async fn create_test_engine(test_name: &str) -> (EmbeddedEngine, TempDir) { +async fn create_test_engine(test_name: &str) -> (DefaultEmbeddedEngine, TempDir) { let temp_dir = tempfile::tempdir().expect("Failed to create temp dir"); let db_path = temp_dir.path().join(test_name); @@ -43,7 +43,7 @@ max_batch_size = 100 ); std::fs::write(&config_path, config_content).expect("Failed to write config"); - let engine = EmbeddedEngine::start_with(config_path.to_str().unwrap()) + let engine = DefaultEmbeddedEngine::start_with(config_path.to_str().unwrap()) .await .expect("Failed to start engine"); diff --git a/d-engine-server/tests/embedded_client/embedded_client_operations.rs b/d-engine-server/tests/embedded_client/embedded_client_operations.rs index 59a82b1f..83c02451 100644 --- a/d-engine-server/tests/embedded_client/embedded_client_operations.rs +++ b/d-engine-server/tests/embedded_client/embedded_client_operations.rs @@ -9,17 +9,17 @@ use std::time::Duration; use bytes::Bytes; use d_engine_core::ClientApi; -use d_engine_server::EmbeddedEngine; +use d_engine_server::DefaultEmbeddedEngine; use tempfile::TempDir; use crate::common::get_available_ports; -/// Helper to create a test EmbeddedEngine without any lease config section. +/// Helper to create a test DefaultEmbeddedEngine without any lease config section. /// /// Used by regression tests that must verify behaviour with default (zero-config) /// settings — notably #398 where omitting [raft.state_machine.lease] previously /// caused a fatal crash when put_with_ttl was called. -async fn create_test_engine_default_config(test_name: &str) -> (EmbeddedEngine, TempDir) { +async fn create_test_engine_default_config(test_name: &str) -> (DefaultEmbeddedEngine, TempDir) { let temp_dir = tempfile::tempdir().expect("Failed to create temp dir"); let db_path = temp_dir.path().join(test_name); let config_path = temp_dir.path().join("d-engine.toml"); @@ -37,7 +37,7 @@ single_node = true db_path.display() ); std::fs::write(&config_path, config_content).expect("Failed to write config"); - let engine = EmbeddedEngine::start_with(config_path.to_str().unwrap()) + let engine = DefaultEmbeddedEngine::start_with(config_path.to_str().unwrap()) .await .expect("Failed to start engine"); engine.wait_ready(Duration::from_secs(5)).await.expect("Engine not ready"); @@ -45,7 +45,7 @@ single_node = true } /// Helper to create a test EmbeddedEngine -async fn create_test_engine(test_name: &str) -> (EmbeddedEngine, TempDir) { +async fn create_test_engine(test_name: &str) -> (DefaultEmbeddedEngine, TempDir) { let temp_dir = tempfile::tempdir().expect("Failed to create temp dir"); let db_path = temp_dir.path().join(test_name); @@ -65,7 +65,7 @@ single_node = true ); std::fs::write(&config_path, config_content).expect("Failed to write config"); - let engine = EmbeddedEngine::start_with(config_path.to_str().unwrap()) + let engine = DefaultEmbeddedEngine::start_with(config_path.to_str().unwrap()) .await .expect("Failed to start engine"); @@ -709,14 +709,14 @@ async fn test_get_multi_linearizable_empty_keys() { // Watch operations (requires `watch` feature) // ============================================================================= -/// Helper: create an EmbeddedEngine with watch feature enabled. +/// Helper: create an DefaultEmbeddedEngine with watch feature enabled. /// /// Separate from create_test_engine because watch requires the /// `[raft.watch]` config section to activate the WatchRegistry. #[cfg(feature = "watch")] async fn create_watch_engine( test_name: &str -) -> (d_engine_server::EmbeddedEngine, tempfile::TempDir) { +) -> (d_engine_server::DefaultEmbeddedEngine, tempfile::TempDir) { let temp_dir = tempfile::tempdir().expect("Failed to create temp dir"); let db_path = temp_dir.path().join(test_name); @@ -752,7 +752,7 @@ watcher_buffer_size = 10 ); std::fs::write(&config_path, config_content).expect("Failed to write watch config"); - let engine = d_engine_server::EmbeddedEngine::start_with(config_path.to_str().unwrap()) + let engine = d_engine_server::DefaultEmbeddedEngine::start_with(config_path.to_str().unwrap()) .await .expect("Failed to start watch engine"); diff --git a/d-engine-server/tests/failover_and_recovery/leader_failover_embedded.rs b/d-engine-server/tests/failover_and_recovery/leader_failover_embedded.rs index da5eb3a0..6716b83c 100644 --- a/d-engine-server/tests/failover_and_recovery/leader_failover_embedded.rs +++ b/d-engine-server/tests/failover_and_recovery/leader_failover_embedded.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use std::time::Duration; use d_engine_server::RocksDBUnifiedEngine; -use d_engine_server::api::EmbeddedEngine; +use d_engine_server::api::DefaultEmbeddedEngine; use tracing::info; use tracing_test::traced_test; @@ -11,7 +11,7 @@ use crate::common::get_available_ports; use crate::common::node_config; use crate::common::wait_for_new_leader; -/// Test 3-node cluster leader failover with EmbeddedEngine API +/// Test 3-node cluster leader failover with DefaultEmbeddedEngine API /// /// Scenario: /// 1. Start 3-node cluster @@ -59,7 +59,7 @@ async fn test_embedded_leader_failover() -> Result<(), Box Result<(), Box> { configs.push((config_str, config_path)); - let engine = EmbeddedEngine::start_custom( + let engine = DefaultEmbeddedEngine::start_custom( Arc::new(storage), Arc::new(state_machine), Some(&configs[i].1), @@ -293,7 +293,7 @@ async fn test_embedded_node_rejoin() -> Result<(), Box> { opened.ok_or_else(|| last_err.unwrap())? }; - let restarted_engine = EmbeddedEngine::start_custom( + let restarted_engine = DefaultEmbeddedEngine::start_custom( Arc::new(storage), Arc::new(state_machine), Some(&killed_config.1), @@ -368,7 +368,7 @@ async fn test_minority_failure_blocks_writes() -> Result<(), Box Result<(EmbeddedEngine, TempDir), Box> { +/// Helper function to create a test DefaultEmbeddedEngine with RocksDB +async fn setup_engine() -> Result<(DefaultEmbeddedEngine, TempDir), Box> { let temp_dir = TempDir::new()?; let db_path = temp_dir.path().join("db"); @@ -37,7 +37,7 @@ watcher_buffer_size = 10 // Start engine with RocksDB storage let (storage, state_machine) = RocksDBUnifiedEngine::open(&db_path)?; - let engine = EmbeddedEngine::start_custom( + let engine = DefaultEmbeddedEngine::start_custom( Arc::new(storage), Arc::new(state_machine), Some(config_path.to_str().unwrap()), @@ -119,7 +119,7 @@ listen_address = "127.0.0.1:{port}" // Start engine with RocksDB storage let (storage, state_machine) = RocksDBUnifiedEngine::open(&db_path)?; - let engine = EmbeddedEngine::start_custom( + let engine = DefaultEmbeddedEngine::start_custom( Arc::new(storage), Arc::new(state_machine), Some(config_path.to_str().unwrap()), @@ -141,7 +141,7 @@ listen_address = "127.0.0.1:{port}" #[tokio::test] async fn test_watch_node_crash_embedded_mode() -> Result<(), Box> { // Scenario: - // 1. Node1: EmbeddedEngine.watch(key) + // 1. Node1: DefaultEmbeddedEngine.watch(key) // 2. Drop engine (simulate crash) // 3. Verify all watchers are automatically cleaned up // @@ -361,7 +361,7 @@ watcher_buffer_size = 10 // Start engine with RocksDB storage let (storage, state_machine) = RocksDBUnifiedEngine::open(&db_path)?; - let engine = EmbeddedEngine::start_custom( + let engine = DefaultEmbeddedEngine::start_custom( Arc::new(storage), Arc::new(state_machine), Some(config_path.to_str().unwrap()), diff --git a/d-engine-server/tests/watch_and_subscriptions/watch_membership_embedded.rs b/d-engine-server/tests/watch_and_subscriptions/watch_membership_embedded.rs index 6c66499e..ae1d3d84 100644 --- a/d-engine-server/tests/watch_and_subscriptions/watch_membership_embedded.rs +++ b/d-engine-server/tests/watch_and_subscriptions/watch_membership_embedded.rs @@ -1,4 +1,4 @@ -//! Integration tests for `EmbeddedEngine::watch_membership()` — Ticket #327 +//! Integration tests for `DefaultEmbeddedEngine::watch_membership()` — Ticket #327 //! //! Verifies in-process membership change notifications via `watch::Receiver`. //! @@ -49,7 +49,7 @@ use std::sync::Arc; use std::time::Duration; -use d_engine_server::{EmbeddedEngine, RocksDBUnifiedEngine}; +use d_engine_server::{DefaultEmbeddedEngine, RocksDBUnifiedEngine}; use tempfile::TempDir; use tokio::time::timeout; use tracing_test::traced_test; @@ -179,18 +179,21 @@ async fn wait_for_node_in_snapshot( .flatten() } -/// Start one `EmbeddedEngine` from a TOML string written to a temp file. +/// Start one `DefaultEmbeddedEngine` from a TOML string written to a temp file. async fn start_engine( toml: &str, node_id: u32, db_root: &std::path::Path, config_path: &str, -) -> Result> { +) -> Result> { tokio::fs::write(config_path, toml).await?; let db_path = db_root.join(format!("node{node_id}/db")); tokio::fs::create_dir_all(&db_path).await?; let (storage, sm) = RocksDBUnifiedEngine::open(&db_path)?; - Ok(EmbeddedEngine::start_custom(Arc::new(storage), Arc::new(sm), Some(config_path)).await?) + Ok( + DefaultEmbeddedEngine::start_custom(Arc::new(storage), Arc::new(sm), Some(config_path)) + .await?, + ) } // ── Tests ────────────────────────────────────────────────────────────────────── diff --git a/d-engine-server/tests/watch_and_subscriptions/watch_performance_gate_embedded.rs b/d-engine-server/tests/watch_and_subscriptions/watch_performance_gate_embedded.rs index 6bd04ccf..a58dfdc2 100644 --- a/d-engine-server/tests/watch_and_subscriptions/watch_performance_gate_embedded.rs +++ b/d-engine-server/tests/watch_and_subscriptions/watch_performance_gate_embedded.rs @@ -17,13 +17,13 @@ use std::sync::Arc; use std::time::Duration; use std::time::Instant; -use d_engine_server::api::EmbeddedEngine; +use d_engine_server::api::DefaultEmbeddedEngine; use tempfile::TempDir; use crate::common::get_available_ports; /// Helper: Create a test engine -async fn create_test_engine() -> (EmbeddedEngine, TempDir) { +async fn create_test_engine() -> (DefaultEmbeddedEngine, TempDir) { let temp_dir = TempDir::new().expect("Failed to create temp dir"); let db_path = temp_dir.path().join("db"); @@ -48,7 +48,7 @@ watcher_buffer_size = 100 let (storage, state_machine) = RocksDBUnifiedEngine::open(&db_path).expect("Failed to open unified DB"); - let engine = EmbeddedEngine::start_custom( + let engine = DefaultEmbeddedEngine::start_custom( Arc::new(storage), Arc::new(state_machine), Some(config_path.to_str().unwrap()), diff --git a/d-engine/src/docs/client_guide/read-consistency.md b/d-engine/src/docs/client_guide/read-consistency.md index 8d977dd0..8abe9386 100644 --- a/d-engine/src/docs/client_guide/read-consistency.md +++ b/d-engine/src/docs/client_guide/read-consistency.md @@ -73,7 +73,7 @@ let user = client.get_lease(b"user:123").await?; **API**: `client.get_eventual(key)` -**Guarantee**: Returns valid committed state, may be ~100ms behind leader. +**Guarantee**: Returns valid committed state, may be slightly behind the latest write. **Performance**: ~0.1ms (20x faster than LinearizableRead) @@ -83,7 +83,7 @@ let user = client.get_lease(b"user:123").await?; - Analytics queries - Monitoring data -**Bonus**: Can read from **any node** (leader/follower), not just leader. +> **Any node**: EventualConsistency is served by the local ReadActor on whichever node the client connects to — leader or follower. No lease is required. Follower reads may return slightly stale data (expected for this policy). ```rust,ignore let stats = client.get_eventual(b"dashboard_stats").await?; @@ -132,6 +132,8 @@ allow_client_override = true # Allow client to override (default: true _Measured on 3-node cluster, AWS same-region_ +**Why 0 RTT for Eventual/LeaseRead (v0.2.5+)**: These policies are served by a dedicated ReadActor that reads directly from the state machine — bypassing the Raft command channel entirely. The leader's read lease is validated atomically; if it expires or the node steps down, ReadActor falls back to the Raft path automatically. + **Important**: All policies return **correct** committed data. Difference is whether you get the **latest** or **slightly older** committed state. --- diff --git a/d-engine/src/docs/examples/ha-deployment-load-balancing.md b/d-engine/src/docs/examples/ha-deployment-load-balancing.md index 088c154a..12bb8d9b 100644 --- a/d-engine/src/docs/examples/ha-deployment-load-balancing.md +++ b/d-engine/src/docs/examples/ha-deployment-load-balancing.md @@ -159,7 +159,7 @@ Each node exposes two health endpoints: **Implementation** (from example code): ```rust,ignore -async fn health_primary(State(engine): State>) -> StatusCode { +async fn health_primary(State(engine): State>) -> StatusCode { if engine.is_leader() { StatusCode::OK } else { diff --git a/d-engine/src/docs/integration-modes.md b/d-engine/src/docs/integration-modes.md index 9e29e5d8..e35255c6 100644 --- a/d-engine/src/docs/integration-modes.md +++ b/d-engine/src/docs/integration-modes.md @@ -63,12 +63,12 @@ d-engine runs **inside your Rust application process**: ### Quick Start ```rust,ignore -use d_engine::EmbeddedEngine; +use d_engine::prelude::*; #[tokio::main] async fn main() -> Result<(), Box> { // Start with config file - let engine = EmbeddedEngine::start_with("d-engine.toml").await?; + let engine = DefaultEmbeddedEngine::start_with("d-engine.toml").await?; // Wait for leader election engine.wait_ready(Duration::from_secs(5)).await?; diff --git a/d-engine/src/docs/overview.md b/d-engine/src/docs/overview.md index 1c4c1b9f..d974c60c 100644 --- a/d-engine/src/docs/overview.md +++ b/d-engine/src/docs/overview.md @@ -21,7 +21,7 @@ use std::time::Duration; #[tokio::main] async fn main() -> Result<(), Box> { - let engine = EmbeddedEngine::start("./data").await?; + let engine = DefaultEmbeddedEngine::start("./data").await?; engine.wait_ready(Duration::from_secs(5)).await?; let client = engine.client(); diff --git a/d-engine/src/docs/performance/throughput-optimization-guide.md b/d-engine/src/docs/performance/throughput-optimization-guide.md index 5234d212..6cc9f83c 100644 --- a/d-engine/src/docs/performance/throughput-optimization-guide.md +++ b/d-engine/src/docs/performance/throughput-optimization-guide.md @@ -56,14 +56,33 @@ The `max_batch_size` controls how many commands are drained per Raft loop iterat max_batch_size = 200 # default, suitable for most deployments ``` -| Deployment | Recommended | Rationale | -|---|---|---| -| Embedded 3-node | **200** | Matches typical concurrent client counts; higher values yield diminishing returns | -| Standalone 3-node | **200** | Network RTT dominates; batch size has limited impact | -| High concurrency (500+ clients) | **500** | Increase if HC Write throughput plateaus | +| Deployment | Recommended | Rationale | +| ------------------------------- | ----------- | --------------------------------------------------------------------------------- | +| Embedded 3-node | **200** | Matches typical concurrent client counts; higher values yield diminishing returns | +| Standalone 3-node | **200** | Network RTT dominates; batch size has limited impact | +| High concurrency (500+ clients) | **500** | Increase if HC Write throughput plateaus | > **Rule of thumb**: For embedded mode, optimal `max_batch_size ≈ concurrent_client_count`. For standalone, keep at 200 unless profiling shows cmd_rx consistently saturated. +## Read Fast Path Tuning (v0.2.5+) + +`Eventual` and `LeaseRead` bypass `cmd_rx` entirely — they never touch `max_batch_size` or `time_threshold_ms`. Two dedicated knobs: + +```toml +[raft.read_actor] +channel_capacity = 512 # default +max_drain = 100 # default +``` + +| Parameter | What it controls | Tune up when… | +|---|---|---| +| `channel_capacity` | mpsc buffer depth between client and ReadActor | read latency spikes under high concurrent Eventual/LeaseRead (channel full → backpressure) | +| `max_drain` | reads batched per ReadActor wakeup (mirrors `max_batch_size`) | read throughput plateaus but CPU is not saturated | + +Same tradeoff as write batching: higher `max_drain` → better throughput, higher tail latency. Start with defaults; only tune if profiling shows a bottleneck here. + +--- + ## Configuration Tuning ### Control Plane (`[network.control]`) diff --git a/d-engine/src/docs/quick-start-5min.md b/d-engine/src/docs/quick-start-5min.md index 501ef50c..bdd8f261 100644 --- a/d-engine/src/docs/quick-start-5min.md +++ b/d-engine/src/docs/quick-start-5min.md @@ -61,7 +61,7 @@ async fn main() -> Result<(), Box> { println!("Starting d-engine...\n"); // Start embedded engine with config file - let engine = EmbeddedEngine::start_with("d-engine.toml").await?; + let engine = DefaultEmbeddedEngine::start_with("d-engine.toml").await?; // Wait for leader election (single-node: instant) let leader = engine.wait_ready(Duration::from_secs(5)).await?; @@ -116,7 +116,7 @@ Done! ### Behind the Scenes ```rust,ignore -EmbeddedEngine::start_with("d-engine.toml").await? +DefaultEmbeddedEngine::start_with("d-engine.toml").await? ``` This one line: @@ -141,7 +141,7 @@ Waits for leader election (combines node initialization + leader election): let client = engine.client(); ``` -Returns `Arc` for zero-overhead KV operations. +Returns `Arc` for zero-overhead KV operations. --- @@ -174,7 +174,8 @@ No gRPC overhead - direct function calls to embedded Raft core. - `put(key, value)` - Write with Raft consensus - `get_linearizable(key)` - Strong consistency read -- `get_eventual(key)` - Fast local read (may be stale) +- `get_lease(key)` - Fast read, no Raft round-trip (recommended default) +- `get_eventual(key)` - Fastest local read (may be slightly stale) - `delete(key)` - Delete key See docs for TTL, multi-key operations, and advanced consistency control. @@ -184,7 +185,7 @@ See docs for TTL, multi-key operations, and advanced consistency control. ### 3. Automatic Lifecycle Management ```rust,ignore -let engine = EmbeddedEngine::start_with("d-engine.toml").await?; +let engine = DefaultEmbeddedEngine::start_with("d-engine.toml").await?; // ↑ Internally spawns node.run() in background engine.stop().await?; @@ -197,17 +198,17 @@ No manual `tokio::spawn()`, no leaked tasks. ## API Reference -### EmbeddedEngine +### DefaultEmbeddedEngine ```rust,ignore // Explicit data directory (highest priority) -EmbeddedEngine::start(data_dir: impl AsRef) -> Result +DefaultEmbeddedEngine::start(data_dir: impl AsRef) -> Result // Use explicit config file -EmbeddedEngine::start_with(config_path: &str) -> Result +DefaultEmbeddedEngine::start_with(config_path: &str) -> Result -// Advanced (custom storage) -EmbeddedEngine::start_custom( +// Advanced (custom storage + state machine — define your own TypeConfig) +EmbeddedEngine::::start_custom( storage: Arc, state_machine: Arc, config_path: Option<&str> @@ -217,7 +218,7 @@ EmbeddedEngine::start_custom( engine.wait_ready(timeout: Duration) -> Result // Get KV client -engine.client() -> Arc +engine.client() -> Arc // Subscribe to leader changes (optional, for monitoring) engine.leader_change_notifier() -> watch::Receiver> @@ -226,14 +227,20 @@ engine.leader_change_notifier() -> watch::Receiver> engine.stop().await -> Result<()> ``` -### EmbeddedClient +### DefaultEmbeddedClient ```rust,ignore // Write (replicates to majority) client.put(key: Vec, value: Vec) -> Result -// Read (local, no network) -client.get(key: Vec) -> Result>> +// Read — linearizable (always latest, 1 Raft round-trip) +client.get_linearizable(key: impl Into>) -> Result> + +// Read — lease (no Raft round-trip, recommended default) +client.get_lease(key: impl Into>) -> Result> + +// Read — eventual (fastest, served directly from state machine) +client.get_eventual(key: impl Into>) -> Result> // Delete client.delete(key: Vec) -> Result @@ -249,14 +256,14 @@ See [complete API documentation](https://docs.rs/d-engine/latest/d_engine/prelud ```rust,ignore // Pass the data directory directly — works in debug and release -let engine = EmbeddedEngine::start("./data").await?; +let engine = DefaultEmbeddedEngine::start("./data").await?; ``` ### Pattern 2: Explicit config file ```rust,ignore // Use specific config file -let engine = EmbeddedEngine::start_with("d-engine.toml").await?; +let engine = DefaultEmbeddedEngine::start_with("d-engine.toml").await?; ``` ### Pattern 3: Monitor Leader Changes diff --git a/d-engine/src/docs/quick-start-standalone.md b/d-engine/src/docs/quick-start-standalone.md index f86e6abd..cf4b5770 100644 --- a/d-engine/src/docs/quick-start-standalone.md +++ b/d-engine/src/docs/quick-start-standalone.md @@ -155,7 +155,7 @@ tail -f logs/2/demo.log # Watch new leader election | **Serialization** | None | Protobuf | | **Language** | Rust only | Any (Go, Python, Java) | | **Deployment** | 1 binary | 3 server processes | -| **Setup** | `EmbeddedEngine::start_with()` | `make start-cluster` | +| **Setup** | `DefaultEmbeddedEngine::start_with()` | `make start-cluster` | --- diff --git a/d-engine/src/docs/server_guide/consistency-tuning.md b/d-engine/src/docs/server_guide/consistency-tuning.md index f7c88bb1..35dbfb63 100644 --- a/d-engine/src/docs/server_guide/consistency-tuning.md +++ b/d-engine/src/docs/server_guide/consistency-tuning.md @@ -9,7 +9,7 @@ All read consistency settings are in your server's `config.toml`: ```toml [raft.read_consistency] default_policy = "LeaseRead" -lease_duration_ms = 500 +lease_duration_ms = 250 allow_client_override = true ``` @@ -33,27 +33,18 @@ The consistency policy applied when clients don't specify one explicitly. **Type**: Integer (milliseconds) -**Default**: `500` +**Default**: `250` -**Range**: `100` to `5000` - -How long the leader considers its lease valid after successful heartbeat. +How long the leader considers its lease valid after a successful heartbeat quorum ACK. **Trade-offs**: -- **Higher values** (e.g., 1000ms): - - Better read performance (fewer verification round-trips) - - Larger window for stale reads if clock drift exists - - Recommended for stable cloud environments with NTP - -- **Lower values** (e.g., 200ms): - - Stronger consistency guarantees - - More sensitive to network latency spikes - - Recommended for environments with unreliable clock sync - -**Formula**: `lease_duration_ms ≥ 2 × heartbeat_interval` +- **Higher values**: better LeaseRead performance; larger staleness window on clock drift +- **Lower values**: more conservative; more sensitive to heartbeat jitter -Default heartbeat interval is 100ms, so minimum recommended lease is 200ms. +**Valid range** (both constraints are enforced by `RaftConfig::validate()`): +- Lower bound: `lease_duration_ms ≥ 2 × heartbeat_interval` (with default 100ms heartbeat: ≥ 200ms) +- Upper bound: `lease_duration_ms < election_timeout_min` (default: < 500ms) ### `allow_client_override` @@ -61,12 +52,10 @@ Default heartbeat interval is 100ms, so minimum recommended lease is 200ms. **Default**: `true` -Whether clients can specify per-request consistency policies. +Whether clients can specify per-request consistency policies that override `default_policy`. -**Use cases**: - -- `true`: Allow mixed workloads (critical reads + analytics reads) -- `false`: Enforce uniform consistency across all operations +- `true`: mixed workloads (e.g., most reads use LeaseRead, critical ops use LinearizableRead) +- `false`: uniform consistency enforced — client-specified policies are ignored ## Tuning by Deployment Scenario @@ -77,8 +66,8 @@ Whether clients can specify per-request consistency policies. ```toml [raft.read_consistency] default_policy = "LinearizableRead" -lease_duration_ms = 500 # Unused for LinearizableRead, but keep default -allow_client_override = false # Enforce strict policy +lease_duration_ms = 250 # unused for LinearizableRead +allow_client_override = false # enforce strict policy ``` **Result**: @@ -94,14 +83,13 @@ allow_client_override = false # Enforce strict policy ```toml [raft.read_consistency] default_policy = "EventualConsistency" -lease_duration_ms = 500 # Unused for EventualConsistency -allow_client_override = true # Allow critical reads to override +lease_duration_ms = 250 # unused for EventualConsistency +allow_client_override = true # allow critical reads to override ``` **Result**: -- Reads served immediately from any node -- Maximum throughput (~20x baseline) +- Reads served from local state machine on any node (leader or follower) - Sub-millisecond latency ### Scenario 3: Production Web Service (Recommended) @@ -111,15 +99,14 @@ allow_client_override = true # Allow critical reads to override ```toml [raft.read_consistency] default_policy = "LeaseRead" -lease_duration_ms = 500 # 5x heartbeat interval -allow_client_override = true # Mixed workload support +lease_duration_ms = 250 # default; increase if read latency matters more than freshness +allow_client_override = true # mixed workload support ``` **Result**: - Strong consistency with low latency - Clients can use LinearizableRead for critical operations -- 7x better performance than LinearizableRead ### Scenario 4: Unstable Network Environment @@ -128,49 +115,14 @@ allow_client_override = true # Mixed workload support ```toml [raft.read_consistency] default_policy = "LeaseRead" -lease_duration_ms = 200 # Shorter lease for safety +lease_duration_ms = 200 # shorter lease — faster fallback on partition allow_client_override = true ``` **Result**: -- Faster fallback to LinearizableRead if lease expires -- Better handling of network partitions -- Slightly more verification overhead - -## Monitoring Recommendations - -Track these metrics to validate your configuration: - -### Key Metrics - -```bash -# Lease-related -raft.lease_renewal.success # Should match heartbeat frequency -raft.lease_renewal.failed # Should be near zero - -# Linearizable reads -raft.linearizable_read.success # Total count -raft.leadership_verification.duration_us # Should be <2000 (2ms) - -# Read coalescing effectiveness -raft.linearizable_read.coalesced # Higher = better optimization -raft.pending_reads.queue_depth # Should spike during high concurrency -``` - -### Health Indicators - -**Healthy LeaseRead configuration**: - -- `lease_renewal.success` rate: ~10/sec (100ms heartbeat interval) -- `lease_renewal.failed` < 1% of attempts -- `leadership_verification.duration_us` p99 < 5ms - -**Signs of misconfiguration**: - -- Frequent lease expirations → Increase `lease_duration_ms` -- High `linearizable_read.coalesced` but low throughput → Network bottleneck -- `leadership_verification.duration_us` p99 > 10ms → Check network latency +- Faster fallback to Raft path when lease expires under partition +- Slightly higher fallback frequency under heartbeat jitter ## Clock Synchronization Requirements @@ -206,10 +158,12 @@ Most cloud providers offer time synchronization services: | Configuration | Read Latency (p50) | Throughput | Clock Dependency | | ---------------------------------------------------------- | ------------------ | ---------- | ---------------- | | `default_policy = "LinearizableRead"` | 2.1ms | Baseline | None | -| `default_policy = "LeaseRead"`, `lease_duration_ms = 500` | 0.3ms | ~7x | Low (NTP) | -| `default_policy = "LeaseRead"`, `lease_duration_ms = 1000` | 0.2ms | ~8x | Medium | +| `default_policy = "LeaseRead"`, `lease_duration_ms = 250` | 0.3ms | ~7x | Low (NTP) | +| `default_policy = "LeaseRead"`, `lease_duration_ms = 400` | 0.2ms | ~8x | Medium | | `default_policy = "EventualConsistency"` | 0.1ms | ~20x | None | +_Measured on 3-node cluster, AWS same-region._ + ## Further Reading - Client usage guide: [Read Consistency Guide](crate::docs::client_guide::read_consistency) diff --git a/d-engine/src/docs/server_guide/customize-state-machine.md b/d-engine/src/docs/server_guide/customize-state-machine.md index a732d907..004ba578 100644 --- a/d-engine/src/docs/server_guide/customize-state-machine.md +++ b/d-engine/src/docs/server_guide/customize-state-machine.md @@ -30,10 +30,21 @@ impl StateMachine for CustomStateMachine { } fn stop(&self) -> Result<(), Error> { - // Cleanup resources + // Graceful shutdown: stop accepting requests. + // Called by the Raft loop on shutdown. May be called multiple times. Ok(()) } + // close_storage() has a default no-op implementation. + // Override ONLY if your backend holds an exclusive OS resource (e.g. a lock + // file or an open DB handle) that must be released immediately when + // EmbeddedEngine::stop() is called — without waiting for all Arc clones + // held by EmbeddedClient to drop. + // + // fn close_storage(&self) { + // self.backend.close(); // release lock file / DB handle + // } + fn is_running(&self) -> bool { // Return running status true @@ -104,7 +115,8 @@ impl StateMachine for CustomStateMachine { | Method | Purpose | Sync/Async | Criticality | | ---------------------------------- | ---------------------------------- | ---------- | ----------- | | `start()` | Initialize state machine service | Sync | High | -| `stop()` | Graceful shutdown | Sync | High | +| `stop()` | Graceful shutdown (reversible) | Sync | High | +| `close_storage()` | Release exclusive OS resources (e.g. DB lock file). Default no-op — override only if needed. | Sync | Medium | | `is_running()` | Check service status | Sync | Medium | | `get()` | Read value by key | Sync | High | | `entry_term()` | Get term for log index | Sync | Medium | diff --git a/d-engine/src/docs/server_guide/watch-feature.md b/d-engine/src/docs/server_guide/watch-feature.md index 46107dcc..64a302af 100644 --- a/d-engine/src/docs/server_guide/watch-feature.md +++ b/d-engine/src/docs/server_guide/watch-feature.md @@ -63,10 +63,10 @@ while let Some(event) = stream.next().await { For in-process usage (e.g., `EmbeddedEngine`), use the client's `watch()` method: ```rust,ignore -use d_engine::EmbeddedEngine; +use d_engine::prelude::*; // Initialize engine with config enabling watch -let engine = EmbeddedEngine::start_with("d-engine.toml").await?; +let engine = DefaultEmbeddedEngine::start_with("d-engine.toml").await?; // Watch via client (unified API) let mut watcher = engine.client().watch("my_key")?; @@ -386,7 +386,7 @@ Write → Raft Consensus → StateMachine.apply() → broadcast::send() When `EmbeddedEngine` is dropped or crashes: ```rust,ignore -let engine = EmbeddedEngine::start_with(config).await?; +let engine = DefaultEmbeddedEngine::start_with(config).await?; let mut watcher = engine.client().watch(b"key")?; // If engine crashes or is dropped: @@ -399,7 +399,7 @@ let mut watcher = engine.client().watch(b"key")?; ```rust,ignore loop { - let engine = EmbeddedEngine::start_with(config).await?; + let engine = DefaultEmbeddedEngine::start_with(config).await?; let mut watcher = engine.watch(b"key").await?; while let Some(event) = watcher.recv().await { diff --git a/d-engine/src/lib.rs b/d-engine/src/lib.rs index 596803cd..079eb251 100644 --- a/d-engine/src/lib.rs +++ b/d-engine/src/lib.rs @@ -22,7 +22,9 @@ pub mod prelude { }; #[cfg(feature = "rocksdb")] - pub use d_engine_server::{RocksDBStateMachine, RocksDBStorageEngine}; + pub use d_engine_server::{ + DefaultEmbeddedClient, DefaultEmbeddedEngine, RocksDBStateMachine, RocksDBStorageEngine, + }; #[cfg(feature = "client")] pub use d_engine_client::{Client, ClientApi, ClientBuilder}; diff --git a/examples/quick-start-embedded/README.md b/examples/quick-start-embedded/README.md index ba16edcf..13eaee5d 100644 --- a/examples/quick-start-embedded/README.md +++ b/examples/quick-start-embedded/README.md @@ -85,7 +85,7 @@ The example demonstrates 3 simple steps: ```rust // 1. Start embedded engine with config file -let engine = EmbeddedEngine::start_with("d-engine.toml").await?; +let engine = DefaultEmbeddedEngine::start_with("d-engine.toml").await?; // 2. Wait for leader election (single-node: instant) let leader = engine.wait_ready(Duration::from_secs(5)).await?; @@ -131,7 +131,7 @@ Copy this example and modify `src/main.rs`: Example: ```rust -async fn my_app(client: &EmbeddedClient) -> Result<(), Box> { +async fn my_app(client: &DefaultEmbeddedClient) -> Result<(), Box> { // Your code here client.put("user:1:name".as_bytes().to_vec(), b"alice".to_vec()).await?; client.put("user:1:email".as_bytes().to_vec(), b"alice@example.com".to_vec()).await?; diff --git a/examples/quick-start-embedded/src/main.rs b/examples/quick-start-embedded/src/main.rs index 4e1c3191..9b8be264 100644 --- a/examples/quick-start-embedded/src/main.rs +++ b/examples/quick-start-embedded/src/main.rs @@ -22,7 +22,7 @@ async fn main() -> std::result::Result<(), Box> { // Start embedded engine with explicit config // Config file specifies data directory and other settings - let engine = EmbeddedEngine::start_with("d-engine.toml").await?; + let engine = DefaultEmbeddedEngine::start_with("d-engine.toml").await?; // Wait for leader election (single-node: instant) let leader = engine.wait_ready(Duration::from_secs(5)).await?; @@ -47,7 +47,7 @@ async fn main() -> std::result::Result<(), Box> { Ok(()) } -async fn run_demo(client: Arc) -> std::result::Result<(), Box> { +async fn run_demo(client: Arc) -> std::result::Result<(), Box> { println!("=== Quick Start Demo ==="); // Store workflow state diff --git a/examples/service-discovery-embedded/server.rs b/examples/service-discovery-embedded/server.rs index 8f1adb3a..6c774308 100644 --- a/examples/service-discovery-embedded/server.rs +++ b/examples/service-discovery-embedded/server.rs @@ -12,7 +12,7 @@ use std::collections::HashMap; use std::sync::{Arc, Mutex}; use std::time::Duration; -use d_engine::EmbeddedEngine; +use d_engine::DefaultEmbeddedEngine; use d_engine::WatchEventType; use tokio::signal; @@ -27,7 +27,7 @@ async fn main() -> Result<(), Box> { println!("Starting embedded d-engine for service discovery...\n"); - let engine = EmbeddedEngine::start_with("d-engine.toml").await?; + let engine = DefaultEmbeddedEngine::start_with("d-engine.toml").await?; let leader = engine.wait_ready(Duration::from_secs(5)).await?; println!( diff --git a/examples/sled-cluster/src/sled_state_machine.rs b/examples/sled-cluster/src/sled_state_machine.rs index 613794e9..14f3909a 100644 --- a/examples/sled-cluster/src/sled_state_machine.rs +++ b/examples/sled-cluster/src/sled_state_machine.rs @@ -93,6 +93,12 @@ impl StateMachine for SledStateMachine { Ok(()) } + fn close_storage(&self) { + self.is_serving.store(false, Ordering::Release); + // Flush pending sled writes. The sled::Db is released when SledStateMachine drops. + let _ = self.db.load().flush(); + } + fn stop(&self) -> Result<()> { debug!("stop state machine"); self.is_serving.store(false, Ordering::Release); diff --git a/examples/three-nodes-embedded/src/main.rs b/examples/three-nodes-embedded/src/main.rs index 2e3ebeb0..81c07d24 100644 --- a/examples/three-nodes-embedded/src/main.rs +++ b/examples/three-nodes-embedded/src/main.rs @@ -8,7 +8,7 @@ use axum::{ http::StatusCode, routing::{get, post}, }; -use d_engine::{ClientApiError, EmbeddedEngine, ErrorCode}; +use d_engine::{ClientApiError, DefaultEmbeddedEngine, ErrorCode}; use serde::{Deserialize, Serialize}; use tokio::sync::watch; @@ -60,7 +60,7 @@ async fn main() { // d-engine Integration (Start & Wait for Leader Election) // ============================================================ let engine = Arc::new( - EmbeddedEngine::start_with(&cli.config_path) + DefaultEmbeddedEngine::start_with(&cli.config_path) .await .expect("Failed to start engine"), ); @@ -135,7 +135,7 @@ async fn main() { } async fn start_health_check_server( - engine: Arc, + engine: Arc, port: u16, mut shutdown_rx: watch::Receiver<()>, ) { @@ -158,7 +158,7 @@ async fn start_health_check_server( .expect("Health check server failed"); } -async fn health_primary(State(engine): State>) -> StatusCode { +async fn health_primary(State(engine): State>) -> StatusCode { if engine.is_leader() { StatusCode::OK } else { @@ -166,7 +166,7 @@ async fn health_primary(State(engine): State>) -> StatusCode } } -async fn health_replica(State(engine): State>) -> StatusCode { +async fn health_replica(State(engine): State>) -> StatusCode { if !engine.is_leader() { StatusCode::OK } else { @@ -175,7 +175,7 @@ async fn health_replica(State(engine): State>) -> StatusCode } async fn start_business_server( - engine: Arc, + engine: Arc, port: u16, mut shutdown_rx: watch::Receiver<()>, ) { @@ -203,7 +203,7 @@ fn is_not_leader(e: &ClientApiError) -> bool { } async fn handle_put( - State(engine): State>, + State(engine): State>, Json(req): Json, ) -> (StatusCode, Json) { match engine.client().put(req.key.into_bytes(), req.value.into_bytes()).await { @@ -223,7 +223,7 @@ async fn handle_put( } async fn handle_get( - State(engine): State>, + State(engine): State>, Path(key): Path, ) -> Result, StatusCode> { match engine.client().get_eventual(key.into_bytes()).await { diff --git a/examples/three-nodes-standalone/config/n1.toml b/examples/three-nodes-standalone/config/n1.toml index 0d47e2f5..54ec4a0a 100644 --- a/examples/three-nodes-standalone/config/n1.toml +++ b/examples/three-nodes-standalone/config/n1.toml @@ -15,7 +15,6 @@ general_raft_timeout_duration_in_ms = 100 cmd_channel_capacity = 1024 ordered_channel_capacity = 1024 - [raft.election] election_timeout_min = 1000 election_timeout_max = 2000 @@ -24,6 +23,10 @@ election_timeout_max = 2000 default_policy = "LeaseRead" lease_duration_ms = 500 +[raft.read_actor] +channel_capacity = 10240 +max_drain = 2000 + [raft.batching] # Maximum number of commands to accumulate in a single batch during drain operations max_batch_size = 200 diff --git a/examples/three-nodes-standalone/config/n2.toml b/examples/three-nodes-standalone/config/n2.toml index 83342a3d..3ce38ad5 100644 --- a/examples/three-nodes-standalone/config/n2.toml +++ b/examples/three-nodes-standalone/config/n2.toml @@ -23,6 +23,10 @@ election_timeout_max = 2000 default_policy = "LeaseRead" lease_duration_ms = 100 +[raft.read_actor] +channel_capacity = 10240 +max_drain = 2000 + [raft.batching] # Maximum number of commands to accumulate in a single batch during drain operations max_batch_size = 200 diff --git a/examples/three-nodes-standalone/config/n3.toml b/examples/three-nodes-standalone/config/n3.toml index a6f7c9ee..e71a9b42 100644 --- a/examples/three-nodes-standalone/config/n3.toml +++ b/examples/three-nodes-standalone/config/n3.toml @@ -23,6 +23,10 @@ election_timeout_max = 2000 default_policy = "LeaseRead" lease_duration_ms = 500 +[raft.read_actor] +channel_capacity = 10240 +max_drain = 2000 + [raft.batching] # Maximum number of commands to accumulate in a single batch during drain operations max_batch_size = 200