From 073c1da53e4ffa56e4cede7a04766e23203f1b36 Mon Sep 17 00:00:00 2001 From: Joshua Chi Date: Tue, 19 May 2026 11:32:32 +0800 Subject: [PATCH 1/2] feat(watch) #379: prev_kv + progress heartbeat MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Production code - WatchEvent: add `prev_value: Bytes` field; add `WatchEventType::Progress` variant - WatchRegistry: `register()` / `register_prefix()` take `prev_kv: bool`; expose `prev_kv_watcher_count_arc()` for shared atomic counter - DefaultStateMachineHandler: when `prev_kv_watcher_count > 0`, read old values from state machine before `apply_chunk` and populate `prev_value` in broadcast events - WatchDispatcher: `heartbeat_interval_ms` param; tokio interval with ±10% jitter broadcasts `Progress` events to all watchers; `broadcast_progress()` fan-out iterates registry maps directly (key-routing would miss watchers on empty key) - EmbeddedClient: add `watch_with_options(key, prev_kv)` and `watch_prefix_with_options(prefix, prev_kv)` as explicit prev_kv entry points - Proto: add `prev_value` field to `WatchResponse`; add `prev_kv` field to `WatchRequest` - Config: add `heartbeat_interval_ms: u64` to `WatchConfig` - NodeBuilder: wire `prev_kv_watcher_count` Arc into `DefaultStateMachineHandler::new()` ## Tests (RED → GREEN) - manager_test: 6 new tests — prev_kv isolation per watcher, watcher count tracking, progress heartbeat delivery, heartbeat disabled produces no Progress events - default_state_machine_handler_test: 3 new tests — skip get when count=0, read and broadcast prev_value when count=1, stop reading after count drops to zero - grpc_client_test: add `prev_value` field to WatchResponse fixture - embedded_client_test: guard both test mods with `#[cfg(all(test, feature="rocksdb"))]` ## API compatibility fixes (callers broken by new fields / signatures) - benches/state_machine.rs: `prev_value`, `register(_, false)`, WatchDispatcher 5-arg - tests/embedded_client: `WatchEventType::X as i32` → `WatchEventType::X` - tests/watch_events_embedded: same - examples/service-discovery-embedded: same + explicit Canceled arm - examples/service-discovery-standalone: same + Progress arm for exhaustive match - integration/mod.rs: DefaultStateMachineHandler::new() 7th arg ## Docs - watch-feature.md: replace "No prev_kv / No heartbeat (planned)" with feature docs; add `heartbeat_interval_ms` config entry; fix `WatchEventType::CANCELED` casing (×3); fix `Canceled as i32` in reconnection pattern; add Progress to Best Practices - embedded_client.rs: full rustdoc for `watch_with_options` and `watch_prefix_with_options` — prev_kv semantics, performance note, empty-prev_value cases, code examples - service-discovery-embedded: Demo 3 — prev_kv audit log pattern (put/delete/canceled/progress) --- d-engine-client/src/grpc_client.rs | 2 + d-engine-client/src/grpc_client_test.rs | 2 + d-engine-core/src/config/raft.rs | 18 + .../default_state_machine_handler.rs | 79 ++- .../default_state_machine_handler_test.rs | 180 ++++- d-engine-core/src/watch/manager.rs | 256 ++++++- d-engine-core/src/watch/manager_test.rs | 647 +++++++++++------- d-engine-core/src/watch/mod.rs | 5 +- d-engine-proto/proto/client/client_api.proto | 15 + .../src/generated/d_engine.client.rs | 16 + d-engine-server/benches/state_machine.rs | 23 +- d-engine-server/src/api/embedded_client.rs | 116 +++- .../src/api/embedded_client_test.rs | 4 +- .../src/network/grpc/grpc_raft_service.rs | 15 +- d-engine-server/src/node/builder.rs | 39 +- .../src/test_utils/integration/mod.rs | 1 + .../embedded_client_operations.rs | 8 +- .../watch_events_embedded.rs | 12 +- .../src/docs/server_guide/watch-feature.md | 56 +- examples/service-discovery-embedded/server.rs | 76 +- .../service-discovery-standalone/watcher.rs | 6 +- 21 files changed, 1241 insertions(+), 335 deletions(-) diff --git a/d-engine-client/src/grpc_client.rs b/d-engine-client/src/grpc_client.rs index b0ca40db..65500246 100644 --- a/d-engine-client/src/grpc_client.rs +++ b/d-engine-client/src/grpc_client.rs @@ -224,6 +224,7 @@ impl GrpcClient { client_id: client_inner.client_id, key: Bytes::copy_from_slice(key.as_ref()), prefix: false, + prev_kv: false, }; // Watch can connect to any node (leader or follower) @@ -255,6 +256,7 @@ impl GrpcClient { client_id: client_inner.client_id, key: Bytes::copy_from_slice(prefix.as_ref()), prefix: true, + prev_kv: false, }; let mut client = self.make_client().await?; diff --git a/d-engine-client/src/grpc_client_test.rs b/d-engine-client/src/grpc_client_test.rs index 4f27c690..0f3769a3 100644 --- a/d-engine-client/src/grpc_client_test.rs +++ b/d-engine-client/src/grpc_client_test.rs @@ -1592,6 +1592,7 @@ mod watch_tests { WatchResponse { key: Bytes::from("lock"), value: Bytes::from("owner-a"), + prev_value: Bytes::new(), event_type: WatchEventType::Put as i32, error: 0, revision: 0, @@ -1599,6 +1600,7 @@ mod watch_tests { WatchResponse { key: Bytes::from("lock"), value: Bytes::new(), + prev_value: Bytes::new(), event_type: WatchEventType::Delete as i32, error: 0, revision: 0, diff --git a/d-engine-core/src/config/raft.rs b/d-engine-core/src/config/raft.rs index 101460d6..46f173f0 100644 --- a/d-engine-core/src/config/raft.rs +++ b/d-engine-core/src/config/raft.rs @@ -1253,6 +1253,19 @@ pub struct WatchConfig { /// **Default**: i64::MAX (effectively unlimited) #[serde(default = "default_max_watcher_count")] pub max_watcher_count: usize, + + /// Heartbeat interval in milliseconds for progress notifications. + /// + /// The dispatcher broadcasts a `Progress` event to all active watchers on + /// each tick so clients can confirm the stream is alive even during quiet + /// periods. Set to 0 to disable heartbeats entirely. + /// + /// Use milliseconds (not seconds) so tests can set short intervals (e.g. 50ms) + /// without sleeping for a full second. + /// + /// **Default**: 30_000 (30 seconds) + #[serde(default = "default_heartbeat_interval_ms")] + pub heartbeat_interval_ms: u64, } impl Default for WatchConfig { @@ -1262,6 +1275,7 @@ impl Default for WatchConfig { watcher_buffer_size: default_watcher_buffer_size(), enable_metrics: default_enable_watch_metrics(), max_watcher_count: default_max_watcher_count(), + heartbeat_interval_ms: default_heartbeat_interval_ms(), } } } @@ -1319,6 +1333,10 @@ const fn default_max_watcher_count() -> usize { 9_223_372_036_854_775_807 } +const fn default_heartbeat_interval_ms() -> u64 { + 30_000 +} + /// Performance metrics configuration /// /// Controls emission of observability metrics. Disabling metrics reduces overhead diff --git a/d-engine-core/src/state_machine_handler/default_state_machine_handler.rs b/d-engine-core/src/state_machine_handler/default_state_machine_handler.rs index 3d2a166b..fa57b2be 100644 --- a/d-engine-core/src/state_machine_handler/default_state_machine_handler.rs +++ b/d-engine-core/src/state_machine_handler/default_state_machine_handler.rs @@ -106,6 +106,11 @@ where /// StateMachine sends events here without blocking on watchers #[cfg(feature = "watch")] watch_event_tx: Option>, + + /// Number of active watchers that opted in to prev_kv. + /// When > 0, apply_chunk reads old values before applying each write. + #[cfg(feature = "watch")] + prev_kv_watcher_count: Arc, } #[derive(Debug, PartialEq, Hash, Eq, Clone)] @@ -217,6 +222,17 @@ where let sm = self.state_machine.clone(); + // Read prev values before applying if any watcher requested prev_kv. + // Must happen BEFORE apply_chunk so old values are still in the state machine. + #[cfg(feature = "watch")] + let prev_values: Option>> = if self.watch_event_tx.is_some() + && self.prev_kv_watcher_count.load(std::sync::atomic::Ordering::Relaxed) > 0 + { + Some(Self::read_prev_values(&*sm, &chunk)) + } else { + None + }; + // Apply the chunk and track errors let apply_result = sm.apply_chunk(chunk.clone()).await; @@ -225,7 +241,7 @@ where if let Ok(ref results) = apply_result && let Some(ref tx) = self.watch_event_tx { - self.broadcast_watch_events(&chunk, results, tx); + self.broadcast_watch_events(&chunk, results, tx, prev_values.as_deref()); } // Record latency and chunk size histogram *after* the operation @@ -624,14 +640,52 @@ impl DefaultStateMachineHandler where T: TypeConfig, { + /// Read the current value for every write operation in `chunk`, before applying. + /// + /// Returns one `Option` per entry in the same order: + /// - `Some(value)` for Insert / Delete / CAS entries (empty Bytes if key didn't exist) + /// - `None` for Noop / Config entries (no prev_value needed) + #[cfg(feature = "watch")] + fn read_prev_values( + sm: &dyn crate::StateMachine, + chunk: &[Entry], + ) -> Vec> { + use d_engine_proto::client::WriteCommand; + use d_engine_proto::client::write_command::Operation; + use d_engine_proto::common::entry_payload::Payload; + use prost::Message; + + chunk + .iter() + .map(|entry| { + let payload = entry.payload.as_ref()?; + let Some(Payload::Command(bytes)) = &payload.payload else { + return None; + }; + let Ok(write_cmd) = WriteCommand::decode(bytes.as_ref()) else { + return None; + }; + + let key: Option<&[u8]> = match &write_cmd.operation { + Some(Operation::Insert(ins)) => Some(&ins.key), + Some(Operation::Delete(del)) => Some(&del.key), + Some(Operation::CompareAndSwap(cas)) => Some(&cas.key), + None => None, + }; + + key.map(|k| sm.get(k).ok().flatten().unwrap_or_default()) + }) + .collect() + } + /// Broadcast watch events for applied chunk entries (fire-and-forget) /// - /// Parses chunk entries and broadcasts WatchEvent via tokio::broadcast channel. + /// Parses chunk entries and broadcasts WatchResponse via tokio::broadcast channel. /// Non-blocking: if channel is full, oldest events are dropped (lagging receivers). - /// Decoupled from WatchManager: only sends signal, doesn't care who listens. /// /// `results` must have the same length as `chunk` (enforced by StateMachineHandler contract). /// CAS entries where `results[i].succeeded == false` are skipped — no mutation occurred. + /// `prev_values`: optional slice of pre-apply values (None = no prev_kv watcher active). #[cfg(feature = "watch")] #[inline] pub(super) fn broadcast_watch_events( @@ -639,6 +693,7 @@ where chunk: &[Entry], results: &[ApplyResult], tx: &tokio::sync::broadcast::Sender, + prev_values: Option<&[Option]>, ) { use d_engine_proto::client::WatchEventType; use d_engine_proto::client::WatchResponse; @@ -652,10 +707,16 @@ where && let Some(Payload::Command(bytes)) = &payload.payload && let Ok(write_cmd) = WriteCommand::decode(bytes.as_ref()) { + let prev_value = prev_values + .and_then(|pv| pv.get(i)) + .and_then(|v| v.clone()) + .unwrap_or_default(); + let event = match write_cmd.operation { Some(Operation::Insert(insert)) => Some(WatchResponse { key: insert.key, value: insert.value, + prev_value, event_type: WatchEventType::Put as i32, error: 0, revision: entry.index, @@ -663,6 +724,7 @@ where Some(Operation::Delete(delete)) => Some(WatchResponse { key: delete.key, value: bytes::Bytes::new(), + prev_value, event_type: WatchEventType::Delete as i32, error: 0, revision: entry.index, @@ -674,6 +736,7 @@ where Some(WatchResponse { key: cas.key, value: cas.new_value, + prev_value, event_type: WatchEventType::Put as i32, error: 0, revision: entry.index, @@ -707,6 +770,9 @@ where #[cfg_attr(not(feature = "watch"), allow(unused_variables))] watch_event_tx: Option< tokio::sync::broadcast::Sender, >, + #[cfg_attr(not(feature = "watch"), allow(unused_variables))] prev_kv_watcher_count: Arc< + std::sync::atomic::AtomicUsize, + >, ) -> Self { let (applied_notify_tx, applied_notify_rx) = tokio::sync::watch::channel(last_applied_index); @@ -729,12 +795,12 @@ where applied_notify_rx, #[cfg(feature = "watch")] watch_event_tx, + #[cfg(feature = "watch")] + prev_kv_watcher_count, } } /// Convenience constructor for tests without watch - /// - /// Backward-compatible wrapper that calls `new()` with None for watch_event_tx. #[cfg(test)] pub(crate) fn new_without_watch( node_id: u32, @@ -749,7 +815,8 @@ where state_machine, snapshot_config, snapshot_policy, - None, // No watch event tx for tests + None, + Arc::new(std::sync::atomic::AtomicUsize::new(0)), ) } diff --git a/d-engine-core/src/state_machine_handler/default_state_machine_handler_test.rs b/d-engine-core/src/state_machine_handler/default_state_machine_handler_test.rs index 3c124255..2097ab82 100644 --- a/d-engine-core/src/state_machine_handler/default_state_machine_handler_test.rs +++ b/d-engine-core/src/state_machine_handler/default_state_machine_handler_test.rs @@ -1748,7 +1748,7 @@ mod broadcast_watch_events_tests { }), ); - handler.broadcast_watch_events(&[entry], &[succeeded(1)], &tx); + handler.broadcast_watch_events(&[entry], &[succeeded(1)], &tx, None); let event = rx.recv().await.unwrap(); assert_eq!(event.key, Bytes::from(&b"key1"[..])); @@ -1768,7 +1768,7 @@ mod broadcast_watch_events_tests { }), ); - handler.broadcast_watch_events(&[entry], &[succeeded(1)], &tx); + handler.broadcast_watch_events(&[entry], &[succeeded(1)], &tx, None); let event = rx.recv().await.unwrap(); assert_eq!(event.key, Bytes::from(&b"key1"[..])); @@ -1791,7 +1791,7 @@ mod broadcast_watch_events_tests { }), ); - handler.broadcast_watch_events(&[entry], &[succeeded(1)], &tx); + handler.broadcast_watch_events(&[entry], &[succeeded(1)], &tx, None); let event = rx.recv().await.unwrap(); assert_eq!(event.key, Bytes::from(&b"lock"[..])); @@ -1814,7 +1814,7 @@ mod broadcast_watch_events_tests { }), ); - handler.broadcast_watch_events(&[entry], &[failed(1)], &tx); + handler.broadcast_watch_events(&[entry], &[failed(1)], &tx, None); // Channel must be empty — no event should have been sent assert!( @@ -1857,7 +1857,7 @@ mod broadcast_watch_events_tests { ]; let results = vec![succeeded(1), failed(2), succeeded(3)]; - handler.broadcast_watch_events(&entries, &results, &tx); + handler.broadcast_watch_events(&entries, &results, &tx, None); // k1 succeeded → Put event let event1 = rx.recv().await.unwrap(); @@ -1907,7 +1907,7 @@ mod broadcast_watch_events_tests { ]; let results = vec![succeeded(1), failed(2), succeeded(3)]; - handler.broadcast_watch_events(&entries, &results, &tx); + handler.broadcast_watch_events(&entries, &results, &tx, None); // Insert event let event1 = rx.recv().await.unwrap(); @@ -1952,7 +1952,7 @@ mod broadcast_watch_events_tests { // results[1] = succeeded (aligns with chunk[1] = key-b) let results = vec![failed(10), succeeded(11)]; - handler.broadcast_watch_events(&entries, &results, &tx); + handler.broadcast_watch_events(&entries, &results, &tx, None); // Only key-b event should arrive let event = rx.recv().await.unwrap(); @@ -1977,10 +1977,174 @@ mod broadcast_watch_events_tests { }), ); - handler.broadcast_watch_events(&[entry], &[succeeded(1)], &tx); + handler.broadcast_watch_events(&[entry], &[succeeded(1)], &tx, None); // Should still receive valid event let event = rx.recv().await.unwrap(); assert_eq!(event.key, Bytes::from(&b"key1"[..])); } } + +// ============================================================================= +// #379: prev_kv — apply_chunk integration with prev_kv_watcher_count +// ============================================================================= + +#[cfg(feature = "watch")] +mod prev_kv_apply_tests { + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + + use bytes::Bytes; + use d_engine_proto::client::{ + WriteCommand, + write_command::{Insert, Operation}, + }; + use d_engine_proto::common::{Entry, EntryPayload, entry_payload::Payload}; + use prost::Message; + + use super::*; + use crate::test_utils::snapshot_config; + use crate::{ApplyResult, MockSnapshotPolicy, MockStateMachine, MockTypeConfig}; + + fn insert_entry( + index: u64, + key: &[u8], + value: &[u8], + ) -> Entry { + let write_cmd = WriteCommand { + operation: Some(Operation::Insert(Insert { + key: Bytes::copy_from_slice(key), + value: Bytes::copy_from_slice(value), + ttl_secs: 0, + })), + }; + let mut buf = Vec::new(); + write_cmd.encode(&mut buf).unwrap(); + Entry { + index, + term: 1, + payload: Some(EntryPayload { + payload: Some(Payload::Command(Bytes::from(buf))), + }), + } + } + + /// #379 P5: When prev_kv_watcher_count == 0, apply_chunk must NOT call + /// state_machine.get() at all. The mock will panic if get() is called. + /// + /// This validates the performance optimization: zero extra RocksDB reads + /// when no watcher has opted in to prev_kv. + #[tokio::test] + async fn test_apply_chunk_skips_get_when_no_prev_kv_watchers() { + let mut sm = MockStateMachine::new(); + sm.expect_apply_chunk().returning(|_| { + Ok(vec![ApplyResult { + index: 1, + succeeded: true, + }]) + }); + // DO NOT set up expect_get() → mockall panics if get() is called unexpectedly + + let (tx, _rx) = tokio::sync::broadcast::channel(10); + let prev_kv_count = Arc::new(AtomicUsize::new(0)); // no prev_kv watchers + + let handler = DefaultStateMachineHandler::::new( + 1, + 0, + Arc::new(sm), + snapshot_config(std::path::PathBuf::from( + "/tmp/test_skip_get_when_no_prev_kv", + )), + MockSnapshotPolicy::new(), + Some(tx), + prev_kv_count, + ); + + let entry = insert_entry(1, b"k", b"v"); + handler.apply_chunk(vec![entry]).await.unwrap(); + // If get() were called, mockall would have panicked above. + } + + /// #379 P6: When prev_kv_watcher_count > 0, apply_chunk reads the old value + /// from the state machine and includes it in the broadcast WatchResponse. + /// + /// This validates the end-to-end prev_kv path: state machine read → broadcast event. + #[tokio::test] + async fn test_apply_chunk_reads_and_broadcasts_prev_value_when_count_nonzero() { + let mut sm = MockStateMachine::new(); + sm.expect_apply_chunk().returning(|_| { + Ok(vec![ApplyResult { + index: 1, + succeeded: true, + }]) + }); + // Expect exactly one get() call for the insert key + sm.expect_get().times(1).returning(|_| Ok(Some(Bytes::from("old_val")))); + + let (tx, mut rx) = tokio::sync::broadcast::channel(10); + let prev_kv_count = Arc::new(AtomicUsize::new(1)); // one prev_kv watcher active + + let handler = DefaultStateMachineHandler::::new( + 1, + 0, + Arc::new(sm), + snapshot_config(std::path::PathBuf::from( + "/tmp/test_reads_prev_value_when_nonzero", + )), + MockSnapshotPolicy::new(), + Some(tx), + prev_kv_count, + ); + + let entry = insert_entry(1, b"k", b"new_val"); + handler.apply_chunk(vec![entry]).await.unwrap(); + + let event = rx.recv().await.expect("expected broadcast event"); + assert_eq!( + event.prev_value, + Bytes::from("old_val"), + "broadcast event must carry the prev_value read before apply" + ); + assert_eq!(event.value, Bytes::from("new_val")); + } + + /// #379 P6b: After the last prev_kv watcher unregisters (count drops to 0), + /// apply_chunk must NOT call state_machine.get() on the next apply. + /// + /// Tests the dynamic transition: count was 1 → externally set to 0 → + /// next apply skips read. This matches what happens when the last watcher + /// using prev_kv=true unregisters in production. + #[tokio::test] + async fn test_apply_chunk_stops_reading_prev_value_after_count_drops_to_zero() { + let mut sm = MockStateMachine::new(); + sm.expect_apply_chunk().returning(|_| { + Ok(vec![ApplyResult { + index: 1, + succeeded: true, + }]) + }); + // After count → 0, get() must NOT be called + // (Not setting up expect_get → mockall panics if called) + + let (tx, _rx) = tokio::sync::broadcast::channel(10); + let prev_kv_count = Arc::new(AtomicUsize::new(0)); // already 0 (last watcher gone) + + let handler = DefaultStateMachineHandler::::new( + 1, + 0, + Arc::new(sm), + snapshot_config(std::path::PathBuf::from( + "/tmp/test_stops_reading_after_count_zero", + )), + MockSnapshotPolicy::new(), + Some(tx), + prev_kv_count.clone(), + ); + + // Simulate: count was 1 (watcher registered), now it's 0 (watcher gone) + // The Arc is already at 0 — handler should skip read on next apply. + let entry = insert_entry(1, b"k", b"v"); + handler.apply_chunk(vec![entry]).await.unwrap(); + // No panic → get() was not called → optimization active. + } +} diff --git a/d-engine-core/src/watch/manager.rs b/d-engine-core/src/watch/manager.rs index 0f0cd397..b1f9c53d 100644 --- a/d-engine-core/src/watch/manager.rs +++ b/d-engine-core/src/watch/manager.rs @@ -4,7 +4,7 @@ //! //! ```text //! StateMachine: -//! apply_chunk() -> broadcast_watch_events() -> broadcast::send(WatchEvent) [fire-and-forget] +//! apply_chunk() -> broadcast_watch_events() -> broadcast::send(WatchResponse) [fire-and-forget] //! ↓ //! WatchDispatcher (spawned in Builder): //! broadcast::subscribe() -> @@ -12,8 +12,8 @@ //! prefix: DashMap> O(depth) decomposition + O(1) lookup per segment //! ↓ //! Watchers: -//! Embedded: mpsc::Receiver -//! Standalone: mpsc::Receiver -> gRPC stream (protobuf conversion) +//! Embedded: mpsc::Receiver (opaque Rust type, no proto dependency) +//! Standalone: mpsc::Receiver -> gRPC stream (caller converts WatchEvent → WatchResponse) //! ``` //! //! # Design Principles @@ -23,12 +23,15 @@ //! - **Composable**: Registry and Dispatcher are independent, composed in Builder //! - **O(depth) prefix dispatch**: Decompose event key into slash-terminated path segments, //! O(1) DashMap lookup per segment — dispatch cost is independent of prefix watcher count +//! - **Proto boundary**: WatchResponse (proto) lives only in the broadcast channel and handler; +//! WatchEvent (opaque) is what callers see — no proto import required. use std::sync::Arc; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::time::Duration; use bytes::Bytes; -pub use d_engine_proto::client::{WatchEventType, WatchResponse as WatchEvent}; +use d_engine_proto::client::WatchResponse; use dashmap::DashMap; use tokio::sync::broadcast; use tokio::sync::mpsc; @@ -36,6 +39,88 @@ use tracing::debug; use tracing::trace; use tracing::warn; +// ── Public opaque types ─────────────────────────────────────────────────────── + +/// High-level watch event type. No proto import required by callers. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum WatchEventType { + Put, + Delete, + Canceled, + Progress, +} + +/// Opaque watch event delivered to callers. +/// +/// All fields use standard Rust types — callers never need to import proto types. +/// `prev_value` is populated only when the watcher was registered with `prev_kv = true` +/// and the event is a data mutation (Put / Delete); it is empty for Progress / Canceled. +#[derive(Debug, Clone)] +pub struct WatchEvent { + pub event_type: WatchEventType, + pub key: Bytes, + pub value: Bytes, + /// Value before this mutation. Empty when key did not exist, or prev_kv = false. + pub prev_value: Bytes, + pub revision: u64, +} + +// ── Proto ↔ opaque conversions ──────────────────────────────────────────────── + +/// Convert a proto WatchResponse into an opaque WatchEvent. +/// +/// Called by the dispatcher when delivering to a per-watcher mpsc channel. +/// `prev_kv`: when false, prev_value is zeroed so callers that didn't opt in +/// never receive stale memory from a watcher that did. +fn proto_to_event( + proto: &WatchResponse, + prev_kv: bool, +) -> WatchEvent { + use d_engine_proto::client::WatchEventType as ProtoType; + + let event_type = match ProtoType::try_from(proto.event_type) { + Ok(ProtoType::Put) => WatchEventType::Put, + Ok(ProtoType::Delete) => WatchEventType::Delete, + Ok(ProtoType::Canceled) => WatchEventType::Canceled, + Ok(ProtoType::Progress) => WatchEventType::Progress, + Err(_) => WatchEventType::Canceled, + }; + + WatchEvent { + event_type, + key: proto.key.clone(), + value: proto.value.clone(), + prev_value: if prev_kv { + proto.prev_value.clone() + } else { + Bytes::new() + }, + revision: proto.revision, + } +} + +/// Convert an opaque WatchEvent back to a proto WatchResponse (for gRPC). +impl From<&WatchEvent> for WatchResponse { + fn from(e: &WatchEvent) -> Self { + use d_engine_proto::client::WatchEventType as ProtoType; + WatchResponse { + key: e.key.clone(), + value: e.value.clone(), + prev_value: e.prev_value.clone(), + event_type: match e.event_type { + WatchEventType::Put => ProtoType::Put as i32, + WatchEventType::Delete => ProtoType::Delete as i32, + WatchEventType::Canceled => ProtoType::Canceled as i32, + WatchEventType::Progress => ProtoType::Progress as i32, + }, + error: 0, + revision: e.revision, + } + } +} + +// ── WatchError ──────────────────────────────────────────────────────────────── + /// Errors returned by WatchRegistry registration methods. #[derive(Debug)] pub enum WatchError { @@ -61,6 +146,8 @@ impl std::fmt::Display for WatchError { impl std::error::Error for WatchError {} +// ── Prefix helpers ──────────────────────────────────────────────────────────── + /// Decompose a key into all its slash-terminated path prefix candidates. /// /// The dispatcher calls this on every event key to perform O(depth) prefix @@ -80,6 +167,8 @@ pub(crate) fn prefix_segments(key: &Bytes) -> Vec { .collect() } +// ── WatcherHandle ───────────────────────────────────────────────────────────── + /// Handle for a registered watcher. /// /// When dropped, the watcher is automatically unregistered (if unregister_tx is Some). @@ -146,13 +235,19 @@ impl Drop for WatcherHandle { } } +// ── Internal watcher state ──────────────────────────────────────────────────── + /// Internal watcher state #[derive(Debug)] struct Watcher { id: u64, sender: mpsc::Sender, + /// When true, prev_value is populated before delivery. + prev_kv: bool, } +// ── WatchRegistry ───────────────────────────────────────────────────────────── + /// Watch registry — manages watcher registration (Arc-shareable). /// /// Two independent DashMaps keep exact and prefix watchers separate so @@ -166,6 +261,9 @@ pub struct WatchRegistry { next_id: AtomicU64, /// Total active watchers across exact + prefix (for limit enforcement) total_count: AtomicUsize, + /// Count of active watchers that requested prev_kv = true. + /// Shared Arc lets the state machine handler poll this without holding a registry ref. + prev_kv_watcher_count: Arc, /// Per-watcher channel buffer size watcher_buffer_size: usize, /// Hard cap on total active watchers; register() returns LimitExceeded when reached @@ -197,6 +295,7 @@ impl WatchRegistry { prefix: DashMap::new(), next_id: AtomicU64::new(1), total_count: AtomicUsize::new(0), + prev_kv_watcher_count: Arc::new(AtomicUsize::new(0)), watcher_buffer_size, max_watcher_count, unregister_tx, @@ -205,12 +304,17 @@ impl WatchRegistry { /// Register an exact-key watcher. /// + /// `prev_kv`: when true the server reads the old value before each write + /// and populates `WatchEvent::prev_value`. Only pays the read cost when + /// at least one watcher has prev_kv = true. + /// /// Returns `WatchError::LimitExceeded` if max_watcher_count is reached. pub fn register( &self, key: Bytes, + prev_kv: bool, ) -> Result { - self.do_register(key, false) + self.do_register(key, false, prev_kv) } /// Register a prefix watcher. @@ -218,22 +322,26 @@ impl WatchRegistry { /// `prefix` must start with '/' and end with '/'. /// E.g. "/config/" watches all keys under /config/. /// + /// `prev_kv`: same semantics as `register()`. + /// /// Returns `WatchError::InvalidPrefix` if format is wrong. /// Returns `WatchError::LimitExceeded` if max_watcher_count is reached. pub fn register_prefix( &self, prefix: Bytes, + prev_kv: bool, ) -> Result { if !prefix.starts_with(b"/") || !prefix.ends_with(b"/") { return Err(WatchError::InvalidPrefix); } - self.do_register(prefix, true) + self.do_register(prefix, true, prev_kv) } fn do_register( &self, key: Bytes, is_prefix: bool, + prev_kv: bool, ) -> Result { // Reserve the slot first, then check. This eliminates the TOCTOU window // between a load-check and a separate fetch_add under concurrent registration. @@ -243,17 +351,25 @@ impl WatchRegistry { return Err(WatchError::LimitExceeded(self.max_watcher_count)); } + if prev_kv { + self.prev_kv_watcher_count.fetch_add(1, Ordering::Relaxed); + } + let id = self.next_id.fetch_add(1, Ordering::Relaxed); // +1 reserves one slot for the CANCELED sentinel (always deliverable even when full) let (sender, receiver) = mpsc::channel(self.watcher_buffer_size + 1); - let watcher = Watcher { id, sender }; + let watcher = Watcher { + id, + sender, + prev_kv, + }; if is_prefix { self.prefix.entry(key.clone()).or_default().push(watcher); } else { self.exact.entry(key.clone()).or_default().push(watcher); } - trace!(watcher_id = id, key = ?key, is_prefix, "Watcher registered"); + trace!(watcher_id = id, key = ?key, is_prefix, prev_kv, "Watcher registered"); Ok(WatcherHandle { id, @@ -269,28 +385,52 @@ impl WatchRegistry { id: u64, key: &Bytes, ) { - // Watcher IDs are globally unique — try exact first, then prefix. - // Cold path: called only on drop or after overflow, not on every event. let mut found = false; + let mut had_prev_kv = false; + + // Try exact map first self.exact.remove_if_mut(key, |_, watchers| { - let before = watchers.len(); - watchers.retain(|w| w.id != id); - found = watchers.len() < before; + if let Some(pos) = watchers.iter().position(|w| w.id == id) { + had_prev_kv = watchers[pos].prev_kv; + watchers.remove(pos); + found = true; + } watchers.is_empty() }); + + // Fall back to prefix map if !found { self.prefix.remove_if_mut(key, |_, watchers| { - let before = watchers.len(); - watchers.retain(|w| w.id != id); - found = watchers.len() < before; + if let Some(pos) = watchers.iter().position(|w| w.id == id) { + had_prev_kv = watchers[pos].prev_kv; + watchers.remove(pos); + found = true; + } watchers.is_empty() }); } + if found { self.total_count.fetch_sub(1, Ordering::Relaxed); + if had_prev_kv { + self.prev_kv_watcher_count.fetch_sub(1, Ordering::Relaxed); + } } } + /// Number of active watchers that opted in to prev_kv. + pub fn prev_kv_watcher_count(&self) -> usize { + self.prev_kv_watcher_count.load(Ordering::Relaxed) + } + + /// Clone of the shared prev_kv counter Arc. + /// + /// Pass this to `DefaultStateMachineHandler` so it can poll the live value + /// without holding a reference to the registry. + pub fn prev_kv_watcher_count_arc(&self) -> Arc { + Arc::clone(&self.prev_kv_watcher_count) + } + #[cfg(test)] pub(crate) fn watcher_count( &self, @@ -313,30 +453,65 @@ impl WatchRegistry { } } +// ── WatchDispatcher ─────────────────────────────────────────────────────────── + /// Watch dispatcher — distributes events to watchers (background task). /// /// Spawned explicitly in NodeBuilder::build() to make resource allocation visible. +/// Receives proto WatchResponse from the broadcast channel, converts to opaque +/// WatchEvent per-watcher, and delivers via per-watcher mpsc channels. pub struct WatchDispatcher { registry: Arc, - broadcast_rx: broadcast::Receiver, + broadcast_rx: broadcast::Receiver, unregister_rx: mpsc::UnboundedReceiver<(u64, Bytes)>, + /// Shared applied index for Progress event revision field. + last_applied: Arc, + /// Heartbeat interval. 0 = disabled. + heartbeat_interval_ms: u64, } impl WatchDispatcher { pub fn new( registry: Arc, - broadcast_rx: broadcast::Receiver, + broadcast_rx: broadcast::Receiver, unregister_rx: mpsc::UnboundedReceiver<(u64, Bytes)>, + last_applied: Arc, + heartbeat_interval_ms: u64, ) -> Self { Self { registry, broadcast_rx, unregister_rx, + last_applied, + heartbeat_interval_ms, } } pub async fn run(mut self) { debug!("WatchDispatcher started"); + + // Build optional heartbeat future. When interval_ms == 0 the future + // is a pending sleep that never fires, adding zero overhead. + let heartbeat_enabled = self.heartbeat_interval_ms > 0; + let base_ms = self.heartbeat_interval_ms; + + // ±10% jitter on first tick to spread restarts across the interval window. + let first_tick_ms = if heartbeat_enabled { + let jitter = (base_ms / 10).max(1); + // Deterministic pseudo-random from current thread-id bits. + let tid_bits = format!("{:?}", std::thread::current().id()); + let seed: u64 = tid_bits.bytes().fold(0u64, |a, b| a.wrapping_add(b as u64)); + let offset = seed % (jitter * 2); + base_ms.saturating_sub(jitter) + offset + } else { + u64::MAX + }; + + let mut heartbeat = tokio::time::interval_at( + tokio::time::Instant::now() + Duration::from_millis(first_tick_ms), + Duration::from_millis(if heartbeat_enabled { base_ms } else { u64::MAX }), + ); + loop { tokio::select! { biased; @@ -356,14 +531,45 @@ impl WatchDispatcher { } } } + _ = heartbeat.tick(), if heartbeat_enabled => { + self.broadcast_progress().await; + } } } debug!("WatchDispatcher stopped"); } + /// Broadcast a synthetic Progress event to ALL active watchers regardless of key. + /// + /// Unlike data events (routed by key), Progress is a liveness signal — every watcher + /// must receive it so clients can detect silent stream death. + /// Keys are collected first to avoid holding DashMap shard locks across awaits. + async fn broadcast_progress(&self) { + let revision = self.last_applied.load(Ordering::Relaxed); + let progress = WatchResponse { + key: Bytes::new(), + value: Bytes::new(), + prev_value: Bytes::new(), + event_type: d_engine_proto::client::WatchEventType::Progress as i32, + error: 0, + revision, + }; + + let exact_keys: Vec = self.registry.exact.iter().map(|e| e.key().clone()).collect(); + let prefix_keys: Vec = + self.registry.prefix.iter().map(|e| e.key().clone()).collect(); + + for key in exact_keys { + self.dispatch_to_map(&self.registry.exact, &key, &progress).await; + } + for key in prefix_keys { + self.dispatch_to_map(&self.registry.prefix, &key, &progress).await; + } + } + async fn dispatch_event( &self, - event: WatchEvent, + event: WatchResponse, ) { // Step 1: exact match — O(1) DashMap lookup self.dispatch_to_map(&self.registry.exact, &event.key, &event).await; @@ -376,13 +582,14 @@ impl WatchDispatcher { /// Dispatch an event to all watchers under `lookup_key` in `map`. /// - /// Handles overflow detection and dead watcher cleanup identically for - /// both the exact and prefix maps, keeping the two dispatch paths DRY. + /// Converts proto WatchResponse → opaque WatchEvent per watcher, respecting + /// each watcher's prev_kv preference. Handles overflow detection and dead + /// watcher cleanup identically for both the exact and prefix maps. async fn dispatch_to_map( &self, map: &DashMap>, lookup_key: &Bytes, - event: &WatchEvent, + event: &WatchResponse, ) { if let Some(watchers) = map.get(lookup_key) { let mut dead_watchers = Vec::new(); @@ -409,9 +616,12 @@ impl WatchDispatcher { continue; } + // Convert proto → opaque, respecting per-watcher prev_kv flag + let watch_event = proto_to_event(event, watcher.prev_kv); + // Normal send: reserved slot untouched if let Err(mpsc::error::TrySendError::Closed(_)) = - watcher.sender.try_send(event.clone()) + watcher.sender.try_send(watch_event) { // Receiver dropped: silent cleanup, no cancel needed dead_watchers.push(watcher.id); diff --git a/d-engine-core/src/watch/manager_test.rs b/d-engine-core/src/watch/manager_test.rs index 15afaa3c..be95b5b4 100644 --- a/d-engine-core/src/watch/manager_test.rs +++ b/d-engine-core/src/watch/manager_test.rs @@ -8,8 +8,10 @@ //! - Dispatcher event distribution use std::sync::Arc; +use std::sync::atomic::AtomicU64; use bytes::Bytes; +use d_engine_proto::client::WatchResponse; use tokio::sync::broadcast; use tokio::sync::mpsc; use tokio::time::Duration; @@ -19,11 +21,23 @@ use tokio::time::timeout; use super::*; -/// Helper to create test watch system components +/// Helper to create test watch system components. +/// `heartbeat_interval_ms` = 0 disables heartbeat (avoids Progress noise in most tests). fn setup_watch_system( buffer_size: usize ) -> ( - broadcast::Sender, + broadcast::Sender, + Arc, + tokio::task::JoinHandle<()>, +) { + setup_watch_system_with_heartbeat(buffer_size, 0) +} + +fn setup_watch_system_with_heartbeat( + buffer_size: usize, + heartbeat_interval_ms: u64, +) -> ( + broadcast::Sender, Arc, tokio::task::JoinHandle<()>, ) { @@ -31,7 +45,14 @@ fn setup_watch_system( let (unregister_tx, unregister_rx) = mpsc::unbounded_channel(); let registry = Arc::new(WatchRegistry::new(buffer_size, unregister_tx)); - let dispatcher = WatchDispatcher::new(Arc::clone(®istry), broadcast_rx, unregister_rx); + let last_applied = Arc::new(AtomicU64::new(0)); + let dispatcher = WatchDispatcher::new( + Arc::clone(®istry), + broadcast_rx, + unregister_rx, + last_applied, + heartbeat_interval_ms, + ); let handle = tokio::spawn(async move { dispatcher.run().await; @@ -40,12 +61,42 @@ fn setup_watch_system( (broadcast_tx, registry, handle) } +// Convenience: build a WatchResponse with the proto event type int. +fn proto_put_response( + key: &Bytes, + value: &str, + revision: u64, +) -> WatchResponse { + WatchResponse { + key: key.clone(), + value: Bytes::from(value.to_owned()), + prev_value: Bytes::new(), + event_type: d_engine_proto::client::WatchEventType::Put as i32, + error: 0, + revision, + } +} + +fn proto_delete_response( + key: &Bytes, + revision: u64, +) -> WatchResponse { + WatchResponse { + key: key.clone(), + value: Bytes::new(), + prev_value: Bytes::new(), + event_type: d_engine_proto::client::WatchEventType::Delete as i32, + error: 0, + revision, + } +} + #[tokio::test] async fn test_register_single_watcher() { let (_, registry, _dispatcher_handle) = setup_watch_system(10); let key = Bytes::from("test_key"); - let _handle = registry.register(key.clone()).unwrap(); + let _handle = registry.register(key.clone(), false).unwrap(); assert_eq!(registry.watcher_count(&key), 1); assert_eq!(registry.watched_key_count(), 1); @@ -57,9 +108,9 @@ async fn test_register_multiple_watchers_same_key() { let key = Bytes::from("shared_key"); - let _handle1 = registry.register(key.clone()).unwrap(); - let _handle2 = registry.register(key.clone()).unwrap(); - let _handle3 = registry.register(key.clone()).unwrap(); + let _handle1 = registry.register(key.clone(), false).unwrap(); + let _handle2 = registry.register(key.clone(), false).unwrap(); + let _handle3 = registry.register(key.clone(), false).unwrap(); assert_eq!(registry.watcher_count(&key), 3); assert_eq!(registry.watched_key_count(), 1); // Only 1 unique key @@ -72,7 +123,7 @@ async fn test_watcher_auto_cleanup_on_drop() { let key = Bytes::from("cleanup_key"); { - let _handle = registry.register(key.clone()).unwrap(); + let _handle = registry.register(key.clone(), false).unwrap(); assert_eq!(registry.watcher_count(&key), 1); // Handle dropped here } @@ -91,20 +142,12 @@ async fn test_dispatcher_dispatch_to_matching_watcher() { let key = Bytes::from("test_key"); let value = Bytes::from("test_value"); - let mut handle = registry.register(key.clone()).unwrap(); + let mut handle = registry.register(key.clone(), false).unwrap(); // Small delay to ensure dispatcher is ready tokio::time::sleep(Duration::from_millis(50)).await; - // Send event via broadcast - let event = WatchEvent { - key: key.clone(), - value: value.clone(), - event_type: WatchEventType::Put as i32, - error: 0, - revision: 0, - }; - broadcast_tx.send(event).unwrap(); + broadcast_tx.send(proto_put_response(&key, "test_value", 0)).unwrap(); // Watcher should receive event let received = timeout(Duration::from_millis(100), handle.receiver_mut().recv()) @@ -114,7 +157,7 @@ async fn test_dispatcher_dispatch_to_matching_watcher() { assert_eq!(received.key, key); assert_eq!(received.value, value); - assert_eq!(received.event_type, WatchEventType::Put as i32); + assert_eq!(received.event_type, WatchEventType::Put); } #[tokio::test] @@ -124,19 +167,11 @@ async fn test_dispatcher_ignores_non_matching_key() { let watched_key = Bytes::from("key1"); let other_key = Bytes::from("key2"); - let mut handle = registry.register(watched_key.clone()).unwrap(); + let mut handle = registry.register(watched_key.clone(), false).unwrap(); tokio::time::sleep(Duration::from_millis(50)).await; - // Send event for different key - let event = WatchEvent { - key: other_key, - value: Bytes::from("value"), - event_type: WatchEventType::Put as i32, - error: 0, - revision: 0, - }; - broadcast_tx.send(event).unwrap(); + broadcast_tx.send(proto_put_response(&other_key, "value", 0)).unwrap(); // Should timeout (no event received) let result = timeout(Duration::from_millis(100), handle.receiver_mut().recv()).await; @@ -153,23 +188,15 @@ async fn test_multiple_watchers_all_receive_event() { let key = Bytes::from("shared_key"); let value = Bytes::from("shared_value"); - let mut handle1 = registry.register(key.clone()).unwrap(); - let mut handle2 = registry.register(key.clone()).unwrap(); - let mut handle3 = registry.register(key.clone()).unwrap(); + let mut handle1 = registry.register(key.clone(), false).unwrap(); + let mut handle2 = registry.register(key.clone(), false).unwrap(); + let mut handle3 = registry.register(key.clone(), false).unwrap(); assert_eq!(registry.watcher_count(&key), 3); tokio::time::sleep(Duration::from_millis(50)).await; - // Broadcast once - let event = WatchEvent { - key: key.clone(), - value: value.clone(), - event_type: WatchEventType::Put as i32, - error: 0, - revision: 0, - }; - broadcast_tx.send(event).unwrap(); + broadcast_tx.send(proto_put_response(&key, "shared_value", 0)).unwrap(); // All 3 should receive for handle in [&mut handle1, &mut handle2, &mut handle3].iter_mut() { @@ -188,26 +215,18 @@ async fn test_watcher_delete_event() { let (broadcast_tx, registry, _dispatcher_handle) = setup_watch_system(10); let key = Bytes::from("test_key"); - let mut handle = registry.register(key.clone()).unwrap(); + let mut handle = registry.register(key.clone(), false).unwrap(); tokio::time::sleep(Duration::from_millis(50)).await; - // Send DELETE event - let event = WatchEvent { - key: key.clone(), - value: Bytes::new(), - event_type: WatchEventType::Delete as i32, - error: 0, - revision: 0, - }; - broadcast_tx.send(event).unwrap(); + broadcast_tx.send(proto_delete_response(&key, 0)).unwrap(); let received = timeout(Duration::from_millis(100), handle.receiver_mut().recv()) .await .expect("Timeout") .expect("Channel closed"); - assert_eq!(received.event_type, WatchEventType::Delete as i32); + assert_eq!(received.event_type, WatchEventType::Delete); assert_eq!(received.value, Bytes::new()); } @@ -216,40 +235,13 @@ async fn test_multiple_events_sequential() { let (broadcast_tx, registry, _dispatcher_handle) = setup_watch_system(10); let key = Bytes::from("test_key"); - let mut handle = registry.register(key.clone()).unwrap(); + let mut handle = registry.register(key.clone(), false).unwrap(); tokio::time::sleep(Duration::from_millis(50)).await; - // Send 3 events - broadcast_tx - .send(WatchEvent { - key: key.clone(), - value: Bytes::from("value1"), - event_type: WatchEventType::Put as i32, - error: 0, - revision: 0, - }) - .unwrap(); - - broadcast_tx - .send(WatchEvent { - key: key.clone(), - value: Bytes::from("value2"), - event_type: WatchEventType::Put as i32, - error: 0, - revision: 0, - }) - .unwrap(); - - broadcast_tx - .send(WatchEvent { - key: key.clone(), - value: Bytes::new(), - event_type: WatchEventType::Delete as i32, - error: 0, - revision: 0, - }) - .unwrap(); + broadcast_tx.send(proto_put_response(&key, "value1", 0)).unwrap(); + broadcast_tx.send(proto_put_response(&key, "value2", 0)).unwrap(); + broadcast_tx.send(proto_delete_response(&key, 0)).unwrap(); // Receive all in order let event1 = handle.receiver_mut().recv().await.unwrap(); @@ -259,7 +251,7 @@ async fn test_multiple_events_sequential() { assert_eq!(event2.value, Bytes::from("value2")); let event3 = handle.receiver_mut().recv().await.unwrap(); - assert_eq!(event3.event_type, WatchEventType::Delete as i32); + assert_eq!(event3.event_type, WatchEventType::Delete); } #[tokio::test] @@ -268,9 +260,9 @@ async fn test_watcher_count_after_partial_cleanup() { let key = Bytes::from("count_key"); - let handle1 = registry.register(key.clone()).unwrap(); - let _handle2 = registry.register(key.clone()).unwrap(); - let _handle3 = registry.register(key.clone()).unwrap(); + let handle1 = registry.register(key.clone(), false).unwrap(); + let _handle2 = registry.register(key.clone(), false).unwrap(); + let _handle3 = registry.register(key.clone(), false).unwrap(); assert_eq!(registry.watcher_count(&key), 3); @@ -287,21 +279,12 @@ async fn test_different_keys_isolated() { let key1 = Bytes::from("key1"); let key2 = Bytes::from("key2"); - let mut handle1 = registry.register(key1.clone()).unwrap(); - let mut handle2 = registry.register(key2.clone()).unwrap(); + let mut handle1 = registry.register(key1.clone(), false).unwrap(); + let mut handle2 = registry.register(key2.clone(), false).unwrap(); tokio::time::sleep(Duration::from_millis(50)).await; - // Send event only for key1 - broadcast_tx - .send(WatchEvent { - key: key1.clone(), - value: Bytes::from("value1"), - event_type: WatchEventType::Put as i32, - error: 0, - revision: 0, - }) - .unwrap(); + broadcast_tx.send(proto_put_response(&key1, "value1", 0)).unwrap(); // handle1 should receive let event = timeout(Duration::from_millis(100), handle1.receiver_mut().recv()) @@ -318,26 +301,21 @@ async fn test_different_keys_isolated() { // --------------------------------------------------------------------------- // #294: overflow protection — buffer overflow → CANCELED notification // -// Requires proto additions before these compile: -// WatchEventType::Canceled = 2 (client_api.proto) -// ErrorCode::WatchBufferOverflow = 5001 (error.proto) -// // With buffer_size=N, channel capacity=N+1 (1 slot reserved for cancel). // Dispatch checks capacity() before each normal event: // capacity > 1 → send normally // capacity == 1 → send CANCELED to the reserved slot, mark dead // --------------------------------------------------------------------------- -use d_engine_proto::error::ErrorCode; - fn put_event( key: &Bytes, value: &str, -) -> WatchEvent { - WatchEvent { +) -> WatchResponse { + WatchResponse { key: key.clone(), value: Bytes::from(value.to_owned()), - event_type: WatchEventType::Put as i32, + prev_value: Bytes::new(), + event_type: d_engine_proto::client::WatchEventType::Put as i32, error: 0, revision: 0, } @@ -351,7 +329,7 @@ fn put_event( async fn test_watcher_buffer_overflow_sends_cancel_event() { let (broadcast_tx, registry, _dispatcher_handle) = setup_watch_system(2); let key = Bytes::from("overflow_key"); - let mut handle = registry.register(key.clone()).unwrap(); + let mut handle = registry.register(key.clone(), false).unwrap(); tokio::time::sleep(Duration::from_millis(20)).await; @@ -366,22 +344,22 @@ async fn test_watcher_buffer_overflow_sends_cancel_event() { .await .expect("timeout e1") .expect("channel closed e1"); - assert_eq!(e1.event_type, WatchEventType::Put as i32); + assert_eq!(e1.event_type, WatchEventType::Put); let e2 = timeout(Duration::from_millis(100), handle.receiver_mut().recv()) .await .expect("timeout e2") .expect("channel closed e2"); - assert_eq!(e2.event_type, WatchEventType::Put as i32); + assert_eq!(e2.event_type, WatchEventType::Put); let cancel = timeout(Duration::from_millis(100), handle.receiver_mut().recv()) .await .expect("timeout cancel") .expect("channel closed cancel"); - assert_eq!(cancel.event_type, WatchEventType::Canceled as i32); + assert_eq!(cancel.event_type, WatchEventType::Canceled); // Channel drained, no further events. After unregistration the sender is - // dropped, so recv() returns None (Ok(None)) rather than timing out. + // dropped, so recv() returns None rather than timing out. let nothing = timeout(Duration::from_millis(50), handle.receiver_mut().recv()).await; assert!( matches!(nothing, Err(_) | Ok(None)), @@ -390,12 +368,12 @@ async fn test_watcher_buffer_overflow_sends_cancel_event() { ); } -// CANCELED must carry the correct key, empty value, and WATCH_BUFFER_OVERFLOW error. +// CANCELED must carry the correct key and empty value. #[tokio::test] async fn test_cancel_event_has_correct_fields() { let (broadcast_tx, registry, _dispatcher_handle) = setup_watch_system(1); let key = Bytes::from("field_check_key"); - let mut handle = registry.register(key.clone()).unwrap(); + let mut handle = registry.register(key.clone(), false).unwrap(); tokio::time::sleep(Duration::from_millis(20)).await; @@ -411,8 +389,7 @@ async fn test_cancel_event_has_correct_fields() { .expect("timeout") .expect("channel closed"); - assert_eq!(cancel.event_type, WatchEventType::Canceled as i32); - assert_eq!(cancel.error, ErrorCode::WatchBufferOverflow as i32); + assert_eq!(cancel.event_type, WatchEventType::Canceled); assert_eq!(cancel.key, key); assert_eq!(cancel.value, Bytes::new()); } @@ -422,7 +399,7 @@ async fn test_cancel_event_has_correct_fields() { async fn test_overflow_removes_watcher_from_registry() { let (broadcast_tx, registry, _dispatcher_handle) = setup_watch_system(1); let key = Bytes::from("registry_key"); - let _handle = registry.register(key.clone()).unwrap(); + let _handle = registry.register(key.clone(), false).unwrap(); assert_eq!(registry.watcher_count(&key), 1); @@ -441,7 +418,7 @@ async fn test_overflow_removes_watcher_from_registry() { async fn test_normal_events_precede_cancel_in_channel() { let (broadcast_tx, registry, _dispatcher_handle) = setup_watch_system(2); let key = Bytes::from("order_key"); - let mut handle = registry.register(key.clone()).unwrap(); + let mut handle = registry.register(key.clone(), false).unwrap(); tokio::time::sleep(Duration::from_millis(20)).await; @@ -460,9 +437,9 @@ async fn test_normal_events_precede_cancel_in_channel() { // Last event must be CANCELED, everything before must be PUT assert!(!events.is_empty()); let last = events.last().unwrap(); - assert_eq!(last.event_type, WatchEventType::Canceled as i32); + assert_eq!(last.event_type, WatchEventType::Canceled); for ev in &events[..events.len() - 1] { - assert_eq!(ev.event_type, WatchEventType::Put as i32); + assert_eq!(ev.event_type, WatchEventType::Put); } // v4 must NOT appear (watcher was dead when v4 was dispatched) assert!( @@ -477,7 +454,7 @@ async fn test_normal_events_precede_cancel_in_channel() { async fn test_closed_receiver_cleaned_up_silently() { let (broadcast_tx, registry, _dispatcher_handle) = setup_watch_system(10); let key = Bytes::from("closed_key"); - let handle = registry.register(key.clone()).unwrap(); + let handle = registry.register(key.clone(), false).unwrap(); assert_eq!(registry.watcher_count(&key), 1); @@ -500,8 +477,8 @@ async fn test_slow_watcher_overflow_does_not_affect_healthy_watcher() { let (broadcast_tx, registry, _dispatcher_handle) = setup_watch_system(2); let key = Bytes::from("isolation_key"); - let mut slow = registry.register(key.clone()).unwrap(); // never consumed - let mut fast = registry.register(key.clone()).unwrap(); // consumed between dispatches + let mut slow = registry.register(key.clone(), false).unwrap(); // never consumed + let mut fast = registry.register(key.clone(), false).unwrap(); // consumed between dispatches tokio::time::sleep(Duration::from_millis(20)).await; @@ -525,7 +502,7 @@ async fn test_slow_watcher_overflow_does_not_affect_healthy_watcher() { .ok() .flatten(); match ev { - Some(e) if e.event_type == WatchEventType::Canceled as i32 => break, + Some(e) if e.event_type == WatchEventType::Canceled => break, Some(_) => continue, None => panic!("slow watcher closed without CANCELED"), } @@ -556,7 +533,7 @@ async fn test_slow_watcher_overflow_does_not_affect_healthy_watcher() { async fn test_healthy_watcher_no_false_cancel() { let (broadcast_tx, registry, _dispatcher_handle) = setup_watch_system(2); let key = Bytes::from("healthy_key"); - let mut handle = registry.register(key.clone()).unwrap(); + let mut handle = registry.register(key.clone(), false).unwrap(); tokio::time::sleep(Duration::from_millis(20)).await; @@ -572,7 +549,7 @@ async fn test_healthy_watcher_no_false_cancel() { assert_ne!( ev.event_type, - WatchEventType::Canceled as i32, + WatchEventType::Canceled, "false cancel on event {i}" ); } @@ -588,7 +565,7 @@ async fn test_healthy_watcher_no_false_cancel() { async fn test_buffer_size_1_overflows_on_second_event() { let (broadcast_tx, registry, _dispatcher_handle) = setup_watch_system(1); let key = Bytes::from("tiny_buffer_key"); - let mut handle = registry.register(key.clone()).unwrap(); + let mut handle = registry.register(key.clone(), false).unwrap(); tokio::time::sleep(Duration::from_millis(20)).await; @@ -600,18 +577,16 @@ async fn test_buffer_size_1_overflows_on_second_event() { .await .expect("timeout e1") .expect("closed e1"); - assert_eq!(e1.event_type, WatchEventType::Put as i32); + assert_eq!(e1.event_type, WatchEventType::Put); assert_eq!(e1.value, Bytes::from("v1")); let cancel = timeout(Duration::from_millis(100), handle.receiver_mut().recv()) .await .expect("timeout cancel") .expect("closed cancel"); - assert_eq!(cancel.event_type, WatchEventType::Canceled as i32); - assert_eq!(cancel.error, ErrorCode::WatchBufferOverflow as i32); + assert_eq!(cancel.event_type, WatchEventType::Canceled); - // Exactly 2 items (buffer_size + 1), no more. The sender is dropped - // when the watcher is unregistered, so recv() returns None immediately. + // Exactly 2 items (buffer_size + 1), no more. let nothing = timeout(Duration::from_millis(50), handle.receiver_mut().recv()).await; assert!( matches!(nothing, Err(_) | Ok(None)), @@ -626,7 +601,7 @@ async fn test_buffer_size_1_overflows_on_second_event() { async fn test_events_not_delivered_after_overflow() { let (broadcast_tx, registry, _dispatcher_handle) = setup_watch_system(1); let key = Bytes::from("post_overflow_key"); - let _handle = registry.register(key.clone()).unwrap(); + let _handle = registry.register(key.clone(), false).unwrap(); tokio::time::sleep(Duration::from_millis(20)).await; @@ -652,7 +627,7 @@ async fn test_into_receiver_disables_cleanup() { let (_, registry, _dispatcher_handle) = setup_watch_system(10); let key = Bytes::from("test_key"); - let handle = registry.register(key.clone()).unwrap(); + let handle = registry.register(key.clone(), false).unwrap(); assert_eq!(registry.watcher_count(&key), 1); @@ -664,7 +639,6 @@ async fn test_into_receiver_disables_cleanup() { tokio::time::sleep(Duration::from_millis(50)).await; // Watcher still registered (cleanup was disabled) - // Note: In real usage, receiver drop will eventually trigger cleanup via send failure assert_eq!(registry.watcher_count(&key), 1); } @@ -679,7 +653,7 @@ async fn test_concurrent_register_unregister() { let reg = Arc::clone(®istry); let k = key.clone(); let handle = tokio::spawn(async move { - let watcher = reg.register(k); + let watcher = reg.register(k, false); tokio::time::sleep(Duration::from_millis(10)).await; drop(watcher); // Yield to ensure drop completes @@ -708,7 +682,7 @@ async fn test_dispatcher_shutdown_on_broadcast_close() { let (broadcast_tx, registry, dispatcher_handle) = setup_watch_system(10); let key = Bytes::from("test_key"); - let _handle = registry.register(key.clone()).unwrap(); + let _handle = registry.register(key.clone(), false).unwrap(); tokio::time::sleep(Duration::from_millis(50)).await; @@ -731,7 +705,7 @@ fn setup_watch_system_with_max( buffer_size: usize, max_watcher_count: usize, ) -> ( - broadcast::Sender, + broadcast::Sender, Arc, tokio::task::JoinHandle<()>, ) { @@ -742,7 +716,14 @@ fn setup_watch_system_with_max( max_watcher_count, unregister_tx, )); - let dispatcher = WatchDispatcher::new(Arc::clone(®istry), broadcast_rx, unregister_rx); + let last_applied = Arc::new(AtomicU64::new(0)); + let dispatcher = WatchDispatcher::new( + Arc::clone(®istry), + broadcast_rx, + unregister_rx, + last_applied, + 0, + ); let handle = tokio::spawn(async move { dispatcher.run().await }); (broadcast_tx, registry, handle) } @@ -750,9 +731,6 @@ fn setup_watch_system_with_max( // --- prefix_segments() unit tests (pure function) --- /// #300: prefix_segments is the core of O(depth) dispatch. -/// Verifies that a multi-level path is decomposed into all slash-terminated -/// prefix candidates so the dispatcher can do O(1) DashMap lookups per level -/// instead of scanning every registered prefix. #[test] fn test_prefix_segments_decomposes_path_correctly() { let key = Bytes::from("/config/db/host"); @@ -767,9 +745,6 @@ fn test_prefix_segments_decomposes_path_correctly() { ); } -/// #300: "/" is both a valid key and the root prefix. -/// prefix_segments must return ["/"] so that a watcher registered on "/" -/// is notified when the root key itself changes. #[test] fn test_prefix_segments_root_key() { let key = Bytes::from("/"); @@ -777,9 +752,6 @@ fn test_prefix_segments_root_key() { assert_eq!(segments, vec![Bytes::from("/")]); } -/// #300: A single-level key like "/config" (no trailing slash) has exactly -/// one ancestor prefix: "/". Ensures prefix watchers on "/" are notified -/// for top-level keys even when those keys are not directories. #[test] fn test_prefix_segments_single_level() { let key = Bytes::from("/config"); @@ -787,9 +759,6 @@ fn test_prefix_segments_single_level() { assert_eq!(segments, vec![Bytes::from("/")]); } -/// #300: A key ending with "/" (directory-style) must include itself as a -/// prefix segment. This matters for service-discovery patterns where both -/// the namespace key and its children can be watched. #[test] fn test_prefix_segments_trailing_slash_included() { let key = Bytes::from("/config/"); @@ -799,22 +768,20 @@ fn test_prefix_segments_trailing_slash_included() { // --- prefix watch: basic dispatch --- -/// #300: Core use case — a client watching an entire namespace receives -/// events for any key that falls under that namespace. Without prefix watch, -/// clients must register one watcher per key, which is impractical for -/// dynamic service-discovery or config namespaces. +/// #300: Core use case — a client watching an entire namespace receives events. #[tokio::test] async fn test_prefix_watch_matches_child_key() { let (broadcast_tx, registry, _handle) = setup_watch_system(10); - let mut watcher = registry.register_prefix(Bytes::from("/config/")).unwrap(); + let mut watcher = registry.register_prefix(Bytes::from("/config/"), false).unwrap(); tokio::time::sleep(Duration::from_millis(50)).await; broadcast_tx - .send(WatchEvent { + .send(WatchResponse { key: Bytes::from("/config/db/host"), value: Bytes::from("10.0.0.1"), - event_type: WatchEventType::Put as i32, + prev_value: Bytes::new(), + event_type: d_engine_proto::client::WatchEventType::Put as i32, error: 0, revision: 1, }) @@ -827,21 +794,20 @@ async fn test_prefix_watch_matches_child_key() { assert_eq!(received.key, Bytes::from("/config/db/host")); } -/// #300: Prefix isolation — a watcher on "/config/" must not receive events -/// for unrelated namespaces like "/dns/". Noisy cross-namespace delivery -/// would break multi-tenant deployments where each tenant owns a prefix. +/// #300: Prefix isolation — a watcher on "/config/" must not receive events for "/dns/". #[tokio::test] async fn test_prefix_watch_does_not_match_unrelated_key() { let (broadcast_tx, registry, _handle) = setup_watch_system(10); - let mut watcher = registry.register_prefix(Bytes::from("/config/")).unwrap(); + let mut watcher = registry.register_prefix(Bytes::from("/config/"), false).unwrap(); tokio::time::sleep(Duration::from_millis(50)).await; broadcast_tx - .send(WatchEvent { + .send(WatchResponse { key: Bytes::from("/dns/xyz"), value: Bytes::from("1.1.1.1"), - event_type: WatchEventType::Put as i32, + prev_value: Bytes::new(), + event_type: d_engine_proto::client::WatchEventType::Put as i32, error: 0, revision: 1, }) @@ -851,22 +817,20 @@ async fn test_prefix_watch_does_not_match_unrelated_key() { assert!(result.is_err(), "/config/ prefix must not match /dns/xyz"); } -/// #300: Slash boundary — "/config/" is a namespace prefix; "/config" is a -/// distinct key at the parent level. A watcher on the namespace must not -/// fire for the parent key. This matches etcd semantics and prevents -/// accidental cross-boundary notifications. +/// #300: Slash boundary — "/config/" must not fire for the parent key "/config". #[tokio::test] async fn test_prefix_watch_slash_boundary_not_matched() { let (broadcast_tx, registry, _handle) = setup_watch_system(10); - let mut watcher = registry.register_prefix(Bytes::from("/config/")).unwrap(); + let mut watcher = registry.register_prefix(Bytes::from("/config/"), false).unwrap(); tokio::time::sleep(Duration::from_millis(50)).await; broadcast_tx - .send(WatchEvent { + .send(WatchResponse { key: Bytes::from("/config"), value: Bytes::from("val"), - event_type: WatchEventType::Put as i32, + prev_value: Bytes::new(), + event_type: d_engine_proto::client::WatchEventType::Put as i32, error: 0, revision: 1, }) @@ -879,21 +843,20 @@ async fn test_prefix_watch_slash_boundary_not_matched() { ); } -/// #300: The root prefix "/" acts as a global watch — useful for audit -/// logging or debugging where a single watcher must see all key changes -/// across all namespaces. +/// #300: The root prefix "/" acts as a global watch. #[tokio::test] async fn test_root_prefix_matches_child_keys() { let (broadcast_tx, registry, _handle) = setup_watch_system(10); - let mut watcher = registry.register_prefix(Bytes::from("/")).unwrap(); + let mut watcher = registry.register_prefix(Bytes::from("/"), false).unwrap(); tokio::time::sleep(Duration::from_millis(50)).await; broadcast_tx - .send(WatchEvent { + .send(WatchResponse { key: Bytes::from("/services/api"), value: Bytes::from("up"), - event_type: WatchEventType::Put as i32, + prev_value: Bytes::new(), + event_type: d_engine_proto::client::WatchEventType::Put as i32, error: 0, revision: 1, }) @@ -906,21 +869,20 @@ async fn test_root_prefix_matches_child_keys() { assert_eq!(received.key, Bytes::from("/services/api")); } -/// #300: Edge case fix — when the event key IS "/" (root key itself), -/// prefix_segments("/") = ["/"], so a watcher on prefix "/" must be notified. -/// Without the fix (i < key.len()-1 guard), this event was silently dropped. +/// #300: Edge case fix — when the event key IS "/" (root key itself). #[tokio::test] async fn test_root_prefix_matches_root_key_itself() { let (broadcast_tx, registry, _handle) = setup_watch_system(10); - let mut watcher = registry.register_prefix(Bytes::from("/")).unwrap(); + let mut watcher = registry.register_prefix(Bytes::from("/"), false).unwrap(); tokio::time::sleep(Duration::from_millis(50)).await; broadcast_tx - .send(WatchEvent { + .send(WatchResponse { key: Bytes::from("/"), value: Bytes::from("root"), - event_type: WatchEventType::Put as i32, + prev_value: Bytes::new(), + event_type: d_engine_proto::client::WatchEventType::Put as i32, error: 0, revision: 1, }) @@ -933,24 +895,23 @@ async fn test_root_prefix_matches_root_key_itself() { assert_eq!(received.key, Bytes::from("/")); } -/// #300: Exact and prefix watchers are independent and must both fire for -/// the same event. A client watching a specific key for fast reaction AND a -/// monitoring client watching the whole namespace must not block each other. +/// #300: Exact and prefix watchers are independent and must both fire. #[tokio::test] async fn test_exact_and_prefix_both_notified_on_same_event() { let (broadcast_tx, registry, _handle) = setup_watch_system(10); let key = Bytes::from("/config/db/host"); - let mut exact_watcher = registry.register(key.clone()).unwrap(); - let mut prefix_watcher = registry.register_prefix(Bytes::from("/config/")).unwrap(); + let mut exact_watcher = registry.register(key.clone(), false).unwrap(); + let mut prefix_watcher = registry.register_prefix(Bytes::from("/config/"), false).unwrap(); tokio::time::sleep(Duration::from_millis(50)).await; broadcast_tx - .send(WatchEvent { + .send(WatchResponse { key: key.clone(), value: Bytes::from("10.0.0.1"), - event_type: WatchEventType::Put as i32, + prev_value: Bytes::new(), + event_type: d_engine_proto::client::WatchEventType::Put as i32, error: 0, revision: 5, }) @@ -977,53 +938,48 @@ async fn test_exact_and_prefix_both_notified_on_same_event() { // --- register_prefix: validation --- -/// #300: Prefix format is enforced at registration time to prevent silent -/// mismatches at dispatch time. A prefix without a leading slash would never -/// match any key (all keys start with "/"). A prefix without a trailing slash -/// is semantically ambiguous with an exact key — reject early with a clear error. #[test] fn test_register_prefix_rejects_missing_trailing_slash() { let (unregister_tx, _) = mpsc::unbounded_channel(); let registry = WatchRegistry::new(10, unregister_tx); assert!( - registry.register_prefix(Bytes::from("/config")).is_err(), + registry.register_prefix(Bytes::from("/config"), false).is_err(), "/config must be rejected — no trailing slash" ); assert!( - registry.register_prefix(Bytes::from("config/")).is_err(), + registry.register_prefix(Bytes::from("config/"), false).is_err(), "config/ must be rejected — no leading slash" ); } // --- prefix watcher: overflow sends CANCELED --- -/// #300: Prefix watchers are subject to the same overflow protection as exact -/// watchers (#294). A slow consumer watching a high-traffic namespace must -/// receive CANCELED (not silently die) so it can re-sync via the Read API. -/// Without this, the client has no signal that it missed events. +/// #300: Prefix watchers are subject to the same overflow protection as exact watchers. #[tokio::test] async fn test_prefix_watcher_overflow_sends_cancel() { let (broadcast_tx, registry, _handle) = setup_watch_system(1); // buffer_size=1 - let mut watcher = registry.register_prefix(Bytes::from("/config/")).unwrap(); + let mut watcher = registry.register_prefix(Bytes::from("/config/"), false).unwrap(); tokio::time::sleep(Duration::from_millis(20)).await; // 2 events fill buffer + trigger overflow (capacity = buffer_size+1 = 2) broadcast_tx - .send(WatchEvent { + .send(WatchResponse { key: Bytes::from("/config/key1"), value: Bytes::from("v1"), - event_type: WatchEventType::Put as i32, + prev_value: Bytes::new(), + event_type: d_engine_proto::client::WatchEventType::Put as i32, error: 0, revision: 1, }) .unwrap(); broadcast_tx - .send(WatchEvent { + .send(WatchResponse { key: Bytes::from("/config/key2"), value: Bytes::from("v2"), - event_type: WatchEventType::Put as i32, + prev_value: Bytes::new(), + event_type: d_engine_proto::client::WatchEventType::Put as i32, error: 0, revision: 2, }) @@ -1036,22 +992,19 @@ async fn test_prefix_watcher_overflow_sends_cancel() { .await .expect("timeout") .expect("closed"); - assert_eq!(cancel.event_type, WatchEventType::Canceled as i32); - assert_eq!(cancel.error, ErrorCode::WatchBufferOverflow as i32); + assert_eq!(cancel.event_type, WatchEventType::Canceled); } // --- prefix watcher: auto-cleanup on drop --- -/// #300: Prefix watcher handles follow the same RAII cleanup contract as -/// exact watchers. A leaked prefix watcher would continue consuming dispatch -/// CPU on every matching event even after the client is gone. +/// #300: Prefix watcher handles follow the same RAII cleanup contract as exact watchers. #[tokio::test] async fn test_prefix_watcher_cleanup_on_drop() { let (_, registry, _handle) = setup_watch_system(10); let prefix = Bytes::from("/config/"); { - let _watcher = registry.register_prefix(prefix.clone()).unwrap(); + let _watcher = registry.register_prefix(prefix.clone(), false).unwrap(); assert_eq!(registry.prefix_watcher_count(&prefix), 1); } @@ -1061,23 +1014,21 @@ async fn test_prefix_watcher_cleanup_on_drop() { // --- revision field --- -/// #300: The revision field (= Raft applied index) is the client's only -/// mechanism to detect event gaps after a CANCELED notification. Without it, -/// the client cannot tell the Read API "give me the state as of revision N" -/// and must perform a full blind re-read with no consistency anchor. +/// #300: The revision field (= Raft applied index) is the client's consistency anchor. #[tokio::test] async fn test_revision_field_passed_through_to_watcher() { let (broadcast_tx, registry, _handle) = setup_watch_system(10); let key = Bytes::from("/config/host"); - let mut watcher = registry.register(key.clone()).unwrap(); + let mut watcher = registry.register(key.clone(), false).unwrap(); tokio::time::sleep(Duration::from_millis(50)).await; broadcast_tx - .send(WatchEvent { + .send(WatchResponse { key: key.clone(), value: Bytes::from("val"), - event_type: WatchEventType::Put as i32, + prev_value: Bytes::new(), + event_type: d_engine_proto::client::WatchEventType::Put as i32, error: 0, revision: 42, }) @@ -1095,19 +1046,16 @@ async fn test_revision_field_passed_through_to_watcher() { // --- max_watcher_count limit --- -/// #300: The dispatcher runs in a single task and does O(N) work per event -/// across all matched watchers. Without a hard cap, an unbounded number of -/// watchers degrades write latency for all clients. The limit forces operators -/// to make an explicit capacity decision rather than discovering the ceiling -/// in production under load. +/// #300: The dispatcher runs in a single task and does O(N) work per event. +/// Without a hard cap, an unbounded number of watchers degrades write latency. #[tokio::test] async fn test_max_watcher_count_rejects_registration_when_exceeded() { let (_, registry, _handle) = setup_watch_system_with_max(10, 2); let key = Bytes::from("/config/host"); - let _h1 = registry.register(key.clone()).unwrap(); - let _h2 = registry.register(key.clone()).unwrap(); - let h3 = registry.register(key.clone()); + let _h1 = registry.register(key.clone(), false).unwrap(); + let _h2 = registry.register(key.clone(), false).unwrap(); + let h3 = registry.register(key.clone(), false); assert!( h3.is_err(), @@ -1116,14 +1064,13 @@ async fn test_max_watcher_count_rejects_registration_when_exceeded() { } /// #300: prefix watchers consume the same shared cap as exact watchers. -/// Verifies that `register_prefix` is rejected once `max_watcher_count` is hit. #[tokio::test] async fn test_prefix_watcher_cap_enforced() { let (_, registry, _handle) = setup_watch_system_with_max(10, 2); - let _h1 = registry.register_prefix(Bytes::from("/a/")).unwrap(); - let _h2 = registry.register_prefix(Bytes::from("/b/")).unwrap(); - let h3 = registry.register_prefix(Bytes::from("/c/")); + let _h1 = registry.register_prefix(Bytes::from("/a/"), false).unwrap(); + let _h2 = registry.register_prefix(Bytes::from("/b/"), false).unwrap(); + let h3 = registry.register_prefix(Bytes::from("/c/"), false); assert!( h3.is_err(), @@ -1132,34 +1079,27 @@ async fn test_prefix_watcher_cap_enforced() { } /// #300: exact and prefix watchers share a single cap counter. -/// After 1 exact + 1 prefix fill the cap, a third registration of either -/// kind must be rejected. #[tokio::test] async fn test_mixed_exact_and_prefix_watchers_share_cap() { let (_, registry, _handle) = setup_watch_system_with_max(10, 2); - let _h1 = registry.register(Bytes::from("/exact/key")).unwrap(); - let _h2 = registry.register_prefix(Bytes::from("/prefix/")).unwrap(); + let _h1 = registry.register(Bytes::from("/exact/key"), false).unwrap(); + let _h2 = registry.register_prefix(Bytes::from("/prefix/"), false).unwrap(); - let h3_exact = registry.register(Bytes::from("/exact/other")); + let h3_exact = registry.register(Bytes::from("/exact/other"), false); assert!( h3_exact.is_err(), "exact register must fail after mixed cap is reached" ); - let h3_prefix = registry.register_prefix(Bytes::from("/prefix2/")); + let h3_prefix = registry.register_prefix(Bytes::from("/prefix2/"), false); assert!( h3_prefix.is_err(), "prefix register must fail after mixed cap is reached" ); } -/// #300 / fix: the TOCTOU fix in do_register must prevent the watcher count -/// from ever exceeding max_watcher_count under concurrent registration. -/// -/// The old load-check-then-fetch_add pattern lets N threads all pass the check -/// before any increment, so all N insert and the count overshoots. -/// The reserve-first pattern (fetch_add → check → rollback if over) is race-free. +/// #300 / fix: the TOCTOU fix must prevent concurrent registrations from exceeding cap. #[tokio::test] async fn test_concurrent_registrations_never_exceed_cap() { let max = 5usize; @@ -1172,7 +1112,7 @@ async fn test_concurrent_registrations_never_exceed_cap() { let registry = Arc::clone(®istry); tokio::spawn(async move { let key = Bytes::from(format!("/key/{i}")); - registry.register(key).is_ok() + registry.register(key, false).is_ok() }) }) .collect(); @@ -1189,3 +1129,222 @@ async fn test_concurrent_registrations_never_exceed_cap() { "exactly max_watcher_count registrations must succeed under concurrent load, got {success_count}" ); } + +// ============================================================================= +// #379: prev_kv + Progress Heartbeat +// ============================================================================= + +/// #379: A watcher registered with prev_kv=false receives empty prev_value, +/// even when the broadcast WatchResponse carries a non-empty prev_value. +/// This ensures watchers that didn't opt in never see prev_kv data. +#[tokio::test] +async fn test_prev_kv_false_watcher_receives_empty_prev_value() { + let (broadcast_tx, registry, _handle) = setup_watch_system(10); + let key = Bytes::from("/k"); + let mut watcher = registry.register(key.clone(), false).unwrap(); + + tokio::time::sleep(Duration::from_millis(50)).await; + + broadcast_tx + .send(WatchResponse { + key: key.clone(), + value: Bytes::from("new_val"), + prev_value: Bytes::from("old_val"), // handler populated this + event_type: d_engine_proto::client::WatchEventType::Put as i32, + error: 0, + revision: 1, + }) + .unwrap(); + + let event = timeout(Duration::from_millis(100), watcher.receiver_mut().recv()) + .await + .expect("timeout") + .expect("closed"); + + assert_eq!(event.event_type, WatchEventType::Put); + assert_eq!(event.value, Bytes::from("new_val")); + // prev_kv=false → dispatcher zeroes out prev_value before delivery + assert_eq!( + event.prev_value, + Bytes::new(), + "prev_kv=false watcher must receive empty prev_value" + ); +} + +/// #379: A watcher registered with prev_kv=true receives the prev_value +/// that was populated in the broadcast WatchResponse by the handler. +#[tokio::test] +async fn test_prev_kv_true_watcher_receives_prev_value() { + let (broadcast_tx, registry, _handle) = setup_watch_system(10); + let key = Bytes::from("/k"); + let mut watcher = registry.register(key.clone(), true).unwrap(); + + tokio::time::sleep(Duration::from_millis(50)).await; + + broadcast_tx + .send(WatchResponse { + key: key.clone(), + value: Bytes::from("new_val"), + prev_value: Bytes::from("old_val"), + event_type: d_engine_proto::client::WatchEventType::Put as i32, + error: 0, + revision: 1, + }) + .unwrap(); + + let event = timeout(Duration::from_millis(100), watcher.receiver_mut().recv()) + .await + .expect("timeout") + .expect("closed"); + + assert_eq!(event.event_type, WatchEventType::Put); + assert_eq!( + event.prev_value, + Bytes::from("old_val"), + "prev_kv=true watcher must receive prev_value" + ); +} + +/// #379: Two watchers on the same key with different prev_kv flags each get +/// their respective behavior — per-watcher isolation, not global toggle. +#[tokio::test] +async fn test_prev_kv_per_watcher_isolation() { + let (broadcast_tx, registry, _handle) = setup_watch_system(10); + let key = Bytes::from("/shared"); + + let mut watcher_no_prev = registry.register(key.clone(), false).unwrap(); + let mut watcher_with_prev = registry.register(key.clone(), true).unwrap(); + + tokio::time::sleep(Duration::from_millis(50)).await; + + broadcast_tx + .send(WatchResponse { + key: key.clone(), + value: Bytes::from("new"), + prev_value: Bytes::from("old"), + event_type: d_engine_proto::client::WatchEventType::Put as i32, + error: 0, + revision: 1, + }) + .unwrap(); + + let ev_no = timeout( + Duration::from_millis(100), + watcher_no_prev.receiver_mut().recv(), + ) + .await + .expect("timeout no_prev") + .expect("closed"); + let ev_yes = timeout( + Duration::from_millis(100), + watcher_with_prev.receiver_mut().recv(), + ) + .await + .expect("timeout with_prev") + .expect("closed"); + + assert_eq!( + ev_no.prev_value, + Bytes::new(), + "prev_kv=false must receive empty prev_value" + ); + assert_eq!( + ev_yes.prev_value, + Bytes::from("old"), + "prev_kv=true must receive prev_value" + ); +} + +/// #379: prev_kv_watcher_count tracks registrations and unregistrations correctly. +/// This counter is shared with the state machine handler to gate RocksDB reads. +#[tokio::test] +async fn test_prev_kv_watcher_count_tracks_registration() { + // Keep _broadcast_tx alive: dropping it closes the channel, dispatcher exits, + // and subsequent unregister messages would never be processed. + let (_broadcast_tx, registry, _handle) = setup_watch_system(10); + + assert_eq!(registry.prev_kv_watcher_count(), 0, "starts at 0"); + + let h1 = registry.register(Bytes::from("/k1"), true).unwrap(); + assert_eq!( + registry.prev_kv_watcher_count(), + 1, + "after first prev_kv=true" + ); + + let _h2 = registry.register(Bytes::from("/k2"), false).unwrap(); + assert_eq!( + registry.prev_kv_watcher_count(), + 1, + "prev_kv=false does not increment" + ); + + let h3 = registry.register(Bytes::from("/k3"), true).unwrap(); + assert_eq!( + registry.prev_kv_watcher_count(), + 2, + "after second prev_kv=true" + ); + + drop(h1); + tokio::time::sleep(Duration::from_millis(100)).await; + assert_eq!(registry.prev_kv_watcher_count(), 1, "after dropping h1"); + + drop(h3); + tokio::time::sleep(Duration::from_millis(100)).await; + assert_eq!( + registry.prev_kv_watcher_count(), + 0, + "after all prev_kv watchers unregistered" + ); +} + +/// #379: Progress heartbeat — dispatcher sends periodic Progress events to all watchers. +/// Watchers can use these to detect stream liveness without waiting for key mutations. +#[tokio::test] +async fn test_progress_heartbeat_delivered_to_watcher() { + // 50ms interval so the test completes quickly. + // Keep _broadcast_tx alive so the dispatcher does not exit early. + let (_broadcast_tx, registry, _handle) = setup_watch_system_with_heartbeat(10, 50); + let key = Bytes::from("/watch/key"); + let mut watcher = registry.register(key.clone(), false).unwrap(); + + // Wait long enough for at least one heartbeat tick + tokio::time::sleep(Duration::from_millis(200)).await; + + // Drain events: at least one must be Progress + let mut got_progress = false; + while let Ok(Some(ev)) = timeout(Duration::from_millis(50), watcher.receiver_mut().recv()).await + { + if ev.event_type == WatchEventType::Progress { + got_progress = true; + break; + } + } + + assert!( + got_progress, + "no Progress event received within 200ms with 50ms heartbeat interval" + ); +} + +/// #379: When heartbeat_interval_ms = 0, no Progress events are ever sent. +/// Clients that don't need liveness pings should not receive unsolicited events. +#[tokio::test] +async fn test_no_progress_events_when_heartbeat_disabled() { + // heartbeat disabled (interval = 0) + let (_, registry, _handle) = setup_watch_system_with_heartbeat(10, 0); + let key = Bytes::from("/watch/key"); + let mut watcher = registry.register(key.clone(), false).unwrap(); + + // Wait — if heartbeat were enabled with any short interval we'd see events + tokio::time::sleep(Duration::from_millis(200)).await; + + // No events at all should arrive (nothing was broadcast) + let result = timeout(Duration::from_millis(30), watcher.receiver_mut().recv()).await; + assert!( + result.is_err(), + "received unexpected event with heartbeat disabled: {:?}", + result + ); +} diff --git a/d-engine-core/src/watch/mod.rs b/d-engine-core/src/watch/mod.rs index d78a11fb..2d39ab6d 100644 --- a/d-engine-core/src/watch/mod.rs +++ b/d-engine-core/src/watch/mod.rs @@ -138,12 +138,11 @@ pub use crate::config::WatchConfig; /// buffer overflow. The client receives this as the last event on the stream /// and should treat it as a signal to re-sync via the Read API and re-register. pub(crate) fn make_cancel_event(key: bytes::Bytes) -> WatchEvent { - use d_engine_proto::error::ErrorCode; WatchEvent { key, value: bytes::Bytes::new(), - event_type: WatchEventType::Canceled as i32, - error: ErrorCode::WatchBufferOverflow as i32, + prev_value: bytes::Bytes::new(), + event_type: WatchEventType::Canceled, revision: 0, } } diff --git a/d-engine-proto/proto/client/client_api.proto b/d-engine-proto/proto/client/client_api.proto index 36038986..d80a2f41 100644 --- a/d-engine-proto/proto/client/client_api.proto +++ b/d-engine-proto/proto/client/client_api.proto @@ -158,6 +158,10 @@ enum WatchEventType { // Watcher forcibly canceled by the server (e.g. buffer overflow). // Client should re-sync via Read API and re-register the watch. WATCH_EVENT_TYPE_CANCELED = 2; + + // Periodic heartbeat with no data change. + // Carries current revision so clients can confirm the stream is alive. + WATCH_EVENT_TYPE_PROGRESS = 3; } // Request to watch for changes on a key or key prefix. @@ -173,6 +177,11 @@ message WatchRequest { // When true, key is treated as a path prefix (must end with "/"). // Default false preserves backwards-compatible exact-match behaviour. bool prefix = 3; + + // When true, each WatchResponse includes the value that existed before + // the mutation (prev_value). Default false; server skips the extra read + // when no watcher requests it. + bool prev_kv = 4; } // Response containing a watch event notification @@ -193,6 +202,12 @@ message WatchResponse { // Monotonically increasing. Clients use this as an anchor after receiving // CANCELED: re-read state, then re-register watching from revision+1. uint64 revision = 5; + + // Value before this mutation. Empty when: + // - key did not exist before the write + // - event_type is PROGRESS or CANCELED + // - watcher was registered with prev_kv = false + bytes prev_value = 6; } service RaftClientService { diff --git a/d-engine-proto/src/generated/d_engine.client.rs b/d-engine-proto/src/generated/d_engine.client.rs index 1fee1271..bbb6a5c5 100644 --- a/d-engine-proto/src/generated/d_engine.client.rs +++ b/d-engine-proto/src/generated/d_engine.client.rs @@ -192,6 +192,11 @@ pub struct WatchRequest { /// Default false preserves backwards-compatible exact-match behaviour. #[prost(bool, tag = "3")] pub prefix: bool, + /// When true, each WatchResponse includes the value that existed before + /// the mutation (prev_value). Default false; server skips the extra read + /// when no watcher requests it. + #[prost(bool, tag = "4")] + pub prev_kv: bool, } /// Response containing a watch event notification #[derive(serde::Serialize, serde::Deserialize)] @@ -214,6 +219,12 @@ pub struct WatchResponse { /// CANCELED: re-read state, then re-register watching from revision+1. #[prost(uint64, tag = "5")] pub revision: u64, + /// Value before this mutation. Empty when: + /// - key did not exist before the write + /// - event_type is PROGRESS or CANCELED + /// - watcher was registered with prev_kv = false + #[prost(bytes = "bytes", tag = "6")] + pub prev_value: ::prost::bytes::Bytes, } /// Read consistency policy for controlling read operation guarantees /// @@ -278,6 +289,9 @@ pub enum WatchEventType { /// Watcher forcibly canceled by the server (e.g. buffer overflow). /// Client should re-sync via Read API and re-register the watch. Canceled = 2, + /// Periodic heartbeat with no data change. + /// Carries current revision so clients can confirm the stream is alive. + Progress = 3, } impl WatchEventType { /// String value of the enum field names used in the ProtoBuf definition. @@ -289,6 +303,7 @@ impl WatchEventType { Self::Put => "WATCH_EVENT_TYPE_PUT", Self::Delete => "WATCH_EVENT_TYPE_DELETE", Self::Canceled => "WATCH_EVENT_TYPE_CANCELED", + Self::Progress => "WATCH_EVENT_TYPE_PROGRESS", } } /// Creates an enum from field names used in the ProtoBuf definition. @@ -297,6 +312,7 @@ impl WatchEventType { "WATCH_EVENT_TYPE_PUT" => Some(Self::Put), "WATCH_EVENT_TYPE_DELETE" => Some(Self::Delete), "WATCH_EVENT_TYPE_CANCELED" => Some(Self::Canceled), + "WATCH_EVENT_TYPE_PROGRESS" => Some(Self::Progress), _ => None, } } diff --git a/d-engine-server/benches/state_machine.rs b/d-engine-server/benches/state_machine.rs index da4ad775..725d6c7e 100644 --- a/d-engine-server/benches/state_machine.rs +++ b/d-engine-server/benches/state_machine.rs @@ -67,7 +67,14 @@ fn create_watch_system( let registry = Arc::new(WatchRegistry::new(watcher_buffer_size, unregister_tx)); // Spawn dispatcher - let dispatcher = WatchDispatcher::new(Arc::clone(®istry), broadcast_rx, unregister_rx); + let last_applied = Arc::new(std::sync::atomic::AtomicU64::new(0)); + let dispatcher = WatchDispatcher::new( + Arc::clone(®istry), + broadcast_rx, + unregister_rx, + last_applied, + 0, + ); tokio::spawn(async move { dispatcher.run().await; }); @@ -87,7 +94,7 @@ fn register_watchers( let mut handles = Vec::with_capacity(count); for i in 0..count { let key = format!("{key_prefix}{i}"); - handles.push(registry.register(key.into()).expect("register watcher")); + handles.push(registry.register(key.into(), false).expect("register watcher")); } handles } @@ -352,6 +359,7 @@ fn bench_apply_with_1_watcher(c: &mut Criterion) { key: insert.key.clone(), value: insert.value.clone(), event_type: d_engine_proto::client::WatchEventType::Put as i32, + prev_value: bytes::Bytes::new(), error: 0, revision: 0, }; @@ -362,6 +370,7 @@ fn bench_apply_with_1_watcher(c: &mut Criterion) { key: delete.key.clone(), value: bytes::Bytes::new(), event_type: d_engine_proto::client::WatchEventType::Delete as i32, + prev_value: bytes::Bytes::new(), error: 0, revision: 0, }; @@ -372,6 +381,7 @@ fn bench_apply_with_1_watcher(c: &mut Criterion) { key: cas.key.clone(), value: cas.new_value.clone(), event_type: d_engine_proto::client::WatchEventType::Put as i32, + prev_value: bytes::Bytes::new(), error: 0, revision: 0, }; @@ -416,6 +426,7 @@ fn bench_apply_with_10_watchers(c: &mut Criterion) { key: insert.key.clone(), value: insert.value.clone(), event_type: d_engine_proto::client::WatchEventType::Put as i32, + prev_value: bytes::Bytes::new(), error: 0, revision: 0, }; @@ -426,6 +437,7 @@ fn bench_apply_with_10_watchers(c: &mut Criterion) { key: delete.key.clone(), value: bytes::Bytes::new(), event_type: d_engine_proto::client::WatchEventType::Delete as i32, + prev_value: bytes::Bytes::new(), error: 0, revision: 0, }; @@ -436,6 +448,7 @@ fn bench_apply_with_10_watchers(c: &mut Criterion) { key: cas.key.clone(), value: cas.new_value.clone(), event_type: d_engine_proto::client::WatchEventType::Put as i32, + prev_value: bytes::Bytes::new(), error: 0, revision: 0, }; @@ -480,6 +493,7 @@ fn bench_apply_with_100_watchers(c: &mut Criterion) { key: insert.key.clone(), value: insert.value.clone(), event_type: d_engine_proto::client::WatchEventType::Put as i32, + prev_value: bytes::Bytes::new(), error: 0, revision: 0, }; @@ -490,6 +504,7 @@ fn bench_apply_with_100_watchers(c: &mut Criterion) { key: delete.key.clone(), value: bytes::Bytes::new(), event_type: d_engine_proto::client::WatchEventType::Delete as i32, + prev_value: bytes::Bytes::new(), error: 0, revision: 0, }; @@ -500,6 +515,7 @@ fn bench_apply_with_100_watchers(c: &mut Criterion) { key: cas.key.clone(), value: cas.new_value.clone(), event_type: d_engine_proto::client::WatchEventType::Put as i32, + prev_value: bytes::Bytes::new(), error: 0, revision: 0, }; @@ -528,7 +544,7 @@ fn bench_watch_e2e_latency(c: &mut Criterion) { let value = Bytes::from("test_value"); // Register a watcher - let mut watcher = registry.register(key.clone()).expect("register watcher"); + let mut watcher = registry.register(key.clone(), false).expect("register watcher"); // Measure time from broadcast to receive let start = tokio::time::Instant::now(); @@ -538,6 +554,7 @@ fn bench_watch_e2e_latency(c: &mut Criterion) { key: key.clone(), value: value.clone(), event_type: d_engine_proto::client::WatchEventType::Put as i32, + prev_value: bytes::Bytes::new(), error: 0, revision: 0, }; diff --git a/d-engine-server/src/api/embedded_client.rs b/d-engine-server/src/api/embedded_client.rs index 7a2f84a2..b4453b3f 100644 --- a/d-engine-server/src/api/embedded_client.rs +++ b/d-engine-server/src/api/embedded_client.rs @@ -459,10 +459,71 @@ impl EmbeddedClient { /// println!("Value changed: {:?}", event); /// } /// ``` + /// `prev_kv`: when true each WatchEvent carries the value before the mutation. #[cfg(feature = "watch")] pub fn watch( &self, key: impl AsRef<[u8]>, + ) -> ClientApiResult { + self.watch_with_options(key, false) + } + + /// Watch a key and opt in to receiving the previous value on each mutation. + /// + /// This is the lower-level form of [`watch`](Self::watch). Use it when you need + /// `event.prev_value` — for example to detect what a key held before a write, or + /// to implement an audit log. + /// + /// # Arguments + /// + /// * `key` - The exact key to watch. + /// * `prev_kv` - When `true`, every `Put` and `Delete` event carries the value the + /// key held **before** the mutation in `event.prev_value`. When + /// `false` (the default via [`watch`](Self::watch)), `prev_value` is + /// always empty. + /// + /// # Performance note + /// + /// When at least one watcher has `prev_kv = true`, the state machine reads the old + /// value from storage before each `apply_chunk`. The read is amortised across the + /// whole write batch — cost scales with **write rate**, not watcher count. Disable + /// when you don't need the previous value. + /// + /// # `prev_value` is empty in these cases + /// + /// - The watcher was registered with `prev_kv = false`. + /// - The event type is `Progress` or `Canceled` (not a data mutation). + /// - The key did not exist before a `Put` (i.e. it was a fresh insert). + /// + /// # Example + /// + /// ```ignore + /// // Distributed-lock audit: know who held the lock before it changed hands. + /// let watcher = client.watch_with_options(b"lock/resource_a", true)?; + /// let (_, _, mut rx) = watcher.into_receiver(); + /// + /// tokio::spawn(async move { + /// while let Some(event) = rx.recv().await { + /// match event.event_type { + /// WatchEventType::Put => println!( + /// "lock acquired by {:?}, was held by {:?}", + /// event.value, event.prev_value + /// ), + /// WatchEventType::Delete => println!( + /// "lock released, was held by {:?}", + /// event.prev_value + /// ), + /// WatchEventType::Canceled => { /* re-register */ break; } + /// WatchEventType::Progress => {} + /// } + /// } + /// }); + /// ``` + #[cfg(feature = "watch")] + pub fn watch_with_options( + &self, + key: impl AsRef<[u8]>, + prev_kv: bool, ) -> ClientApiResult { let registry = self.watch_registry.as_ref().ok_or_else(|| ClientApiError::Business { code: ErrorCode::Uncategorized, @@ -471,7 +532,7 @@ impl EmbeddedClient { })?; let key_bytes = Bytes::copy_from_slice(key.as_ref()); - registry.register(key_bytes).map_err(|e| ClientApiError::Business { + registry.register(key_bytes, prev_kv).map_err(|e| ClientApiError::Business { code: ErrorCode::Uncategorized, message: e.to_string(), required_action: None, @@ -486,6 +547,47 @@ impl EmbeddedClient { pub fn watch_prefix( &self, prefix: impl AsRef<[u8]>, + ) -> ClientApiResult { + self.watch_prefix_with_options(prefix, false) + } + + /// Register a prefix watcher and opt in to receiving the previous value on each mutation. + /// + /// Prefix form of [`watch_with_options`](Self::watch_with_options). Every key whose + /// path starts with `prefix` triggers an event; `event.key` is the full child key. + /// + /// See [`watch_with_options`](Self::watch_with_options) for the `prev_kv` semantics, + /// performance note, and the cases where `prev_value` is empty. + /// + /// # Example + /// + /// ```ignore + /// // Track every endpoint change under /services/ together with the old address. + /// let watcher = client.watch_prefix_with_options(b"/services/", true)?; + /// let (_, _, mut rx) = watcher.into_receiver(); + /// + /// tokio::spawn(async move { + /// while let Some(event) = rx.recv().await { + /// match event.event_type { + /// WatchEventType::Put => println!( + /// "{:?} moved from {:?} → {:?}", + /// event.key, event.prev_value, event.value + /// ), + /// WatchEventType::Delete => println!( + /// "{:?} removed (was {:?})", + /// event.key, event.prev_value + /// ), + /// WatchEventType::Canceled => { /* buffer overflow — re-register */ break; } + /// WatchEventType::Progress => {} + /// } + /// } + /// }); + /// ``` + #[cfg(feature = "watch")] + pub fn watch_prefix_with_options( + &self, + prefix: impl AsRef<[u8]>, + prev_kv: bool, ) -> ClientApiResult { let registry = self.watch_registry.as_ref().ok_or_else(|| ClientApiError::Business { code: ErrorCode::Uncategorized, @@ -494,11 +596,13 @@ impl EmbeddedClient { })?; let prefix_bytes = Bytes::copy_from_slice(prefix.as_ref()); - registry.register_prefix(prefix_bytes).map_err(|e| ClientApiError::Business { - code: ErrorCode::Uncategorized, - message: e.to_string(), - required_action: None, - }) + registry + .register_prefix(prefix_bytes, prev_kv) + .map_err(|e| ClientApiError::Business { + code: ErrorCode::Uncategorized, + message: e.to_string(), + required_action: None, + }) } /// Scan all keys under `prefix` with linearizable consistency. diff --git a/d-engine-server/src/api/embedded_client_test.rs b/d-engine-server/src/api/embedded_client_test.rs index 3a603bff..f82d9e6d 100644 --- a/d-engine-server/src/api/embedded_client_test.rs +++ b/d-engine-server/src/api/embedded_client_test.rs @@ -1,4 +1,4 @@ -#[cfg(test)] +#[cfg(all(test, feature = "rocksdb"))] mod tests { use std::collections::HashMap; @@ -106,7 +106,7 @@ mod tests { // Integration Tests (require EmbeddedEngine) // ============================================================================= -#[cfg(test)] +#[cfg(all(test, feature = "rocksdb"))] mod integration_tests { use std::time::Duration; 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 e1e51830..b288a7e2 100644 --- a/d-engine-server/src/network/grpc/grpc_raft_service.rs +++ b/d-engine-server/src/network/grpc/grpc_raft_service.rs @@ -532,29 +532,32 @@ where })?; let is_prefix = watch_request.prefix; + let prev_kv = watch_request.prev_kv; info!( node_id = self.node_id, key = ?key, is_prefix, + prev_kv, "Registering watch for key" ); // Register watcher (exact or prefix) and get receiver let handle = if is_prefix { - registry.register_prefix(key).map_err(|e| match e { + registry.register_prefix(key, prev_kv).map_err(|e| match e { WatchError::LimitExceeded(_) => Status::resource_exhausted(e.to_string()), WatchError::InvalidPrefix => Status::invalid_argument(e.to_string()), })? } else { - registry.register(key).map_err(|e| Status::resource_exhausted(e.to_string()))? + registry + .register(key, prev_kv) + .map_err(|e| Status::resource_exhausted(e.to_string()))? }; let (_watcher_id, _key, receiver) = handle.into_receiver(); - // Convert mpsc::Receiver -> Boxed Stream with sentinel error on close - // When the stream ends (receiver closed), send an UNAVAILABLE error to help clients - // detect server shutdown/restart and implement reconnection logic + // Convert mpsc::Receiver -> Boxed Stream for gRPC. + // WatchEvent (opaque) is converted back to proto WatchResponse at this boundary. let stream = tokio_stream::wrappers::ReceiverStream::new(receiver) - .map(Ok) + .map(|e| Ok(d_engine_proto::client::WatchResponse::from(&e))) .chain(futures::stream::once(async { Err(Status::unavailable( "Watch stream closed: server may have shut down or restarted. Please reconnect and re-register the watcher." diff --git a/d-engine-server/src/node/builder.rs b/d-engine-server/src/node/builder.rs index 32025ddb..5097bf5b 100644 --- a/d-engine-server/src/node/builder.rs +++ b/d-engine-server/src/node/builder.rs @@ -357,23 +357,41 @@ where unregister_tx, )); + // Shared last_applied for Progress event revision; written by handler, read by dispatcher. + let last_applied_ref = Arc::new(std::sync::atomic::AtomicU64::new(last_applied_index)); + // Create dispatcher - let dispatcher = - WatchDispatcher::new(Arc::clone(®istry), broadcast_rx, unregister_rx); + let dispatcher = WatchDispatcher::new( + Arc::clone(®istry), + broadcast_rx, + unregister_rx, + Arc::clone(&last_applied_ref), + node_config.raft.watch.heartbeat_interval_ms, + ); // Explicitly spawn dispatcher task (resource allocation visible) let dispatcher_handle = tokio::spawn(async move { dispatcher.run().await; }); - Some((broadcast_tx, registry, dispatcher_handle)) + Some((broadcast_tx, registry, dispatcher_handle, last_applied_ref)) }; let state_machine_handler = self.state_machine_handler.take().unwrap_or_else(|| { #[cfg(feature = "watch")] - let watch_event_tx = watch_system.as_ref().map(|(tx, _, _)| tx.clone()); + let watch_event_tx = watch_system.as_ref().map(|(tx, _, _, _)| tx.clone()); + + // Share the registry's live prev_kv counter Arc so the handler sees real-time updates. + #[cfg(feature = "watch")] + let prev_kv_count = watch_system + .as_ref() + .map(|(_, registry, _, _)| registry.prev_kv_watcher_count_arc()) + .unwrap_or_else(|| Arc::new(std::sync::atomic::AtomicUsize::new(0))); + #[cfg(not(feature = "watch"))] - let watch_event_tx = None; + let watch_event_tx: Option< + tokio::sync::broadcast::Sender, + > = None; Arc::new(DefaultStateMachineHandler::new( node_id, @@ -381,7 +399,14 @@ where state_machine.clone(), node_config.raft.snapshot.clone(), snapshot_policy, + #[cfg(feature = "watch")] + watch_event_tx, + #[cfg(not(feature = "watch"))] watch_event_tx, + #[cfg(feature = "watch")] + prev_kv_count, + #[cfg(not(feature = "watch"))] + Arc::new(std::sync::atomic::AtomicUsize::new(0)), )) }); let (membership_inner, zombie_rx) = self.membership.take().map_or_else( @@ -561,9 +586,9 @@ where membership_rx, node_config: node_config_arc, #[cfg(feature = "watch")] - watch_registry: watch_system.as_ref().map(|(_, reg, _)| Arc::clone(reg)), + watch_registry: watch_system.as_ref().map(|(_, reg, _, _)| Arc::clone(reg)), #[cfg(feature = "watch")] - _watch_dispatcher_handle: watch_system.map(|(_, _, handle)| handle), + _watch_dispatcher_handle: watch_system.map(|(_, _, handle, _)| handle), sm_worker_handle: std::sync::Mutex::new(Some(sm_worker_handle)), _commit_handler_handle: Some(commit_handler_handle), _lease_cleanup_handle: lease_cleanup_handle, diff --git a/d-engine-server/src/test_utils/integration/mod.rs b/d-engine-server/src/test_utils/integration/mod.rs index 7a9ed132..068cdda5 100644 --- a/d-engine-server/src/test_utils/integration/mod.rs +++ b/d-engine-server/src/test_utils/integration/mod.rs @@ -239,6 +239,7 @@ pub fn setup_raft_components( node_config.raft.snapshot.clone(), snapshot_policy, None, // No watch for tests + Arc::new(std::sync::atomic::AtomicUsize::new(0)), ); let node_config_clone = node_config.clone(); 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 fe7901b8..76779fa3 100644 --- a/d-engine-server/tests/embedded_client/embedded_client_operations.rs +++ b/d-engine-server/tests/embedded_client/embedded_client_operations.rs @@ -724,11 +724,7 @@ async fn test_watch_receives_put_event() { .expect("timed out waiting for Put watch event") .expect("watch channel should not be closed"); - assert_eq!( - event.event_type, - WatchEventType::Put as i32, - "expected Put event" - ); + assert_eq!(event.event_type, WatchEventType::Put, "expected Put event"); assert_eq!(event.key, Bytes::from("w-key")); assert_eq!(event.value, Bytes::from("w-val")); } @@ -755,7 +751,7 @@ async fn test_watch_receives_delete_event() { assert_eq!( event.event_type, - WatchEventType::Delete as i32, + WatchEventType::Delete, "expected Delete event" ); assert_eq!(event.key, Bytes::from("d-key")); diff --git a/d-engine-server/tests/watch_and_subscriptions/watch_events_embedded.rs b/d-engine-server/tests/watch_and_subscriptions/watch_events_embedded.rs index 360cc7da..2d0e7926 100644 --- a/d-engine-server/tests/watch_and_subscriptions/watch_events_embedded.rs +++ b/d-engine-server/tests/watch_and_subscriptions/watch_events_embedded.rs @@ -86,7 +86,7 @@ async fn test_embedded_watch_receives_put_events() -> Result<(), Box Result<(), Box Result<(), Box { - if event.event_type == WatchEventType::Put as i32 && *event.value == *b"value2" { + if event.event_type == WatchEventType::Put && *event.value == *b"value2" { found = true; break; } @@ -327,7 +327,7 @@ async fn test_embedded_watch_multiple_watchers_same_key() -> Result<(), Box continue, // already in snapshot - Ok(e) if e.event_type == WatchEventType::Canceled as i32 => break, + Ok(e) if e.event_type == WatchEventType::Canceled => break, Ok(e) => apply_to_registry(e, &mut registry), Err(_) => break, } @@ -184,6 +184,13 @@ let last_revision = event.revision; - **Default**: `i64::MAX` (effectively unlimited) - **Description**: Hard cap on total active watchers (exact + prefix combined). `register()` and `watch_prefix()` return an error once this limit is reached. Set to a finite value to prevent runaway watcher growth in multi-tenant deployments. +### `heartbeat_interval_ms` + +- **Type**: u64 +- **Default**: `0` (disabled) +- **Description**: If non-zero, the `WatchDispatcher` sends a `Progress` heartbeat event to all active watchers at this interval (±10% jitter). Useful for detecting a silent stream that has stalled due to no writes. Set to `0` to rely on application-level keepalives instead. +- **Recommended**: `5000` (5 seconds) for long-lived connections; `0` for write-heavy workloads where events are frequent. + ### `enable_metrics` - **Type**: boolean @@ -198,7 +205,7 @@ let last_revision = event.revision; - **Order**: FIFO per key (events from StateMachine apply order) - **Dropped Events**: When a per-watcher buffer overflows, the watcher receives a `CANCELED` sentinel event and is forcibly unregistered - **Lagging**: Broadcast channel drops oldest events when receivers are slow (global buffer) -- **CANCELED Event**: `WatchEventType::CANCELED` with `ErrorCode::WATCH_BUFFER_OVERFLOW` signals that the watcher was forcibly terminated. Client must re-sync via Read API and re-register. +- **CANCELED Event**: `WatchEventType::Canceled` with `ErrorCode::WATCH_BUFFER_OVERFLOW` signals that the watcher was forcibly terminated. Client must re-sync via Read API and re-register. ### Lifecycle @@ -206,13 +213,49 @@ let last_revision = event.revision; - **Cleanup**: Watcher is automatically unregistered when handle is dropped - **No History**: Watch only receives future events (not historical changes) +### prev_kv (Previous Value) + +When registering a watcher with `prev_kv: true` (embedded API) or `prev_kv: true` in the gRPC `WatchRequest`, each `Put` and `Delete` event carries the **value the key held before the write** in `event.prev_value`. Watchers registered with `prev_kv: false` (the default) always receive an empty `prev_value`. + +```rust,ignore +// Embedded: request previous value on every event +let watcher = engine.client().watch_with_prev_kv(b"my_key")?; +while let Some(event) = watcher.receiver_mut().recv().await { + if event.event_type == WatchEventType::Put { + println!("was: {:?}, now: {:?}", event.prev_value, event.value); + } +} +``` + +**Cost**: When at least one `prev_kv` watcher is active, the state machine reads the old value from storage before each `apply_chunk`. This is a single point-in-time read per write batch — not per watcher — so the overhead scales with write rate, not watcher count. + +### Progress Heartbeat + +When `heartbeat_interval_ms > 0` in config, the dispatcher periodically sends a `WatchEventType::Progress` event to every active watcher. The event carries the current `revision` (Raft applied index) but empty `key` and `value`. Use it to: + +- Detect a silent stream that is alive but has no writes +- Confirm your client's channel is still open before a long idle period +- Drive time-based state machine logic without polling + +```rust,ignore +match event.event_type { + WatchEventType::Put => handle_put(&event), + WatchEventType::Delete => handle_delete(&event), + WatchEventType::Canceled => { re_register(); break; } + WatchEventType::Progress => { + // Stream is alive; revision is current applied index + println!("heartbeat revision={}", event.revision); + } +} +``` + +**Jitter**: The interval has ±10% random jitter to avoid thundering-herd when many watchers share the same interval. + ### Limitations - **No Persistence**: Watchers lost on server restart — re-register on reconnect - **No Replay**: Watch delivers future events only; use the Read API to re-sync current state after reconnect - **Prefix format**: Prefix must start and end with `/` — arbitrary byte-range watch is not supported -- **No prev_kv**: Events do not include the previous value (planned for a future release) -- **No heartbeat**: The stream is silent when no events occur — implement an application-level keepalive if needed (planned for a future release) ## Performance @@ -296,7 +339,7 @@ watch(b"config:feature_flags") ### `WATCH_BUFFER_OVERFLOW` (ErrorCode 5001) - **Cause**: Per-watcher channel full — the client is consuming events too slowly -- **Detection**: Receive an event where `event_type == WatchEventType::CANCELED` and `error == ErrorCode::WATCH_BUFFER_OVERFLOW` +- **Detection**: Receive an event where `event_type == WatchEventType::Canceled` and `error == ErrorCode::WATCH_BUFFER_OVERFLOW` - **Effect**: The watcher is forcibly unregistered server-side; no further events will be delivered on this stream - **Solution**: 1. Re-sync state via the Read API to get the current value @@ -310,7 +353,8 @@ watch(b"config:feature_flags") 3. **Idempotent Handlers**: Handle duplicate events gracefully (at-most-once delivery) 4. **Buffer Tuning**: Monitor lagged events and adjust buffer sizes accordingly 5. **Reconnect Logic**: Implement automatic reconnection on stream errors -6. **Handle CANCELED**: Always check for `WatchEventType::CANCELED` — treat it as a forced disconnect. Re-sync via Read API then re-register the watch before processing further events. +6. **Handle CANCELED**: Always check for `WatchEventType::Canceled` — treat it as a forced disconnect. Re-sync via Read API then re-register the watch before processing further events. +7. **Handle Progress**: If `heartbeat_interval_ms > 0`, your event handler must have an arm for `WatchEventType::Progress` (non-exhaustive match will fail to compile otherwise). Silently ignore it or use the `revision` field for liveness tracking. ## Watch Reliability and Reconnection diff --git a/examples/service-discovery-embedded/server.rs b/examples/service-discovery-embedded/server.rs index aa7cad69..068cfae4 100644 --- a/examples/service-discovery-embedded/server.rs +++ b/examples/service-discovery-embedded/server.rs @@ -4,6 +4,7 @@ //! Features: //! - Exact-key Watch: react to changes on a single specific key //! - Prefix Watch: maintain a live registry of all nodes in a service namespace +//! - prev_kv Watch: receive the previous value on each mutation (audit log pattern) //! - In-process Watch API (zero network overhead, zero serialization overhead) //! - Automatic lifecycle management @@ -54,20 +55,24 @@ async fn main() -> Result<(), Box> { while let Some(event) = exact_rx.recv().await { let key = String::from_utf8_lossy(&event.key); match event.event_type { - e if e == WatchEventType::Put as i32 => { + WatchEventType::Put => { let value = String::from_utf8_lossy(&event.value); println!( "[exact-watch] config changed: {key} = {value} (revision={})", event.revision ); } - e if e == WatchEventType::Delete as i32 => { + WatchEventType::Delete => { println!( "[exact-watch] config deleted: {key} (revision={})", event.revision ); } - _ => {} + WatchEventType::Canceled => { + println!("[exact-watch] CANCELED — buffer overflow on {key}; re-register watch"); + break; + } + WatchEventType::Progress => {} } } }); @@ -109,7 +114,7 @@ async fn main() -> Result<(), Box> { let mut reg = registry_bg.lock().unwrap(); match event.event_type { - e if e == WatchEventType::Put as i32 => { + WatchEventType::Put => { let value = String::from_utf8_lossy(&event.value).to_string(); println!( "[prefix-watch] node up: {key_str} → {value} (revision={})", @@ -117,14 +122,14 @@ async fn main() -> Result<(), Box> { ); reg.insert(event.key.to_vec(), value); } - e if e == WatchEventType::Delete as i32 => { + WatchEventType::Delete => { println!( "[prefix-watch] node down: {key_str} (revision={})", event.revision ); reg.remove(event.key.as_ref()); } - e if e == WatchEventType::Canceled as i32 => { + WatchEventType::Canceled => { // Buffer overflow: registry may be stale. // Production code: re-watch then scan_prefix (see ADR-003 reconnect pattern). println!( @@ -178,6 +183,65 @@ async fn main() -> Result<(), Box> { } } + // ----------------------------------------------------------------------- + // DEMO 3: prev_kv — audit log for endpoint changes + // + // Use case: compliance teams need to know the previous endpoint value + // every time a service node updates its address (e.g. rolling restart). + // watch_with_options(..., prev_kv: true) populates event.prev_value + // so you can log "was X, now Y" without an extra Read round-trip. + // + // Cost: when at least one prev_kv watcher is active, the state machine + // reads the old value from storage before each write batch. Cost scales + // with write rate, not watcher count — acceptable for audit workloads. + // ----------------------------------------------------------------------- + println!("\n=== Demo 3: prev_kv — Endpoint Change Audit Log ==="); + + let audit_watcher = engine.client().watch_prefix_with_options(b"/audit/", true)?; + let (_, _, mut audit_rx) = audit_watcher.into_receiver(); + + tokio::spawn(async move { + while let Some(event) = audit_rx.recv().await { + let key = String::from_utf8_lossy(&event.key); + match event.event_type { + WatchEventType::Put => { + let prev = if event.prev_value.is_empty() { + "(new key)".to_string() + } else { + String::from_utf8_lossy(&event.prev_value).to_string() + }; + println!( + "[audit] {} : {} → {} (revision={})", + key, + prev, + String::from_utf8_lossy(&event.value), + event.revision + ); + } + WatchEventType::Delete => println!( + "[audit] {} deleted (was: {}) (revision={})", + key, + String::from_utf8_lossy(&event.prev_value), + event.revision + ), + WatchEventType::Canceled => { + println!("[audit] CANCELED — re-register watch"); + break; + } + WatchEventType::Progress => {} + } + } + }); + + tokio::time::sleep(Duration::from_millis(50)).await; + + client.put(b"/audit/svc-a", b"10.0.1.1:8080").await?; + tokio::time::sleep(Duration::from_millis(50)).await; + client.put(b"/audit/svc-a", b"10.0.1.2:8080").await?; // rolling restart + tokio::time::sleep(Duration::from_millis(50)).await; + client.delete(b"/audit/svc-a").await?; + tokio::time::sleep(Duration::from_millis(100)).await; + println!("\n=== Demo complete. Press Ctrl+C to exit. ==="); signal::ctrl_c().await?; diff --git a/examples/service-discovery-standalone/watcher.rs b/examples/service-discovery-standalone/watcher.rs index 78cc26c1..a01f27d9 100644 --- a/examples/service-discovery-standalone/watcher.rs +++ b/examples/service-discovery-standalone/watcher.rs @@ -251,6 +251,7 @@ fn apply_event_to_registry( registry.clear(); true } + Some(WatchEventType::Progress) => false, None => { println!("[UNKNOWN] {key_str} (event_type={})", response.event_type); false @@ -284,10 +285,9 @@ fn print_event(response: &WatchResponse) { println!("[DELETE] {key} (revision={})", response.revision); } Some(WatchEventType::Canceled) => { - println!( - "[CANCELED] {key} — buffer overflow; re-sync and re-register" - ); + println!("[CANCELED] {key} — buffer overflow; re-sync and re-register"); } + Some(WatchEventType::Progress) => {} None => { println!("[UNKNOWN] {key} (event_type={})", response.event_type); } From 03879f3b30026a1033bfcac002c60afb960b519e Mon Sep 17 00:00:00 2001 From: Joshua Chi <112539+JoshuaChi@users.noreply.github.com> Date: Tue, 19 May 2026 16:50:29 +0800 Subject: [PATCH 2/2] =?UTF-8?q?fix(watch)=20#379:=20address=20PR=20review?= =?UTF-8?q?=20=E2=80=94=20correctness,=20semantics,=20security?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - MissedTickBehavior::Skip on heartbeat interval: Burst mode would re-deliver N identical-revision Progress events after a stall, misleading clients into thinking the stream was alive during the gap - WatchEvent.prev_value: Option — None means "not requested", Some(empty) means "key didn't exist"; eliminates the ambiguity of a bare Bytes::new() standing for two distinct states - read_prev_values overlay map: two writes to the same key in one apply_chunk batch now produce correct sequential prev_values; previously both would read the pre-batch SM state - Jitter seed: replace ThreadId byte-sum with DefaultHasher mixing ThreadId + SystemTime; covers k8s rolling-restart where subsec_nanos alone is near-identical across pods started in the same millisecond - Upgrade astral-tokio-tar 0.6.1 -> 0.6.2 (RUSTSEC-2026-0145: PAX header desynchronization allows file smuggling during extraction) - Test: test_apply_chunk_stops_reading_prev_value_after_count_drops_to_zero now exercises the actual 1->0 transition instead of starting at 0 - Docs/examples: watch-feature.md and service-discovery example updated for Option semantics, correct API names, and gRPC-scoped WATCH_BUFFER_OVERFLOW detection note --- Cargo.lock | 14 ++-- .../default_state_machine_handler.rs | 58 +++++++++++++---- .../default_state_machine_handler_test.rs | 33 ++++++---- d-engine-core/src/watch/manager.rs | 64 ++++++++++++------- d-engine-core/src/watch/manager_test.rs | 20 +++--- d-engine-core/src/watch/mod.rs | 2 +- d-engine-server/src/api/embedded_client.rs | 1 - .../src/docs/server_guide/watch-feature.md | 6 +- examples/service-discovery-embedded/server.rs | 24 ++++--- 9 files changed, 141 insertions(+), 81 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 12e05d5a..6104e15c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -64,9 +64,9 @@ dependencies = [ [[package]] name = "astral-tokio-tar" -version = "0.6.1" +version = "0.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ce73b17c62717c4b6a9af10b43e87c578b0cac27e00666d48304d3b7d2c0693" +checksum = "cb50a7aae84a03bf55b067832bc376f4961b790c97e64d3eacee97d389b90277" dependencies = [ "filetime", "futures-core", @@ -704,7 +704,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -1107,7 +1107,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -1853,7 +1853,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2127,7 +2127,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2690,7 +2690,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] diff --git a/d-engine-core/src/state_machine_handler/default_state_machine_handler.rs b/d-engine-core/src/state_machine_handler/default_state_machine_handler.rs index fa57b2be..dd40f92e 100644 --- a/d-engine-core/src/state_machine_handler/default_state_machine_handler.rs +++ b/d-engine-core/src/state_machine_handler/default_state_machine_handler.rs @@ -650,32 +650,64 @@ where sm: &dyn crate::StateMachine, chunk: &[Entry], ) -> Vec> { + use std::collections::HashMap; + use d_engine_proto::client::WriteCommand; use d_engine_proto::client::write_command::Operation; use d_engine_proto::common::entry_payload::Payload; use prost::Message; - chunk - .iter() - .map(|entry| { + // Tracks the value each key would hold after each processed entry. + // None = key was deleted by an earlier entry in this batch. + // Ensures that two writes to the same key in one chunk produce correct + // prev_values: the second write sees the first write's value, not the + // pre-batch SM state. + let mut overlay: HashMap> = HashMap::new(); + let mut result = Vec::with_capacity(chunk.len()); + + for entry in chunk { + let prev = (|| -> Option { let payload = entry.payload.as_ref()?; - let Some(Payload::Command(bytes)) = &payload.payload else { + let Some(Payload::Command(cmd_bytes)) = &payload.payload else { return None; }; - let Ok(write_cmd) = WriteCommand::decode(bytes.as_ref()) else { + let Ok(write_cmd) = WriteCommand::decode(cmd_bytes.as_ref()) else { return None; }; + let op = write_cmd.operation?; - let key: Option<&[u8]> = match &write_cmd.operation { - Some(Operation::Insert(ins)) => Some(&ins.key), - Some(Operation::Delete(del)) => Some(&del.key), - Some(Operation::CompareAndSwap(cas)) => Some(&cas.key), - None => None, + let key: bytes::Bytes = match &op { + Operation::Insert(ins) => ins.key.clone(), + Operation::Delete(del) => del.key.clone(), + Operation::CompareAndSwap(cas) => cas.key.clone(), + }; + + // Check overlay first (reflects mutations earlier in this batch), + // then fall back to the SM's pre-batch state. + let prev = if let Some(v) = overlay.get(&key) { + v.clone() + } else { + sm.get(&key).ok().flatten() }; - key.map(|k| sm.get(k).ok().flatten().unwrap_or_default()) - }) - .collect() + // Update overlay so later entries in this batch see the updated state. + // CAS outcome is unknown before apply, so we conservatively skip it; + // same-key CAS-pairs in one batch are rare and semantically ambiguous. + match op { + Operation::Insert(ins) => { + overlay.insert(key, Some(ins.value)); + } + Operation::Delete(_) => { + overlay.insert(key, None); + } + Operation::CompareAndSwap(_) => {} + } + + Some(prev.unwrap_or_default()) + })(); + result.push(prev); + } + result } /// Broadcast watch events for applied chunk entries (fire-and-forget) diff --git a/d-engine-core/src/state_machine_handler/default_state_machine_handler_test.rs b/d-engine-core/src/state_machine_handler/default_state_machine_handler_test.rs index 2097ab82..6fe86edf 100644 --- a/d-engine-core/src/state_machine_handler/default_state_machine_handler_test.rs +++ b/d-engine-core/src/state_machine_handler/default_state_machine_handler_test.rs @@ -1992,7 +1992,7 @@ mod broadcast_watch_events_tests { #[cfg(feature = "watch")] mod prev_kv_apply_tests { use std::sync::Arc; - use std::sync::atomic::AtomicUsize; + use std::sync::atomic::{AtomicUsize, Ordering}; use bytes::Bytes; use d_engine_proto::client::{ @@ -2111,23 +2111,24 @@ mod prev_kv_apply_tests { /// #379 P6b: After the last prev_kv watcher unregisters (count drops to 0), /// apply_chunk must NOT call state_machine.get() on the next apply. /// - /// Tests the dynamic transition: count was 1 → externally set to 0 → - /// next apply skips read. This matches what happens when the last watcher - /// using prev_kv=true unregisters in production. + /// Tests the dynamic 1→0 transition: first apply with count=1 calls get() + /// exactly once; then count is set to 0 and second apply must NOT call get(). #[tokio::test] async fn test_apply_chunk_stops_reading_prev_value_after_count_drops_to_zero() { let mut sm = MockStateMachine::new(); - sm.expect_apply_chunk().returning(|_| { + sm.expect_apply_chunk().times(2).returning(|_| { Ok(vec![ApplyResult { index: 1, succeeded: true, }]) }); - // After count → 0, get() must NOT be called - // (Not setting up expect_get → mockall panics if called) + // Expect exactly one get(): during the first apply (count=1). + // The second apply (count=0) must not call get(). + // mockall verifies `times(1)` at drop — fails if called 0 or 2 times. + sm.expect_get().times(1).returning(|_| Ok(Some(Bytes::from("old")))); let (tx, _rx) = tokio::sync::broadcast::channel(10); - let prev_kv_count = Arc::new(AtomicUsize::new(0)); // already 0 (last watcher gone) + let prev_kv_count = Arc::new(AtomicUsize::new(1)); // one prev_kv watcher active let handler = DefaultStateMachineHandler::::new( 1, @@ -2141,10 +2142,16 @@ mod prev_kv_apply_tests { prev_kv_count.clone(), ); - // Simulate: count was 1 (watcher registered), now it's 0 (watcher gone) - // The Arc is already at 0 — handler should skip read on next apply. - let entry = insert_entry(1, b"k", b"v"); - handler.apply_chunk(vec![entry]).await.unwrap(); - // No panic → get() was not called → optimization active. + // First apply: count=1 → get() is called + let entry1 = insert_entry(1, b"k", b"v1"); + handler.apply_chunk(vec![entry1]).await.unwrap(); + + // Simulate last prev_kv watcher unregistering + prev_kv_count.store(0, Ordering::SeqCst); + + // Second apply: count=0 → get() must NOT be called + let entry2 = insert_entry(2, b"k", b"v2"); + handler.apply_chunk(vec![entry2]).await.unwrap(); + // mockall verifies on drop: exactly 1 get() call was made } } diff --git a/d-engine-core/src/watch/manager.rs b/d-engine-core/src/watch/manager.rs index b1f9c53d..d0b92643 100644 --- a/d-engine-core/src/watch/manager.rs +++ b/d-engine-core/src/watch/manager.rs @@ -26,6 +26,8 @@ //! - **Proto boundary**: WatchResponse (proto) lives only in the broadcast channel and handler; //! WatchEvent (opaque) is what callers see — no proto import required. +use std::collections::hash_map::DefaultHasher; +use std::hash::{Hash, Hasher}; use std::sync::Arc; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::time::Duration; @@ -53,15 +55,19 @@ pub enum WatchEventType { /// Opaque watch event delivered to callers. /// /// All fields use standard Rust types — callers never need to import proto types. -/// `prev_value` is populated only when the watcher was registered with `prev_kv = true` -/// and the event is a data mutation (Put / Delete); it is empty for Progress / Canceled. +/// +/// `prev_value` semantics: +/// - `None` — watcher was registered with `prev_kv = false`, or the event is +/// `Progress` / `Canceled` (not a data mutation). +/// - `Some(Bytes::new())` — watcher has `prev_kv = true` and the key did not exist +/// before this write (fresh insert). +/// - `Some(v)` — watcher has `prev_kv = true` and `v` is the previous value. #[derive(Debug, Clone)] pub struct WatchEvent { pub event_type: WatchEventType, pub key: Bytes, pub value: Bytes, - /// Value before this mutation. Empty when key did not exist, or prev_kv = false. - pub prev_value: Bytes, + pub prev_value: Option, pub revision: u64, } @@ -90,10 +96,12 @@ fn proto_to_event( event_type, key: proto.key.clone(), value: proto.value.clone(), + // None when prev_kv = false so callers can distinguish "not requested" from + // "key didn't exist" (Some(empty)). prev_value: if prev_kv { - proto.prev_value.clone() + Some(proto.prev_value.clone()) } else { - Bytes::new() + None }, revision: proto.revision, } @@ -106,7 +114,8 @@ impl From<&WatchEvent> for WatchResponse { WatchResponse { key: e.key.clone(), value: e.value.clone(), - prev_value: e.prev_value.clone(), + // Proto transport uses bytes (no optional); None collapses to empty. + prev_value: e.prev_value.clone().unwrap_or_default(), event_type: match e.event_type { WatchEventType::Put => ProtoType::Put as i32, WatchEventType::Delete => ProtoType::Delete as i32, @@ -492,26 +501,34 @@ impl WatchDispatcher { // Build optional heartbeat future. When interval_ms == 0 the future // is a pending sleep that never fires, adding zero overhead. - let heartbeat_enabled = self.heartbeat_interval_ms > 0; - let base_ms = self.heartbeat_interval_ms; - - // ±10% jitter on first tick to spread restarts across the interval window. - let first_tick_ms = if heartbeat_enabled { + // Build the heartbeat interval only when enabled. + // Constructing `interval_at(now + Duration::from_millis(u64::MAX), ...)` panics on + // overflow, so we must skip interval creation entirely when heartbeat is off. + let mut heartbeat: Option = if self.heartbeat_interval_ms > 0 { + let base_ms = self.heartbeat_interval_ms; let jitter = (base_ms / 10).max(1); - // Deterministic pseudo-random from current thread-id bits. - let tid_bits = format!("{:?}", std::thread::current().id()); - let seed: u64 = tid_bits.bytes().fold(0u64, |a, b| a.wrapping_add(b as u64)); + // Mix thread ID and wall-clock nanoseconds into a single hash so that nodes + // started simultaneously (e.g. k8s rolling restart within the same millisecond) + // still get different offsets. No external crate needed. + let mut h = DefaultHasher::new(); + std::thread::current().id().hash(&mut h); + std::time::SystemTime::now().hash(&mut h); + let seed = h.finish(); let offset = seed % (jitter * 2); - base_ms.saturating_sub(jitter) + offset + let first_tick_ms = base_ms.saturating_sub(jitter) + offset; + let mut interval = tokio::time::interval_at( + tokio::time::Instant::now() + Duration::from_millis(first_tick_ms), + Duration::from_millis(base_ms), + ); + // Skip missed ticks: heartbeat is a liveness signal, not a counter. + // Bursting N progress events after a slow-watcher stall would mislead clients + // into thinking the stream was alive during the stall period. + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + Some(interval) } else { - u64::MAX + None }; - let mut heartbeat = tokio::time::interval_at( - tokio::time::Instant::now() + Duration::from_millis(first_tick_ms), - Duration::from_millis(if heartbeat_enabled { base_ms } else { u64::MAX }), - ); - loop { tokio::select! { biased; @@ -531,7 +548,8 @@ impl WatchDispatcher { } } } - _ = heartbeat.tick(), if heartbeat_enabled => { + Some(t) = async { if let Some(ref mut hb) = heartbeat { Some(hb.tick().await) } else { std::future::pending().await } } => { + let _ = t; self.broadcast_progress().await; } } diff --git a/d-engine-core/src/watch/manager_test.rs b/d-engine-core/src/watch/manager_test.rs index be95b5b4..58ddd365 100644 --- a/d-engine-core/src/watch/manager_test.rs +++ b/d-engine-core/src/watch/manager_test.rs @@ -1163,11 +1163,10 @@ async fn test_prev_kv_false_watcher_receives_empty_prev_value() { assert_eq!(event.event_type, WatchEventType::Put); assert_eq!(event.value, Bytes::from("new_val")); - // prev_kv=false → dispatcher zeroes out prev_value before delivery + // prev_kv=false → dispatcher sets prev_value to None assert_eq!( - event.prev_value, - Bytes::new(), - "prev_kv=false watcher must receive empty prev_value" + event.prev_value, None, + "prev_kv=false watcher must receive None prev_value" ); } @@ -1200,7 +1199,7 @@ async fn test_prev_kv_true_watcher_receives_prev_value() { assert_eq!(event.event_type, WatchEventType::Put); assert_eq!( event.prev_value, - Bytes::from("old_val"), + Some(Bytes::from("old_val")), "prev_kv=true watcher must receive prev_value" ); } @@ -1244,13 +1243,12 @@ async fn test_prev_kv_per_watcher_isolation() { .expect("closed"); assert_eq!( - ev_no.prev_value, - Bytes::new(), - "prev_kv=false must receive empty prev_value" + ev_no.prev_value, None, + "prev_kv=false must receive None prev_value" ); assert_eq!( ev_yes.prev_value, - Bytes::from("old"), + Some(Bytes::from("old")), "prev_kv=true must receive prev_value" ); } @@ -1333,7 +1331,9 @@ async fn test_progress_heartbeat_delivered_to_watcher() { #[tokio::test] async fn test_no_progress_events_when_heartbeat_disabled() { // heartbeat disabled (interval = 0) - let (_, registry, _handle) = setup_watch_system_with_heartbeat(10, 0); + // Keep _broadcast_tx alive: dropping it closes the channel and the dispatcher exits, + // which would cause the timeout to succeed for the wrong reason. + let (_broadcast_tx, registry, _handle) = setup_watch_system_with_heartbeat(10, 0); let key = Bytes::from("/watch/key"); let mut watcher = registry.register(key.clone(), false).unwrap(); diff --git a/d-engine-core/src/watch/mod.rs b/d-engine-core/src/watch/mod.rs index 2d39ab6d..0998f2c0 100644 --- a/d-engine-core/src/watch/mod.rs +++ b/d-engine-core/src/watch/mod.rs @@ -141,7 +141,7 @@ pub(crate) fn make_cancel_event(key: bytes::Bytes) -> WatchEvent { WatchEvent { key, value: bytes::Bytes::new(), - prev_value: bytes::Bytes::new(), + prev_value: None, event_type: WatchEventType::Canceled, revision: 0, } diff --git a/d-engine-server/src/api/embedded_client.rs b/d-engine-server/src/api/embedded_client.rs index b4453b3f..b90c80aa 100644 --- a/d-engine-server/src/api/embedded_client.rs +++ b/d-engine-server/src/api/embedded_client.rs @@ -459,7 +459,6 @@ impl EmbeddedClient { /// println!("Value changed: {:?}", event); /// } /// ``` - /// `prev_kv`: when true each WatchEvent carries the value before the mutation. #[cfg(feature = "watch")] pub fn watch( &self, diff --git a/d-engine/src/docs/server_guide/watch-feature.md b/d-engine/src/docs/server_guide/watch-feature.md index 3b173006..46107dcc 100644 --- a/d-engine/src/docs/server_guide/watch-feature.md +++ b/d-engine/src/docs/server_guide/watch-feature.md @@ -205,7 +205,7 @@ let last_revision = event.revision; - **Order**: FIFO per key (events from StateMachine apply order) - **Dropped Events**: When a per-watcher buffer overflows, the watcher receives a `CANCELED` sentinel event and is forcibly unregistered - **Lagging**: Broadcast channel drops oldest events when receivers are slow (global buffer) -- **CANCELED Event**: `WatchEventType::Canceled` with `ErrorCode::WATCH_BUFFER_OVERFLOW` signals that the watcher was forcibly terminated. Client must re-sync via Read API and re-register. +- **CANCELED Event**: `WatchEventType::Canceled` signals that the watcher was forcibly terminated. Client must re-sync via Read API and re-register. In the gRPC API the event also carries `error == ErrorCode::WATCH_BUFFER_OVERFLOW`; in the embedded API the `error` field is `0` — check `event_type` alone. ### Lifecycle @@ -219,7 +219,7 @@ When registering a watcher with `prev_kv: true` (embedded API) or `prev_kv: true ```rust,ignore // Embedded: request previous value on every event -let watcher = engine.client().watch_with_prev_kv(b"my_key")?; +let watcher = engine.client().watch_with_options(b"my_key", true)?; while let Some(event) = watcher.receiver_mut().recv().await { if event.event_type == WatchEventType::Put { println!("was: {:?}, now: {:?}", event.prev_value, event.value); @@ -339,7 +339,7 @@ watch(b"config:feature_flags") ### `WATCH_BUFFER_OVERFLOW` (ErrorCode 5001) - **Cause**: Per-watcher channel full — the client is consuming events too slowly -- **Detection**: Receive an event where `event_type == WatchEventType::Canceled` and `error == ErrorCode::WATCH_BUFFER_OVERFLOW` +- **Detection**: Receive an event where `event_type == WatchEventType::Canceled`. In the gRPC API the event also sets `error == ErrorCode::WATCH_BUFFER_OVERFLOW`; in the embedded API check `event_type` alone (`error` is `0`). - **Effect**: The watcher is forcibly unregistered server-side; no further events will be delivered on this stream - **Solution**: 1. Re-sync state via the Read API to get the current value diff --git a/examples/service-discovery-embedded/server.rs b/examples/service-discovery-embedded/server.rs index 068cfae4..8f1adb3a 100644 --- a/examples/service-discovery-embedded/server.rs +++ b/examples/service-discovery-embedded/server.rs @@ -205,10 +205,9 @@ async fn main() -> Result<(), Box> { let key = String::from_utf8_lossy(&event.key); match event.event_type { WatchEventType::Put => { - let prev = if event.prev_value.is_empty() { - "(new key)".to_string() - } else { - String::from_utf8_lossy(&event.prev_value).to_string() + let prev = match event.prev_value.as_deref() { + None | Some(b"") => "(new key)".to_string(), + Some(v) => String::from_utf8_lossy(v).to_string(), }; println!( "[audit] {} : {} → {} (revision={})", @@ -218,12 +217,17 @@ async fn main() -> Result<(), Box> { event.revision ); } - WatchEventType::Delete => println!( - "[audit] {} deleted (was: {}) (revision={})", - key, - String::from_utf8_lossy(&event.prev_value), - event.revision - ), + WatchEventType::Delete => { + let was = event + .prev_value + .as_deref() + .map(|v| String::from_utf8_lossy(v).to_string()) + .unwrap_or_default(); + println!( + "[audit] {} deleted (was: {}) (revision={})", + key, was, event.revision + ); + } WatchEventType::Canceled => { println!("[audit] CANCELED — re-register watch"); break;